diff --git a/docs/kernel-contract.md b/docs/kernel-contract.md index 1ed43ac146..8f3a676fc4 100644 --- a/docs/kernel-contract.md +++ b/docs/kernel-contract.md @@ -215,6 +215,8 @@ shared memory, the device's limits and the planted controls are checked the same A restored 13.0 receipt can arrive after the later compiler sets. Such records mark `later_pending` so 13.0 and 13.1 use their base usage, while 13.2 and later keep their independently measured or pending selection. A missing later receipt remains pending; a 13.0 measurement never substitutes for it. +Flash Next's EXL3 launch (`flashnext_exl3_pending.zig`) carries only an nvcc 13.4 receipt: it stays pending under +13.0 to 13.3 and takes its 13.4 usage from `cuda_134`, which `pending_132.lookup` serves for that release. Gemma 4 CUDA tower QKV norms and nine-term pooling register FP32 sums and BF16 outputs. Their contract keeps the exact FP32 mirror and adds independent FP64 ranges at the BF16 nodes. diff --git a/tools/exl3_len_abba.py b/tools/exl3_len_abba.py new file mode 100644 index 0000000000..fe4cafce50 --- /dev/null +++ b/tools/exl3_len_abba.py @@ -0,0 +1,97 @@ +"""Concurrent two-request HTTP bench (repeated-ID prompts) used for the EXL3 ABBA runs: PORT IDS.json OUT.json LENGTH DRAFT.""" +import concurrent.futures +import json +import statistics +import sys +import time +import urllib.request + + +def one(port, prompt, tokens=128, draft=False): + body = { + "prompt": prompt, + "max_tokens": tokens, + "temperature": 0, + "ignore_eos": True, + "stream": True, + "stream_options": {"include_usage": True}, + "draft": draft, + } + request = urllib.request.Request( + f"http://127.0.0.1:{port}/v1/completions", + json.dumps(body).encode(), + {"Content-Type": "application/json"}, + ) + started = time.perf_counter() + first = last = None + usage = {} + runtime = {} + text_events = 0 + complete = False + with urllib.request.urlopen(request, timeout=1800) as response: + for raw in response: + line = raw.decode().strip() + if line == "data: [DONE]": + complete = True + break + if not line.startswith("data:"): + continue + event = json.loads(line[5:]) + usage = event.get("usage") or usage + runtime = event.get("tensorfold") or runtime + for choice in event.get("choices", []): + piece = choice.get("text") or (choice.get("delta") or {}).get("content") or "" + if piece: + now = time.perf_counter() + first = first or now + last = now + text_events += 1 + generated = int(usage.get("completion_tokens", 0)) + span = (last - first) if first is not None and last is not None else 0.0 + return { + "complete": complete, + "prompt_tokens": usage.get("prompt_tokens"), + "tokens": generated, + "ttft_s": (first - started) if first is not None else None, + "first_token_at": first, + "last_token_at": last, + "prefill_s": runtime.get("prefill_seconds"), + "decode_tps_server": runtime.get("tokens_per_second"), + "decode_tps_client": ((generated - 1) / span) if generated > 1 and span > 0 else None, + "itl_ms_client": (1000 * span / (generated - 1)) if generated > 1 and span > 0 else None, + "token_sha256": runtime.get("token_sha256"), + "text_events": text_events, + } + + +def batch(port, prompt, tokens, draft=False): + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + results = list(pool.map(lambda _: one(port, prompt, tokens, draft), range(2))) + end = time.perf_counter() + total = sum(row["tokens"] for row in results) + first_times = [row["first_token_at"] for row in results if row["first_token_at"] is not None] + last_times = [row["last_token_at"] for row in results if row["last_token_at"] is not None] + decode_span = max(last_times) - min(first_times) if first_times and last_times else 0.0 + return { + "requests": results, + "batch_wall_s": end - start, + "aggregate_decode_tps": total / decode_span if decode_span > 0 else None, + "end_to_end_aggregate_tps": total / (end - start) if end > start else None, + "mean_prefill_s": statistics.mean(row["prefill_s"] for row in results if row["prefill_s"] is not None), + } + + +if __name__ == "__main__": + port = int(sys.argv[1]) + with open(sys.argv[2], encoding="utf-8") as source: + base = json.load(source) + length = int(sys.argv[4]); draft = sys.argv[5] == "1" + ids = (base * (length // len(base) + 1))[:length] + report = {"prompt_tokens": len(ids), "tokens": 128, "drafts": draft, "batches": []} + for phase in ("warmup-1", "measure-1", "measure-2"): + result = batch(port, ids, int(__import__("os").environ.get("TOK", "128")), draft) + report["batches"].append({"phase": phase, **result}) + print(json.dumps(report["batches"][-1]), flush=True) + with open(sys.argv[3], "w", encoding="utf-8") as target: + json.dump(report, target, indent=2) diff --git a/tools/exl3_longctx_client.py b/tools/exl3_longctx_client.py new file mode 100644 index 0000000000..8ef9a845bb --- /dev/null +++ b/tools/exl3_longctx_client.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +"""Long-context checks against tensorfold-native serve (EXL3 Flash Next, docs/plans/exl3-kv-memory.md). + +usage: + exl3_longctx_client.py PORT needle CTX[,CTX..] CONC TEXT + A haystack of real text (the file TEXT) cut to about CTX prompt tokens (counted by the + server's /tokenize on the rendered chat), one needle sentence at a depth, a question at the end; greedy chat + reply (thinking off). CONC requests at once, each with its own needle, depth and text offset. Every request + runs drafts off, then all again drafts on: the answer must hold the needle's value and the drafted reply must + equal the undrafted one (token_sha256). + exl3_longctx_client.py PORT cont N[,N..] GEN IDS + The first N real-text ids (the little-endian u32 file IDS), GEN greedy tokens, drafts off: prints the reply's + sha and ids (for comparing two servers' continuations).""" +import hashlib, json, os, struct, sys, threading, time, urllib.request + +port, mode = int(sys.argv[1]), sys.argv[2] +base = f"http://127.0.0.1:{port}" + + +def post(path, body, timeout=7200): + req = urllib.request.Request(base + path, json.dumps(body).encode(), {"Content-Type": "application/json"}) + return json.load(urllib.request.urlopen(req, timeout=timeout)) + + +NEEDLES = [ + ("The secret code for the blue vault is 7391-4826.", "What is the secret code for the blue vault?", "7391-4826"), + ("The password for the red door is amber-heron-55.", "What is the password for the red door?", "amber-heron-55"), + ("The launch date written on the green folder is 14 March 2031.", "What launch date is written on the green folder?", "14 March 2031"), + ("The name of the night watchman at the north gate is Oskar Lindqvist.", "What is the name of the night watchman at the north gate?", "Lindqvist"), +] +DEPTHS = [0.5, 0.25, 0.75, 0.1] + + +def messages(hay, needle, question): + content = ("Read the following document carefully. Somewhere in it is one sentence that does not belong to it.\n\n" + + hay + "\n\nQuestion: " + question + " Answer with the exact value only.") + return [{"role": "user", "content": content}] + + +def build(text, ctx, i): + needle, question, answer = NEEDLES[i % len(NEEDLES)] + depth = DEPTHS[i % len(DEPTHS)] + off = (i * 97_000) % max(1, len(text) // 2) + chars = int(ctx * 3.3) + for _ in range(4): + hay = text[off:off + chars] + if len(hay) < chars: + hay = (text[off:] + "\n" + text)[:chars] + cut = int(len(hay) * depth) + cut = hay.rfind("\n", 0, cut) + 1 or cut + doc = hay[:cut] + needle + "\n" + hay[cut:] + msgs = messages(doc, needle, question) + n = post("/v1/tokenize", {"messages": msgs, "chat_template_kwargs": {"enable_thinking": False}})["count"] + if abs(n - ctx) <= max(64, ctx // 200) and n <= ctx: + break + chars = int(chars * (ctx - 48) / n) + return msgs, answer, n + + +def chat(msgs, draft, max_tokens=32): + body = {"messages": msgs, "max_tokens": max_tokens, "temperature": 0, "chat_template_kwargs": {"enable_thinking": False}} + if not draft: + body["draft"] = False + t0 = time.time() + r = post("/v1/chat/completions", body) + r["_wall"] = time.time() - t0 + return r + + +def run_all(reqs, draft): + out = [None] * len(reqs) + + def one(k): + try: + out[k] = chat(reqs[k][0], draft) + except Exception as e: # noqa: BLE001 - a failed request is reported, not raised + out[k] = {"error": repr(e)} + + ts = [threading.Thread(target=one, args=(k,)) for k in range(len(reqs))] + t0 = time.time() + for t in ts: + t.start() + for t in ts: + t.join() + return out, time.time() - t0 + + +if mode == "needle": + ctxs = [int(x) for x in sys.argv[3].split(",")] + conc = int(sys.argv[4]) + text = open(sys.argv[5], encoding="utf-8", errors="replace").read() + all_ok = True + for ctx in ctxs: + reqs = [build(text, ctx, i) for i in range(conc)] + res = {} + for draft in (False, True): + out, wall = run_all(reqs, draft) + res[draft] = out + for k, r in enumerate(out): + if "error" in r: + print(f"NEEDLE ctx={ctx} req={k} drafts={'on' if draft else 'off'}: ERROR {r['error']}", flush=True) + all_ok = False + continue + tf, sp = r.get("tensorfold", {}), r.get("speculative", {}) + ans = r["choices"][0]["message"]["content"] + ok = reqs[k][1].lower() in ans.lower() + all_ok &= ok + print(f"NEEDLE ctx={ctx} req={k} prompt_tokens={r['usage']['prompt_tokens']} drafts={'on' if draft else 'off'}: " + f"{'CORRECT' if ok else 'WRONG'} answer={ans.strip()[:80]!r} want={reqs[k][1]!r} prefill {tf.get('prefill_seconds', 0):.1f} s " + f"decode {tf.get('tokens_per_second', 0):.1f} tok/s accepted {sp.get('accepted')}/{sp.get('drafted')} " + f"sha {tf.get('token_sha256', '')[:12]} wall {r['_wall']:.1f} s", flush=True) + print(f"NEEDLE ctx={ctx} conc={conc} drafts={'on' if draft else 'off'}: all requests done in {wall:.1f} s", flush=True) + for k in range(conc): + a, b = res[False][k], res[True][k] + same = "error" not in a and "error" not in b and a["tensorfold"]["token_sha256"] == b["tensorfold"]["token_sha256"] + all_ok &= same + print(f"NEEDLE ctx={ctx} req={k}: drafted {'==' if same else '!='} undrafted", flush=True) + print("NEEDLE ALL OK" if all_ok else "NEEDLE NOT ALL OK", flush=True) +elif mode == "cont": + ns = [int(x) for x in sys.argv[3].split(",")] + gen = int(sys.argv[4]) + raw = open(sys.argv[5], "rb").read() + ids = list(struct.unpack("<%dI" % (len(raw) // 4), raw)) + for n in ns: + body = {"prompt": ids[:n], "max_tokens": gen, "temperature": 0, "ignore_eos": True, "draft": False} + r = post("/v1/completions", body) + tf = r["tensorfold"] + print(f"CONT n={n} gen={gen}: sha {tf['token_sha256'][:16]} prefill {tf['prefill_seconds']:.1f} s decode {tf['tokens_per_second']:.1f} tok/s " + f"text {r['choices'][0]['text'][:100]!r}", flush=True) +else: + sys.exit(__doc__) diff --git a/tools/exl3_pack_sizes.py b/tools/exl3_pack_sizes.py new file mode 100644 index 0000000000..ccaadbccb0 --- /dev/null +++ b/tools/exl3_pack_sizes.py @@ -0,0 +1,26 @@ +"""Group an EXL3 pack's safetensors tensors by role and print their sizes (reads headers only). usage: exl3_pack_sizes.py PACK_DIR.""" +import json, struct, glob, collections, sys +P = sys.argv[1] # the pack directory +g = collections.Counter() +def cat(n): + for k in ("ngram_embedding", "ple_", ".ple."): + if k in n and "ngram" in n: return "ple n-gram embedding shards" + if ".ple." in n: return "ple other (projections)" + if "mtp." in n: return "mtp (head)" + if "experts." in n and "shared" not in n: return "routed experts" + if "shared_expert" in n: return "shared expert" + if "hyper_connection" in n: return "hyper-connection mixes" + if "linear_attn" in n: return "linear (DeltaNet) attention" + if "self_attn" in n: return "full attention" + if "embed_tokens" in n: return "embed_tokens" + if "lm_head" in n: return "lm_head" + if "router" in n or "gate." in n: return "router" + return "other: " + n.split(".")[-2] if "." in n else "other" +for f in sorted(glob.glob(P + "/*.safetensors")): + with open(f, "rb") as h: + n = struct.unpack(" 1 else 8080 +TOK = int(os.environ.get("TOK", "192")) +root = os.environ.get("TF_SRC") or os.path.dirname(os.path.dirname(os.path.abspath(__file__))) # the checkout whose sources are the corpus +a = t.corpus([root + "/zig/src/families/flashnext/*.zig", root + "/zig/src/core/**/*.zig"], 50000) +b = t.corpus([root + "/docs/**/*.md", root + "/*.md", root + "/zig/tests/**/*.zig"], 50000) + + +def one(prompt): + body = { + "model": os.environ.get("MODEL", "qwen3.8-flash-next-iq3_xxs"), + "messages": [{"role": "user", "content": ("%s\n%s" % (uuid.uuid4().hex if not os.environ.get("NOUUID") else "fixed", prompt))}], + "max_tokens": TOK, "temperature": 0, "ignore_eos": True, "stream": True, + "stream_options": {"include_usage": True}, + "chat_template_kwargs": {"enable_thinking": False}, + } + req = urllib.request.Request("http://127.0.0.1:%d/v1/chat/completions" % PORT, json.dumps(body).encode(), {"Content-Type": "application/json"}) + start = time.perf_counter(); first = last = None; usage = {}; pieces = 0; digest = hashlib.sha256() + with urllib.request.urlopen(req, timeout=1800) as r: + for raw in r: + line = raw.decode().strip() + if line == "data: [DONE]": + break + if not line.startswith("data:"): + continue + ev = json.loads(line[5:]) + usage = ev.get("usage") or usage + for c in ev.get("choices", []): + d = c.get("delta") or {} + if d.get("content") or d.get("reasoning_content"): + now = time.perf_counter(); first = first or now; last = now; pieces += 1; digest.update((d.get("content") or d.get("reasoning_content")).encode()) + n = int(usage.get("completion_tokens", 0)) + span = (last - first) if first and last else 0 + return {"prompt_tokens": usage.get("prompt_tokens"), "completion_tokens": n, + "ttft_s": round(first - start, 2) if first else None, + "decode_tps": round((n - 1) / span, 2) if n > 1 and span > 0 else None, + "text_sha": digest.hexdigest()[:12]} + + +short = "Write a long, detailed story about a lighthouse keeper who finds a message in a bottle. Keep writing until told to stop." +for name, prompt in (("short", short), ("zig", a), ("docs", b)): + for rep in range(3): + r = one(prompt); r["prompt"] = name; r["rep"] = rep + print(json.dumps(r), flush=True) diff --git a/tools/exl3_snapshot_gate.py b/tools/exl3_snapshot_gate.py new file mode 100644 index 0000000000..eb1ef83481 --- /dev/null +++ b/tools/exl3_snapshot_gate.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Strict two-lane serving gate for BF16 paged prompt-cache restore.""" + +import argparse +import concurrent.futures +import json +import threading +import time +import urllib.request +from pathlib import Path + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("url") + parser.add_argument("--tokens", type=int, default=4099) + parser.add_argument("--decode", type=int, default=32) + parser.add_argument("--drafts", action="store_true") + parser.add_argument("--reference", type=Path) + parser.add_argument("--record-reference", action="store_true", help="record a cold baseline; use a fresh server with prompt cache disabled") + parser.add_argument("--output", type=Path) + args = parser.parse_args(argv) + if args.tokens < 32 or args.decode < 1: + parser.error("require at least 32 prompt tokens and one decode token") + if args.record_reference and not args.output: + parser.error("--record-reference requires --output") + if args.record_reference and args.reference: + parser.error("--record-reference and --reference are mutually exclusive") + shared = " cache" * max(args.tokens, 1) + history = [{"role": "user", "content": shared}, {"role": "assistant", "content": "I have reviewed the shared document."}] + prompts = [history + [{"role": "user", "content": f"Original ending for lane {lane}."}] for lane in range(2)] + reference = json.loads(args.reference.read_text()) if args.reference else None + report = {"prompt_tokens": args.tokens, "decode_tokens": args.decode, "drafts": args.drafts, "phases": {}} + if reference: + for field in ("prompt_tokens", "decode_tokens", "drafts"): + if reference[field] != report[field]: + raise AssertionError(f"reference configuration mismatch: {field}") + + def pair(request_prompts): + barrier = threading.Barrier(2) + + def ask(ids): + started = time.perf_counter() + body = {"messages": ids, "max_tokens": args.decode, "temperature": 0, "ignore_eos": True, "draft": args.drafts} + request = urllib.request.Request(args.url.rstrip("/") + "/v1/chat/completions", json.dumps(body).encode(), {"Content-Type": "application/json"}) + barrier.wait(timeout=10) + with urllib.request.urlopen(request, timeout=900) as reply: + response = json.load(reply) + usage = response["usage"] + if usage["prompt_tokens"] < args.tokens or usage["completion_tokens"] != args.decode: + raise AssertionError(f"inexact token counts: {usage}") + runtime = response["tensorfold"] + digest = runtime.get("token_sha256") or runtime["token_sha"] + if len(digest) < 12: + raise AssertionError("missing output token digest") + return { + "sha": digest, + "cached_tokens": usage["prompt_tokens_details"]["cached_tokens"], + "prompt_tokens": usage["prompt_tokens"], + "elapsed_seconds": time.perf_counter() - started, + "prefill_seconds": runtime.get("prefill_seconds"), + "tokens_per_second": runtime.get("tokens_per_second"), + } + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + return list(pool.map(ask, request_prompts)) + + branches = [history + [{"role": "user", "content": f"Edited ending for branch {lane}."}] for lane in range(2)] + for phase, ids in (("cold", prompts), ("warm", prompts), ("branch", branches), ("rewind", prompts)): + results = pair(ids) + for lane, result in enumerate(results): + if args.record_reference and result["cached_tokens"] != 0: + raise AssertionError(f"reference phase {phase} lane {lane} was not cold: {result}") + if not args.record_reference and phase != "cold" and not 0 < result["cached_tokens"] < result["prompt_tokens"]: + raise AssertionError(f"{phase} lane {lane} did not restore a cached prefix: {result}") + if phase in ("warm", "rewind") and result["sha"] != report["phases"]["cold"][lane]["sha"]: + raise AssertionError(f"{phase} lane {lane} differs from its cold continuation") + if reference and result["sha"] != reference["phases"][phase][lane]["sha"]: + raise AssertionError(f"{phase} lane {lane} differs from reference continuation") + report["phases"][phase] = results + print(json.dumps({"phase": phase, "lanes": results}), flush=True) + if args.output: + args.output.write_text(json.dumps(report, indent=2) + "\n") + if args.record_reference: + print("PASS: recorded no-cache cold continuations for cold, warm, branch and rewind inputs") + else: + print("PASS: exact prompt/reply lengths, two submitted lanes, warm/branch/rewind cache hits and cold-baseline parity") + + +if __name__ == "__main__": + main() diff --git a/tools/exl3_text_abba.py b/tools/exl3_text_abba.py new file mode 100644 index 0000000000..ff773c5405 --- /dev/null +++ b/tools/exl3_text_abba.py @@ -0,0 +1,45 @@ +"""Concurrent two-request bench with two different real-text prompts built from a checkout: PORT OUT.json CHARS.""" +import concurrent.futures, glob, json, os, statistics, sys, time, urllib.request +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import exl3_len_abba as base + + +def corpus(patterns, chars): + out, total = [], 0 + for pattern in patterns: + for path in sorted(glob.glob(os.path.expanduser(pattern), recursive=True)): + try: + text = open(path, encoding="utf-8").read() + except Exception: + continue + out.append("// file: %s\n%s\n" % (os.path.basename(path), text)) + total += len(out[-1]) + if total >= chars: + return "".join(out)[:chars] + return "".join(out)[:chars] + + +def post(port, prompt, tokens): + return base.one(port, prompt, tokens, False) + + +if __name__ == "__main__": + port, out_path, chars, tokens = int(sys.argv[1]), sys.argv[2], int(sys.argv[3]), int(os.environ.get("TOK", "128")) + root = os.environ.get("TF_SRC") or os.path.dirname(os.path.dirname(os.path.abspath(__file__))) # the checkout whose sources are the corpus + a = corpus([root + "/zig/src/families/flashnext/*.zig", root + "/zig/src/core/**/*.zig"], chars) + b = corpus([root + "/docs/**/*.md", root + "/*.md", root + "/zig/tests/**/*.zig"], chars) + prompts = [a, b] + report = {"chars": [len(a), len(b)], "batches": []} + for phase in ("warmup-1", "measure-1", "measure-2"): + start = time.perf_counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + rows = list(pool.map(lambda p: post(port, p, tokens), prompts)) + batch = { + "phase": phase, + "requests": rows, + "mean_prefill_s": statistics.mean(r["prefill_s"] for r in rows), + "batch_wall_s": time.perf_counter() - start, + } + report["batches"].append(batch) + print(json.dumps(batch), flush=True) + json.dump(report, open(out_path, "w"), indent=2) diff --git a/tools/tests/test_exl3_snapshot_gate.py b/tools/tests/test_exl3_snapshot_gate.py new file mode 100644 index 0000000000..7363c7c345 --- /dev/null +++ b/tools/tests/test_exl3_snapshot_gate.py @@ -0,0 +1,86 @@ +"""Model-free checks for the strict chat-history snapshot gate.""" + +from contextlib import redirect_stderr, redirect_stdout +import hashlib +import importlib.util +import io +import json +from pathlib import Path +import tempfile +import unittest +from unittest.mock import patch + + +SCRIPT = Path(__file__).resolve().parents[1] / "exl3_snapshot_gate.py" +SPEC = importlib.util.spec_from_file_location("exl3_snapshot_gate", SCRIPT) +gate = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(gate) + + +class Response: + def __init__(self, payload): + self.payload = payload + + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def read(self): + return json.dumps(self.payload).encode() + + +class SnapshotGateTests(unittest.TestCase): + def invoke(self, args, cached): + requests = [] + + def open_request(request, timeout): + body = json.loads(request.data) + requests.append((request.full_url, body)) + digest = hashlib.sha256(json.dumps(body["messages"]).encode()).hexdigest() + return Response({ + "usage": {"prompt_tokens": 80, "completion_tokens": 2, + "prompt_tokens_details": {"cached_tokens": cached}}, + "tensorfold": {"token_sha256": digest}, + }) + + output, errors = io.StringIO(), io.StringIO() + with patch.object(gate.urllib.request, "urlopen", open_request), redirect_stdout(output), redirect_stderr(errors): + try: + gate.main(["http://example.invalid", "--tokens", "40", "--decode", "2", *args]) + code = 0 + except AssertionError as error: + errors.write(str(error)) + code = 1 + return code, output.getvalue(), errors.getvalue(), requests + + def test_records_cold_reference_then_requires_exact_chat_cache_reuse(self): + with tempfile.TemporaryDirectory() as directory: + reference = Path(directory) / "reference.json" + receipt = Path(directory) / "receipt.json" + code, _, _, cold_requests = self.invoke(["--record-reference", "--output", str(reference)], 0) + self.assertEqual(code, 0) + self.assertEqual(len(cold_requests), 8) + self.assertTrue(all(url.endswith("/v1/chat/completions") for url, _ in cold_requests)) + self.assertTrue(all(len(body["messages"]) == 3 for _, body in cold_requests)) + code, output, _, warm_requests = self.invoke(["--reference", str(reference), "--output", str(receipt)], 40) + self.assertEqual(code, 0) + self.assertEqual(len(warm_requests), 8) + self.assertIn("PASS:", output) + saved = json.loads(receipt.read_text()) + for phase in ("warm", "branch", "rewind"): + self.assertTrue(all(0 < lane["cached_tokens"] < lane["prompt_tokens"] + for lane in saved["phases"][phase])) + + def test_missing_warm_cache_hit_fails(self): + with tempfile.TemporaryDirectory() as directory: + reference = Path(directory) / "reference.json" + self.invoke(["--record-reference", "--output", str(reference)], 0) + code, output, errors, _ = self.invoke(["--reference", str(reference)], 0) + self.assertEqual(code, 1) + self.assertIn("did not restore a cached prefix", errors) + + +if __name__ == "__main__": + unittest.main() diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 054d85922b..e53e74b715 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -433,7 +433,8 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st golden_build.step(b, host, server_tests, native_modules.api, mods.tokenizer, template, native_modules.engines); @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } - // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them + // the tokenizer's own tests run here: no other test binary reaches them; the EXL3 hooks alone run as "exl3: ..." + b.step("test-flashnext-exl3-host", "Flash Next EXL3 admission, store, sizing and loader host tests").dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision, mods.api), .filters = &.{"exl3: "} })).step); for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.qwen_vision, flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision, mods.api), lfm2Module(b, host, .debug, cuda, mods.core), mods.radio, mods.heat, mods.tokenizer, native_modules.api, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const selected = b.option([]const u8, "modelopt-host-filter", "One portable ModelOpt regression") orelse "modelopt host:"; const run = b.addRunArtifact(b.addTest(.{ .root_module = mods.nemotron, .filters = &.{selected} })); diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 887b2d84ba..3dc4c2a816 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -120,6 +120,7 @@ pub const kernels = [_]Kernel{ .{ .name = "drama3", .flags = &.{"-O3"} }, .{ .name = "drama3_fp8", .flags = &.{"-O3"} }, .{ .name = "drama3_codec", .flags = &.{"-O3"} }, + .{ .name = "flashnext_exl3", .src = "flashnext_exl3/flashnext_exl3", .flags = glue }, // ours: Flash Next's EXL3 glue (an fp16 row split exactly into two bf16) }; /// torch.utils.cpp_extension's own nvcc flags (torch 2.13): C++20 and which half/bf16 operators the headers define. diff --git a/zig/kernels/cuda/flashnext_exl3/flashnext_exl3.cu b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3.cu new file mode 100644 index 0000000000..ae58b291de --- /dev/null +++ b/zig/kernels/cuda/flashnext_exl3/flashnext_exl3.cu @@ -0,0 +1,15 @@ +// The Flash Next family's EXL3 glue (cuda_exl3_pass.zig): an fp16 row as two bf16 rows whose sum is it exactly. + +#include +#include + +// hi = bf16(x), lo = bf16(x - hi): an fp16 has 11 significant bits, hi keeps 8 and the remainder has at most 3. +extern "C" __global__ void fn_exl3_split_f16(const __half* __restrict__ x, __nv_bfloat16* __restrict__ hi, + __nv_bfloat16* __restrict__ lo, long long n) { + const long long i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) return; + const float v = __half2float(x[i]); + const __nv_bfloat16 h = __float2bfloat16_rn(v); + hi[i] = h; + lo[i] = __float2bfloat16_rn(v - __bfloat162float(h)); +} diff --git a/zig/src/core/contract/flashnext_exl3_pending.zig b/zig/src/core/contract/flashnext_exl3_pending.zig new file mode 100644 index 0000000000..e8d0f82fec --- /dev/null +++ b/zig/src/core/contract/flashnext_exl3_pending.zig @@ -0,0 +1,33 @@ +//! Flash Next EXL3 launches (the fp16 row split); nvcc 13.4 measured, older pending. +const std = @import("std"); +const limits = @import("limits.zig"); +const registry = @import("registry.zig"); + +/// Launches with no nvcc 13.0 to 13.3 receipt yet: the geometry is recorded and usage stays pending until measured. +pub const records = [_]registry.CudaRecord{ + .{ .launch = .{ .name = "flashnext_exl3 fn_exl3_split_f16", .symbol = "fn_exl3_split_f16", .image = "flashnext_exl3", .threads = 256 }, .usage = registry.none }, +}; + +/// nvcc 13.4's measured registers, static shared bytes and launch bounds for sm_80, sm_86, sm_89, sm_120, sm_121. +pub const cuda_134 = [_]registry.Override{ + .{ .name = "flashnext_exl3 fn_exl3_split_f16", .usage = .{ .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 10, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 }, .{ .regs = 12, .static = 0, .bound = 0 } } }, +}; + +test "Flash Next EXL3 launches: pending below nvcc 13.4, measured 13.4 resources fit every SM launch limit" { + const t = std.testing; + try t.expectEqual(records.len, cuda_134.len); + for (records, cuda_134) |r, o| { + try t.expectEqualStrings(r.launch.name, o.name); + try t.expectEqual(limits.sms.len, try registry.checkCuda(t.allocator, &.{r}, true)); + const measured = registry.usageFor(r, .{ .major = 13, .minor = 4 }); + try t.expect(std.meta.eql(measured, o.usage)); + var current = r; + current.usage = measured; + try t.expectEqual(@as(usize, 0), try registry.checkCuda(t.allocator, &.{current}, true)); + for (measured) |u| { + try t.expect(u.regs > 0); + try t.expect(u.bound == 0 or u.bound == r.launch.threads); + } + for (registry.cuda) |known| try t.expect(!std.mem.eql(u8, r.launch.symbol, known.launch.symbol)); + } +} diff --git a/zig/src/core/contract/pending_132.zig b/zig/src/core/contract/pending_132.zig index 09f0d0023a..11e92b7807 100644 --- a/zig/src/core/contract/pending_132.zig +++ b/zig/src/core/contract/pending_132.zig @@ -8,6 +8,10 @@ pub fn lookup(name: []const u8, release: limits.Nvcc) ?[limits.sms.len]limits.Us for (@import("measurements/qwen_lanes_133.zig").records ++ @import("measurements/pending_133.zig").records) |r| if (std.mem.eql(u8, r.name, name)) return r.usage; } + if (release.major == 13 and release.minor == 4) { + for (@import("flashnext_exl3_pending.zig").cuda_134) |r| + if (std.mem.eql(u8, r.name, name)) return r.usage; + } if (release.major < 13 or (release.major == 13 and release.minor < 2)) return null; for (records) |r| if (std.mem.eql(u8, r.name, name)) return r.usage; return null; diff --git a/zig/src/core/contract/registry.zig b/zig/src/core/contract/registry.zig index 7d83d71249..493c9d5a02 100644 --- a/zig/src/core/contract/registry.zig +++ b/zig/src/core/contract/registry.zig @@ -324,6 +324,7 @@ pub const cuda_sets = [_]CudaSet{ .{ .name = "router tf_nemo_topk", .usage = .{ .{ .regs = 94, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 95, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 }, .{ .regs = 96, .static = 0, .bound = 32 } } }, .{ .name = "flashnext_batch fn_qsa_select", .usage = .{ .{ .regs = 32, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 50, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 }, .{ .regs = 44, .static = 1168, .bound = 1024 } } }, } ++ @import("flashnext_tile_records.zig").cuda_133 ++ @import("glm_cuda_records.zig").cuda_133 ++ @import("nvcc_133_records.zig").cuda_133) }, + .{ .release = .{ .major = 13, .minor = 4 }, .usage = &.{.{ .name = "exl3_candidate tf_exl3_prompt_2", .usage = .{ .{ .regs = 80, .static = 32768, .bound = 256 }, .{ .regs = 80, .static = 32768, .bound = 256 }, .{ .regs = 80, .static = 32768, .bound = 256 }, .{ .regs = 77, .static = 32768, .bound = 256 }, .{ .regs = 77, .static = 32768, .bound = 256 } } }} }, }; pub fn newer(x: limits.Nvcc, y: limits.Nvcc) bool { return x.major > y.major or (x.major == y.major and x.minor > y.minor); @@ -578,7 +579,7 @@ const learning = @import("registry_learning.zig"); pub const gradient_pending = learning.gradient_pending; pub const qwen_learning_pending = learning.qwen_learning_pending; pub const glm_pending = @import("glm_sampling_pending.zig"); -pub const landed_cuda = cuda ++ @import("dory_cuda_records.zig").cuda ++ gradient_pending ++ qwen_learning_pending ++ paging.extra ++ @import("qwen_served_records.zig").records ++ cuda_pending ++ @import("qwen_decode_pending.zig").records ++ glm_pending.sampling ++ @import("glm_kv_pending.zig").records ++ @import("proposal_sampling_records.zig").records ++ @import("flashnext_pages_pending.zig").records ++ @import("glm_f32w_pending.zig").records ++ @import("penalty_sampling_records.zig").records; +pub const landed_cuda = cuda ++ @import("dory_cuda_records.zig").cuda ++ gradient_pending ++ qwen_learning_pending ++ paging.extra ++ @import("qwen_served_records.zig").records ++ cuda_pending ++ @import("qwen_decode_pending.zig").records ++ glm_pending.sampling ++ @import("glm_kv_pending.zig").records ++ @import("proposal_sampling_records.zig").records ++ @import("flashnext_pages_pending.zig").records ++ @import("flashnext_exl3_pending.zig").records ++ @import("glm_f32w_pending.zig").records ++ @import("penalty_sampling_records.zig").records; test { _ = learning; } diff --git a/zig/src/core/contract/registry_pending_test.zig b/zig/src/core/contract/registry_pending_test.zig index c708e15bbc..2453c70030 100644 --- a/zig/src/core/contract/registry_pending_test.zig +++ b/zig/src/core/contract/registry_pending_test.zig @@ -9,6 +9,23 @@ const cuda = registry.cuda; const checkCuda = registry.checkCuda; const targetDecision = registry.targetDecision; const TargetDecision = registry.TargetDecision; + +test { + _ = @import("flashnext_exl3_pending.zig"); +} +test "EXL3 prompt resource records select measured CUDA 13.4 usage without changing older compiler records" { + const record = for (cuda) |candidate| { + if (std.mem.eql(u8, candidate.launch.name, "exl3_candidate tf_exl3_prompt_2")) break candidate; + } else return error.MissingExl3PromptRecord; + const older = registry.usageFor(record, .{ .major = 13, .minor = 3 }); + const current = registry.usageFor(record, .{ .major = 13, .minor = 4 }); + for (current, 0..) |usage, index| { + try std.testing.expectEqual(@as(u32, if (index < 3) 80 else 77), usage.regs); + try std.testing.expectEqual(@as(u32, if (index < 3) 80 else 79), older[index].regs); + try std.testing.expectEqual(@as(u32, 32768), usage.static); + try std.testing.expectEqual(@as(u32, 256), usage.bound); + } +} test "CUDA lane resources distinguish recorded launches from missing receipts" { for (cuda_pending) |r| { try std.testing.expect(r.launch.threads > 0); diff --git a/zig/src/core/contract/registry_sets_test.zig b/zig/src/core/contract/registry_sets_test.zig index a494f045cb..d964985766 100644 --- a/zig/src/core/contract/registry_sets_test.zig +++ b/zig/src/core/contract/registry_sets_test.zig @@ -39,7 +39,8 @@ test "a compiler's records: its own set, else the newest below it, else 13.0's" .{ .{ .major = 13, .minor = 1 }, .{ .major = 13, .minor = 0 }, 91, 45 }, .{ .{ .major = 13, .minor = 2 }, .{ .major = 13, .minor = 2 }, 91, 43 }, .{ .{ .major = 13, .minor = 3 }, .{ .major = 13, .minor = 3 }, 96, 43 }, - .{ .{ .major = 14, .minor = 0 }, .{ .major = 13, .minor = 3 }, 96, 43 }, + .{ .{ .major = 13, .minor = 4 }, .{ .major = 13, .minor = 4 }, 96, 43 }, + .{ .{ .major = 14, .minor = 0 }, .{ .major = 13, .minor = 4 }, 96, 43 }, }; for (cases) |c| { try std.testing.expectEqual(c[1], setFor(c[0])); diff --git a/zig/src/cuda/exl3_dma.zig b/zig/src/cuda/exl3_dma.zig new file mode 100644 index 0000000000..f2b78ffb6c --- /dev/null +++ b/zig/src/cuda/exl3_dma.zig @@ -0,0 +1,38 @@ +const std = @import("std"); +const Stream = @import("stream.zig").Stream; + +pub fn directLines(stream: Stream, dst: u64, capacity: usize, offsets: [3]usize, src: [3][]const u8, pinned_bank: ?[]const u8) !usize { + const Span = struct { dst: u64, bytes: []const u8, banked: bool }; + var spans: [3]Span = undefined; + var count: usize = 0; + var moved: usize = 0; + _ = std.math.add(u64, dst, capacity) catch return error.Invalid; + const bank_start = if (pinned_bank) |bank| @intFromPtr(bank.ptr) else 0; + const bank_end = if (pinned_bank) |bank| std.math.add(usize, bank_start, bank.len) catch return error.Invalid else 0; + for (src, offsets) |bytes, offset| { + if (offset > capacity or bytes.len > capacity - offset) return error.Invalid; + const start = @intFromPtr(bytes.ptr); + const end = std.math.add(usize, start, bytes.len) catch return error.Invalid; + moved = std.math.add(usize, moved, bytes.len) catch return error.Invalid; + const address = dst + offset; + const banked = pinned_bank != null and start >= bank_start and end <= bank_end; + if (count > 0 and bytes.len > 0 and banked) { + const previous = &spans[count - 1]; + if (previous.banked and previous.bytes.len > 0 and + @intFromPtr(previous.bytes.ptr) + previous.bytes.len == start and + previous.dst + previous.bytes.len == address) + { + previous.bytes = pinned_bank.?[@intFromPtr(previous.bytes.ptr) - bank_start ..][0 .. previous.bytes.len + bytes.len]; + continue; + } + } + spans[count] = .{ .dst = address, .bytes = bytes, .banked = banked }; + count += 1; + } + for (spans[0..count]) |span| try stream.d.check(stream.d.api.cuMemcpyHtoDAsync_v2(span.dst, span.bytes.ptr, span.bytes.len, stream.handle), "cuMemcpyHtoDAsync"); + return moved; +} + +test { + _ = @import("exl3_dma_test.zig"); +} diff --git a/zig/src/cuda/exl3_dma_test.zig b/zig/src/cuda/exl3_dma_test.zig new file mode 100644 index 0000000000..8cbafd61cb --- /dev/null +++ b/zig/src/cuda/exl3_dma_test.zig @@ -0,0 +1,216 @@ +const std = @import("std"); +const dma = @import("exl3_dma.zig"); +const pool = @import("exl3_pool.zig"); +const abi = @import("abi.zig"); +const Driver = @import("driver.zig").Driver; +const Stream = @import("stream.zig").Stream; +const Event = @import("stream.zig").Event; + +const Fake = struct { + const Call = enum { query, wait, copy, table, event }; + const Copy = struct { dst: u64, src: usize, len: usize, stream: abi.Stream }; + const Word = struct { dst: u64, value: u32 }; + var calls: [32]Call = undefined; + var count: usize = 0; + var copies: [3]Copy = undefined; + var copied: usize = 0; + var fail_copy: usize = 0; + var words: [18]Word = undefined; + var written: usize = 0; + + fn note(call: Call) void { + calls[count] = call; + count += 1; + } + + fn copy(dst: u64, src: ?*const anyopaque, len: usize, stream: abi.Stream) callconv(.c) abi.Result { + note(.copy); + copies[copied] = .{ .dst = dst, .src = @intFromPtr(src), .len = len, .stream = stream }; + copied += 1; + return if (copied == fail_copy) abi.error_not_ready else abi.success; + } + + fn table(dst: u64, value: c_uint, len: usize, stream: abi.Stream) callconv(.c) abi.Result { + std.debug.assert(len == 1 and stream == @as(abi.Stream, @ptrFromInt(7))); + note(.table); + words[written] = .{ .dst = dst, .value = value }; + written += 1; + return abi.success; + } + + fn event(_: abi.Event, _: abi.Stream) callconv(.c) abi.Result { + note(.event); + return abi.success; + } + + fn query(_: abi.Event) callconv(.c) abi.Result { + note(.query); + return abi.error_not_ready; + } + + fn wait(_: abi.Stream, _: abi.Event, _: c_uint) callconv(.c) abi.Result { + note(.wait); + return abi.success; + } + + fn driver() Driver { + count = 0; + copied = 0; + fail_copy = 0; + written = 0; + var result: Driver = undefined; + result.api.cuMemcpyHtoDAsync_v2 = copy; + result.api.cuMemsetD32Async = table; + result.api.cuEventRecord = event; + result.api.cuEventQuery = query; + result.api.cuStreamWaitEvent = wait; + return result; + } +}; + +fn run(src: [3][]const u8, offsets: [3]usize, bank: ?[]const u8, lengths: []const usize) !void { + var driver = Fake.driver(); + const stream: Stream = .{ .d = &driver, .handle = @ptrFromInt(7) }; + const moved = try dma.directLines(stream, 4096, 32, offsets, src, bank); + try std.testing.expectEqual(src[0].len + src[1].len + src[2].len, moved); + try std.testing.expectEqual(lengths.len, Fake.copied); + var projection: usize = 0; + for (Fake.copies[0..Fake.copied], lengths) |copy, len| { + try std.testing.expectEqual(len, copy.len); + try std.testing.expectEqual(4096 + offsets[projection], copy.dst); + try std.testing.expectEqual(@intFromPtr(src[projection].ptr), copy.src); + try std.testing.expectEqual(stream.handle, copy.stream); + var consumed: usize = 0; + while (projection < 3) { + const bytes = src[projection]; + const uploaded: [*]const u8 = @ptrFromInt(copy.src); + try std.testing.expectEqualSlices(u8, bytes, uploaded[consumed..][0..bytes.len]); + consumed += bytes.len; + projection += 1; + if (consumed == len) break; + } + try std.testing.expectEqual(len, consumed); + } + try std.testing.expectEqual(@as(usize, 3), projection); +} + +test "DMA coalescing: exact unequal widths merge three projections or either pair" { + const bank = [_]u8{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 }; + const src: [3][]const u8 = .{ bank[0..2], bank[2..5], bank[5..9] }; + try run(src, .{ 0, 2, 5 }, &bank, &.{9}); + try run(src, .{ 0, 2, 8 }, &bank, &.{ 5, 4 }); + try run(src, .{ 0, 4, 7 }, &bank, &.{ 2, 7 }); + try run(.{ bank[0..2], bank[2..5], bank[6..10] }, .{ 0, 2, 5 }, &bank, &.{ 5, 4 }); + try run(.{ bank[0..2], bank[3..6], bank[6..10] }, .{ 0, 2, 5 }, &bank, &.{ 2, 7 }); +} + +test "DMA coalescing: absent provenance and source or destination gaps retain exact fallback copies" { + const bank: [32]u8 = @splat(1); + const src: [3][]const u8 = .{ bank[0..2], bank[2..5], bank[5..9] }; + try run(src, .{ 0, 2, 5 }, null, &.{ 2, 3, 4 }); + try run(src, .{ 0, 4, 9 }, &bank, &.{ 2, 3, 4 }); + try run(.{ bank[0..2], bank[4..7], bank[9..13] }, .{ 0, 2, 5 }, &bank, &.{ 2, 3, 4 }); + try run(src, .{ 0, 2, 5 }, bank[10..20], &.{ 2, 3, 4 }); + try run(src, .{ 0, 2, 5 }, bank[0..4], &.{ 2, 3, 4 }); + try run(src, .{ 0, 2, 5 }, bank[2..9], &.{ 2, 7 }); + try run(src, .{ 0, 2, 5 }, bank[0..0], &.{ 2, 3, 4 }); + const foreign: [9]u8 = @splat(3); + try run(.{ foreign[0..2], foreign[2..5], foreign[5..9] }, .{ 0, 2, 5 }, &bank, &.{ 2, 3, 4 }); +} + +test "DMA coalescing: zero lengths preserve fallback and never bridge a gap" { + const bank: [9]u8 = @splat(2); + try run(.{ bank[0..0], bank[0..3], bank[3..9] }, .{ 0, 0, 3 }, &bank, &.{ 0, 9 }); + try run(.{ bank[0..2], bank[2..2], bank[2..9] }, .{ 0, 2, 2 }, &bank, &.{ 2, 0, 7 }); + try run(.{ bank[0..2], bank[2..2], bank[3..9] }, .{ 0, 2, 3 }, &bank, &.{ 2, 0, 6 }); + try run(.{ bank[0..0], bank[0..0], bank[0..0] }, .{ 0, 0, 0 }, &bank, &.{ 0, 0, 0 }); +} + +test "DMA coalescing: bounds and overflowing source bank or destination fail before copying" { + var driver = Fake.driver(); + const stream: Stream = .{ .d = &driver, .handle = null }; + const bank: [9]u8 = @splat(1); + const src: [3][]const u8 = .{ bank[0..2], bank[2..5], bank[5..9] }; + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 8, .{ 0, 2, 5 }, src, &bank)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 10, 5 }, src, &bank)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 2, std.math.maxInt(usize) }, src, &bank)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, std.math.maxInt(u64) - 8, 9, .{ 0, 2, 5 }, src, &bank)); + const high: [*]const u8 = @ptrFromInt(std.math.maxInt(usize) - 1); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 2, 5 }, .{ src[0], src[1], high[0..4] }, null)); + try std.testing.expectError(error.Invalid, dma.directLines(stream, 4096, 9, .{ 0, 2, 5 }, src, high[0..4])); + try std.testing.expectEqual(@as(usize, 0), Fake.count); +} + +test "DMA coalescing: production three-bit sizes issue one exact transfer" { + const width = 614400; + const bank = try std.testing.allocator.alloc(u8, 3 * width); + defer std.testing.allocator.free(bank); + @memset(bank, 1); + var driver = Fake.driver(); + const stream: Stream = .{ .d = &driver, .handle = null }; + const layout = pool.Layout.init(.{ width, width, width }); + const src: [3][]const u8 = .{ bank[0..width], bank[width .. 2 * width], bank[2 * width ..] }; + try std.testing.expectEqual(@as(usize, 3 * width), try dma.directLines(stream, 4096, layout.line_bytes, layout.off, src, bank)); + try std.testing.expectEqual(@as(usize, 1), Fake.copied); + try std.testing.expectEqual(@as(usize, 3 * width), Fake.copies[0].len); +} + +test "DMA coalescing: pool preserves reader waits table writes event and expert byte statistics" { + const bank: [1536]u8 = @splat(1); + const src: [3][]const u8 = .{ bank[0..256], bank[256..768], bank[768..1536] }; + for ([_]?[]const u8{ null, &bank, bank[256..] }, [_]usize{ 3, 1, 2 }) |provenance, copies| { + for (0..copies + 1) |failure| { + var driver = Fake.driver(); + Fake.fail_copy = failure; + var epochs = [_]u64{1}; + var pending = [_]bool{false}; + var events = [_]Event{.{ .d = &driver, .handle = @ptrFromInt(2) }}; + var lines: pool.Lines = undefined; + lines.d = &driver; + lines.spec = .{ .experts = 2, .bytes = .{ 256, 512, 768 }, .lines = 1, .staging = 0 }; + lines.layout = pool.Layout.init(lines.spec.bytes); + lines.tlay = pool.Tables.init(2); + lines.mem = .{ .d = &driver, .ptr = 4096, .len = 1536 }; + lines.tables = .{ .d = &driver, .ptr = 8192, .len = lines.tlay.bytes }; + lines.staging = null; + lines.pinned_bank = provenance; + lines.copy = .{ .d = &driver, .handle = @ptrFromInt(7) }; + lines.line_ev = &events; + lines.line_pending = &pending; + lines.line_epoch = &epochs; + lines.rel_ev[1] = .{ .d = &driver, .handle = @ptrFromInt(3) }; + lines.done_epoch = 0; + lines.stats = .{}; + const result = lines.load(0, 1, src, .{ 3, 6, 9 }, 0, true, undefined); + try std.testing.expectEqualSlices(Fake.Call, &.{ .query, .wait }, Fake.calls[0..2]); + if (failure > 0) { + try std.testing.expectError(error.NotReady, result); + try std.testing.expectEqual(2 + failure, Fake.count); + try std.testing.expectEqual(failure, Fake.copied); + try std.testing.expect(!pending[0]); + try std.testing.expectEqual(@as(u64, 0), lines.stats.copies); + try std.testing.expectEqual(@as(u64, 0), lines.stats.bytes); + } else { + try result; + try std.testing.expectEqual(2 + copies + 18 + 1, Fake.count); + for (Fake.calls[2 + copies .. Fake.count - 1]) |call| try std.testing.expectEqual(Fake.Call.table, call); + try std.testing.expectEqual(Fake.Call.event, Fake.calls[Fake.count - 1]); + try std.testing.expect(pending[0]); + try std.testing.expectEqual(@as(u64, 1), lines.stats.copies); + try std.testing.expectEqual(@as(u64, bank.len), lines.stats.bytes); + try std.testing.expectEqual(@as(u64, 1), lines.stats.release_waits); + for (0..3) |which| { + const old = Fake.words[3 * which ..][0..3]; + const installed = Fake.words[9 + 3 * which ..][0..3]; + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which), .value = 0 }, old[0]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which) + 4, .value = 0 }, old[1]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.widths(which), .value = 0 }, old[2]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which) + 8, .value = @intCast(4096 + lines.layout.off[which]) }, installed[0]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.table(which) + 12, .value = 0 }, installed[1]); + try std.testing.expectEqualDeep(Fake.Word{ .dst = lines.widths(which) + 4, .value = @intCast(3 * (which + 1)) }, installed[2]); + } + } + for (Fake.calls[2 .. 2 + Fake.copied]) |call| try std.testing.expectEqual(Fake.Call.copy, call); + } + } +} diff --git a/zig/src/cuda/exl3_pool.zig b/zig/src/cuda/exl3_pool.zig new file mode 100644 index 0000000000..4ac452ff2e --- /dev/null +++ b/zig/src/cuda/exl3_pool.zig @@ -0,0 +1,331 @@ +//! Device lines for one layer's EXL3 routed-expert offload on upstream's `exl3_experts.Experts`. + +const std = @import("std"); +const abi = @import("abi.zig"); +const memory = @import("memory.zig"); +const stream_ = @import("stream.zig"); +pub const dma = @import("exl3_dma.zig"); + +const Driver = @import("driver.zig").Driver; +const Error = @import("driver.zig").Error; +const DeviceBuffer = memory.DeviceBuffer; +const HostBuffer = memory.HostBuffer; +const Stream = stream_.Stream; +const Event = stream_.Event; + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +pub const Spec = struct { + /// Table length E (513: 512 routed experts and the shared expert as id 512). + experts: usize, + /// Largest trellis bytes of any expert, per projection (gate, up, down): a line fits every expert of the layer. + bytes: [3]usize, + lines: usize, + /// Pinned staging lines (a ring); 0 when every copy is `direct`. + staging: usize = 8, +}; + +/// Byte offsets of gate, up and down inside a line, each 256-aligned. +pub const Layout = struct { + off: [3]usize, + line_bytes: usize, + + pub fn init(bytes: [3]usize) Layout { + var l: Layout = .{ .off = undefined, .line_bytes = 0 }; + for (bytes, &l.off) |b, *o| { + o.* = l.line_bytes; + l.line_bytes += align256(b); + } + return l; + } +}; + +/// Where the six tables sit inside one allocation: three int64 [E] pointer tables, then three int32 [E] width tables. +pub const Tables = struct { + ptr: [3]usize, + k2: [3]usize, + bytes: usize, + + pub fn init(experts: usize) Tables { + var t: Tables = .{ .ptr = undefined, .k2 = undefined, .bytes = 0 }; + for (&t.ptr) |*o| { + o.* = t.bytes; + t.bytes += align256(experts * 8); + } + for (&t.k2) |*o| { + o.* = t.bytes; + t.bytes += align256(experts * 4); + } + return t; + } +}; + +pub const Stats = struct { + copies: u64 = 0, + bytes: u64 = 0, + staged_bytes: u64 = 0, // bytes the host memcpy'd into pinned staging + stage_ns: u64 = 0, // host time in those memcpys + staging_waits: u64 = 0, // a staging slot was still being copied from (the host blocked) + release_waits: u64 = 0, // a copy waited on the device for a pass that read its line + gate_waits: u64 = 0, // a pass waited on the device for a line copy +}; + +const release_ring = 64; + +pub const Lines = struct { + d: *const Driver, + spec: Spec, + layout: Layout, + tlay: Tables, + mem: DeviceBuffer, + tables: DeviceBuffer, + shared: ?DeviceBuffer = null, + pinned_bank: ?[]const u8 = null, + staging: ?HostBuffer, + stage_ev: []Event, + stage_busy: []bool, + stage_next: usize = 0, + copy: Stream, + line_ev: []Event, + line_pending: []bool, + /// Epoch of the last compute pass that read the line (0: never read). + line_epoch: []u64, + rel_ev: [release_ring]Event, + epoch: u64 = 0, // last epoch released + done_epoch: u64 = 0, // highest epoch known finished + stats: Stats = .{}, + gpa: std.mem.Allocator, + + pub fn init(gpa: std.mem.Allocator, d: *const Driver, spec: Spec) !Lines { + if (spec.lines == 0 or spec.experts == 0) return error.Invalid; + const layout = Layout.init(spec.bytes); + const tlay = Tables.init(spec.experts); + var mem = try DeviceBuffer.alloc(d, spec.lines * layout.line_bytes); + errdefer mem.free(); + var tables = try DeviceBuffer.alloc(d, tlay.bytes); + errdefer tables.free(); + try tables.fill8(0, null); // every expert absent; synchronous, before any stream uses the tables + var staging: ?HostBuffer = if (spec.staging > 0) try HostBuffer.alloc(d, spec.staging * layout.line_bytes) else null; + errdefer if (staging) |*h| h.free(); + var copy = try Stream.init(d, true); + errdefer copy.deinit(); + const stage_ev = try makeEvents(gpa, d, spec.staging); + errdefer freeEvents(gpa, stage_ev); + const stage_busy = try gpa.alloc(bool, spec.staging); + errdefer gpa.free(stage_busy); + @memset(stage_busy, false); + const line_ev = try makeEvents(gpa, d, spec.lines); + errdefer freeEvents(gpa, line_ev); + const line_pending = try gpa.alloc(bool, spec.lines); + errdefer gpa.free(line_pending); + @memset(line_pending, false); + const line_epoch = try gpa.alloc(u64, spec.lines); + errdefer gpa.free(line_epoch); + @memset(line_epoch, 0); + var rel: [release_ring]Event = undefined; + var made: usize = 0; + errdefer for (rel[0..made]) |*e| e.deinit(); + for (&rel) |*e| { + e.* = try Event.init(d, false); + made += 1; + } + return .{ + .d = d, + .spec = spec, + .layout = layout, + .tlay = tlay, + .mem = mem, + .tables = tables, + .staging = staging, + .stage_ev = stage_ev, + .stage_busy = stage_busy, + .copy = copy, + .line_ev = line_ev, + .line_pending = line_pending, + .line_epoch = line_epoch, + .rel_ev = rel, + .gpa = gpa, + }; + } + + fn makeEvents(gpa: std.mem.Allocator, d: *const Driver, n: usize) ![]Event { + const evs = try gpa.alloc(Event, n); + var made: usize = 0; + errdefer { + for (evs[0..made]) |*e| e.deinit(); + gpa.free(evs); + } + for (evs) |*e| { + e.* = try Event.init(d, false); + made += 1; + } + return evs; + } + + fn freeEvents(gpa: std.mem.Allocator, evs: []Event) void { + for (evs) |*e| e.deinit(); + gpa.free(evs); + } + + pub fn deinit(l: *Lines) void { + l.copy.synchronize() catch {}; + for (&l.rel_ev) |*e| e.deinit(); + freeEvents(l.gpa, l.line_ev); + freeEvents(l.gpa, l.stage_ev); + l.gpa.free(l.stage_busy); + l.gpa.free(l.line_pending); + l.gpa.free(l.line_epoch); + l.copy.deinit(); + if (l.staging) |*h| h.free(); + if (l.shared) |*b| b.free(); + l.tables.free(); + l.mem.free(); + l.* = undefined; + } + + /// Device bytes: the lines, the tables and the resident shared expert. + pub fn deviceBytes(l: Lines) usize { + return l.mem.len + l.tables.len + if (l.shared) |b| b.len else 0; + } + + pub fn table(l: Lines, which: usize) u64 { + return l.tables.ptr + l.tlay.ptr[which]; + } + + pub fn widths(l: Lines, which: usize) u64 { + return l.tables.ptr + l.tlay.k2[which]; + } + + /// Device address of projection `which` (0 gate, 1 up, 2 down) of line `line`. + pub fn at(l: Lines, line: usize, which: usize) u64 { + return l.mem.ptr + line * l.layout.line_bytes + l.layout.off[which]; + } + + /// Uploads the never-offloaded expert `id` (the shared expert) once into its own region (it may be wider than a + pub fn setResident(l: *Lines, id: usize, src: [3][]const u8, k2: [3]i32) !void { + if (id >= l.spec.experts or l.shared != null) return error.Invalid; + const lay = Layout.init(.{ src[0].len, src[1].len, src[2].len }); + var b = try DeviceBuffer.alloc(l.d, lay.line_bytes); + errdefer b.free(); + for (src, lay.off) |s, o| try b.upload(o, s); + l.shared = b; + for (0..3) |w| try l.writeEntry(id, b.ptr + lay.off[w], k2[w], w); + try l.copy.synchronize(); + } + + fn memset32(l: *Lines, addr: u64, value: u32, st: Stream) !void { + try l.d.check(l.d.api.cuMemsetD32Async(addr, value, 1, st.handle), "cuMemsetD32Async"); + } + + /// One expert's projection `which`: its pointer (two 32-bit words) and width, on the copy stream. + fn writeEntry(l: *Lines, id: usize, ptr: u64, k2: i32, which: usize) !void { + const st = l.copy; + const p = l.table(which) + id * 8; + try l.memset32(p, @truncate(ptr), st); + try l.memset32(p + 4, @truncate(ptr >> 32), st); + try l.memset32(l.widths(which) + id * 4, @bitCast(k2), st); + } + + /// The copy stream waits until the last compute pass that read `line` has finished (no wait if known done). + fn afterReaders(l: *Lines, line: usize) !void { + const e = l.line_epoch[line]; + if (e == 0 or e <= l.done_epoch) return; + const slot = e % release_ring; + const ev = l.rel_ev[slot]; + if (try ev.done()) { + // the slot holds epoch e or a newer one; either way e has finished + l.done_epoch = @max(l.done_epoch, e); + return; + } + try l.copy.wait(ev); // a newer epoch in the slot is a conservative wait + l.stats.release_waits += 1; + } + + /// Makes `id` resident in `line`: the copy waits for the line's readers, then H2D (staged or direct), then `id`'s + pub fn load(l: *Lines, line: usize, id: usize, src: [3][]const u8, k2: [3]i32, evicted: ?usize, direct: bool, io: std.Io) !void { + if (line >= l.spec.lines or id >= l.spec.experts) return error.Invalid; + for (src, l.spec.bytes) |s, cap| if (s.len > cap) return error.Invalid; + try l.afterReaders(line); + const dst = l.mem.ptr + line * l.layout.line_bytes; + var moved: usize = 0; + if (direct or l.staging == null) { + moved = try dma.directLines(l.copy, dst, l.layout.line_bytes, l.layout.off, src, l.pinned_bank); + } else { + const slot = l.stage_next; + l.stage_next = (l.stage_next + 1) % l.spec.staging; + if (l.stage_busy[slot]) { + if (!try l.stage_ev[slot].done()) l.stats.staging_waits += 1; + try l.stage_ev[slot].synchronize(); + } + const buf = l.staging.?.bytes[slot * l.layout.line_bytes ..][0..l.layout.line_bytes]; + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + for (src, l.layout.off) |s, o| @memcpy(buf[o..][0..s.len], s); + l.stats.stage_ns += @intCast(std.Io.Clock.awake.now(io).toNanoseconds() - t0); + // one copy for the line: up to the end of down's bytes (alignment padding rides along) + const n = l.layout.off[2] + src[2].len; + try l.d.check(l.d.api.cuMemcpyHtoDAsync_v2(dst, buf.ptr, n, l.copy.handle), "cuMemcpyHtoDAsync"); + try l.stage_ev[slot].record(l.copy); + l.stage_busy[slot] = true; + l.stats.staged_bytes += src[0].len + src[1].len + src[2].len; + moved = n; + } + if (evicted) |old| for (0..3) |w| try l.writeEntry(old, 0, 0, w); + for (0..3) |w| try l.writeEntry(id, dst + l.layout.off[w], k2[w], w); + try l.line_ev[line].record(l.copy); + l.line_pending[line] = true; + l.stats.copies += 1; + l.stats.bytes += moved; + } + + /// True once the line's last copy finished (or it never had one); clears the pending flag. + pub fn landed(l: *Lines, line: usize) !bool { + if (!l.line_pending[line]) return true; + if (!try l.line_ev[line].done()) return false; + l.line_pending[line] = false; + return true; + } + + /// `compute` waits on the device for the line's copy, if one may still be running. The host does not block. + pub fn gate(l: *Lines, compute: Stream, line: usize) !void { + if (try l.landed(line)) return; + try compute.wait(l.line_ev[line]); + l.stats.gate_waits += 1; + } + + /// After a compute pass that read `lines`: a release event on `compute`; a later copy into any of those lines + pub fn release(l: *Lines, compute: Stream, lines: []const u32) !void { + l.epoch += 1; + const slot = l.epoch % release_ring; + try l.rel_ev[slot].record(compute); + for (lines) |i| l.line_epoch[i] = l.epoch; + } + + /// Blocks until every queued copy and table write finished. + pub fn drain(l: *Lines) !void { + try l.copy.synchronize(); + for (l.line_pending) |*p| p.* = false; + @memset(l.stage_busy, false); + } +}; + +test { + _ = dma; +} + +test "a line packs gate, up and down 256-aligned; the tables sit in one allocation" { + // Flash Next layer 0: 3-bit experts, K=2560 N=640 -> 160 x 40 tiles of 48 words + const tb = 160 * 40 * 48 * 2; + const l = Layout.init(.{ tb, tb, tb }); + try std.testing.expectEqual([3]usize{ 0, tb, 2 * tb }, l.off); + try std.testing.expectEqual(@as(usize, 3 * tb), l.line_bytes); + const odd = Layout.init(.{ 100, 300, 5 }); + try std.testing.expectEqual([3]usize{ 0, 256, 768 }, odd.off); + try std.testing.expectEqual(@as(usize, 1024), odd.line_bytes); + const t = Tables.init(513); + try std.testing.expectEqual([3]usize{ 0, 4352, 8704 }, t.ptr); + try std.testing.expectEqual(@as(usize, 13056), t.k2[0]); + try std.testing.expectEqual(@as(usize, 13056 + 3 * 2304), t.bytes); + for (t.ptr ++ t.k2) |o| try std.testing.expectEqual(@as(usize, 0), o % 256); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 33a3bbfe55..0594bdb4ce 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -18,6 +18,7 @@ pub const qwen4_gdn: []const u8 = if (available) &Blob("fatbin_qwen4_gdn").bytes pub const qwen4_projection: []const u8 = if (available) &Blob("fatbin_qwen4_projection").bytes else &.{}; pub const gdn: []const u8 = if (available) &Blob("fatbin_gdn").bytes else &.{}; pub const flashnext_gdn: []const u8 = if (available) &Blob("fatbin_flashnext_gdn").bytes else &.{}; +pub const flashnext_exl3: []const u8 = if (available) &Blob("fatbin_flashnext_exl3").bytes else &.{}; pub const flashnext_batch: []const u8 = if (available) &Blob("fatbin_flashnext_batch").bytes else &.{}; pub const flashnext_mtp: []const u8 = if (available) &Blob("fatbin_flashnext_mtp").bytes else &.{}; pub const probe: []const u8 = if (available) &Blob("fatbin_probe").bytes else &.{}; diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 5489b019b8..98da1f7efa 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -69,6 +69,7 @@ pub const exl3_experts = @import("exl3_experts.zig"); pub const exl3_upload = @import("exl3_upload.zig"); pub const exl3_candidate = @import("exl3_candidate.zig"); pub const exl3_grouped = @import("exl3_grouped.zig"); +pub const exl3_pool = @import("exl3_pool.zig"); pub const arena = @import("arena.zig"); pub const pack = @import("pack.zig"); @@ -118,6 +119,7 @@ test { _ = exl3; _ = exl3_experts; _ = exl3_upload; + _ = exl3_pool; _ = @import("modelopt.zig"); _ = @import("ck.zig"); _ = @import("tp_watch.zig"); diff --git a/zig/src/families/flashnext/cuda_batch.zig b/zig/src/families/flashnext/cuda_batch.zig index 6066b1ea63..8ebdea1f63 100644 --- a/zig/src/families/flashnext/cuda_batch.zig +++ b/zig/src/families/flashnext/cuda_batch.zig @@ -8,6 +8,7 @@ const ple = @import("cuda_ple.zig"); const wt = @import("cuda_prompt_wt.zig"); const load = @import("cuda_prompt_load.zig"); const parts = @import("cuda_batch_state.zig"); +const exl3_pass = @import("exl3/model/cuda_exl3_pass.zig"); const batch_lanes = @import("cuda_batch_lanes.zig"); const tiles = @import("cuda_tiles.zig"); const tk = @import("cuda_tile_kernels.zig"); @@ -123,6 +124,10 @@ pub fn Pass(comptime Tri: type) type { /// qmm.cu's arithmetic over m rows into out (bf16) or part (fp32 K slices); returns the slices (1 when folded). pub fn lane(self: *Self, x: u64, xs: u64, w: Q4, out: u64, part: u64, m: usize, sk: usize) !usize { + if (w.exl3) |e| { + try exl3_pass.lane(self, x, e, out, part, m, sk); + return sk; // the first slice written, the other sk - 1 zero-filled (cuda_exl3_pass.lane) + } if (try zs.lane(self.k, self.stream, self.s, self.zero_slices, x, w, part, m, sk)) |one| return one; return tiles.matmul(self.k, self.stream, x, xs, w, out, part, self.s.native_part, m, sk); } @@ -228,27 +233,7 @@ pub fn Pass(comptime Tri: type) type { } pub fn pleBlock(self: *Self, lw: *const Layer, ids: []const u32) !void { - const s = self.s; - const n = self.rows; - // the rows' n-gram embeddings (host gathers from the mapped table) are already uploaded by `stage` - _ = ids; - _ = try self.lane(s.ple_emb, s.ple_xs, lw.key, s.ple_keys, s.dummy, n, 1); - const vsk = try self.lane(s.ple_emb, s.ple_xs, lw.val, s.dummy, s.ple_part, n, 4); - try self.sumSlices(s.ple_part, s.ple_vals, true, n * dims, vsk); - try self.tri.pleGate(s.ple_keys, s.ple_vals, s.h, lw.nk.ptr, lw.nq.ptr, s.ple_gated, s.ple_pss, eps, n, dims, streams); - if (self.lanes.len == 0) { - try self.tri.pleConv(s.ple_gated, s.ple_pss, lw.nc.ptr, self.seq.tail.ptr, lw.ple_cw.ptr, s.h, s.h, s.ple_nrow, eps, n, dims, streams, 4, 3); - try self.shiftWindow(self.seq.tail.ptr, s.ple_nrow, wide, n, ple_tail_n, wide); - return; - } - // shared rounds: each lane's rows convolved over its own stream's tail, the launch its lone pass makes - var row0: usize = 0; - for (self.lanes) |l| { - const g = row0 * wide * 2; - try self.tri.pleConv(s.ple_gated + g, s.ple_pss + row0 * streams * 4, lw.nc.ptr, l.seq.tail.ptr, lw.ple_cw.ptr, s.h + g, s.h + g, s.ple_nrow + g, eps, l.rows, dims, streams, 4, 3); - try self.shiftWindow(l.seq.tail.ptr, s.ple_nrow + g, wide, l.rows, ple_tail_n, wide); - row0 += l.rows; - } + return @import("cuda_batch_ple.zig").pleBlock(self, lw, ids); } pub fn gdnBlock(self: *Self, lw: *const Layer, layer: usize, li: usize, keep: LayerKeep) !void { @@ -427,6 +412,14 @@ pub fn Pass(comptime Tri: type) type { const s = self.s; const n = self.rows; const pairs = n * slots; + if (lw.exl3) |x| { + if (x.needsPool()) return error.Exl3PoolNotAttached; + if (keep.moe_x != 0) try self.copy(keep.moe_x, s.mixed, n * dims * 2); + try exl3_pass.moe(self, lw, x); + if (keep.moe_ids != 0) try self.copy(keep.moe_ids, s.picks, pairs * 4); + if (keep.moe_wts != 0) try self.copy(keep.moe_wts, s.wts, pairs * 4); + return; + } try self.copy(keep.moe_x, s.mixed, n * dims * 2); try self.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); try self.tri.topkMany(s.router_l, s.picks, s.wts, n); diff --git a/zig/src/families/flashnext/cuda_batch_gather.zig b/zig/src/families/flashnext/cuda_batch_gather.zig new file mode 100644 index 0000000000..b2a1ab8773 --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch_gather.zig @@ -0,0 +1,40 @@ +//! The batched pass's n-gram gather: the rows' embeddings on a host thread while the GPU runs layer 0. +const std = @import("std"); +const ple = @import("cuda_ple.zig"); +const wt = @import("cuda_prompt_wt.zig"); +const Lane = @import("cuda_batch_state.zig").Lane; + +const dims = wt.dims; + +/// The rows' n-gram embeddings, gathered on a host thread while the GPU embeds and runs layer 0. +pub const Gather = struct { + ngram: *const ple.Ngram, + history: []i64, + ids: []const u32, + emb: []u16, + xs: []f32, + failed: ?anyerror = null, + ns: i96 = 0, + io: std.Io, + /// A round's rows, each gathered after its own stream's history (empty: every row after `history`). + lanes: []const Lane = &.{}, + + pub fn main(g: *Gather) void { + const t0 = std.Io.Clock.awake.now(g.io).toNanoseconds(); + defer g.ns = std.Io.Clock.awake.now(g.io).toNanoseconds() - t0; + if (g.lanes.len == 0) { + g.ngram.rows(g.history, g.ids, g.emb, g.xs) catch |err| { + g.failed = err; + }; + return; + } + var row0: usize = 0; + for (g.lanes) |l| { + g.ngram.rows(l.seq.history[0..], g.ids[row0..][0..l.rows], g.emb[row0 * dims ..][0 .. l.rows * dims], g.xs[row0 * (dims / 32) ..][0 .. l.rows * (dims / 32)]) catch |err| { + g.failed = err; + return; + }; + row0 += l.rows; + } + } +}; diff --git a/zig/src/families/flashnext/cuda_batch_ple.zig b/zig/src/families/flashnext/cuda_batch_ple.zig new file mode 100644 index 0000000000..4ff2ffc346 --- /dev/null +++ b/zig/src/families/flashnext/cuda_batch_ple.zig @@ -0,0 +1,32 @@ +//! The shared batched pass's per-layer embedding block: the n-gram keys and values, the gate and the convolution. +const wt = @import("cuda_prompt_wt.zig"); + +const dims = wt.dims; +const streams = wt.streams; +const wide = wt.wide; +const eps = wt.eps; +const ple_tail_n = wt.ple_tail_n; + +pub fn pleBlock(self: anytype, lw: anytype, ids: []const u32) !void { + const s = self.s; + const n = self.rows; + // the rows' n-gram embeddings (host gathers from the mapped table) are already uploaded by `stage` + _ = ids; + _ = try self.lane(s.ple_emb, s.ple_xs, lw.key, s.ple_keys, s.dummy, n, 1); + const vsk = try self.lane(s.ple_emb, s.ple_xs, lw.val, s.dummy, s.ple_part, n, 4); + try self.sumSlices(s.ple_part, s.ple_vals, true, n * dims, vsk); + try self.tri.pleGate(s.ple_keys, s.ple_vals, s.h, lw.nk.ptr, lw.nq.ptr, s.ple_gated, s.ple_pss, eps, n, dims, streams); + if (self.lanes.len == 0) { + try self.tri.pleConv(s.ple_gated, s.ple_pss, lw.nc.ptr, self.seq.tail.ptr, lw.ple_cw.ptr, s.h, s.h, s.ple_nrow, eps, n, dims, streams, 4, 3); + try self.shiftWindow(self.seq.tail.ptr, s.ple_nrow, wide, n, ple_tail_n, wide); + return; + } + // shared rounds: each lane's rows convolved over its own stream's tail, the launch its lone pass makes + var row0: usize = 0; + for (self.lanes) |l| { + const g = row0 * wide * 2; + try self.tri.pleConv(s.ple_gated + g, s.ple_pss + row0 * streams * 4, lw.nc.ptr, l.seq.tail.ptr, lw.ple_cw.ptr, s.h + g, s.h + g, s.ple_nrow + g, eps, l.rows, dims, streams, 4, 3); + try self.shiftWindow(l.seq.tail.ptr, s.ple_nrow + g, wide, l.rows, ple_tail_n, wide); + row0 += l.rows; + } +} diff --git a/zig/src/families/flashnext/cuda_batch_state.zig b/zig/src/families/flashnext/cuda_batch_state.zig index 98faec84be..14035aa5d9 100644 --- a/zig/src/families/flashnext/cuda_batch_state.zig +++ b/zig/src/families/flashnext/cuda_batch_state.zig @@ -192,8 +192,11 @@ pub fn Staging(comptime Q: type) type { } /// Device scratch for chunks of up to `cap` rows. +pub const ScratchFormat = enum { classic, exl3 }; + pub const Scratch = struct { cap: usize, + format: ScratchFormat = .classic, bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, ids: u64 = 0, pos: u64 = 0, @@ -312,110 +315,127 @@ pub const Scratch = struct { return @alignCast(std.mem.bytesAsSlice(f32, s.host.?.bytes[hostAt(s.cap)[3]..][0 .. n * (dims / 32) * 4])); } - /// Bytes `alloc(cap)` takes on the device. - pub fn bytes(cap: usize) usize { + /// The buffers in the order `allocationSizes` lists their sizes. + const order = [_][]const u8{ + "ids", "pos", "h", "pss", "norm", "rinv", "dn", "hc_part", "hc_up", "native_part", "act", + "xs_act", "inj", "mixed", "xsm", "proj", "part", "gout", "gxs", "ys", "gq", "gk", + "gv", "ggb", "q", "iq", "po", "pm", "pl", "qids", "qnk", "qspr", "ao", + "gated", "gxs_attn", "router_l", "picks", "wts", "members", "items", "counts", "rank", "hist", "shared_gate", + "shared_up", "shared_act", "shared_out", "moe_act", "y", "branch", "moe_f32", "pocket_act", "ple_emb", "ple_xs", "ple_keys", + "ple_vals", "ple_part", "ple_gated", "ple_pss", "ple_nrow", "logits", "head_local", "inv", "dummy", + }; + + pub fn allocationSizes(cap: usize, format: ScratchFormat) [64]?usize { const pairs = cap * slots; const nblk = cdiv(pairs, plan_small); const items = cdiv(pairs, tile) + experts_n + 1; - var n: usize = 0; - n += cap * 4 + (cap + 1) * 16; - n += cap * streams * 4 + cap * wide * 2 * 2 + cap * pss_n * 4 + cap * (low + 4) * 2 + hc_sk * cap * (low + 4) * 4; - n += cap * wide * 2 + 8 * cap * dims * 4; - n += cap * low * 2 + cap * (low / 32) * 4 + cap * streams * 2 + cap * dims * 2 + cap * (dims / 32) * 4; - n += cap * proj_row * 2 + out_sk * cap * dims * 4 + cap * gdn.value_dim * (2 + 4 + 4) + cap * (gdn.value_dim / 32) * 4; - n += cap * (2 * gdn.nk * gdn.dk + 2 * gdn.nv) * 4; - n += cap * q_heads * head_dim * 2 + cap * index_heads * index_dim * 2; - n += att_rows * nch * q_heads * (head_dim + 2) * 4 + cap * out_k * 2 * 2 + cap * (out_k / 32) * 4; - n += fused_rows * (idw + 2) * 4; - n += cap * (experts_n + 1) * 4 + pairs * 4 * 2 + pairs * 4 + items * 3 * 4 + 8 + pairs * 4 + nblk * (experts_n + 1) * 4; - n += cap * moe_w * 6 + cap * dims * 4; - n += pairs * moe_w * 2 + pairs * dims * 4 + cap * dims * 2; - n += cap * dims * 2 + cap * (dims / 32) * 4 + cap * wide * 2 + cap * dims * 2 + 4 * cap * dims * 4; - n += cap * wide * 2 * 2 + cap * streams * 4 + vocab * 2 + cap * vocab + half * 4 + 4096; - n += cap * dims * 4 + cap * pocket_stride * 4; - return n; + return .{ + cap * 4, + (cap + 1) * 16, + cap * wide * 2, + cap * pss_n * 4, + cap * wide * 2, + cap * streams * 4, + cap * (low + 4) * 2, + hc_sk * cap * (low + 4) * 4, + cap * wide * 2, + 8 * cap * dims * 4, + cap * low * 2, + cap * (low / 32) * 4, + cap * streams * 2, + cap * dims * 2, + cap * (dims / 32) * 4, + cap * proj_row * 2, + out_sk * cap * dims * 4, + cap * gdn.value_dim * 2, + cap * (gdn.value_dim / 32) * 4, + cap * gdn.value_dim * 4, + cap * gdn.nk * gdn.dk * 4, + cap * gdn.nk * gdn.dk * 4, + cap * gdn.value_dim * 4, + cap * gdn.nv * 2 * 4, + cap * q_heads * head_dim * 2, + cap * index_heads * index_dim * 2, + att_rows * nch * q_heads * head_dim * 4, + att_rows * nch * q_heads * 4, + att_rows * nch * q_heads * 4, + fused_rows * idw * 4, + fused_rows * 4, + fused_rows * 4, + cap * out_k * 2, + cap * out_k * 2, + cap * (out_k / 32) * 4, + cap * (experts_n + 1) * 4, + pairs * 4, + pairs * 4, + if (format == .classic) pairs * 4 else null, + if (format == .classic) items * 3 * 4 else null, + if (format == .classic) 8 else null, + if (format == .classic) pairs * 4 else null, + if (format == .classic) nblk * (experts_n + 1) * 4 else null, + if (format == .classic) cap * moe_w * 2 else null, + if (format == .classic) cap * moe_w * 2 else null, + if (format == .classic) cap * moe_w * 2 else null, + if (format == .classic) cap * dims * 4 else null, + if (format == .classic) pairs * moe_w * 2 else null, + pairs * dims * 4, + cap * dims * 2, + cap * dims * 4, + cap * pocket_stride * 4, + cap * dims * 2, + cap * (dims / 32) * 4, + cap * wide * 2, + cap * dims * 2, + 4 * cap * dims * 4, + cap * wide * 2, + cap * streams * 4, + cap * wide * 2, + vocab * 2, + if (format == .classic) cap * vocab else null, + half * 4, + 4096, + }; + } + + pub fn bytes(cap: usize) usize { + return bytesForFormat(cap, .classic); + } + + pub fn bytesForFormat(cap: usize, format: ScratchFormat) usize { + var total: usize = 0; + for (allocationSizes(cap, format)) |optional| if (optional) |len| { + total += @max(len, 256); + }; + return total; } pub fn alloc(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize) !Scratch { - var s: Scratch = .{ .cap = cap }; + return allocForFormat(gpa, d, stream, cap, .classic); + } + + pub fn allocForFormat(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, cap: usize, format: ScratchFormat) !Scratch { + var s: Scratch = .{ .cap = cap, .format = format }; errdefer { cuda.completion.drain(d, stream, "Flash Next scratch work did not drain before rollback"); s.free(gpa); } - const pairs = cap * slots; - const nblk = cdiv(pairs, plan_small); - const items = cdiv(pairs, tile) + experts_n + 1; - s.ids = try s.add(gpa, d, cap * 4); - // a position every 16 bytes: one a block of att_rows rows, or one a row of a decode round - s.pos = try s.add(gpa, d, (cap + 1) * 16); - s.h = try s.add(gpa, d, cap * wide * 2); - s.pss = try s.add(gpa, d, cap * pss_n * 4); - s.norm = try s.add(gpa, d, cap * wide * 2); - s.rinv = try s.add(gpa, d, cap * streams * 4); - s.dn = try s.add(gpa, d, cap * (low + 4) * 2); - s.hc_part = try s.add(gpa, d, hc_sk * cap * (low + 4) * 4); - s.hc_up = try s.add(gpa, d, cap * wide * 2); - s.native_part = try s.add(gpa, d, 8 * cap * dims * 4); - s.act = try s.add(gpa, d, cap * low * 2); - s.xs_act = try s.add(gpa, d, cap * (low / 32) * 4); - s.inj = try s.add(gpa, d, cap * streams * 2); - s.mixed = try s.add(gpa, d, cap * dims * 2); - s.xsm = try s.add(gpa, d, cap * (dims / 32) * 4); - s.proj = try s.add(gpa, d, cap * proj_row * 2); - s.part = try s.add(gpa, d, out_sk * cap * dims * 4); - try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); - s.part_zero = true; - s.gout = try s.add(gpa, d, cap * gdn.value_dim * 2); - s.gxs = try s.add(gpa, d, cap * (gdn.value_dim / 32) * 4); - s.ys = try s.add(gpa, d, cap * gdn.value_dim * 4); - s.gq = try s.add(gpa, d, cap * gdn.nk * gdn.dk * 4); - s.gk = try s.add(gpa, d, cap * gdn.nk * gdn.dk * 4); - s.gv = try s.add(gpa, d, cap * gdn.value_dim * 4); - s.ggb = try s.add(gpa, d, cap * gdn.nv * 2 * 4); - s.q = try s.add(gpa, d, cap * q_heads * head_dim * 2); - s.iq = try s.add(gpa, d, cap * index_heads * index_dim * 2); - s.po = try s.add(gpa, d, att_rows * nch * q_heads * head_dim * 4); - s.pm = try s.add(gpa, d, att_rows * nch * q_heads * 4); - s.pl = try s.add(gpa, d, att_rows * nch * q_heads * 4); - s.qids = try s.add(gpa, d, fused_rows * idw * 4); - s.qnk = try s.add(gpa, d, fused_rows * 4); - s.qspr = try s.add(gpa, d, fused_rows * 4); - s.ao = try s.add(gpa, d, cap * out_k * 2); - s.gated = try s.add(gpa, d, cap * out_k * 2); - s.gxs_attn = try s.add(gpa, d, cap * (out_k / 32) * 4); - s.router_l = try s.add(gpa, d, cap * (experts_n + 1) * 4); - s.picks = try s.add(gpa, d, pairs * 4); - s.wts = try s.add(gpa, d, pairs * 4); - s.members = try s.add(gpa, d, pairs * 4); - s.items = try s.add(gpa, d, items * 3 * 4); - s.counts = try s.add(gpa, d, 8); - s.rank = try s.add(gpa, d, pairs * 4); - s.hist = try s.add(gpa, d, nblk * (experts_n + 1) * 4); - s.shared_gate = try s.add(gpa, d, cap * moe_w * 2); - s.shared_up = try s.add(gpa, d, cap * moe_w * 2); - s.shared_act = try s.add(gpa, d, cap * moe_w * 2); - s.shared_out = try s.add(gpa, d, cap * dims * 4); - s.moe_act = try s.add(gpa, d, pairs * moe_w * 2); - s.y = try s.add(gpa, d, pairs * dims * 4); - s.branch = try s.add(gpa, d, cap * dims * 2); - s.moe_f32 = try s.add(gpa, d, cap * dims * 4); - s.pocket_act = try s.add(gpa, d, cap * pocket_stride * 4); - s.ple_emb = try s.add(gpa, d, cap * dims * 2); - s.ple_xs = try s.add(gpa, d, cap * (dims / 32) * 4); - s.ple_keys = try s.add(gpa, d, cap * wide * 2); - s.ple_vals = try s.add(gpa, d, cap * dims * 2); - s.ple_part = try s.add(gpa, d, 4 * cap * dims * 4); - s.ple_gated = try s.add(gpa, d, cap * wide * 2); - s.ple_pss = try s.add(gpa, d, cap * streams * 4); - s.ple_nrow = try s.add(gpa, d, cap * wide * 2); - s.logits = try s.add(gpa, d, vocab * 2); - s.head_local = try s.add(gpa, d, cap * vocab); + const sizes = allocationSizes(cap, format); + // a row of a decode round or a block of att_rows rows takes `pos`'s 16 bytes var inv: [half]f32 = undefined; for (&inv, 0..) |*o, i| o.* = @floatCast(std.math.pow(f64, 10_000_000.0, -@as(f64, @floatFromInt(i)) / @as(f64, half))); - s.inv = try s.add(gpa, d, half * 4); - try s.bufs.items[s.bufs.items.len - 1].uploadAsync(0, std.mem.sliceAsBytes(&inv), stream.handle); - s.dummy = try s.add(gpa, d, 4096); - try s.bufs.items[s.bufs.items.len - 1].fill8(0, stream.handle); + inline for (order, 0..) |name, i| { + @field(s, name) = if (sizes[i]) |len| try s.add(gpa, d, len) else 0; + if (sizes[i] != null) { + const last = &s.bufs.items[s.bufs.items.len - 1]; + // part's second cap-stride slice stays zeros (native lanes write only the first); dummy is all zeros + if (comptime std.mem.eql(u8, name, "part")) { + try last.fill8(0, stream.handle); + s.part_zero = true; + } + if (comptime std.mem.eql(u8, name, "inv")) try last.uploadAsync(0, std.mem.sliceAsBytes(&inv), stream.handle); + if (comptime std.mem.eql(u8, name, "dummy")) try last.fill8(0, stream.handle); + } + } try stream.synchronize(); s.host = try cuda.HostBuffer.alloc(d, hostAt(cap)[4]); s.staging = .{ .q = .{ .stream = stream, .event = try cuda.Event.init(d, false) } }; @@ -483,42 +503,12 @@ pub const Model = struct { pocket: ?MoeAdd = null, }; -/// The rows' n-gram embeddings, gathered on a host thread while the GPU embeds and runs layer 0. -pub const Gather = struct { - ngram: *const ple.Ngram, - history: []i64, - ids: []const u32, - emb: []u16, - xs: []f32, - failed: ?anyerror = null, - ns: i96 = 0, - io: std.Io, - /// A round's rows, each gathered after its own stream's history (empty: every row after `history`). - lanes: []const Lane = &.{}, - - pub fn main(g: *Gather) void { - const t0 = std.Io.Clock.awake.now(g.io).toNanoseconds(); - defer g.ns = std.Io.Clock.awake.now(g.io).toNanoseconds() - t0; - if (g.lanes.len == 0) { - g.ngram.rows(g.history, g.ids, g.emb, g.xs) catch |err| { - g.failed = err; - }; - return; - } - var row0: usize = 0; - for (g.lanes) |l| { - g.ngram.rows(l.seq.history[0..], g.ids[row0..][0..l.rows], g.emb[row0 * dims ..][0 .. l.rows * dims], g.xs[row0 * (dims / 32) ..][0 .. l.rows * (dims / 32)]) catch |err| { - g.failed = err; - return; - }; - row0 += l.rows; - } - } -}; +pub const Gather = @import("cuda_batch_gather.zig").Gather; test { _ = @import("cuda_batch_state_test.zig"); _ = @import("cuda_nvfp4_tile.zig"); + _ = @import("cuda_scratch_layout_test.zig"); } pub const Pg = @import("cuda_glue.zig").Pg; diff --git a/zig/src/families/flashnext/cuda_config.zig b/zig/src/families/flashnext/cuda_config.zig index 41b4099220..659df5c70b 100644 --- a/zig/src/families/flashnext/cuda_config.zig +++ b/zig/src/families/flashnext/cuda_config.zig @@ -3,6 +3,7 @@ const std = @import("std"); const wt = @import("cuda_prompt_wt.zig"); const gdn = @import("cuda_gdn.zig"); +const exl3 = @import("exl3/pack/cuda_exl3_config.zig"); /// One config field and the value the kernels and buffers assume. const Want = struct { name: []const u8, value: i64 }; @@ -100,7 +101,7 @@ pub fn mismatch(text: std.json.ObjectMap, quant: ?std.json.ObjectMap) ?[]const u return geometry(text); } -fn geometry(text: std.json.ObjectMap) ?[]const u8 { +pub fn geometry(text: std.json.ObjectMap) ?[]const u8 { // a folder Living Weights taught widens the shared expert by its pocket (a whole number of groups, 128 at most) const learned = int(text, learned_key) orelse 0; if (learned < 0 or learned > 128 or @mod(learned, 32) != 0) return learned_key; @@ -111,7 +112,7 @@ fn geometry(text: std.json.ObjectMap) ?[]const u8 { return semantics(text); } -pub const Mode = enum { affine, modelopt }; +pub const Mode = enum { affine, modelopt, exl3 }; fn projection(q: std.json.ObjectMap, name: []const u8, algo: []const u8, group: i64) bool { const p = object(q, name) orelse return false; @@ -201,8 +202,8 @@ pub fn mode(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Mode { const root = parsed.value.object; const text = if (root.get("text_config")) |v| (if (v == .object) v.object else root) else root; const quant = if (root.get("quantization")) |v| (if (v == .object) v.object else null) else null; - const selected: Mode = if (quant != null) .affine else .modelopt; - var field = if (selected == .affine) mismatch(text, quant) else if (object(root, "quantization_config")) |q| modeloptMismatch(text, q) else "quantization"; + const selected: Mode = if (quant != null) .affine else if (exl3.named(root)) .exl3 else .modelopt; + var field = if (selected == .affine) mismatch(text, quant) else if (selected == .exl3) exl3.mismatch(root, text, geometry) else if (object(root, "quantization_config")) |q| modeloptMismatch(text, q) else "quantization"; if (field == null) if (given(root, "tie_word_embeddings")) |v| if (!holds(v, .{ .flag = false })) { field = "tie_word_embeddings"; }; diff --git a/zig/src/families/flashnext/cuda_embed.zig b/zig/src/families/flashnext/cuda_embed.zig index 453ba86271..d79a64f84c 100644 --- a/zig/src/families/flashnext/cuda_embed.zig +++ b/zig/src/families/flashnext/cuda_embed.zig @@ -83,6 +83,8 @@ pub const Table = struct { rows: usize, dims: usize, dense: bool = false, + /// A dense table kept in pinned host memory the kernels read in place (`w` borrows its device address). + host: ?cuda.HostBuffer = null, /// Copies the three tensors. `dims` is hidden size (eight values per stored word). pub fn upload(d: *const cuda.Driver, weight: Tensor, scales: Tensor, biases: Tensor) !Table { @@ -111,7 +113,27 @@ pub const Table = struct { return .{ .w = try cuda.DeviceBuffer.fromHost(d, weight.bytes), .s = zero, .b = zero, .rows = weight.dim(0), .dims = weight.dim(1), .dense = true }; } + /// A dense table in mapped host memory when `host`, else in VRAM as main has it. + pub fn uploadDenseIn(d: *const cuda.Driver, weight: Tensor, host: bool) !Table { + return if (host) uploadDenseHost(d, weight) else uploadDense(d, weight); + } + + /// Own a dense table in mapped host memory instead of VRAM: the gather reads its rows over PCIe. + pub fn uploadDenseHost(d: *const cuda.Driver, weight: Tensor) !Table { + if (weight.dtype != .bf16 or weight.rank != 2 or weight.dim(0) == 0 or weight.dim(1) == 0 or weight.dim(1) % 8 != 0) + return error.UnexpectedTensor; + const count = try std.math.mul(usize, weight.dim(0), weight.dim(1)); + if (weight.bytes.len != try std.math.mul(usize, count, 2)) return error.UnexpectedTensor; + var host = try cuda.HostBuffer.allocMapped(d, weight.bytes.len); + errdefer host.free(); + @memcpy(host.bytes, weight.bytes); + const zero: cuda.DeviceBuffer = .{ .d = d, .ptr = 0, .len = 0 }; + const w: cuda.DeviceBuffer = .{ .d = d, .ptr = try host.device(), .len = weight.bytes.len, .borrowed = true }; + return .{ .w = w, .s = zero, .b = zero, .rows = weight.dim(0), .dims = weight.dim(1), .dense = true, .host = host }; + } + pub fn deinit(self: *Table) void { + if (self.host) |*h| h.free(); self.w.free(); self.s.free(); self.b.free(); diff --git a/zig/src/families/flashnext/cuda_lanes.zig b/zig/src/families/flashnext/cuda_lanes.zig index 5fe5814d7d..fc28b9dcb3 100644 --- a/zig/src/families/flashnext/cuda_lanes.zig +++ b/zig/src/families/flashnext/cuda_lanes.zig @@ -27,18 +27,18 @@ pub fn Namespace(comptime Tri: type) type { return struct { pub const model_type = "qwen4_exp"; pub const tp_draft_snapshots = false; - pub const formats: []const []const u8 = &.{ "mlx-q4g32", "modelopt-mixed" }; + pub const formats: []const []const u8 = &.{ "mlx-q4g32", "modelopt-mixed", @import("exl3/pack/cuda_exl3_config.zig").format }; // the model's window: past 2,048 keys the block indexer picks each row's 512 best blocks of 4 (cuda_batch.zig) pub const default_context: i64 = @intCast(prompt.max_tokens); // a prompt runs on the batched pass in chunks of this many rows, one a prefill_step: no staggered segments pub const max_segments: u32 = 1; pub const prompt_rows: u32 = @intCast(prompt.batch.default_chunk); /// The prompt cache's copies of a stream's state (cuda_snap.zig), called with `Loaded.target`. - pub const snap_bytes = snapBytes; + pub const snap_bytes = copies.bytes; pub const snap_repeat_min: u32 = 16; - pub const snap_save = snapSave; - pub const snap_restore = snapRestore; - pub const snap_drop = snapDrop; + pub const snap_save = copies.save; + pub const snap_restore = copies.restore; + pub const snap_drop = copies.drop; pub const kept_disk = @import("cuda_disk.zig").Kept(@This()); // their disk image, as qwen3_5 declares its own /// TF_FLASHNEXT_KV_PAGES: the pool's lanes, the pages of a request's rows and a drafted window (kvl). pub const window_rows = kvl.window_rows; @@ -57,7 +57,10 @@ pub fn Namespace(comptime Tri: type) type { pub const min_gap: u32 = 256; pub const min_prompt: u32 = 512; }; - pub const Options = struct { vision: ?@import("cuda_images.zig").Limits = null, context: usize, drafts: bool, segments: usize = 1, slide: bool = false, tp: usize = 1, rank: usize = 0, master: ?[]const u8 = null, master_port: u16 = 29551 }; + + pub const Options = struct { vision: ?@import("cuda_images.zig").Limits = null, context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 1, tp: usize = 1, rank: usize = 0, master: ?[]const u8 = null, master_port: u16 = 29551 }; + pub const preallocates_lanes = false; + /// Living Weights' learner (--slide): lessons put neurons in each layer's pocket beside its shared expert. pub const slide = slide_mod.Slide(Tri); /// The first file a lesson would rewrite through a link, or null: a teaching copy may hard-link the others. @@ -79,6 +82,10 @@ pub fn Namespace(comptime Tri: type) type { stream_bytes: usize, /// Every stream allocates its own sequence; none rides on the engine's. free_streams: u32 = 0, + /// Kept prompt states in pinned host memory (cuda_snap.on_host): the native host budgets them there. + host_snaps: bool = false, + /// The prefill step when the prompt arithmetic depends on the chunk cuts (EXL3): states keep on its starts + snap_grid: u32 = 0, ctx: *anyopaque, deinit: *const fn (*anyopaque) void, lone: ?LoneRun = null, @@ -134,6 +141,7 @@ pub fn Namespace(comptime Tri: type) type { }; const Ops = mtp.Ops(Tri); + const copies = @import("cuda_lanes_snap.zig").Copies(of); pub const Cuda = struct { gpa: std.mem.Allocator, @@ -142,6 +150,8 @@ pub fn Namespace(comptime Tri: type) type { vision: ?*@import("cuda_images.zig").Vision = null, ranks: ?*ranks.Session = null, head: ?*mtp.Head = null, + profile: ?@import("exl3/offload/cuda_exl3_frequency_state.zig").State = null, + profile_counts: ?@import("exl3/offload/cuda_exl3_frequency.zig").Counts = null, depth: usize = mtp.default_depth, confidence: f64 = mtp.default_confidence, /// A captured Triton set (TENSORFOLD_CUDA_KERNELS) for comparison, else null: our own glue. @@ -202,7 +212,7 @@ pub fn Namespace(comptime Tri: type) type { .score = @import("cuda_score.zig").Hook(Cuda, Tri).call, .first_masked = firstMaskedFn, }; - const ranked_table = @import("cuda_score.zig").withoutScore(scoring_table); + const ranked_table = comptime @import("cuda_score.zig").withoutScore(scoring_table); // static: backend() returns its address return .{ .ptr = self, .proposal_head = if (self.head) |h| &h.service else null, .vtable = if (self.ranks == null) &scoring_table else &ranked_table }; } @@ -320,6 +330,9 @@ pub fn Namespace(comptime Tri: type) type { }; const last = stop == ids.len; self.engine.ngram.warmAhead(ids, stop, @min(ids.len, stop + self.engine.chunk)); + const previous_telemetry_phase = self.engine.telemetry.forward_phase; + self.engine.telemetry.forward_phase = .prefill; + defer self.engine.telemetry.forward_phase = previous_telemetry_phase; if (l.draft) |*d| { try self.ops().prefillChunk(&l.seq, d, ids, at, stop); } else if (self.ranks) |r| { @@ -354,6 +367,7 @@ pub fn Namespace(comptime Tri: type) type { pub fn release(p: *anyopaque) void { const own: *Cuda = @ptrCast(@alignCast(p)); cuda.completion.drain(own.engine.driver, own.stream, "flashnext: release did not drain"); + if (own.profile) |*profile| profile.save(own); var it = own.lanes.valueIterator(); while (it.next()) |l| dropLane(own, l); own.lanes.deinit(own.gpa); @@ -367,6 +381,7 @@ pub fn Namespace(comptime Tri: type) type { if (own.set) |*x| x.deinit(); own.glue.deinit(); own.stream.deinit(); + if (own.profile) |*profile| profile.deinit(); own.gpa.destroy(own); } @@ -544,7 +559,7 @@ pub fn Namespace(comptime Tri: type) type { test { _ = @import("cuda_images_test.zig").Cases(Cuda, Cuda.open, dropLane); - _ = @import("cuda_lanes_restore_test.zig").Cases(Cuda, Cuda.open, snapRestore, dropLane); + _ = @import("cuda_lanes_restore_test.zig").Cases(Cuda, Cuda.open, copies.restore, dropLane); _ = @import("cuda_lanes_admission_test.zig").Cases(Cuda, Cuda.open, dropLane); _ = @import("cuda_score.zig").CapabilityCases(Cuda); } @@ -560,41 +575,5 @@ pub fn Namespace(comptime Tri: type) type { try std.testing.expectError(bad[2], open(std.testing.allocator, std.testing.io, &ctx, "no-such-checkpoint", null, .{ .context = 4096, .drafts = true })); } } - - // -- the prompt cache's copies ------------------------------------------------------------------------ - - fn snapBytes(ptr: *anyopaque, at: u32) u64 { - return snap.bytes(at, of(ptr).head != null); - } - - /// The owner stream's state, its prompt pass standing at `at` (a drafting stream's with the head's rows). - fn snapSave(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { - const self = of(ptr); - const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); - const l = try self.lane(s); - if (l.seq.fed != at) return error.NotAtMark; - if (self.ranks) |r| { - if (l.draft != null) return error.TensorRankDraftSnapshotUnavailable; - return try r.save(&l.seq); - } - return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); - } - - fn snapRestore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { - const self = of(ptr); - const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); - const l = try self.lane(s); - if (self.ranks) |r| { - if (l.draft != null) return error.NoHeadState; - return r.restore(&l.seq, @ptrCast(@alignCast(saved))); - } - try snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))); - } - - fn snapDrop(ptr: *anyopaque, saved: *anyopaque) void { - const self = of(ptr); - if (self.ranks) |r| return r.dropSaved(@ptrCast(@alignCast(saved))); - snap.drop(self.gpa, self.stream, @ptrCast(@alignCast(saved))); - } }; } diff --git a/zig/src/families/flashnext/cuda_lanes_bytes.zig b/zig/src/families/flashnext/cuda_lanes_bytes.zig index bb959edc89..556ef895a7 100644 --- a/zig/src/families/flashnext/cuda_lanes_bytes.zig +++ b/zig/src/families/flashnext/cuda_lanes_bytes.zig @@ -3,9 +3,11 @@ const std = @import("std"); const weights = @import("cuda_prompt_wt.zig"); const mtp = @import("cuda_mtp.zig"); const config = @import("cuda_config.zig"); +const exl3_memory = @import("exl3/pack/cuda_exl3_memory.zig"); /// Device bytes of the checkpoint: the n-gram tables stay mapped, the MTP head counts; 0 if unread. pub fn deviceWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { + if (exl3_memory.isExl3(gpa, io, dir)) return exl3_memory.weightBytes(gpa, io, dir); const main = weights.deviceBytes(gpa, io, dir) catch return 0; return main + mtp.headBytes(); } @@ -13,6 +15,7 @@ pub fn deviceWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u6 pub fn placementWeightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize, drafts: bool) !u64 { if ((world != 1 and world != 2) or rank >= world) return error.BadRank; const mode = try config.mode(gpa, io, dir); + if (mode == .exl3) return exl3_memory.placement(gpa, io, dir, world, rank); if (mode == .modelopt) { const main = try @import("cuda_rank_memory.zig").estimate(gpa, io, dir, world, rank); // Dense draft projections and vocabulary rows exceed the affine head's estimate. diff --git a/zig/src/families/flashnext/cuda_lanes_load.zig b/zig/src/families/flashnext/cuda_lanes_load.zig index 508ce51c51..3513ebf653 100644 --- a/zig/src/families/flashnext/cuda_lanes_load.zig +++ b/zig/src/families/flashnext/cuda_lanes_load.zig @@ -7,6 +7,8 @@ const ranks = @import("cuda_rank_engine.zig"); const config = @import("cuda_config.zig"); const Glue = @import("cuda_glue.zig").Glue; const images = @import("cuda_images.zig"); +const memory = @import("exl3/pack/exl3_startup_memory.zig"); +const exl3_pass = @import("exl3/model/cuda_exl3_pass.zig"); pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, captured: ?[]const u8, o: F.Options) !F.Loaded { if (o.vision) |v| try @import("qwen_vision").media.checkLimits(v); const Cuda = F.Cuda; @@ -19,11 +21,12 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu const chunk = try prompt.chunkRows(); _ = try @import("cuda_modelopt_threads.zig").configured(); // TF_FLASHNEXT_LOAD_THREADS, refused before any load const mode = try config.mode(gpa, io, dir); + if (mode == .exl3) try @import("exl3/model/cuda_exl3_arch.zig").check(ctx.d); // sm_120 only: refused before any weight loads if (o.tp != 1 and (mode != .modelopt or captured != null)) return error.UnsupportedQuantization; if (o.slide and (mode != .affine or o.tp != 1)) return error.UnsupportedLearningFormat; if (o.slide and captured != null) return error.TrainingOnCapturedGlue; const page_tokens = try @import("cuda_kv_lanes.zig").paging(); - if (page_tokens > 0 and (o.tp != 1 or o.slide or captured != null)) return error.FlashNextPagesUnsupported; + if (page_tokens > 0 and (o.tp != 1 or o.slide or captured != null or mode == .exl3)) return error.FlashNextPagesUnsupported; // the pool's pages are not wired to EXL3's K/V const driver: *cuda.Driver = @constCast(ctx.d); const own = try gpa.create(Cuda); errdefer gpa.destroy(own); @@ -61,11 +64,16 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu own.engine.image_factory = own.vision; } errdefer if (own.vision) |v| v.deinit(); + var persistent_bytes = cuda.usage(false).device; + try @import("exl3/kv/cuda_exl3_kv.zig").setup(own.engine, mode == .exl3); // K10: the scratch format and the chunk of an EXL3 pack // the pass's scratch for a full chunk, once for the engine (every prompt, window, round and head pass) _ = try own.engine.reserve(@max(@min(o.context, own.engine.chunk), if (o.slide) @as(usize, @import("cuda_train.zig").pack_rows) else 0)); // the MTP head unless --no-drafts (TF_FLASHNEXT_DRAFTS and _CONFIDENCE: the chain's rule, 6 and 0.60) - if (o.drafts and (o.tp == 1 or o.rank == 0)) { + if (o.drafts and mode == .exl3) std.log.info("flashnext: an EXL3 pack decodes without the MTP head", .{}); + if (o.drafts and mode != .exl3 and (o.tp == 1 or o.rank == 0)) { + const head_before = cuda.usage(false).device; own.head = try mtp.Head.init(gpa, io, driver, &own.stream, dir); + persistent_bytes = try std.math.add(u64, persistent_bytes, cuda.usage(false).device -| head_before); own.depth = @max(1, @min(mtp.max_chain, F.envUsize("TF_FLASHNEXT_DRAFTS", mtp.default_depth))); if (std.c.getenv("TF_FLASHNEXT_CONFIDENCE")) |v| own.confidence = std.fmt.parseFloat(f64, std.mem.span(v)) catch mtp.default_confidence; } @@ -79,18 +87,38 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu own.lanes = .empty; own.next = 0; if (o.slide) own.lay = try @import("cuda_slide.zig").layoutOf(gpa, io, dir); + const scratch_bytes = cuda.usage(false).device -| persistent_bytes; const own_seq_bytes = try prompt.seqBytesPaged(window, try @import("cuda_pages.zig").pageTest()); var stream_bytes = own_seq_bytes; if (own.head) |h| { - stream_bytes += mtp.Draft.bytes(window, h.ids.len, own.engine.chunk); + stream_bytes = try std.math.add(u64, stream_bytes, mtp.Draft.bytes(window, h.ids.len, own.engine.chunk)); } else if (o.drafts and o.tp == 2) stream_bytes += mtp.Draft.bytes(window, 1, own.engine.chunk); - if (o.vision) |v| stream_bytes += @as(usize, v.image_tokens) * @import("cuda_prompt_wt.zig").dims * 2 + 16 * o.context; + if (o.vision) |v| stream_bytes = try std.math.add(u64, stream_bytes, try std.math.add(u64, try std.math.mul(u64, v.image_tokens, @import("cuda_prompt_wt.zig").dims * 2), try std.math.mul(u64, 16, o.context))); if (page_tokens > 0) { // the lane's head runs at the engine's chunk: its cache on the stream's pages, its buffers in the run own.kv.head = if (own.head) |h| .{ .draft_n = h.ids.len, .rows = own.engine.chunk } else null; const head_bytes = if (own.head) |h| mtp.Draft.bytes(window, h.ids.len, own.engine.chunk) else 0; own.kv.layout = @import("cuda_kv_lanes.zig").layout(o.context, page_tokens, own.kv.head, stream_bytes - own_seq_bytes - head_bytes); } + try @import("exl3/kv/cuda_exl3_kv.zig").logStream(own.engine, window, stream_bytes); // K10: what a stream holds, by piece + const pool_budget = try exl3_pass.budgetForStreams(stream_bytes, o.lanes); + const placement = @import("exl3/offload/cuda_exl3_frequency_state.zig"); + if (try placement.prepare(gpa, io, dir, own.engine.layers[0..own.engine.n_layers], @intCast(o.tp), @intCast(o.rank))) |prepared| { + own.profile = prepared.state; + own.profile_counts = prepared.counts; + } + errdefer if (own.profile) |*profile| profile.deinit(); + const profile_counts = if (own.profile_counts) |*counts| &counts.decoder else null; + const attached = try exl3_pass.attach(gpa, io, driver, own.stream, own.engine.layers[0..own.engine.n_layers], own.engine.chunk, pool_budget, profile_counts); + if (own.profile_counts) |*counts| try placement.finish(own, counts); + const ledger: memory.Ledger = .{ + .persistent = persistent_bytes, + .scratch = scratch_bytes, + .pool_scratch = attached.scratch_bytes, + .expert_pools = attached.device_bytes, + .reservation = .{ .stream_bytes = stream_bytes, .streams = if (stream_bytes == 0) 0 else pool_budget.other / stream_bytes }, + }; + std.log.info("flashnext rank {d} memory ledger: persistent={d} scratch={d} pool_scratch={d} expert_pools={d} stream_bytes={d} streams={d} planned_total={d} free_after_pools={d}", .{ o.rank, ledger.persistent, ledger.scratch, ledger.pool_scratch, ledger.expert_pools, stream_bytes, ledger.reservation.streams, try ledger.total(), attached.free_after }); return .{ .layout = own.kv.layout, .vision = own.vision, @@ -98,7 +126,9 @@ pub fn open(comptime F: type, gpa: std.mem.Allocator, io: std.Io, ctx: *const cu .facts = own.facts(), .rows = if (own.head != null) mtp.max_rows else 2, .prompt_rows = @intCast(own.engine.chunk), + .snap_grid = if (mode == .exl3) @intCast(own.engine.chunk) else 0, .stream_bytes = stream_bytes, + .host_snaps = @import("exl3/kv/cuda_exl3_kv.zig").snapHost(o.context, mode == .exl3), .ctx = own, .deinit = F.release, .learns = if (o.slide) own else null, diff --git a/zig/src/families/flashnext/cuda_lanes_snap.zig b/zig/src/families/flashnext/cuda_lanes_snap.zig new file mode 100644 index 0000000000..fe8cc9b71c --- /dev/null +++ b/zig/src/families/flashnext/cuda_lanes_snap.zig @@ -0,0 +1,51 @@ +//! The prompt cache's copies of a Flash Next stream (cuda_snap.zig), called with `Loaded.target`. + +const std = @import("std"); +const lanes = @import("lanes"); +const snap = @import("cuda_snap.zig"); + +/// Size, save, restore and drop over the lane owner `of` finds behind a target pointer. +pub fn Copies(comptime of: anytype) type { + return struct { + pub fn bytes(ptr: *anyopaque, at: u32) u64 { + return snap.bytes(at, of(ptr).head != null); + } + + /// The owner stream's state, its prompt pass standing at `at` (a drafting stream's with the head's rows). + pub fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + const self = of(ptr); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + const l = try self.lane(s); + if (l.seq.fed != at) return error.NotAtMark; + if (self.ranks) |r| { + if (l.draft != null) return error.TensorRankDraftSnapshotUnavailable; + return try r.save(&l.seq); + } + return try snap.save(self.gpa, self.engine.driver, self.stream, &l.seq, if (l.draft) |*d| d else null); + } + + pub fn restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + const self = of(ptr); + const s: *lanes.Stream = @ptrCast(@alignCast(owner orelse return error.NoStream)); + const l = self.lane(s) catch |err| { + std.log.warn("flashnext: a kept state was not restored: no lane ({t})", .{err}); + return err; + }; + if (self.ranks) |r| { + if (l.draft != null) return error.NoHeadState; + return r.restore(&l.seq, @ptrCast(@alignCast(saved))); + } + self.engine.round_graph.clear(); + snap.restore(self.stream, &l.seq, if (l.draft) |*d| d else null, @ptrCast(@alignCast(saved))) catch |err| { + std.log.warn("flashnext: a kept state was not restored ({t})", .{err}); + return err; + }; + } + + pub fn drop(ptr: *anyopaque, saved: *anyopaque) void { + const self = of(ptr); + if (self.ranks) |r| return r.dropSaved(@ptrCast(@alignCast(saved))); + snap.drop(self.gpa, self.stream, @ptrCast(@alignCast(saved))); + } + }; +} diff --git a/zig/src/families/flashnext/cuda_launch.zig b/zig/src/families/flashnext/cuda_launch.zig index fb3fce8d91..4171f2e5fc 100644 --- a/zig/src/families/flashnext/cuda_launch.zig +++ b/zig/src/families/flashnext/cuda_launch.zig @@ -6,6 +6,7 @@ const aot = cuda.aot; pub const embed = @import("cuda_embed.zig"); pub const projection = @import("cuda_projection.zig"); +pub const exl3 = @import("exl3/model/cuda_exl3.zig"); pub const rank_engine = @import("cuda_rank_engine.zig"); pub const collective = @import("cuda_collective.zig"); pub const rank_wire = @import("cuda_rank_wire.zig"); @@ -22,8 +23,11 @@ pub const moe = @import("cuda_moe.zig"); pub const attn = @import("cuda_attn.zig"); pub const ple = @import("cuda_ple.zig"); pub const ple_check = @import("cuda_ple_check.zig"); +/// An EXL3 pack's n-gram table (cuda_ple.Ngram's `.exl3` format) +pub const exl3_ple = @import("exl3/model/cuda_exl3_ple.zig"); pub const mix = @import("cuda_mix.zig"); pub const prompt = @import("cuda_prompt.zig"); +pub const snapshot = @import("cuda_snap.zig"); /// The row-by-row reference path, for flashnext-prompt pub const row_path = @import("cuda_rows.zig"); /// The MTP head's drafts: its forward, the chain, the verify window's rows and their rollback @@ -32,6 +36,7 @@ pub const mtp = @import("cuda_mtp.zig"); pub const pocket = @import("cuda_pocket.zig"); /// Living Weights' trainer over the batched pass pub const train = @import("cuda_train.zig"); +pub const exl3_pool = @import("exl3/offload/cuda_exl3_pool.zig"); /// Decisions' label scores (the lanes backend's `score`) pub const score = @import("cuda_score.zig"); pub const native = @import("cuda_lanes.zig").Namespace(Tri); @@ -431,6 +436,7 @@ test { _ = @import("cuda_pages.zig"); _ = @import("cuda_kv_lanes.zig"); _ = @import("cuda_config.zig"); + _ = @import("exl3/model/cuda_exl3.zig"); _ = @import("cuda_projection_load.zig"); _ = @import("cuda_rank_engine.zig"); _ = @import("cuda_round_graph.zig"); @@ -448,6 +454,10 @@ test { _ = @import("living.zig"); } +test "exl3: dispatch host contracts" { + _ = @import("exl3/model/cuda_exl3_dispatch_test.zig"); +} + test "embed grid is one program per row and per group of 32" { try std.testing.expectEqual([3]u32{ 4, 80, 1 }, embedGrid(4, 2560)); } diff --git a/zig/src/families/flashnext/cuda_paging_batch_test.zig b/zig/src/families/flashnext/cuda_paging_batch_test.zig index 356cf096ac..70ae7ad70b 100644 --- a/zig/src/families/flashnext/cuda_paging_batch_test.zig +++ b/zig/src/families/flashnext/cuda_paging_batch_test.zig @@ -121,12 +121,14 @@ test "paged grouped prompts: every current attention route receives the sequence s.part = 0x1300000000; var lw: Layer = undefined; lw.proj.native = null; + lw.proj.exl3 = null; lw.proj.w.ptr = 0x1000; lw.proj.s.ptr = 0x2000; lw.proj.b.ptr = 0x3000; lw.proj.k = 7680; lw.proj.n = 13952; lw.out.native = null; + lw.out.exl3 = null; lw.out.w.ptr = 0x4000; lw.out.s.ptr = 0x5000; lw.out.b.ptr = 0x6000; diff --git a/zig/src/families/flashnext/cuda_ple.zig b/zig/src/families/flashnext/cuda_ple.zig index e4ba3ca438..f5e4b11b63 100644 --- a/zig/src/families/flashnext/cuda_ple.zig +++ b/zig/src/families/flashnext/cuda_ple.zig @@ -5,6 +5,7 @@ const cuda = @import("cuda"); const core = @import("core"); const embed = @import("cuda_embed.zig"); const hc = @import("cuda_hc.zig"); +const exl3_ple = @import("exl3/model/cuda_exl3_ple.zig"); const ple_pool = @import("cuda_ple_pool.zig"); const ple_reader = @import("cuda_ple_reader.zig"); @@ -343,11 +344,12 @@ pub const Ngram = struct { w: [max_shards][]const u8 = undefined, s: [max_shards][]const u8 = undefined, b: [max_shards][]const u8 = undefined, - format: enum { mlx, nvfp4 } = .mlx, + format: enum { mlx, nvfp4, exl3 } = .mlx, tables: [max_shards]core.nvfp4_rows.Table = undefined, readers: [max_shards]ple_reader.Reader = undefined, read_mode: ple_reader.Mode = .mmap, warm_ahead: bool = false, + x3: ?*exl3_ple.Table = null, pool: ?*ple_pool.Pool = null, // the gather's parked workers (null: every row on the caller's thread) pub fn init(gpa: std.mem.Allocator, io: std.Io, model_dir: []const u8) !Ngram { @@ -365,6 +367,7 @@ pub const Ngram = struct { .ple_index = cfg.ple_index, }), .scale = 1 }; errdefer g.deinit(); + if (try exl3_ple.Table.open(gpa, io, model_dir, g.hash, exl3_ple.optionsFromEnv())) |t| return .{ .gpa = gpa, .hash = g.hash, .scale = 1, .format = .exl3, .x3 = t }; const index_path = try std.fs.path.join(gpa, &.{ model_dir, "model.safetensors.index.json" }); defer gpa.free(index_path); const text = try std.Io.Dir.cwd().readFileAlloc(io, index_path, gpa, .limited(1 << 26)); @@ -424,6 +427,7 @@ pub const Ngram = struct { } pub fn deinit(g: *Ngram) void { + if (g.x3) |t| t.destroy(); if (g.pool) |p| p.destroy(); for (g.files.items) |*o| o.close(g.gpa); g.files.deinit(g.gpa); @@ -446,6 +450,7 @@ pub const Ngram = struct { try g.readers[loc.shard].read(g.gpa, g.tables[loc.shard], loc.local, row) else try g.tables[loc.shard].gather(&.{loc.local}, row), + .exl3 => try g.x3.?.row(loc.shard, loc.local, i % g.hash.heads, row), } } } @@ -481,7 +486,8 @@ pub const Ngram = struct { for (0..history.len - 1) |i| hist[i] = hist[i + 1]; hist[history.len - 1] = token; } - const n_threads = @min(ple_pool.max_parts, items / 4); + // EXL3 rows decode in about a microsecond: a decode step's few rows cost less than one thread spawn + const n_threads = if (g.format == .exl3 and items <= 1024) 1 else @min(ple_pool.max_parts, items / 4); if (g.format == .nvfp4 and items > 1) { try @import("cuda_ple_gather.zig").run(g, ids, out); } else if (n_threads <= 1 or g.pool == null) { @@ -498,7 +504,7 @@ pub const Ngram = struct { v.* = toBf16(f * g.scale); } } - sumGroups(row, xs[t * (w / 32) ..][0 .. w / 32]); + if (g.x3) |x3| x3.sums(row, xs[t * (w / 32) ..][0 .. w / 32]) else sumGroups(row, xs[t * (w / 32) ..][0 .. w / 32]); } @memcpy(history, hist[0..history.len]); } @@ -530,6 +536,7 @@ test "a fresh window hashes token 7 into sixteen rows" { test { _ = @import("cuda_ple_nvfp4_test.zig"); + _ = @import("exl3/model/cuda_exl3_ple.zig"); _ = @import("cuda_ple_pool.zig"); _ = @import("cuda_ple_pool_gather_test.zig"); } diff --git a/zig/src/families/flashnext/cuda_projection.zig b/zig/src/families/flashnext/cuda_projection.zig index c0b96adb3b..0d678b580b 100644 --- a/zig/src/families/flashnext/cuda_projection.zig +++ b/zig/src/families/flashnext/cuda_projection.zig @@ -10,8 +10,11 @@ pub const Projection = struct { n: usize, k: usize, native: ?cuda.qmmf.Weight = null, + /// An EXL3 checkpoint's projection: stored-layout EXL3 parts (and dense BF16 ones) a row (cuda_exl3_load.zig). + exl3: ?*@import("exl3/pack/cuda_exl3_load.zig").Proj = null, pub fn free(self: *Projection) void { + if (self.exl3) |e| e.destroy(); self.w.free(); self.s.free(); self.b.free(); diff --git a/zig/src/families/flashnext/cuda_projection_load.zig b/zig/src/families/flashnext/cuda_projection_load.zig index c0c37f1b79..461c698be0 100644 --- a/zig/src/families/flashnext/cuda_projection_load.zig +++ b/zig/src/families/flashnext/cuda_projection_load.zig @@ -3,6 +3,7 @@ const std = @import("std"); const cuda = @import("cuda"); const wt = @import("cuda_prompt_wt.zig"); const projection = @import("cuda_projection.zig"); +const exl3 = @import("exl3/pack/cuda_exl3_load.zig"); pub const Micro = struct { n: usize, @@ -69,6 +70,7 @@ fn uploadMicro(a: std.mem.Allocator, d: *const cuda.Driver, value: Micro) !wt.Q4 } pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, n: usize, k: usize) !wt.Q4 { + if (exl3.has(store, prefix)) return exl3.load(a, d, store, prefix, n, k); var buf: [512]u8 = undefined; const t = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.weight", .{prefix})); switch (t.dtype) { @@ -96,6 +98,7 @@ pub fn load(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []c pub fn loadStitched(a: std.mem.Allocator, d: *cuda.Driver, store: *wt.Store, prefix: []const u8, names: []const []const u8, n: usize, k: usize) !wt.Q4 { if (names.len == 0) return error.UnexpectedTensor; + if (exl3.hasAny(store, prefix, names)) return exl3.stitched(a, d, store, prefix, names, n, k); var buf: [512]u8 = undefined; const first = try store.tensor(try std.fmt.bufPrint(&buf, "{s}{s}.weight", .{ prefix, names[0] })); if (first.dtype == .f8_e4m3) { diff --git a/zig/src/families/flashnext/cuda_prompt.zig b/zig/src/families/flashnext/cuda_prompt.zig index ed2d82c9b1..93932f8803 100644 --- a/zig/src/families/flashnext/cuda_prompt.zig +++ b/zig/src/families/flashnext/cuda_prompt.zig @@ -115,6 +115,7 @@ pub fn pooledBytes(cap: usize) usize { return (cap + weights.index_ratio - 1) / weights.index_ratio * index_dim * 2; } const round_graph = @import("cuda_round_graph.zig"); +const cache_telemetry = @import("cuda_prompt_telemetry.zig"); const pages = @import("cuda_pages.zig"); /// One loaded checkpoint: the embedding table, every layer, the mixer and head, and the batched pass's scratch. pub const Engine = struct { @@ -135,6 +136,7 @@ pub const Engine = struct { ngram: ple.Ngram, /// The pass's scratch, grown to the widest chunk run. batch_s: ?batch.Scratch = null, + scratch_format: @import("cuda_batch_state.zig").ScratchFormat = .classic, /// Block scores for the sparse rows of an attention block, grown to the longest sequence's need. scores: ?cuda.DeviceBuffer = null, /// Rows a batched chunk takes at most (chunkRows). @@ -142,10 +144,14 @@ pub const Engine = struct { round_graph: round_graph.Cache = .{}, image_factory: ?*@import("cuda_images.zig").Vision = null, collective: ?*@import("cuda_collective.zig").Collective = null, + telemetry: cache_telemetry.State = .{}, /// Each layer's Living Weights pocket, from a taught folder's shards or the learner: every pass adds it. pockets: pocket.Pockets = .{}, /// During a round of taught and untaught streams: its rows' pocket scales (0 on an untaught row); 0: none. taught_rows: u64 = 0, + + const measured = cache_telemetry.Operations(Engine, Seq, batch.Hooks, .{ .forward = forwardUnmeasured, .round = roundUnmeasured }); + pub fn init(gpa: std.mem.Allocator, io: std.Io, driver: *cuda.Driver, stream: *cuda.Stream, model_dir: []const u8) !*Engine { return initRank(gpa, io, driver, stream, model_dir, 1, 0); } @@ -161,8 +167,10 @@ pub const Engine = struct { e.driver = driver; e.stream = stream; e.n_layers = 0; + e.scratch_format = .classic; e.collective = null; e.image_factory = null; + e.telemetry = cache_telemetry.State.fromEnv(rank, world); e.round_graph = .{}; errdefer for (e.layers[0..e.n_layers]) |*l| l.free(gpa); e.model_dir = try gpa.dupe(u8, model_dir); @@ -172,7 +180,8 @@ pub const Engine = struct { try e.store.detectNorms(); const weight = try e.store.tensor(embed.weight_name); if (weight.dtype == .bf16) { - e.table = try embed.Table.uploadDense(driver, weight); + // an EXL3 pack keeps its BF16 table in mapped host memory, every other checkpoint on the device + e.table = try embed.Table.uploadDenseIn(driver, weight, @import("exl3/pack/cuda_exl3_load.zig").has(&e.store, "language_model.lm_head")); } else { const scales = try e.store.tensor(embed.scales_name); const biases = try e.store.tensor(embed.biases_name); @@ -182,8 +191,9 @@ pub const Engine = struct { e.eos = try ple.eosOf(gpa, io, e.model_dir); e.mixer = try loadHc(gpa, driver, &e.store, "language_model.model.hyper_connection_mixer", false); errdefer e.mixer.free(); - const head_tensor = try e.store.tensor("language_model.lm_head.weight"); - if (head_tensor.dtype == .bf16) { + const exl3_head = try @import("exl3/pack/cuda_exl3_load.zig").head(gpa, driver, &e.store, world, rank); + const head_tensor = if (exl3_head == null) try e.store.tensor("language_model.lm_head.weight") else undefined; + if (exl3_head) |q| e.head_q = q else if (head_tensor.dtype == .bf16) { if (!head_tensor.is(.bf16, &.{ vocab, dims })) return error.UnexpectedTensor; e.head_q = try @import("cuda_projection.zig").dense(gpa, driver, if (try @import("cuda_head.zig").enabled()) try @import("cuda_head.zig").shard(head_tensor, world, rank) else head_tensor); } else { @@ -345,10 +355,7 @@ pub const Engine = struct { return e.forward(Tri, tri, seq, ids, .{}); } - pub fn complete(e: *Engine) !void { - if (e.collective) |c| return c.finish(e.stream.*); - try e.stream.synchronize(); - } + pub const complete = measured.complete; /// Block scores for a `cap`-token sequence's chunks of up to `rows` rows, reallocated only to grow. pub fn reserveScores(e: *Engine, cap: usize, rows: usize) !u64 { @@ -372,19 +379,22 @@ pub const Engine = struct { /// Scratch for chunks of `rows` rows, reallocated only to grow. pub fn reserve(e: *Engine, rows: usize) !*batch.Scratch { if (e.batch_s) |*bs| { + if (bs.format != e.scratch_format) return error.ScratchFormatChanged; if (bs.cap >= rows) return bs; try e.complete(); e.round_graph.clear(); bs.free(e.gpa); e.batch_s = null; } - e.batch_s = try batch.Scratch.alloc(e.gpa, e.driver, e.stream.*, rows); + e.batch_s = try batch.Scratch.allocForFormat(e.gpa, e.driver, e.stream.*, rows, e.scratch_format); return &e.batch_s.?; } /// The batched pass over `ids` from `seq.fed` in chunks of at most `chunk` rows; a window's hooks take one chunk. - pub fn forward(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32, hooks: batch.Hooks) !void { - if (ids.len == 0 or seq.fed + ids.len > seq.cap) return error.UnexpectedTensor; + pub const forward = measured.forward; + + fn forwardUnmeasured(e: *Engine, comptime Tri: type, tri: Tri, seq: *Seq, ids: []const u32, hooks: batch.Hooks) !void { + if (ids.len == 0 or seq.fed > seq.cap or ids.len > seq.cap - seq.fed) return error.UnexpectedTensor; if ((hooks.keep != null or hooks.logits == .all) and ids.len > e.chunk) return error.UnexpectedTensor; const s = try e.reserve(@min(ids.len, e.chunk)); const scores = try e.reserveScores(seq.cap, s.cap); @@ -417,7 +427,9 @@ pub const Engine = struct { pub const windows = @import("cuda_prompt_windows.zig").windows; /// One decode round: each of `seqs` takes its token of `ids` in one batched pass; scores as if alone. - pub fn round(e: *Engine, comptime Tri: type, tri: Tri, seqs: []const *Seq, ids: []const u32, out: u64) !void { + pub const round = measured.round; + + fn roundUnmeasured(e: *Engine, comptime Tri: type, tri: Tri, seqs: []const *Seq, ids: []const u32, out: u64) !void { if (seqs.len == 0 or seqs.len != ids.len) return error.UnexpectedTensor; var cap: usize = 0; for (seqs) |q| { @@ -500,23 +512,7 @@ pub const Seq = struct { } pub fn deinit(s: *Seq, gpa: std.mem.Allocator) void { - if (s.stream) |stream| cuda.completion.drain(stream.d, stream, "Flash Next sequence work did not drain before release"); - if (s.media) |*m| m.deinit(); - freeBufs(s.conv); - freeBufs(s.state); - freeBufs(s.kc); - freeBufs(s.vc); - freeBufs(s.ikc); - freeBufs(s.pooled); - gpa.free(s.conv); - gpa.free(s.state); - gpa.free(s.kc); - gpa.free(s.vc); - gpa.free(s.ikc); - gpa.free(s.pooled); - if (s.paged) |*p| p.deinit(gpa); - if (s.tail.ptr != 0) s.tail.free(); - gpa.free(s.logits); + @import("cuda_prompt_lifecycle.zig").deinit(s, gpa); } }; @@ -524,55 +520,21 @@ fn freeBufs(bs: []cuda.DeviceBuffer) void { for (bs) |*b| b.free(); } +fn freeOwned(gpa: std.mem.Allocator, bs: []cuda.DeviceBuffer) void { + freeBufs(bs); + gpa.free(bs); +} + test "linear layers skip every fourth index and attention starts at 3" { - try std.testing.expect(linearLayer(0)); - try std.testing.expect(!linearLayer(3)); - try std.testing.expect(!linearLayer(47)); - try std.testing.expectEqual(@as(usize, 0), linearIndex(0)); - try std.testing.expectEqual(@as(usize, 2), linearIndex(2)); - try std.testing.expectEqual(@as(usize, 35), linearIndex(46)); - try std.testing.expectEqual(@as(usize, 0), attnIndex(3)); - try std.testing.expectEqual(@as(usize, 11), attnIndex(47)); + try @import("cuda_prompt_test.zig").layerIndices(); } test "a sequence's device bytes grow by the twelve attention layers' keys, values, indexer keys and pooled blocks" { - const per_token = attn_n * (2 * kv_heads * head_dim * 2 + index_dim * 2); - try std.testing.expectEqual(per_token, seqBytes(6) - seqBytes(5)); - try std.testing.expectEqual(per_token + attn_n * index_dim * 2, seqBytes(5) - seqBytes(4)); - try std.testing.expect(seqBytes(max_tokens) > max_tokens * per_token); + try @import("cuda_prompt_test.zig").sequenceBytes(seqBytes); } test "a failed fill frees the buffer it was filling and every one before it" { - const Fake = struct { - live: *usize, - pub fn free(b: *@This()) void { - b.live.* -= 1; - } - }; - const Src = struct { - live: *usize, - fills: *usize, - fail_at: usize, - fn make(s: @This()) !Fake { - s.live.* += 1; - return .{ .live = s.live }; - } - fn fill(s: @This(), _: *Fake) !void { - s.fills.* += 1; - if (s.fills.* == s.fail_at) return error.FillFailed; - } - }; - for ([_]usize{ 1, 2, 4 }) |fail_at| { - var live: usize = 0; - var fills: usize = 0; - try std.testing.expectError(error.FillFailed, made(Fake, std.testing.allocator, 4, Src{ .live = &live, .fills = &fills, .fail_at = fail_at })); - try std.testing.expectEqual(@as(usize, 0), live); - } - var live: usize = 0; - var fills: usize = 0; - const all = try made(Fake, std.testing.allocator, 3, Src{ .live = &live, .fills = &fills, .fail_at = 0 }); - defer std.testing.allocator.free(all); - try std.testing.expectEqual(@as(usize, 3), live); + try @import("cuda_prompt_test.zig").failedFill(made); } /// TF_FLASHNEXT_CHUNK's rows a pass: unset, the default; otherwise 1 to batch.max_chunk, anything else refused. diff --git a/zig/src/families/flashnext/cuda_prompt_lifecycle.zig b/zig/src/families/flashnext/cuda_prompt_lifecycle.zig new file mode 100644 index 0000000000..1aa3e5e31e --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_lifecycle.zig @@ -0,0 +1,14 @@ +const std = @import("std"); +const prompt = @import("cuda_prompt.zig"); + +pub fn deinit(sequence: *prompt.Seq, allocator: std.mem.Allocator) void { + if (sequence.stream) |stream| @import("cuda").completion.drain(stream.d, stream, "Flash Next sequence work did not drain before release"); + if (sequence.media) |*media| media.deinit(); + for ([_][]@import("cuda").DeviceBuffer{ sequence.conv, sequence.state, sequence.kc, sequence.vc, sequence.ikc, sequence.pooled }) |buffers| { + for (buffers) |*buffer| buffer.free(); + allocator.free(buffers); + } + if (sequence.tail.ptr != 0) sequence.tail.free(); + if (sequence.paged) |*table| table.deinit(allocator); + allocator.free(sequence.logits); +} diff --git a/zig/src/families/flashnext/cuda_prompt_load.zig b/zig/src/families/flashnext/cuda_prompt_load.zig index 45623e98cc..fc42f0a5df 100644 --- a/zig/src/families/flashnext/cuda_prompt_load.zig +++ b/zig/src/families/flashnext/cuda_prompt_load.zig @@ -7,6 +7,7 @@ const gdn = @import("cuda_gdn.zig"); const wt = @import("cuda_prompt_wt.zig"); const mixed = @import("cuda_mixed_experts.zig"); const projections = @import("cuda_projection_load.zig"); +const exl3_load = @import("exl3/pack/cuda_exl3_load.zig"); const layers_n = wt.layers_n; const dims = wt.dims; @@ -80,9 +81,12 @@ pub const Layer = struct { experts_up: cuda.DeviceBuffer, experts_down: cuda.DeviceBuffer, native_experts: ?mixed.Owned = null, + /// EXL3 routed experts with the shared one as id 512 (cuda_exl3_load.zig). + exl3: ?*exl3_load.Experts = null, pub fn free(self: *Layer, gpa: std.mem.Allocator) void { _ = gpa; + if (self.exl3) |x| x.destroy(); self.attn_hc.free(); self.mlp_hc.free(); self.proj.free(); @@ -343,7 +347,7 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S proj_q = try projections.loadStitched(gpa, driver, store, lp, &parts, gdn.proj_width, dims); out_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.out_proj", .{lp}), dims, gdn.value_dim); const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{lp})); - if (conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; + if (conv.dtype != .bf16 or conv.bytes.len != gdn.conv_dim * gdn.taps * 2) return error.UnexpectedTensor; cw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); const a_log = try asF32(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.A_log", .{lp}))); defer gpa.free(a_log); @@ -384,6 +388,7 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S const pp = pre_buf[0..pre.len]; key_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.key_proj", .{pp}), wide, dims); val_q = try projections.load(gpa, driver, store, try std.fmt.bufPrint(&buf, "{s}.value_proj", .{pp}), dims, dims); + try exl3_load.pleRows(gpa, driver, store, &key_q, &val_q); const nk = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_key.weight", .{pp}))); defer gpa.free(nk); const nq = try store.normScale(gpa, try store.tensor(try std.fmt.bufPrint(&buf, "{s}.norm_query.weight", .{pp}))); @@ -395,7 +400,7 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S nq_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nq)); nc_b = try cuda.DeviceBuffer.fromHost(driver, std.mem.sliceAsBytes(nc)); const conv = try store.tensor(try std.fmt.bufPrint(&buf, "{s}.conv1d.weight", .{pp})); - if (conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; + if (conv.dtype != .bf16 or conv.bytes.len != wide * 4 * 2) return error.UnexpectedTensor; pcw_b = try cuda.DeviceBuffer.fromHost(driver, conv.bytes); } @@ -405,7 +410,11 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S errdefer experts_up.free(); var experts_down = deadBuf(driver); errdefer experts_down.free(); - if (store.has(try std.fmt.bufPrint(&buf, "{s}.experts.0.gate_proj.weight", .{mlp_root}))) { + var exl3_experts: ?*exl3_load.Experts = null; + errdefer if (exl3_experts) |x| x.destroy(); + if (exl3_load.hasExperts(store, mlp_root)) { + exl3_experts = try exl3_load.experts(gpa, driver, store, mlp_root, world, rank); + } else if (store.has(try std.fmt.bufPrint(&buf, "{s}.experts.0.gate_proj.weight", .{mlp_root}))) { native_experts = try mixed.load(gpa, driver, store, mlp_root, world, rank); } else { if (world != 1 or rank != 0) return error.UnsupportedQuantization; @@ -468,5 +477,6 @@ fn loadLayerAtRank(gpa: std.mem.Allocator, driver: *cuda.Driver, stream: *cuda.S .experts_up = experts_up, .experts_down = experts_down, .native_experts = native_experts, + .exl3 = exl3_experts, }; } diff --git a/zig/src/families/flashnext/cuda_prompt_telemetry.zig b/zig/src/families/flashnext/cuda_prompt_telemetry.zig new file mode 100644 index 0000000000..b346b9ee74 --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_telemetry.zig @@ -0,0 +1,233 @@ +const std = @import("std"); +const cache = @import("exl3/model/cuda_exl3_telemetry.zig"); +const counters = cache.counters; + +pub const State = struct { + enabled: bool = false, + layers_enabled: bool = false, + layer_cursor: counters.LayerCursor(@import("cuda_prompt_wt.zig").layers_n) = .{}, + rank: counters.Rank = .{ .rank = 0, .world = 1 }, + forward_phase: counters.Phase = .unknown, + + pub fn fromEnv(rank: usize, world: usize) State { + const enabled = enabledFromEnv(); + const state: State = .{ .enabled = enabled, .layers_enabled = enabled and layersEnabledFromEnv(), .rank = .{ .rank = rank, .world = world } }; + if (state.enabled) std.log.warn("EXL3-CACHE-SAMPLING v=1 rank={d} world={d} optin_boundary_sync=1 bytes=logical_payload_not_pcie", .{ rank, world }); + return state; + } + + pub fn forwardPhase(state: State, fed: usize) counters.Phase { + return if (state.forward_phase != .unknown) state.forward_phase else if (fed == 0) .prefill else .unknown; + } + + pub fn enter(state: *State, phase: counters.Phase) Scope { + const previous = state.forward_phase; + state.forward_phase = phase; + return .{ .state = state, .previous = previous }; + } +}; + +pub const Scope = struct { + state: *State, + previous: counters.Phase, + + pub fn restore(scope: Scope) void { + scope.state.forward_phase = scope.previous; + } +}; + +pub fn Operations(comptime Engine: type, comptime Seq: type, comptime Hooks: type, comptime underlying: anytype) type { + return struct { + pub fn complete(engine: *Engine) !void { + if (engine.collective) |collective| return collective.finish(engine.stream.*); + try engine.stream.synchronize(); + } + + pub fn forward(engine: *Engine, comptime Tri: type, tri: Tri, sequence: *Seq, ids: []const u32, hooks: Hooks) !void { + if (!engine.telemetry.enabled) return underlying.forward(engine, Tri, tri, sequence, ids, hooks); + const phase = engine.telemetry.forwardPhase(sequence.fed); + const sequences = [_]*Seq{sequence}; + var before = try begin(engine, &sequences, phase, complete); + defer if (before) |*sample| sample.deinit(engine.gpa); + underlying.forward(engine, Tri, tri, sequence, ids, hooks) catch |failure| { + failed(engine, phase, failure); + return failure; + }; + if (before) |sample| try finish(engine, &sequences, sample, phase, ids.len, complete); + } + + pub fn round(engine: *Engine, comptime Tri: type, tri: Tri, sequences: []const *Seq, ids: []const u32, out: u64) !void { + if (!engine.telemetry.enabled) return underlying.round(engine, Tri, tri, sequences, ids, out); + var before = try begin(engine, sequences, .decode, complete); + defer if (before) |*sample| sample.deinit(engine.gpa); + underlying.round(engine, Tri, tri, sequences, ids, out) catch |failure| { + failed(engine, .decode, failure); + return failure; + }; + if (before) |sample| try finish(engine, sequences, sample, .decode, sequences.len, complete); + } + }; +} + +pub fn enabledFromEnv() bool { + const value = std.c.getenv("TF_FLASHNEXT_CACHE_TELEMETRY") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} + +pub fn layersEnabledFromEnv() bool { + const value = std.c.getenv("TF_FLASHNEXT_EXL3_LAYER_TELEMETRY") orelse return false; + return std.mem.eql(u8, std.mem.span(value), "1"); +} + +pub fn failed(engine: anytype, phase: counters.Phase, failure: anyerror) void { + std.log.warn("EXL3-CACHE-INVALID v=1 rank={d} world={d} phase={s} reason={s}", .{ engine.telemetry.rank.rank, engine.telemetry.rank.world, @tagName(phase), @errorName(failure) }); +} + +pub fn begin(engine: anytype, sequences: anytype, phase: counters.Phase, comptime complete: anytype) !?counters.Snapshot { + return beginWithSampler(engine, sequences, phase, complete, cache.sampleAfterConsumers); +} + +pub fn beginWithSampler(engine: anytype, sequences: anytype, phase: counters.Phase, comptime complete: anytype, comptime sample: anytype) !?counters.Snapshot { + if (!engine.telemetry.enabled) return null; + complete(engine) catch |failure| { + failed(engine, phase, failure); + return failure; + }; + return sample(engine.gpa, engine, sequences, engine.telemetry.rank, cache.Completion.all_rank_consumers_joined) catch |failure| { + failed(engine, phase, failure); + return null; + }; +} + +pub fn finish(engine: anytype, sequences: anytype, before: counters.Snapshot, phase: counters.Phase, rows: usize, comptime complete: anytype) !void { + return finishWithSampler(engine, sequences, before, phase, rows, complete, cache.sampleAfterConsumers, publish); +} + +pub fn finishWithSampler(engine: anytype, sequences: anytype, before: counters.Snapshot, phase: counters.Phase, rows: usize, comptime complete: anytype, comptime sample: anytype, comptime emit: anytype) !void { + complete(engine) catch |failure| { + failed(engine, phase, failure); + return failure; + }; + var after = sample(engine.gpa, engine, sequences, engine.telemetry.rank, cache.Completion.all_rank_consumers_joined) catch |failure| { + failed(engine, phase, failure); + return; + }; + defer after.deinit(engine.gpa); + var report = counters.Report.between(before, after, .{ .phase = phase, .model_rows = rows, .steps = 1 }) catch |failure| { + failed(engine, phase, failure); + return; + }; + const layers = if (engine.telemetry.layers_enabled) engine.telemetry.layer_cursor.claim(engine.gpa, before, after) catch |failure| { + failed(engine, phase, failure); + return; + } else &.{}; + defer engine.gpa.free(layers); + report.layers = layers; + emit(engine, report) catch |failure| { + failed(engine, phase, failure); + return; + }; +} + +fn publish(engine: anytype, report: counters.Report) !void { + const text = try report.formatAlloc(engine.gpa); + defer engine.gpa.free(text); + std.debug.print("{s}", .{text}); +} + +test "exl3: serving telemetry disabled adds no completion or readback and opt in samples only completed boundaries" { + const Fake = struct { + gpa: std.mem.Allocator, + telemetry: State = .{ .rank = .{ .rank = 0, .world = 1 } }, + completions: u64 = 0, + reads: u64 = 0, + records: u64 = 0, + + fn complete(engine: *@This()) !void { + engine.completions += 1; + } + + fn sample(_: std.mem.Allocator, engine: *@This(), _: @TypeOf(null), rank: counters.Rank, _: cache.Completion) !counters.Snapshot { + try std.testing.expect(engine.completions > engine.reads); + engine.reads += 1; + return .{ .rank = rank, .scope = 1, .experts = .{ .device_copies = engine.reads } }; + } + + fn emit(engine: *@This(), report: counters.Report) !void { + try std.testing.expectEqual(counters.Phase.decode, report.interval.phase); + try std.testing.expectEqual(@as(u64, 2), report.interval.model_rows); + try std.testing.expectEqual(@as(u64, 1), report.experts.device_copies); + try std.testing.expectEqual(@as(usize, 0), report.rank.rank); + engine.records += 1; + } + }; + var engine: Fake = .{ .gpa = std.testing.allocator }; + try std.testing.expect((try beginWithSampler(&engine, null, .decode, Fake.complete, Fake.sample)) == null); + try std.testing.expectEqual(@as(u64, 0), engine.completions); + try std.testing.expectEqual(@as(u64, 0), engine.reads); + engine.telemetry.enabled = true; + var first = (try beginWithSampler(&engine, null, .decode, Fake.complete, Fake.sample)).?; + defer first.deinit(engine.gpa); + try finishWithSampler(&engine, null, first, .decode, 2, Fake.complete, Fake.sample, Fake.emit); + try std.testing.expectEqual(@as(u64, 2), engine.completions); + try std.testing.expectEqual(@as(u64, 2), engine.reads); + try std.testing.expectEqual(@as(u64, 1), engine.records); +} + +test "exl3: forward and verify phase scopes restore nested caller state on error" { + const Call = struct { + fn fail(state: *State, phase: counters.Phase) !void { + const scope = state.enter(phase); + defer scope.restore(); + try std.testing.expectEqual(phase, state.forward_phase); + return error.ScopedOperationFailed; + } + }; + for ([_]counters.Phase{ .unknown, .prefill, .decode }) |previous| { + var state: State = .{ .forward_phase = previous }; + const parent = state.enter(.prefill); + try std.testing.expectError(error.ScopedOperationFailed, Call.fail(&state, .verify)); + try std.testing.expectEqual(counters.Phase.prefill, state.forward_phase); + try std.testing.expectError(error.ScopedOperationFailed, Call.fail(&state, .prefill)); + try std.testing.expectEqual(counters.Phase.prefill, state.forward_phase); + parent.restore(); + try std.testing.expectEqual(previous, state.forward_phase); + try std.testing.expect(!state.enabled); + } +} + +test "exl3: extracted disabled forward and round wrappers call underlying operations without joins" { + const prompt = @import("cuda_prompt.zig"); + const cuda = @import("cuda"); + const Fake = @import("exl3/model/cuda_exl3_fake_driver_test.zig").Fake; + const Calls = struct { + var forwards: usize = 0; + var rounds: usize = 0; + + fn forward(_: *prompt.Engine, comptime Tri: type, _: Tri, _: *prompt.Seq, _: []const u32, _: prompt.batch.Hooks) !void { + forwards += 1; + } + + fn round(_: *prompt.Engine, comptime Tri: type, _: Tri, _: []const *prompt.Seq, _: []const u32, _: u64) !void { + rounds += 1; + } + }; + Calls.forwards = 0; + Calls.rounds = 0; + var driver = Fake.driver(); + var stream: cuda.Stream = .{ .d = &driver, .handle = null }; + var engine: prompt.Engine = undefined; + engine.telemetry = .{}; + engine.gpa = std.testing.allocator; + engine.collective = null; + engine.stream = &stream; + var sequence: prompt.Seq = undefined; + sequence.fed = 9000; + const ops = Operations(prompt.Engine, prompt.Seq, prompt.batch.Hooks, .{ .forward = Calls.forward, .round = Calls.round }); + const syncs = Fake.sync_calls; + try ops.forward(&engine, void, {}, &sequence, &.{123}, .{}); + try ops.round(&engine, void, {}, &.{&sequence}, &.{123}, 0); + try std.testing.expectEqual(@as(usize, 1), Calls.forwards); + try std.testing.expectEqual(@as(usize, 1), Calls.rounds); + try std.testing.expectEqual(syncs, Fake.sync_calls); +} diff --git a/zig/src/families/flashnext/cuda_prompt_test.zig b/zig/src/families/flashnext/cuda_prompt_test.zig new file mode 100644 index 0000000000..ec581b9f4e --- /dev/null +++ b/zig/src/families/flashnext/cuda_prompt_test.zig @@ -0,0 +1,53 @@ +const std = @import("std"); +const weights = @import("cuda_prompt_wt.zig"); + +pub fn failedFill(comptime makeBuffers: anytype) !void { + const Fake = struct { + live: *usize, + pub fn free(buffer: *@This()) void { + buffer.live.* -= 1; + } + }; + const Src = struct { + live: *usize, + fills: *usize, + fail_at: usize, + pub fn make(source: @This()) !Fake { + source.live.* += 1; + return .{ .live = source.live }; + } + pub fn fill(source: @This(), _: *Fake) !void { + source.fills.* += 1; + if (source.fills.* == source.fail_at) return error.FillFailed; + } + }; + for ([_]usize{ 1, 2, 4 }) |fail_at| { + var live: usize = 0; + var fills: usize = 0; + try std.testing.expectError(error.FillFailed, makeBuffers(Fake, std.testing.allocator, 4, Src{ .live = &live, .fills = &fills, .fail_at = fail_at })); + try std.testing.expectEqual(@as(usize, 0), live); + } + var live: usize = 0; + var fills: usize = 0; + const all = try makeBuffers(Fake, std.testing.allocator, 3, Src{ .live = &live, .fills = &fills, .fail_at = 0 }); + defer std.testing.allocator.free(all); + try std.testing.expectEqual(@as(usize, 3), live); +} + +pub fn layerIndices() !void { + try std.testing.expect(weights.linearLayer(0)); + try std.testing.expect(!weights.linearLayer(3)); + try std.testing.expect(!weights.linearLayer(47)); + try std.testing.expectEqual(@as(usize, 0), weights.linearIndex(0)); + try std.testing.expectEqual(@as(usize, 2), weights.linearIndex(2)); + try std.testing.expectEqual(@as(usize, 35), weights.linearIndex(46)); + try std.testing.expectEqual(@as(usize, 0), weights.attnIndex(3)); + try std.testing.expectEqual(@as(usize, 11), weights.attnIndex(47)); +} + +pub fn sequenceBytes(comptime sequenceSize: anytype) !void { + const per_token = weights.attn_n * (2 * weights.kv_heads * weights.head_dim * 2 + weights.index_dim * 2); + try std.testing.expectEqual(per_token, sequenceSize(6) - sequenceSize(5)); + try std.testing.expectEqual(per_token + weights.attn_n * weights.index_dim * 2, sequenceSize(5) - sequenceSize(4)); + try std.testing.expect(sequenceSize(weights.window) > weights.window * per_token); +} diff --git a/zig/src/families/flashnext/cuda_prompt_wt.zig b/zig/src/families/flashnext/cuda_prompt_wt.zig index 732e0ce77e..c351aa42c2 100644 --- a/zig/src/families/flashnext/cuda_prompt_wt.zig +++ b/zig/src/families/flashnext/cuda_prompt_wt.zig @@ -6,6 +6,7 @@ const core = @import("core"); const hc = @import("cuda_hc.zig"); const qmm = @import("cuda_qmm.zig"); const gdn = @import("cuda_gdn.zig"); +const exl3_store = @import("exl3/pack/cuda_exl3_store.zig"); pub const layers_n: usize = 48; pub const linear_n: usize = 36; @@ -101,14 +102,19 @@ pub const Store = struct { released: [2]Held = undefined, nreleased: usize = 0, norm_offset: f32 = 0, + /// Dense F16 weights an EXL3 pack stores, widened to BF16 exactly (cuda_exl3_store.zig), freed at release(). + widened: std.ArrayList([]u16) = .empty, + /// The EXL3 kernels and scratch this checkpoint's projections and experts share (cuda_exl3_load.zig), if any. + exl3_ops: ?*@import("exl3/pack/cuda_exl3_load.zig").Ops = null, pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Store { const index = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); defer gpa.free(index); const text = try std.Io.Dir.cwd().readFileAlloc(io, index, gpa, .limited(1 << 26)); defer gpa.free(text); - const parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); + var parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}); errdefer parsed.deinit(); + try exl3_store.addPatch(io, dir, &parsed); _ = try core.safetensors.index.weightMap(parsed.value); return .{ .gpa = gpa, .io = io, .dir = dir, .parsed = parsed }; } @@ -144,6 +150,7 @@ pub const Store = struct { pub fn deinit(self: *Store) void { self.dropStaging(); + if (self.exl3_ops) |o| o.release(); self.release(); self.held.deinit(self.gpa); self.closeReleased(); @@ -174,6 +181,8 @@ pub const Store = struct { /// Borrowers expire here; two indexed shards cover separate weight and calibration planes. pub fn release(self: *Store) void { + for (self.widened.items) |w| self.gpa.free(w); + self.widened.clearAndFree(self.gpa); for (self.held.items) |h| { if (self.nreleased == self.released.len) { self.released[0].file.close(); @@ -207,7 +216,7 @@ pub const Store = struct { const key = try self.resolveKey(requested, &mapped); const file = try self.shard(key); for (self.held.items) |*h| { - if (std.mem.eql(u8, h.rel, file)) return h.file.get(key) orelse return error.MissingTensor; + if (std.mem.eql(u8, h.rel, file)) return self.dense(key, h.file.get(key) orelse return error.MissingTensor); } try self.held.ensureUnusedCapacity(self.gpa, 1); for (self.released[0..self.nreleased], 0..) |cached, i| { @@ -215,7 +224,7 @@ pub const Store = struct { self.held.appendAssumeCapacity(cached); self.nreleased -= 1; if (i < self.nreleased) self.released[i] = self.released[self.nreleased]; - return self.held.items[self.held.items.len - 1].file.get(key) orelse return error.MissingTensor; + return self.dense(key, self.held.items[self.held.items.len - 1].file.get(key) orelse return error.MissingTensor); } } const rel = try self.gpa.dupe(u8, file); @@ -223,10 +232,10 @@ pub const Store = struct { const path = try std.fs.path.join(self.gpa, &.{ self.dir, file }); defer self.gpa.free(path); self.held.appendAssumeCapacity(.{ .rel = rel, .file = try Shard.load(self.gpa, self.io, path) }); - return self.held.items[self.held.items.len - 1].file.get(key) orelse { + return self.dense(key, self.held.items[self.held.items.len - 1].file.get(key) orelse { std.debug.print("absent {s}\n", .{key}); return error.MissingTensor; - }; + }); } pub fn closeFile(self: *Store) void { @@ -277,10 +286,16 @@ pub const Store = struct { var mapped: [512]u8 = undefined; const key = try self.resolveKey(requested, &mapped); try self.use(key); - return self.file.?.get(key) orelse { + return self.dense(key, self.file.?.get(key) orelse { std.debug.print("absent {s}\n", .{key}); return error.MissingTensor; - }; + }); + } + + /// An EXL3 pack's dense F16 weight as the BF16 the loaders read (exact or refused); any other tensor as stored. + fn dense(self: *Store, key: []const u8, t: Tensor) !Tensor { + if (!exl3_store.widens(key, t)) return t; + return exl3_store.widen(self.gpa, key, t, &self.widened); } /// Check the exact or compatibility name without mapping its shard. diff --git a/zig/src/families/flashnext/cuda_scratch_layout_test.zig b/zig/src/families/flashnext/cuda_scratch_layout_test.zig new file mode 100644 index 0000000000..589a972be2 --- /dev/null +++ b/zig/src/families/flashnext/cuda_scratch_layout_test.zig @@ -0,0 +1,41 @@ +const std = @import("std"); +const Scratch = @import("cuda_batch_state.zig").Scratch; + +test "exl3: scratch omits only classic MoE allocations and the split-head logits and keeps every common view" { + for ([_]usize{ 1, 8, 32, 64, 1024, 1536, 2048 }) |cap| { + const classic = Scratch.allocationSizes(cap, .classic); + const exl3 = Scratch.allocationSizes(cap, .exl3); + var removed: usize = 0; + var retained: usize = 0; + for (classic, exl3, 0..) |before, after, index| { + if ((index >= 38 and index <= 47) or index == 61) { + try std.testing.expectEqual(@as(?usize, null), after); + removed += @max(before.?, 256); + } else { + try std.testing.expectEqual(before, after); + retained += @max(after.?, 256); + } + } + try std.testing.expectEqual(retained, Scratch.bytesForFormat(cap, .exl3)); + try std.testing.expectEqual(removed, Scratch.bytes(cap) - retained); + try std.testing.expectEqual(@as(?usize, cap * 11 * 2560 * 4), exl3[48]); + try std.testing.expectEqual(@as(?usize, cap * 2560 * 4), exl3[50]); + try std.testing.expectEqual(@as(?usize, cap * 128 * 4), exl3[51]); + } +} + +test "exl3: scratch savings include the allocator minimum and grow with requested rows" { + var previous: usize = 0; + for ([_]usize{ 1, 8, 32, 64, 1024, 1536, 2048 }) |cap| { + const pairs = cap * 11; + const items = (pairs + 15) / 16 + 513; + const blocks = (pairs + 1023) / 1024; + const expected = 2 * @max(pairs * 4, 256) + @max(items * 12, 256) + 256 + + @max(blocks * 513 * 4, 256) + 3 * @max(cap * 640 * 2, 256) + + @max(cap * 2560 * 4, 256) + @max(pairs * 640 * 2, 256) + @max(cap * 248320, 256); + const saved = Scratch.bytes(cap) - Scratch.bytesForFormat(cap, .exl3); + try std.testing.expectEqual(expected, saved); + try std.testing.expect(saved > previous); + previous = saved; + } +} diff --git a/zig/src/families/flashnext/cuda_snap.zig b/zig/src/families/flashnext/cuda_snap.zig index 7db8fe4043..c8d03ccb1a 100644 --- a/zig/src/families/flashnext/cuda_snap.zig +++ b/zig/src/families/flashnext/cuda_snap.zig @@ -30,6 +30,8 @@ pub fn bytes(at: usize, head: bool) usize { pub const Snap = struct { rank_id: u64 = 0, buf: cuda.DeviceBuffer, + /// The pinned host memory `buf` borrows (states parked on the host: `on_host`), or null (device memory). + host: ?cuda.HostBuffer = null, at: usize, history: [2]i64, head: bool, @@ -102,20 +104,41 @@ fn move(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, base: u64, at: usize, std.debug.assert(off == bytes(at, d != null)); } +fn discard(gpa: std.mem.Allocator, stream: cuda.Stream, snap: *Snap) void { + cuda.completion.drain(stream.d, stream, "snapshot copies did not drain"); + if (snap.host) |*h| h.free() else snap.buf.free(); + gpa.destroy(snap); +} + +/// Kept states go to mapped pinned host memory, set once at load (TF_FLASHNEXT_SNAP_HOST); the host budgets the cache. +pub var on_host: bool = false; + /// A copy of `seq` (and of the head's cache, with `d`) as it stands after `seq.fed` tokens. pub fn save(gpa: std.mem.Allocator, driver: *const cuda.Driver, stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft) !*Snap { const at = seq.fed; - if (d) |h| if (h.len != at or h.drafted != 0) return error.NotAtMark; - const snap = try gpa.create(Snap); - errdefer gpa.destroy(snap); - snap.* = .{ .buf = try cuda.DeviceBuffer.alloc(driver, bytes(at, d != null)), .at = at, .history = seq.history, .head = d != null }; - errdefer snap.buf.free(); - errdefer cuda.completion.drain(driver, stream, "a failed snapshot save did not drain"); + if (at > seq.cap) return error.SequenceTooShort; + if (d) |h| if (h.len != at or h.drafted != 0 or h.open or at > h.cap) return error.NotAtMark; + const n = bytes(at, d != null); + const snap = try allocate(gpa, driver, n, at, seq.history, d != null); + errdefer discard(gpa, stream, snap); try move(stream, seq, d, snap.buf.ptr, at, .save); try stream.synchronize(); return snap; } +fn allocate(gpa: std.mem.Allocator, driver: *const cuda.Driver, n: usize, at: usize, history: [2]i64, head: bool) !*Snap { + const snap = try gpa.create(Snap); + errdefer gpa.destroy(snap); + snap.* = .{ .buf = undefined, .at = at, .history = history, .head = head }; + if (on_host) { + var host = try cuda.HostBuffer.allocMapped(driver, @max(n, 256)); + errdefer host.free(); + snap.buf = .{ .d = driver, .ptr = try host.device(), .len = n, .borrowed = true }; + snap.host = host; + } else snap.buf = try cuda.DeviceBuffer.alloc(driver, n); + return snap; +} + /// `seq` (and the head's cache, with `d`) becomes `snap`'s state; refused before any copy when it cannot hold it. pub fn restore(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const Snap) !void { if (snap.at > seq.cap) return error.SequenceTooShort; @@ -128,13 +151,13 @@ pub fn restore(stream: cuda.Stream, seq: *prompt.Seq, d: ?*Draft, snap: *const S if (d) |h| { h.reset(); h.len = snap.at; + h.open = false; + h.fed0 = snap.at; } } pub fn drop(gpa: std.mem.Allocator, stream: cuda.Stream, snap: *Snap) void { - cuda.completion.drain(stream.d, stream, "a snapshot drop did not drain"); - snap.buf.free(); - gpa.destroy(snap); + discard(gpa, stream, snap); } test "a state's bytes: about 110 MiB of DeltaNet states and the n-gram tail, then 27 KiB a token" { diff --git a/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig new file mode 100644 index 0000000000..c87faba205 --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv.zig @@ -0,0 +1,83 @@ +//! KV memory of an EXL3 pack's streams (K10): bytes by piece, the scratch format, where kept states live. + +const std = @import("std"); +const wt = @import("../../cuda_prompt_wt.zig"); +const gdn = @import("../../cuda_gdn.zig"); + +const head_dim = wt.head_dim; +const kv_heads = wt.kv_heads; +const index_row: usize = wt.index_dim * 2; + +/// One token's K and V in one attention layer, bf16. +pub fn kvRowBytes() usize { + return 2 * kv_heads * head_dim * 2; +} + +/// Bytes a sequence of `cap` tokens holds, by piece. +pub const SeqBytes = struct { + gdn: u64, + kv: u64, + ikc: u64, + pooled: u64, + tail: u64, + + pub fn device(s: SeqBytes) u64 { + return s.gdn + s.kv + s.ikc + s.pooled + s.tail; + } +}; + +pub fn pooledBytes(cap: usize) u64 { + return (cap + wt.index_ratio - 1) / wt.index_ratio * index_row; +} + +/// cuda_prompt.zig seqBytes split by piece (a host test checks that the pieces sum to upstream's number). +pub fn seqBytes(cap: usize) SeqBytes { + const attn: u64 = wt.attn_n; + return .{ + .gdn = wt.linear_n * ((gdn.taps - 1) * gdn.conv_dim * 2 + wt.state_n * 4), + .kv = attn * cap * kvRowBytes(), + .ikc = attn * cap * index_row, + .pooled = attn * pooledBytes(cap), + .tail = wt.ple_tail_n * wt.wide * 2, + }; +} + +pub const gib: f64 = 1024.0 * 1024.0 * 1024.0; + +/// K10 lever 3: the pass's chunk for an EXL3 checkpoint unless TF_FLASHNEXT_CHUNK says. +pub const exl3_chunk: usize = 1024; + +/// The server's load hook (cuda_lanes_load.zig, right after the engine loads, before any scratch or sequence). +pub fn setup(e: anytype, exl3: bool) !void { + if (!exl3) return; + if (e.batch_s != null) return error.ScratchAlreadyReserved; + const decoder_exl3 = e.n_layers != 0 and for (e.layers[0..e.n_layers]) |layer| { + if (layer.exl3 == null) break false; + } else true; + e.scratch_format = if (decoder_exl3) .exl3 else .classic; + if (std.c.getenv("TF_FLASHNEXT_CHUNK") == null) e.chunk = exl3_chunk; +} + +/// One log line: a stream's device bytes at the window by piece. +pub fn logStream(e: anytype, window: usize, stream_bytes: usize) !void { + const b = seqBytes(window); + std.log.info("flashnext: a stream at a {d}-token window holds {d:.3} GiB on the device (K/V {d:.3} bf16, indexer keys {d:.3}, pooled {d:.3}, DeltaNet {d:.3}, draft and the rest {d:.3}); chunk {d} rows, scratch {d:.3} GiB", .{ + window, + @as(f64, @floatFromInt(stream_bytes)) / gib, + @as(f64, @floatFromInt(b.kv)) / gib, + @as(f64, @floatFromInt(b.ikc)) / gib, + @as(f64, @floatFromInt(b.pooled)) / gib, + @as(f64, @floatFromInt(b.gdn)) / gib, + @as(f64, @floatFromInt(stream_bytes -| b.device())) / gib, + e.chunk, + @as(f64, @floatFromInt(@import("../../cuda_batch_state.zig").Scratch.bytesForFormat(e.chunk, e.scratch_format))) / gib, + }); +} + +/// Whether kept prompt states go to pinned host memory (cuda_snap.on_host, set here): +pub fn snapHost(context: usize, exl3: bool) bool { + var on = exl3 and context > 32768; + if (std.c.getenv("TF_FLASHNEXT_SNAP_HOST")) |v| on = std.mem.eql(u8, std.mem.span(v), "1"); + @import("../../cuda_snap.zig").on_host = on; + return on; +} diff --git a/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv_test.zig b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv_test.zig new file mode 100644 index 0000000000..95e577fd59 --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/cuda_exl3_kv_test.zig @@ -0,0 +1,20 @@ +//! Host tests of the KV accounting: the pieces of a sequence's bytes. + +const std = @import("std"); +const hub = @import("cuda_exl3_kv.zig"); + +const seqBytes = hub.seqBytes; + +test "exl3: kv the pieces of a sequence sum to upstream's seqBytes" { + const prompt = @import("../../cuda_prompt.zig"); + for ([_]usize{ 1, 5, 4096, 200_016, 262_144 }) |cap| { + const s = seqBytes(cap); + try std.testing.expectEqual(@as(u64, prompt.seqBytes(cap)), s.device()); + } +} + +test "exl3: kv a token's K and V are upstream's row twice" { + const prompt = @import("../../cuda_prompt.zig"); + try std.testing.expectEqual(@as(usize, 2048), hub.kvRowBytes()); + try std.testing.expectEqual(2 * prompt.kv_row_bytes, hub.kvRowBytes()); +} diff --git a/zig/src/families/flashnext/exl3/kv/exl3_kv.zig b/zig/src/families/flashnext/exl3/kv/exl3_kv.zig new file mode 100644 index 0000000000..d53973f8bb --- /dev/null +++ b/zig/src/families/flashnext/exl3/kv/exl3_kv.zig @@ -0,0 +1,98 @@ +//! The host oracle of the device K/V table: one layer's slots as the kernel addresses them. + +const std = @import("std"); +const shape = @import("../pack/exl3_shape.zig"); +const cache = @import("../model/exl3_cache.zig"); + +/// Full-attention layers only. GDN layers have no KV table. The size is `contextKv`. +const full_layers: usize = shape.layers / shape.full_attention_interval; + +pub const Table = struct { + heads: usize, + tokens: usize, + bytes: []u8, + + pub fn init(gpa: std.mem.Allocator, tokens: usize) error{OutOfMemory}!Table { + const n = cache.contextKv(tokens); + const bytes = try gpa.alloc(u8, @intCast(n)); + @memset(bytes, 0); + return .{ .heads = shape.kv_heads, .tokens = tokens, .bytes = bytes }; + } + + pub fn deinit(self: *Table, gpa: std.mem.Allocator) void { + gpa.free(self.bytes); + } + + /// One head, one token, one full-attention layer. + pub fn write(self: *Table, head: usize, layer: usize, token: usize, src: []const u8) error{ Head, Slot, Length }!void { + const off = try self.locate(head, layer, token, src.len); + @memcpy(self.bytes[off..][0..src.len], src); + } + + pub fn read(self: *const Table, head: usize, layer: usize, token: usize, dest: []u8) error{ Head, Slot, Length }!void { + const off = try self.locate(head, layer, token, dest.len); + @memcpy(dest, self.bytes[off..][0..dest.len]); + } + + fn locate(self: *const Table, head: usize, layer: usize, token: usize, len: usize) error{ Head, Slot, Length }!usize { + if (head >= self.heads) return error.Head; + if (self.tokens == 0 or full_layers == 0 or layer >= full_layers or token >= self.tokens) return error.Slot; + const denom = full_layers * self.tokens * self.heads; + if (self.bytes.len % denom != 0) return error.Slot; + const one = self.bytes.len / denom; + if (len != one) return error.Length; + return ((layer * self.tokens + token) * self.heads + head) * one; + } +}; + +test "the table is contextKv bytes, both heads" { + const gpa = std.testing.allocator; + const tokens: usize = 8192; + var t = try Table.init(gpa, tokens); + defer t.deinit(gpa); + try std.testing.expectEqual(shape.kv_heads, t.heads); + try std.testing.expectEqual(@as(u64, 201_326_592), @as(u64, t.bytes.len)); + try std.testing.expectEqual(cache.contextKv(tokens), @as(u64, t.bytes.len)); + var none = try Table.init(gpa, 0); + defer none.deinit(gpa); + try std.testing.expectEqual(@as(usize, 0), none.bytes.len); +} + +test "a write lands in its own head's slot and nowhere else" { + const gpa = std.testing.allocator; + const tokens: usize = 2; + const slot_bytes: usize = 1024; + const mark: u8 = 0xa5; + var src: [slot_bytes]u8 = undefined; + var got: [slot_bytes]u8 = undefined; + { + var t = try Table.init(gpa, tokens); + defer t.deinit(gpa); + const slot = src[0..slot_bytes]; + @memset(slot, mark); + try std.testing.expectError(error.Head, t.write(shape.kv_heads, 0, 0, slot)); + try std.testing.expectError(error.Length, t.write(0, 0, 0, slot[0 .. slot_bytes - 1])); + try std.testing.expectError(error.Slot, t.write(0, 0, tokens, slot)); + try std.testing.expectError(error.Slot, t.write(0, full_layers, 0, slot)); + try std.testing.expect(std.mem.allEqual(u8, t.bytes, 0)); + + try t.write(0, 0, 1, slot); + @memset(&got, 0); + try t.read(0, 0, 1, got[0..slot_bytes]); + try std.testing.expectEqualSlices(u8, slot, got[0..slot_bytes]); + try t.read(1, 0, 1, got[0..slot_bytes]); + try std.testing.expect(std.mem.allEqual(u8, got[0..slot_bytes], 0)); + + @memset(slot, mark +% 1); + try t.write(1, 1, 0, slot); + try t.read(1, 1, 0, got[0..slot_bytes]); + try std.testing.expectEqualSlices(u8, slot, got[0..slot_bytes]); + try t.read(0, 0, 1, got[0..slot_bytes]); + try std.testing.expect(got[0] == mark); + } + var none = try Table.init(gpa, 0); + defer none.deinit(gpa); + var scratch: [0]u8 = .{}; + try std.testing.expectError(error.Slot, none.write(0, 0, 0, &scratch)); + try std.testing.expectError(error.Slot, none.read(0, 0, 0, &scratch)); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig new file mode 100644 index 0000000000..5b681cdf05 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3.zig @@ -0,0 +1,31 @@ +//! EXL3 checkpoints on the Flash Next CUDA family: what the upstream files' one-line hooks call. +pub const config = @import("../pack/cuda_exl3_config.zig"); +pub const arch = @import("cuda_exl3_arch.zig"); +pub const store = @import("../pack/cuda_exl3_store.zig"); +pub const memory = @import("../pack/cuda_exl3_memory.zig"); +pub const load = @import("../pack/cuda_exl3_load.zig"); +pub const pass = @import("cuda_exl3_pass.zig"); +pub const kv = @import("../kv/cuda_exl3_kv.zig"); +pub const fork = @import("../tp/cuda_exl3_fork.zig"); +pub const telemetry = @import("cuda_exl3_telemetry.zig"); +pub const frequency = @import("../offload/cuda_exl3_frequency.zig"); +/// The upstream modules an EXL3 check drives (tf-cuda-test flashnext-exl3): the store, loaders and the batched pass. +pub const prompt_wt = @import("../../cuda_prompt_wt.zig"); +pub const prompt_load = @import("../../cuda_prompt_load.zig"); +pub const projection_load = @import("../../cuda_projection_load.zig"); +pub const batch = @import("../../cuda_batch.zig"); + +test { + _ = config; + _ = arch; + _ = store; + _ = memory; + _ = load; + _ = pass; + _ = kv; + _ = @import("../kv/cuda_exl3_kv_test.zig"); + _ = fork; + _ = telemetry; + _ = frequency; + _ = @import("../offload/cuda_exl3_frequency_state.zig"); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_arch.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_arch.zig new file mode 100644 index 0000000000..fee7e05a47 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_arch.zig @@ -0,0 +1,32 @@ +//! EXL3 Flash Next runs on sm_120 only; other compute capabilities are refused until their gates have been run on them. +const std = @import("std"); +const cuda = @import("cuda"); + +/// The one compute capability (10 * major + minor) this family's EXL3 path is qualified on: Blackwell sm_120. +pub const qualified: u32 = 120; + +pub const message = "Flash Next EXL3 (expert offload, one rank) is qualified on sm_120 only; this device is sm_{d}: " ++ + "serve an NVFP4 or BF16 Flash Next checkpoint, or set TF_FLASHNEXT_EXL3_ARCH=unqualified to try it without a claim"; + +/// Whether a device of `capability` may serve an EXL3 pack; `override` is the TF_FLASHNEXT_EXL3_ARCH value for testers. +pub fn admits(capability: u32, override: ?[]const u8) bool { + if (capability == qualified) return true; + const v = override orelse return false; + return std.mem.eql(u8, v, "unqualified"); +} + +/// Refuses an EXL3 pack on an unqualified device before any weight loads. +pub fn check(driver: *const cuda.Driver) !void { + const capability = try driver.capability(0); + const v = std.c.getenv("TF_FLASHNEXT_EXL3_ARCH"); + if (admits(capability, if (v) |p| std.mem.span(p) else null)) return; + std.log.err(message, .{capability}); + return error.Exl3ArchUnqualified; +} + +test "only sm_120 admits an EXL3 pack unless a tester overrides the gate" { + try std.testing.expect(admits(120, null)); + for ([_]u32{ 80, 86, 89, 90, 100, 121 }) |c| try std.testing.expect(!admits(c, null)); + try std.testing.expect(!admits(121, "1")); + try std.testing.expect(admits(121, "unqualified")); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_dispatch_test.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_dispatch_test.zig new file mode 100644 index 0000000000..7c0253f1a3 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_dispatch_test.zig @@ -0,0 +1,74 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const batch = @import("../../cuda_batch.zig"); +const load = @import("../../cuda_prompt_load.zig"); +const exl3_load = @import("../pack/cuda_exl3_load.zig"); + +const FakeDriver = struct { + var copies: usize = 0; + + fn copy(_: u64, _: u64, _: usize, _: cuda.abi.Stream) callconv(.c) cuda.abi.Result { + copies += 1; + return 1; + } + + fn errorName(_: cuda.abi.Result, out: *?[*:0]const u8) callconv(.c) cuda.abi.Result { + out.* = null; + return cuda.abi.success; + } + + fn errorString(_: cuda.abi.Result, out: *?[*:0]const u8) callconv(.c) cuda.abi.Result { + out.* = null; + return cuda.abi.success; + } + + fn make() cuda.Driver { + var driver: cuda.Driver = undefined; + driver.api.cuMemcpyDtoDAsync_v2 = copy; + driver.api.cuGetErrorName = errorName; + driver.api.cuGetErrorString = errorString; + return driver; + } +}; + +const MockTri = struct { + own: ?*const @import("../../cuda_glue.zig").Glue = null, + + pub fn routerMany(_: MockTri, _: u64, _: u64, _: u64, _: usize, _: usize) !void {} + pub fn topkMany(_: MockTri, _: u64, _: u64, _: u64, _: usize) !void {} +}; + +test "exl3: Keep dispatch rejects an unattached expert bank before legacy copy" { + FakeDriver.copies = 0; + var driver = FakeDriver.make(); + var host_bytes: [0]u8 align(16) = .{}; + var experts: exl3_load.Experts = .{ + .gpa = std.testing.allocator, + .ops = undefined, + .layer = undefined, + .bank = .{ .d = &driver, .bytes = &host_bytes }, + }; + var layer: load.Layer = undefined; + layer.router.ptr = 1; + layer.exl3 = &experts; + var scratch: batch.Scratch = undefined; + scratch.mixed = 2; + scratch.router_l = 3; + scratch.picks = 4; + scratch.wts = 5; + var pass: batch.Pass(MockTri) = .{ + .gpa = std.testing.allocator, + .io = undefined, + .stream = .{ .d = &driver, .handle = null }, + .tri = .{}, + .k = undefined, + .s = &scratch, + .m = undefined, + .seq = undefined, + .hooks = .{}, + .rows = 1, + }; + + try std.testing.expectError(error.Exl3PoolNotAttached, pass.moeBlockKeep(&layer, .{ .moe_x = 9 })); + try std.testing.expectEqual(@as(usize, 0), FakeDriver.copies); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_fake_driver_test.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_fake_driver_test.zig new file mode 100644 index 0000000000..f8d01f1dfc --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_fake_driver_test.zig @@ -0,0 +1,23 @@ +const cuda = @import("cuda"); + +/// A driver whose only live entry points are the two synchronizations, which count themselves. +pub const Fake = struct { + pub var sync_calls: usize = 0; + + fn sync() callconv(.c) c_int { + sync_calls += 1; + return 0; + } + + fn streamSync(_: cuda.abi.Stream) callconv(.c) c_int { + return sync(); + } + + pub fn driver() cuda.Driver { + sync_calls = 0; + var result: cuda.Driver = undefined; + result.api.cuCtxSynchronize = sync; + result.api.cuStreamSynchronize = streamSync; + return result; + } +}; diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig new file mode 100644 index 0000000000..71de32b822 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_pass.zig @@ -0,0 +1,374 @@ +//! The batched pass's EXL3 branches (cuda_batch.zig): + +const std = @import("std"); +const cuda = @import("cuda"); +const wt = @import("../../cuda_prompt_wt.zig"); +const load = @import("../pack/cuda_exl3_load.zig"); +const pool_ = @import("../offload/cuda_exl3_pool.zig"); +const memory = @import("../pack/exl3_startup_memory.zig"); +const line_deltas = @import("../offload/exl3_line_deltas.zig"); +const profile_line_budget = @import("../offload/exl3_profile_line_budget.zig"); + +const dims = wt.dims; +const slots = wt.slots; + +/// Rows a routed-experts call takes: its scratch grows with rows * 11 (fp16 inputs, fp32 split sums). +pub const moe_rows: usize = 256; + +fn copy2d(s: cuda.Stream, dst: u64, dpitch: usize, src: u64, spitch: usize, width: usize, height: usize) !void { + const m: cuda.abi.Memcpy2D = .{ .src = src, .src_pitch = spitch, .dst = dst, .dst_pitch = dpitch, .width = width, .height = height }; + try s.d.check(s.d.api.cuMemcpy2DAsync_v2(&m, s.handle), "cuMemcpy2DAsync"); +} + +/// One part's rows [m, part.n] into `y` (row pitch part.n), bf16 or fp32. +fn part(p: anytype, ops: *load.Ops, pt: load.Part, x: u64, k: usize, m: usize, y: u64, f32_out: bool) !void { + switch (pt.kind) { + .exl3 => |w| { + const de = try ops.linear(); + try de.decode(p.stream, w, x, .bf16, m, y, if (f32_out) .f32 else .bf16); + }, + .native => |w| { + if (f32_out) return p.k.native_lane.matmulF32(p.stream, x, k, m, w, y, p.s.native_part); + try p.k.native_lane.matmul(p.stream, x, k, m, w, y, p.s.native_part); + }, + } +} + +/// Pass.lane() for an EXL3 projection: out bf16 [m, n] (sk 1), or fp32 [m. +pub fn lane(p: anytype, x: u64, e: *const load.Proj, out: u64, slices: u64, m: usize, sk: usize) !void { + const s = p.stream; + if (e.x == .f16) return laneF16(p, x, e, out, slices, m, sk); + if (sk > 1) { + if (e.parts.len != 1) return error.UnexpectedTensor; + try part(p, e.ops, e.parts[0], x, e.k, m, slices, true); + const bytes = m * e.n * 4; + try s.d.check(s.d.api.cuMemsetD8Async(slices + bytes, 0, bytes * (sk - 1), s.handle), "cuMemsetD8Async"); + return; + } + if (try @import("../tp/cuda_exl3_fork.zig").lane(p, x, e, out, m, onePart)) return; // the parts on parallel streams + for (e.parts) |pt| try onePart(p, e, pt, x, m, out); +} + +/// One part of a bf16 sk-1 lane() on the pass's stream: one part over the whole row, or a single row, writes its +fn onePart(p: anytype, e: *const load.Proj, pt: load.Part, x: u64, m: usize, out: u64) !void { + const s = p.stream; + if (pt.n == e.n or m == 1) return part(p, e.ops, pt, x, e.k, m, out + pt.col * 2, false); + const tmp = try e.ops.scratch(s, &e.ops.tmp, m * pt.n * 2); + try part(p, e.ops, pt, x, e.k, m, tmp, false); + try copy2d(s, out + pt.col * 2, e.n * 2, tmp, pt.n * 2, pt.n * 2, m); +} + +/// fp16 rows through a dense BF16 lane weight: x = hi + lo exactly (fn_exl3_split_f16), hi . w and lo . w in fp32. +fn laneF16(p: anytype, x: u64, e: *const load.Proj, out: u64, slices: u64, m: usize, sk: usize) !void { + const s = p.stream; + if (e.parts.len != 1 or e.parts[0].kind != .native) return error.UnexpectedTensor; + const w = e.parts[0].kind.native; + const count = m * e.k; + const halves = try e.ops.scratch(s, &e.ops.halves, 2 * count * 2); + var a: cuda.Args = .{}; + a.add(x); + a.add(halves); + a.add(halves + count * 2); + a.add(@as(c_longlong, @intCast(count))); + try cuda.launch.launch(try e.ops.split(), .{ .grid = .{ .x = @intCast((count + 255) / 256) }, .block = .{ .x = 256 } }, s, &a); + const bytes = m * e.n * 4; + const sums = if (sk == 1) try e.ops.scratch(s, &e.ops.sums, 2 * bytes) else slices; + try p.k.native_lane.matmulF32(s, halves, e.k, m, w, sums, p.s.native_part); + try p.k.native_lane.matmulF32(s, halves + count * 2, e.k, m, w, sums + bytes, p.s.native_part); + if (sk == 1) return p.sumSlices(sums, out, true, m * e.n, 2); + if (sk > 2) try s.d.check(s.d.api.cuMemsetD8Async(slices + 2 * bytes, 0, bytes * (sk - 2), s.handle), "cuMemsetD8Async"); +} + +/// Pass.moeBlock() for EXL3 experts: the router's 513 logits (the BF16 rows the Store widened exactly), each row's ten +pub fn moe(p: anytype, lw: anytype, x: *const load.Experts) !void { + const s = p.s; + const n = p.rows; + try p.tri.routerMany(s.mixed, lw.router.ptr, s.router_l, dims, n); + try p.tri.topkMany(s.router_l, s.picks, s.wts, n); + if (x.needsPool()) return error.Exl3PoolNotAttached; + const ex = try x.ops.grouped(); + const block = @min(n, moe_rows); + const bytes = try cuda.exl3_experts.scratchBytes(x.layer, block, slots); + const base = try x.ops.scratch(p.stream, &x.ops.moe, bytes); + var r0: usize = 0; + while (r0 < n) : (r0 += block) { + const rows = @min(block, n - r0); + var sc = try cuda.exl3_experts.scratch(base, x.layer, block, slots); + // runSlots leaves each slot's fp32 output in sc.y [rows * slots, dims]: the pass's y rows are that layout + sc.y = s.y + r0 * slots * dims * 4; + const xin = s.mixed + r0 * dims * 2; + const picks = s.picks + r0 * slots * 4; + if (x.pool) |pool| { + try pooled(p.stream, x.ops, pool, ex, x.layer, xin, picks, rows, sc, 0, 0); + continue; + } + try ex.runSlots(p.stream, x.layer, xin, dims, .bf16, picks, rows, sc); + } +} + +/// Host time of the A/B cut (the picks D2H and its wait) and of Pool.plan, summed over every pooled block (read by +pub var pool_time: struct { cut_ns: u64 = 0, plan_ns: u64 = 0, blocks: u64 = 0 } = .{}; + +/// One block of `rows` rows through `pool`: per-slot outputs in sc.y, and with `wts` != 0 the last pass also combines +pub fn pooled(st: cuda.Stream, ops: *load.Ops, pool: *pool_.Pool, ex: *const cuda.exl3_experts.Experts, layer: cuda.exl3_experts.Layer, xin: u64, picks: u64, rows: usize, sc: cuda.exl3_experts.Scratch, wts: u64, out: u64) !void { + const host = (ops.picks orelse return error.Exl3ScratchNotReserved).slice(i32)[0 .. rows * slots]; + const io = pool.io; + const t0 = std.Io.Clock.awake.now(io).toNanoseconds(); + // end of segment A: the block's picks reach the host (the only host wait of the layer) + try st.d.check(st.d.api.cuMemcpyDtoHAsync_v2(host.ptr, picks, rows * slots * 4, st.handle), "cuMemcpyDtoHAsync"); + try st.synchronize(); + const t1 = std.Io.Clock.awake.now(io).toNanoseconds(); + const r = try pool.plan(host, rows); + pool_time.cut_ns += @intCast(t1 - t0); + pool_time.plan_ns += @intCast(std.Io.Clock.awake.now(io).toNanoseconds() - t1); + pool_time.blocks += 1; + // segment B, a pass at a time: its copies and table rewrites, then the experts on the pool's tables + for (0..r.n_passes) |j| { + try pool.stage(r, j, st); + const last = j + 1 == r.n_passes; + try ex.run(st, layer, xin, dims, .bf16, pool.picksFor(r, j, picks), if (last) wts else 0, rows, if (last) out else 0, sc); + try pool.release(r, j, st); + } +} + +/// Every offloaded experts layer of a loaded engine (the 48 decoder layers). +pub fn offloadedLayers(layers: []const @import("../../cuda_prompt_load.zig").Layer, out: []*load.Experts) usize { + var n: usize = 0; + for (layers) |*l| if (l.exl3) |x| if (x.needsPool()) { + out[n] = x; + n += 1; + }; + return n; +} + +pub const Budget = struct { + /// Device bytes left free after the pools; null: the native server's own reserve (TENSORFOLD_MEMORY_RESERVE_GIB). + reserve: ?u64 = null, + /// What the pools leave beyond the server's reserve: allocator rounding and the copy streams' own memory. + margin: u64 = 256 << 20, + /// Bytes the family still allocates after the pools: each admitted stream's caches. + other: u64 = 0, + /// Lines a layer, or null: from free memory (cuda_exl3_round.linesFor). + lines: ?usize = null, + /// Rows a planned round holds (the pass's block of routed rows). + max_rows: usize = moe_rows, +}; + +/// TF_FLASHNEXT_EXPERT_LINES and TF_FLASHNEXT_EXPERT_RESERVE_MB override pool sizing. +pub fn budgetFromEnv(stream_bytes: u64) !Budget { + return budgetFromEnvWithStreams(stream_bytes, null); +} + +pub fn budgetForStreams(stream_bytes: u64, requested_streams: u32) !Budget { + return budgetFromEnvWithStreams(stream_bytes, requested_streams); +} + +fn budgetFromEnvWithStreams(stream_bytes: u64, requested_streams: ?u32) !Budget { + const streams = requested_streams orelse 1; + var b: Budget = .{ .other = try streamReservation(stream_bytes, streams) }; + if (std.c.getenv("TF_FLASHNEXT_EXPERT_LINES")) |v| b.lines = try std.fmt.parseInt(usize, std.mem.span(v), 10); + if (std.c.getenv("TF_FLASHNEXT_EXPERT_RESERVE_MB")) |v| b.reserve = try memory.mebibytes(std.mem.span(v)); + if (requested_streams == null) if (std.c.getenv("TF_FLASHNEXT_EXL3_STREAMS")) |v| { + const env_streams = try std.fmt.parseInt(u64, std.mem.span(v), 10); + if (env_streams == 0) return error.InvalidStreamCount; + b.other = try std.math.mul(u64, stream_bytes, env_streams); + }; + return b; +} + +fn streamReservation(stream_bytes: u64, streams: u32) !u64 { + if (streams == 0) return error.InvalidStreamCount; + return std.math.mul(u64, stream_bytes, streams); +} + +test "EXL3 pool reserves the requested number of concurrent stream caches" { + try std.testing.expectEqual(@as(u64, 200), try streamReservation(100, 2)); + try std.testing.expectError(error.InvalidStreamCount, streamReservation(100, 0)); +} + +pub const Attached = struct { layers: usize = 0, lines: usize = 0, line_bytes: usize = 0, device_bytes: u64 = 0, free_before: u64 = 0, free_after: u64 = 0, scratch_bytes: u64 = 0, bank_bytes: u64 = 0, bank_ns: u64 = 0 }; + +pub fn freeBytes(d: *const cuda.Driver) !u64 { + var free: usize = 0; + var total: usize = 0; + try d.check(d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + return free; +} + +/// The sizing hook (cuda_lanes_load.zig, once everything else of the family is allocated). +pub fn attach(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, st: cuda.Stream, layers: []const @import("../../cuda_prompt_load.zig").Layer, rows: usize, b: Budget, profile_counts: ?*[pool_.shape.layers][pool_.shape.experts]u32) !Attached { + var all: [wt.layers_n]*load.Experts = undefined; + const n = offloadedLayers(layers, &all); + const delta_text = if (std.c.getenv("TF_FLASHNEXT_EXL3_LINE_DELTAS")) |value| std.mem.span(value) else null; + var out: Attached = .{}; + const scratch_before = try freeBytes(d); + // every distinct Ops (the engine's store) reserves its scratch now + var seen: [4]*load.Ops = undefined; + var n_seen: usize = 0; + for (layers) |*l| if (l.exl3) |x| try reserveOnce(x.ops, &seen, &n_seen, st, rows); + if (n == 0) { + if (delta_text != null) return error.InvalidLineDeltaIndex; + return out; + } + for (all[0..n]) |x| { + out.line_bytes = @max(out.line_bytes, x.lineBytes()); + out.bank_bytes += x.bank.?.bytes.len; + out.bank_ns += x.bank_ns; + } + var free: usize = 0; + var total: usize = 0; + try d.check(d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + out.free_before = free; + out.scratch_bytes = scratch_before -| free; + const reserve_floor = b.reserve orelse serverReserve(total); + const reserve = b.reserve orelse try std.math.add(u64, reserve_floor, b.margin); + // beside its lines a layer holds the shared expert's region, its six tables and its masked picks (sized for the + var fixed: u64 = 0; + for (all[0..n]) |x| { + for (x.shared) |t| fixed = try std.math.add(u64, fixed, try std.math.mul(u64, (try std.math.add(u64, t.len, 255)) / 256, 256)); + const picks_bytes = try std.math.mul(u64, try std.math.mul(u64, pool_.passesFor(pool_.min_lines), b.max_rows), slots * 4); + fixed = try std.math.add(u64, fixed, try std.math.add(u64, cuda.exl3_pool.Tables.init(wt.experts_n + 1).bytes + (64 << 10), picks_bytes)); + } + const future = b.other; + const pool_room = try memory.available(free, reserve, future); + const lines = b.lines orelse try pool_.linesFor(free, reserve, try std.math.add(u64, future, fixed), out.line_bytes, n, pool_.min_lines); + if (lines < pool_.min_lines) return error.PoolTooSmall; + const startup_bytes = try memory.poolBytes(lines, out.line_bytes, n, fixed); + if (startup_bytes > pool_room) return error.InsufficientStartupMemory; + out.lines = lines; + var pool_lines: [wt.layers_n]usize = @splat(lines); + var line_bytes: [wt.layers_n]usize = undefined; + var profile_sized = false; + var profile_min_lines: usize = std.math.maxInt(usize); + var profile_max_lines: usize = 0; + for (all[0..n], 0..) |x, index| line_bytes[index] = x.lineBytes(); + if (delta_text == null and shouldProfileSize(profileSizingEnabled(), b.lines)) if (profile_counts) |prior| { + var count_rows: [wt.layers_n][]const u32 = undefined; + var allocation_costs: [wt.layers_n]usize = undefined; + var allocation_lines: [wt.layers_n]usize = undefined; + var decoder_index: usize = 0; + var has_profile_signal = false; + for (all[0..n], 0..) |_, pool_index| { + for (prior[decoder_index]) |count| { + has_profile_signal = has_profile_signal or count != 0; + } + count_rows[decoder_index] = &prior[decoder_index]; + allocation_costs[decoder_index] = line_bytes[pool_index]; + decoder_index += 1; + } + if (has_profile_signal and decoder_index == wt.layers_n) { + const maybe_plan: ?profile_line_budget.Budget = profile_line_budget.allocate(gpa, .{ + .counts = &count_rows, + .rank = 0, + .world = 1, + .line_bytes = &allocation_costs, + .uniform_lines = lines, + .min_lines = pool_.min_lines, + .max_rows = b.max_rows, + .slots = slots, + }, &allocation_lines) catch |err| blk: { + std.log.info("flashnext exl3: profile-sized pools unavailable ({s}); retaining uniform decoder pools", .{@errorName(err)}); + break :blk null; + }; + if (maybe_plan) |plan| { + decoder_index = 0; + for (all[0..n], 0..) |_, pool_index| { + pool_lines[pool_index] = allocation_lines[decoder_index]; + profile_min_lines = @min(profile_min_lines, allocation_lines[decoder_index]); + profile_max_lines = @max(profile_max_lines, allocation_lines[decoder_index]); + decoder_index += 1; + } + if (plan.uniform_fallback) { + @memset(pool_lines[0..n], lines); + std.log.info("flashnext exl3: profile candidate changed total variable CUDA bytes; retaining uniform decoder pools", .{}); + } else { + profile_sized = true; + } + std.log.info("flashnext exl3: profile-weighted decoder allocation used {d}/{d} variable CUDA bytes (expert lines and device picks); saved lifetime picks are a proxy, not latency", .{ plan.allocated_bytes, plan.baseline_bytes }); + std.log.info("flashnext exl3: profile decoder pool sizes range {d}..{d} lines; uniform-equivalent baseline was {d}", .{ profile_min_lines, profile_max_lines, lines }); + } + } + }; + if (delta_text) |text| { + const plan = try line_deltas.apply(text, lines, n, line_bytes[0..n], pool_.min_lines, fixed, startup_bytes, pool_lines[0..n]); + std.log.info("flashnext exl3: experimental rank-local decoder pool line deltas {s}; base={d} startup_bytes={d}", .{ text, lines, plan.pool_bytes }); + } + // the first decoder layer's pool owns the lifetime use table the other decoder layers share + var counts = profile_counts; + var decoder: u8 = 0; + for (all[0..n], pool_lines[0..n]) |x, count| { + try x.attach(io, d, decoder, count, b.max_rows, counts); + if (counts == null) counts = x.pool.?.counts; + decoder += 1; + out.device_bytes += x.pool.?.deviceBytes(); + out.layers += 1; + } + out.free_after = try freeBytes(d); + _ = memory.available(out.free_after, reserve_floor, future) catch |err| { + std.log.err("EXL3 pool reservation shortfall: free_before={d} free_after={d} pool_bytes={d} floor={d} allocator_margin={d} future={d}", .{ free, out.free_after, out.device_bytes, reserve_floor, b.margin, future }); + return err; + }; + const allocation = if (delta_text != null) "uniform-equivalent lines (experimental deltas applied)" else if (profile_sized) "uniform-equivalent lines (profile-sized decoder pools)" else "lines a layer"; + std.log.info("flashnext exl3: {d} offloaded expert layers, {d} {s} of {d:.2} MiB ({d:.2} GiB on the device, {d:.2} GiB free before); pinned host bank {d:.2} GiB copied in {d:.1} s", .{ out.layers, out.lines, allocation, @as(f64, @floatFromInt(out.line_bytes)) / (1 << 20), @as(f64, @floatFromInt(out.device_bytes)) / (1 << 30), @as(f64, @floatFromInt(out.free_before)) / (1 << 30), @as(f64, @floatFromInt(out.bank_bytes)) / (1 << 30), @as(f64, @floatFromInt(out.bank_ns)) / 1e9 }); + return out; +} + +test { + _ = line_deltas; + _ = profile_line_budget; +} + +fn profileSizingEnabled() bool { + const value = std.c.getenv("TF_FLASHNEXT_EXL3_PROFILE_SIZING") orelse return false; + return shouldProfileSize(std.mem.eql(u8, std.mem.span(value), "1"), null); +} + +fn shouldProfileSize(enabled: bool, explicit_lines: ?usize) bool { + return enabled and explicit_lines == null; +} + +test "exl3: profile sizing requires explicit opt-in and does not override fixed line budgets" { + try std.testing.expect(!shouldProfileSize(false, null)); + try std.testing.expect(!shouldProfileSize(true, 32)); + try std.testing.expect(shouldProfileSize(true, null)); +} + +/// native/cuda_memory.zig reserveBytes for a discrete card: TENSORFOLD_MEMORY_RESERVE_GIB, else max(1 GiB, total / 32). +fn serverReserve(total: u64) u64 { + if (std.c.getenv("TENSORFOLD_MEMORY_RESERVE_GIB")) |v| { + const g = std.fmt.parseFloat(f64, std.mem.trim(u8, std.mem.span(v), " ")) catch 0; + if (g >= 0.5) return @intFromFloat(g * (1 << 30)); + } + return @max(@as(u64, 1 << 30), total / 32); +} + +fn reserveOnce(o: *load.Ops, seen: []*load.Ops, n: *usize, st: cuda.Stream, rows: usize) !void { + for (seen[0..n.*]) |x| if (x == o) return; + if (n.* == seen.len) return error.Invalid; + seen[n.*] = o; + n.* += 1; + try o.reserve(st, rows, moe_rows, slots); +} + +/// Pool statistics summed over every attached layer. +pub fn poolStats(layers: []const @import("../../cuda_prompt_load.zig").Layer) pool_.Stats { + var t: pool_.Stats = .{}; + var all: [wt.layers_n]*load.Experts = undefined; + const n = offloadedPools(layers, &all); + for (all[0..n]) |x| { + const s = x.pool.?.stats; + inline for (@typeInfo(pool_.Stats).@"struct".field_names) |f| { + @field(t, f) += @field(s, f); + } + } + return t; +} + +pub fn offloadedPools(layers: []const @import("../../cuda_prompt_load.zig").Layer, out: []*load.Experts) usize { + var n: usize = 0; + for (layers) |*l| if (l.exl3) |x| if (x.pool != null) { + out[n] = x; + n += 1; + }; + return n; +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_ple.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_ple.zig new file mode 100644 index 0000000000..b49662450e --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_ple.zig @@ -0,0 +1,543 @@ +//! An EXL3 pack's n-gram table (ExLlamaV3 `ngram_embedding.safetensors`) behind `cuda_ple.Ngram` (`format = .exl3`). + +const std = @import("std"); +const core = @import("core"); +const ple = @import("../../cuda_ple.zig"); +const fmt = core.exl3_format; + +pub const file_name = "ngram_embedding.safetensors"; +/// Values a row holds: one head's slice of the 2,560-wide embedding. +pub const row_values: usize = 160; +pub const hot_slots: usize = 4096; +/// Bytes a hot slot holds (a row at 8 bits is 162). +pub const hot_cap: usize = 176; + +pub const Read = enum { map, pread }; +pub const Emb = enum { f16, bf16 }; +pub const Options = struct { read: Read = .map, emb: Emb = .f16 }; + +/// TF_FLASHNEXT_PLE_READ (map | pread); anything else keeps the default. The bf16 rows are a test option (`emb`). +pub fn optionsFromEnv() Options { + var o: Options = .{}; + if (std.c.getenv("TF_FLASHNEXT_PLE_READ")) |v| { + if (std.mem.eql(u8, std.mem.span(v), "pread")) o.read = .pread; + } + return o; +} + +// ---- the row codec (ported from exl3-flashnext-window-wip exl3_ngram.zig) ------------------------------------------ + +/// Int16 words in one row at `bits`: the fp16 scale, then 160 values. `bits` outside 2..8 is `NgramWidth`. +pub fn wordsPerRow(bits: usize) error{NgramWidth}!usize { + if (bits < 2 or bits > 8) return error.NgramWidth; + return 1 + row_values * bits / 16; +} + +/// Mul1 state of value `index`: bit m comes from position `index - m / bits` (plus 160 when negative), bit `m % bits`. +pub fn state(words: []const i16, bits: usize, index: usize) u16 { + var acc: u16 = 0; + for (0..16) |m| { + var pos: isize = @as(isize, @intCast(index)) - @as(isize, @intCast(m / bits)); + if (pos < 0) pos += @as(isize, @intCast(row_values)); + const sb: usize = @as(usize, @intCast(pos)) * bits + (m % bits); + const word: u16 = @bitCast(words[1 + (sb >> 4)]); + acc |= ((word >> @as(u4, @intCast(sb & 15))) & 1) << @as(u4, @intCast(m)); + } + return acc; +} + +/// One row into 160 fp16 values: one rounding of `f32(codebook) * f32(scale) + f32(bias[i])`. +pub fn decodeRow(words: []const i16, bits: usize, bias: []const f16, out: []f16) error{NgramWidth}!void { + const width = try wordsPerRow(bits); + if (words.len != width or bias.len != row_values or out.len != row_values) return error.NgramWidth; + const scale: f32 = @floatCast(@as(f16, @bitCast(words[0]))); + // `state(words, bits, i)` for every i at once: value j is stream bits [j * bits, (j + 1) * bits); state i is + var vals: [row_values]u32 = undefined; + const mask: u32 = (@as(u32, 1) << @intCast(bits)) - 1; + for (&vals, 0..) |*v, j| { + const sb = j * bits; + const wi = 1 + (sb >> 4); + const lo: u32 = @as(u16, @bitCast(words[wi])); + const hi: u32 = if (wi + 1 < words.len) @as(u16, @bitCast(words[wi + 1])) else 0; + v.* = ((lo | (hi << 16)) >> @intCast(sb & 15)) & mask; + } + for (out, bias, 0..) |*o, b, i| { + var acc: u32 = 0; + var k: usize = 0; + while (k * bits < 16) : (k += 1) acc |= vals[(i + row_values - k) % row_values] << @intCast(k * bits); + const cb: f32 = @floatCast(fmt.codebookValue(.mul1, @truncate(acc))); + o.* = @floatCast(cb * scale + @as(f32, @floatCast(b))); + } +} + +test "exl3 ple: the packed decode's states equal `state` for every value at every width" { + var prng = std.Random.DefaultPrng.init(11); + for (2..9) |bits| { + const n = try wordsPerRow(bits); + var words: [1 + row_values * 8 / 16]i16 = undefined; + for (words[0..n]) |*w| w.* = @bitCast(prng.random().int(u16)); + words[0] = @bitCast(@as(f16, 0.5)); + var bias: [row_values]f16 = undefined; + for (&bias) |*b| b.* = @floatCast(prng.random().float(f32) - 0.5); + var got: [row_values]f16 = undefined; + try decodeRow(words[0..n], bits, &bias, &got); + const scale: f32 = @floatCast(@as(f16, @bitCast(words[0]))); + for (got, bias, 0..) |g, b, i| { + const cb: f32 = @floatCast(fmt.codebookValue(.mul1, state(words[0..n], bits, i))); + const want: f16 = @floatCast(cb * scale + @as(f32, @floatCast(b))); + try std.testing.expectEqual(@as(u16, @bitCast(want)), @as(u16, @bitCast(g))); + } + } +} + +fn bf16Of(v: f32) u16 { + return ple.toBf16(v); +} + +/// Group sums of 32 fp16 values (the tree `cuda_ple.sumGroups` adds bf16 values in), for an fp16 row. +pub fn sumGroupsF16(row: []const u16, out: []f32) void { + for (0..row.len / 32) |g| { + var buf: [32]f32 = undefined; + for (0..32) |i| buf[i] = @floatCast(@as(f16, @bitCast(row[g * 32 + i]))); + var n: usize = 32; + while (n > 1) { + n /= 2; + for (0..n) |i| buf[i] = buf[2 * i] + buf[2 * i + 1]; + } + out[g] = buf[0]; + } +} + +// ---- the table ---------------------------------------------------------------------------------------------------- + +const c = std.c; +const is_windows = @import("builtin").os.tag == .windows; + +/// The open table. Heap-allocated (`open`/`destroy`): the pread reader's hot cache is 512 KiB. +pub const Table = struct { + gpa: std.mem.Allocator, + io: std.Io, + opts: Options, + heads: usize, + shards: usize, + bits: usize, + words: usize, + /// File byte offset of each shard's row 0, and its row count. + offset: [128]u64 = undefined, + rows: [128]u64 = undefined, + bias: []f16, + /// `.map`: the whole file's mapping (header parsed, nothing populated). + file: ?core.safetensors.File = null, + /// `.pread`: the descriptor; the file is never mapped past `open`. + fd: c.fd_t = if (is_windows) undefined else -1, + hot: ?*Hot = null, + + pub const Hot = struct { + lock: std.atomic.Value(u32) = .init(0), + lookups: u64 = 0, + misses: u64 = 0, + tags: [hot_slots]u64 = @splat(std.math.maxInt(u64)), + stored: [hot_slots][hot_cap]u8 = undefined, + + fn acquire(h: *Hot) void { + while (h.lock.cmpxchgWeak(0, 1, .acquire, .monotonic) != null) std.atomic.spinLoopHint(); + } + fn release(h: *Hot) void { + h.lock.store(0, .release); + } + }; + + fn prefix(buf: []u8, part: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "model.language_model.layers.{d}.ple.ple_embedding.ngram_embedding.{s}", .{ ple.layer, part }); + } + + /// The pack's table when `dir` holds an EXL3 n-gram file whose constants equal `hash`'s; null. + pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, hash: ple.Hash, opts: Options) !?*Table { + const path = try std.fs.path.join(gpa, &.{ dir, file_name }); + defer gpa.free(path); + var file = core.safetensors.File.open(gpa, io, path) catch |err| switch (err) { + error.FileNotFound => return null, + else => return err, + }; + var keep_file = false; + defer if (!keep_file) file.close(io); + var nb: [200]u8 = undefined; + const first = file.get(try prefix(&nb, "shard_0.trellis")) orelse return null; + if (first.dtype != .i16 or first.rank != 2) return error.UnexpectedTensor; + if (hash.shards > 128 or hash.dh != row_values) return error.UnexpectedTensor; + // the hash constants the pack ships against the ones the config derives + { + var got: [ple.max_heads]i64 = undefined; + try ple.readI64(file.get(try prefix(&nb, "layer_multipliers")) orelse return error.MissingTensor, got[0..hash.n]); + if (!std.mem.eql(i64, got[0..hash.n], hash.mult[0..hash.n])) return error.NgramConstants; + try ple.readI64(file.get(try prefix(&nb, "head_vocab_sizes")) orelse return error.MissingTensor, got[0..hash.heads]); + if (!std.mem.eql(i64, got[0..hash.heads], hash.sizes[0..hash.heads])) return error.NgramConstants; + try ple.readI64(file.get(try prefix(&nb, "head_offsets")) orelse return error.MissingTensor, got[0..hash.heads]); + if (!std.mem.eql(i64, got[0..hash.heads], hash.offsets[0..hash.heads])) return error.NgramConstants; + } + // the reference's EXL3 path applies no table scale: refuse a file that carries one + if (file.get(try prefix(&nb, "weight_scale")) != null) return error.UnexpectedTensor; + const bias_t = file.get(try prefix(&nb, "head_bias")) orelse return error.MissingTensor; + if (!bias_t.is(.f16, &.{ hash.heads, row_values })) return error.UnexpectedTensor; + const words = first.dim(1); + const bits = (words - 1) * 16 / row_values; + if ((wordsPerRow(bits) catch return error.UnexpectedTensor) != words) return error.UnexpectedTensor; + + const t = try gpa.create(Table); + errdefer gpa.destroy(t); + t.* = .{ .gpa = gpa, .io = io, .opts = opts, .heads = hash.heads, .shards = hash.shards, .bits = bits, .words = words, .bias = try gpa.alloc(f16, hash.heads * row_values) }; + errdefer gpa.free(t.bias); + @memcpy(std.mem.sliceAsBytes(t.bias), bias_t.bytes); + for (0..hash.shards) |i| { + var part: [32]u8 = undefined; + const e = file.names.get(try prefix(&nb, try std.fmt.bufPrint(&part, "shard_{d}.trellis", .{i}))) orelse return error.MissingTensor; + const want: usize = @intCast(hash.starts[i + 1] - hash.starts[i]); + if (e.dtype != .i16 or e.rank != 2 or e.shape[0] != want or e.shape[1] != words) return error.UnexpectedTensor; + t.offset[i] = file.data + e.begin; + t.rows[i] = want; + } + switch (opts.read) { + .map => { + if (@import("builtin").os.tag == .linux) { + if (file.map.section != null) _ = std.os.linux.madvise(file.map.memory.ptr, file.map.memory.len, std.os.linux.MADV.RANDOM); + } + t.file = file; + keep_file = true; + }, + .pread => { + if (comptime is_windows) return error.UnsupportedPlatform; + const z = try gpa.dupeSentinel(u8, path, 0); + defer gpa.free(z); + const fd = c.open(z.ptr, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return error.OpenFailed; + t.fd = fd; + errdefer _ = c.close(fd); + if (@import("builtin").os.tag == .linux) _ = std.os.linux.fadvise(fd, 0, 0, std.os.linux.POSIX_FADV.RANDOM); + const hot = try gpa.create(Hot); + hot.* = .{}; + t.hot = hot; + }, + } + return t; + } + + pub fn destroy(t: *Table) void { + if (t.file) |*f| f.close(t.io); + if (comptime !is_windows) { + if (t.fd >= 0) _ = c.close(t.fd); + } + if (t.hot) |h| t.gpa.destroy(h); + t.gpa.free(t.bias); + t.gpa.destroy(t); + } + + pub fn rowBytes(t: *const Table) usize { + return t.words * 2; + } + + fn preadAll(t: *const Table, dest: []u8, at: u64) !void { + if (comptime is_windows) return error.UnsupportedPlatform; + var done: usize = 0; + while (done < dest.len) { + const n = c.pread(t.fd, dest.ptr + done, dest.len - done, @intCast(at + done)); + if (n <= 0) return error.ReadFailed; + done += @intCast(n); + } + } + + /// The stored bytes of row `local` of shard `shard` (`rowBytes` long), through the mapping or the hot cache. + pub fn raw(t: *const Table, shard: usize, local: usize, dest: []u8) !void { + if (shard >= t.shards or local >= t.rows[shard] or dest.len != t.rowBytes()) return error.UnexpectedTensor; + const at = t.offset[shard] + @as(u64, local) * t.rowBytes(); + if (t.file) |f| { + @memcpy(dest, f.map.memory[at..][0..dest.len]); + return; + } + const h = t.hot.?; + const slot: usize = @intCast(((at *% 0x9E3779B97F4A7C15) >> 20) % hot_slots); + h.acquire(); + h.lookups += 1; + if (h.tags[slot] == at) { + @memcpy(dest, h.stored[slot][0..dest.len]); + h.release(); + return; + } + h.release(); + try t.preadAll(dest, at); + h.acquire(); + h.misses += 1; + h.tags[slot] = at; + @memcpy(h.stored[slot][0..dest.len], dest); + h.release(); + } + + /// Row `local` of shard `shard`, decoded with head `head`'s bias into 160 fp16 (or bf16) bit patterns. + pub fn row(t: *const Table, shard: usize, local: usize, head: usize, out: []u16) !void { + if (head >= t.heads or out.len != row_values) return error.UnexpectedTensor; + var bytes: [hot_cap]u8 align(2) = undefined; + const rb = t.rowBytes(); + try t.raw(shard, local, bytes[0..rb]); + const words: []const i16 = @as([*]const i16, @ptrCast(&bytes))[0..t.words]; + var vals: [row_values]f16 = undefined; + try decodeRow(words, t.bits, t.bias[head * row_values ..][0..row_values], &vals); + switch (t.opts.emb) { + .f16 => for (out, vals) |*o, v| { + o.* = @bitCast(v); + }, + .bf16 => for (out, vals) |*o, v| { + o.* = bf16Of(@floatCast(v)); + }, + } + } + + /// Group sums of one decoded row in the table's element type. + pub fn sums(t: *const Table, row_bits: []const u16, out: []f32) void { + switch (t.opts.emb) { + .f16 => sumGroupsF16(row_bits, out), + .bf16 => ple.sumGroups(row_bits, out), + } + } + + pub fn hotStats(t: *const Table) struct { lookups: u64, misses: u64 } { + const h = t.hot orelse return .{ .lookups = 0, .misses = 0 }; + return .{ .lookups = h.lookups, .misses = h.misses }; + } +}; + +// ---- tests -------------------------------------------------------------------------------------------------------- + +/// The Triton `_ple_rows` formulation, a second decoder that shares only the three fp16 constants with `decodeRow`. +fn decodeRowRef(words: []const i16, bits: usize, bias: []const f16, out: []f16) void { + const k_inv: f32 = @floatCast(@as(f16, @bitCast(fmt.Mul1ScaleBits))); + const k_bias: f32 = @floatCast(@as(f16, @bitCast(fmt.Mul1BiasBits))); + const scale: f32 = @floatCast(@as(f16, @bitCast(words[0]))); + for (0..out.len) |i| { + var st: u64 = 0; + for (0..16) |m| { + var pos: isize = @as(isize, @intCast(i)) - @as(isize, @intCast(m / bits)); + if (pos < 0) pos += @as(isize, @intCast(out.len)); + const sb: usize = @as(usize, @intCast(pos)) * bits + (m % bits); + const word: u64 = @as(u16, @bitCast(words[1 + (sb >> 4)])); + st |= ((word >> @intCast(sb & 15)) & 1) << @intCast(m); + } + const prod: u64 = (st * 0x83DCD12D) & 0xFFFFFFFF; + const hs: u64 = (prod & 255) + ((prod >> 8) & 255) + ((prod >> 16) & 255) + ((prod >> 24) & 255); + const cb: f16 = @floatCast(@as(f32, @floatFromInt(1024 + hs)) * k_inv + k_bias); + out[i] = @floatCast(@as(f32, @floatCast(cb)) * scale + @as(f32, @floatCast(bias[i]))); + } +} + +test "exl3 ngram: wordsPerRow is one scale word plus 160 K-bit values" { + try std.testing.expectEqual(@as(usize, 21), try wordsPerRow(2)); + try std.testing.expectEqual(@as(usize, 31), try wordsPerRow(3)); + try std.testing.expectEqual(@as(usize, 51), try wordsPerRow(5)); + try std.testing.expectEqual(@as(usize, 81), try wordsPerRow(8)); + try std.testing.expectError(error.NgramWidth, wordsPerRow(1)); + try std.testing.expectError(error.NgramWidth, wordsPerRow(9)); +} + +test "exl3 ngram: patterned rows read tail-biting states" { + var storage: [81]i16 = undefined; + for (&storage, 0..) |*w, i| w.* = @bitCast(@as(u16, @intCast(i))); + try std.testing.expectEqual(@as(u16, 0x5001), state(storage[0..21], 2, 0)); + try std.testing.expectEqual(@as(u16, 0x9001), state(storage[0..31], 3, 0)); + try std.testing.expectEqual(@as(u16, 0x8001), state(storage[0..51], 5, 0)); + try std.testing.expectEqual(@as(u16, 0x0020), state(storage[0..51], 5, 1)); + try std.testing.expectEqual(@as(u16, 0x0400), state(storage[0..51], 5, 2)); + try std.testing.expectEqual(@as(u16, 0x0001), state(storage[0..81], 8, 0)); + // value 0 at 5 bits: bit 5 wraps to position 159, word 50 bit 11 + var words: [51]i16 = @splat(0); + words[50] = @bitCast(@as(u16, 1) << 11); + try std.testing.expectEqual(@as(u16, 1) << 5, state(&words, 5, 0)); +} + +test "exl3 ngram: scale and bias round in f32, then once to f16" { + var words: [51]i16 = @splat(0); + words[0] = @bitCast(@as(u16, 0x3C01)); + var bias: [160]f16 = @splat(0); + bias[0] = @bitCast(@as(u16, 0x1400)); + var out: [160]f16 = undefined; + try decodeRow(&words, 5, &bias, &out); + // codebook(0) is 0xC2E8; times 0x3C01 plus 0x1400 is 0xC2E9 (rounding the product first would give 0xC2EA) + try std.testing.expectEqual(@as(u16, 0xC2E8), @as(u16, @bitCast(fmt.codebookValue(.mul1, 0)))); + try std.testing.expectEqual(@as(u16, 0xC2E9), @as(u16, @bitCast(out[0]))); + try std.testing.expectError(error.NgramWidth, decodeRow(words[0..50], 5, &bias, &out)); + try std.testing.expectError(error.NgramWidth, decodeRow(&words, 2, &bias, &out)); +} + +test "exl3 ngram: decodeRow equals the Triton formulation on random rows at every width" { + var rng = std.Random.DefaultPrng.init(0x91e); + const r = rng.random(); + var words: [81]i16 = undefined; + var bias: [160]f16 = undefined; + var a: [160]f16 = undefined; + var b: [160]f16 = undefined; + for ([_]usize{ 2, 3, 4, 5, 6, 8 }) |bits| { + const n = try wordsPerRow(bits); + for (0..300) |_| { + for (words[0..n]) |*x| x.* = @bitCast(r.int(u16)); + words[0] = @bitCast(@as(f16, @floatCast((r.float(f32) - 0.5) * 4))); + for (&bias) |*x| x.* = @floatCast((r.float(f32) - 0.5) * 0.25); + try decodeRow(words[0..n], bits, &bias, &a); + decodeRowRef(words[0..n], bits, &bias, &b); + try std.testing.expectEqualSlices(u16, @ptrCast(&b), @ptrCast(&a)); + } + } +} + +/// A small table file plus its config: 2 shards, 4 heads of 160 (n-gram 3, 2 heads an n-gram), 5-bit rows. +const Synth = struct { + const vocab = 100; + const eos = 99; + const base = 11; + + fn spec() ple.Hash { + return ple.Hash.init(.{ .vocab = vocab, .ngram = 3, .per = 2, .base = base, .divisor = 4, .shards = 2, .seed = 1234, .eos = eos, .ple_dim = 4 * row_values, .ple_index = 0 }) catch unreachable; + } + + fn put(list: *std.ArrayList(u8), gpa: std.mem.Allocator, bytes: []const u8) !void { + try list.appendSlice(gpa, bytes); + } + + /// Writes config.json, an empty index and ngram_embedding.safetensors under `dir`; row r's words are a hash of r. + fn write(gpa: std.mem.Allocator, io: std.Io, dir: std.Io.Dir) !void { + const h = spec(); + var data: std.ArrayList(u8) = .empty; + defer data.deinit(gpa); + var header: std.ArrayList(u8) = .empty; + defer header.deinit(gpa); + try header.appendSlice(gpa, "{"); + const p = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding."; + var first = true; + const Entry = struct { + fn add(hd: *std.ArrayList(u8), a: std.mem.Allocator, f: *bool, name: []const u8, dtype: []const u8, shape: []const usize, lo: usize, hi: usize) !void { + if (!f.*) try hd.appendSlice(a, ","); + f.* = false; + var buf: [400]u8 = undefined; + var sb: [64]u8 = undefined; + var sl: usize = 0; + for (shape, 0..) |d, i| { + const s = try std.fmt.bufPrint(sb[sl..], "{s}{d}", .{ if (i == 0) "" else ",", d }); + sl += s.len; + } + try hd.appendSlice(a, try std.fmt.bufPrint(&buf, "\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[{s}],\"data_offsets\":[{d},{d}]}}", .{ name, dtype, sb[0..sl], lo, hi })); + } + }; + var nb: [200]u8 = undefined; + for (0..h.shards) |s| { + const rows: usize = @intCast(h.starts[s + 1] - h.starts[s]); + const lo = data.items.len; + for (0..rows) |r| { + const global: u64 = @intCast(h.starts[s] + @as(i64, @intCast(r))); + var prng = std.Random.DefaultPrng.init(global + 1); + for (0..51) |w| { + var v: u16 = prng.random().int(u16); + if (w == 0) v = @bitCast(@as(f16, @floatCast(0.25 + 0.01 * @as(f32, @floatFromInt(global))))); + try data.appendSlice(gpa, &std.mem.toBytes(v)); + } + } + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}shard_{d}.trellis", .{ p, s }), "I16", &.{ rows, 51 }, lo, data.items.len); + } + var lo = data.items.len; + for (0..h.heads * row_values) |i| try data.appendSlice(gpa, &std.mem.toBytes(@as(f16, @floatCast(@as(f32, @floatFromInt(i % 7)) * 0.01)))); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}head_bias", .{p}), "F16", &.{ h.heads, row_values }, lo, data.items.len); + lo = data.items.len; + for (h.mult[0..h.n]) |v| try data.appendSlice(gpa, &std.mem.toBytes(v)); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}layer_multipliers", .{p}), "I64", &.{h.n}, lo, data.items.len); + lo = data.items.len; + for (h.sizes[0..h.heads]) |v| try data.appendSlice(gpa, &std.mem.toBytes(v)); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}head_vocab_sizes", .{p}), "I64", &.{h.heads}, lo, data.items.len); + lo = data.items.len; + for (h.offsets[0..h.heads]) |v| try data.appendSlice(gpa, &std.mem.toBytes(v)); + try Entry.add(&header, gpa, &first, try std.fmt.bufPrint(&nb, "{s}head_offsets", .{p}), "I64", &.{h.heads}, lo, data.items.len); + try header.appendSlice(gpa, "}"); + var out: std.ArrayList(u8) = .empty; + defer out.deinit(gpa); + try out.appendSlice(gpa, &std.mem.toBytes(@as(u64, header.items.len))); + try out.appendSlice(gpa, header.items); + try out.appendSlice(gpa, data.items); + try dir.writeFile(io, .{ .sub_path = file_name, .data = out.items }); + try dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "{\"weight_map\":{}}" }); + try dir.writeFile(io, .{ .sub_path = "config.json", .data = + \\{"text_config":{"vocab_size":100,"eos_token_id":99,"num_hidden_layers":4,"hidden_size":640,"ple_layer_ids":[2], + \\"ngram_size":3,"heads_per_ngram":2,"ngram_vocab_size_base":11,"make_ngram_vocab_size_divisible_by":4, + \\"split_ngram_parts":2,"ple_embed_dim":640}} + }); + } +}; + +test "exl3 ngram: Ngram.rows on a synthetic EXL3 table: hook, map and pread readers, hot cache" { + if (comptime is_windows) return error.SkipZigTest; // the pread reader is POSIX-only + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try Synth.write(gpa, io, tmp.dir); + const dir = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(dir); + + var g = try ple.Ngram.init(gpa, io, dir); + defer g.deinit(); + try std.testing.expect(g.format == .exl3); + const h = g.hash; + try std.testing.expectEqual(@as(usize, 4), h.heads); + const tokens = [_]u32{ 5, 7, Synth.eos, 7, 7, 3 }; + const w = g.width(); + var out: [tokens.len * 4 * row_values]u16 = undefined; + var xs: [tokens.len * 4 * row_values / 32]f32 = undefined; + var history = [_]i64{ Synth.eos, Synth.eos }; + try g.rows(&history, &tokens, &out, &xs); + try std.testing.expectEqualSlices(i64, &.{ 7, 3 }, &history); + + // the same rows by hand: upstream's ids, the shard's bytes, decodeRow + var pread_t = (try Table.open(gpa, io, dir, h, .{ .read = .pread })).?; + defer pread_t.destroy(); + var hist = [_]i64{ Synth.eos, Synth.eos }; + for (tokens, 0..) |tok, t| { + var ids: [4]i64 = undefined; + try ple.lastIds(h, &hist, tok, &ids); + hist = .{ hist[1], tok }; + for (ids, 0..) |id, head| { + const loc = try h.locate(id); + var bytes: [102]u8 align(2) = undefined; + try pread_t.raw(loc.shard, loc.local, &bytes); + var mapped: [102]u8 = undefined; + try g.x3.?.raw(loc.shard, loc.local, &mapped); + try std.testing.expectEqualSlices(u8, &mapped, &bytes); + var vals: [row_values]f16 = undefined; + try decodeRow(@as([*]const i16, @ptrCast(&bytes))[0..51], 5, pread_t.bias[head * row_values ..][0..row_values], &vals); + const got = out[t * w + head * row_values ..][0..row_values]; + try std.testing.expectEqualSlices(u16, @ptrCast(&vals), got); + var gs: [5]f32 = undefined; + sumGroupsF16(got, &gs); + try std.testing.expectEqualSlices(f32, &gs, xs[(t * w + head * row_values) / 32 ..][0..5]); + } + } + // tokens 2 and 3 after an eos: token 7 twice in a fresh segment reads the same rows; the cache serves the repeat + const before = pread_t.hotStats(); + var again: [102]u8 = undefined; + const loc = try h.locate(h.offsets[0]); + try pread_t.raw(loc.shard, loc.local, &again); + try pread_t.raw(loc.shard, loc.local, &again); + const after = pread_t.hotStats(); + try std.testing.expectEqual(before.lookups + 2, after.lookups); + try std.testing.expect(after.misses <= before.misses + 1); + // bf16 elements: the fp16 values rounded once more + var bf_t = (try Table.open(gpa, io, dir, h, .{ .emb = .bf16 })).?; + defer bf_t.destroy(); + var r16: [row_values]u16 = undefined; + var rbf: [row_values]u16 = undefined; + try g.x3.?.row(loc.shard, loc.local, 0, &r16); + try bf_t.row(loc.shard, loc.local, 0, &rbf); + for (r16, rbf) |a, b| try std.testing.expectEqual(ple.toBf16(@floatCast(@as(f16, @bitCast(a)))), b); +} + +test "exl3 ngram: a directory without the n-gram file is not an EXL3 table" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const dir = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(dir); + try std.testing.expect((try Table.open(gpa, io, dir, Synth.spec(), .{})) == null); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig new file mode 100644 index 0000000000..f6f245a852 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry.zig @@ -0,0 +1,85 @@ +const std = @import("std"); +const prompt = @import("../../cuda_prompt.zig"); +const pool = @import("../offload/cuda_exl3_pool.zig"); +const pass = @import("cuda_exl3_pass.zig"); +const weights = @import("../../cuda_prompt_wt.zig"); +pub const counters = @import("exl3_telemetry.zig"); +pub const Completion = enum { all_rank_consumers_joined }; + +pub fn sampleAfterConsumers(allocator: std.mem.Allocator, engine: *prompt.Engine, sequences: []const *prompt.Seq, rank: counters.Rank, completion: Completion) !counters.Snapshot { + _ = completion; + try rank.validate(); + var hash = std.hash.Wyhash.init(@intFromPtr(engine)); + var totals: counters.Experts = .{}; + var attached: [weights.layers_n]*@import("../pack/cuda_exl3_load.zig").Experts = undefined; + const attached_count = pass.offloadedPools(engine.layers[0..engine.n_layers], &attached); + var layers: []counters.LayerSample = if (engine.telemetry.layers_enabled) try allocator.alloc(counters.LayerSample, attached_count) else @constCast(&.{}); + errdefer allocator.free(layers); + var seen: [weights.layers_n]*pool.Pool = undefined; + var seen_count: usize = 0; + for (attached[0..attached_count]) |experts| { + const resident = experts.pool orelse continue; + var duplicate = false; + for (seen[0..seen_count]) |previous| if (previous == resident) { + duplicate = true; + break; + }; + if (duplicate) continue; + if (resident.telemetry_id == 0) return error.TelemetryOwnerChanged; + if (!resident.telemetry_valid) return error.TelemetryCounterRegression; + seen[seen_count] = resident; + seen_count += 1; + hash.update(std.mem.asBytes(&resident.telemetry_id)); + const device = resident.telemetry_device; + var minimum: u64 = std.math.maxInt(u64); + var maximum: u64 = 0; + for (resident.src.routed, resident.k2) |source, widths| { + if (widths[0] == 0) continue; + const bytes = try std.math.add(u64, try std.math.add(u64, source[0].len, source[1].len), source[2].len); + minimum = @min(minimum, bytes); + maximum = @max(maximum, bytes); + } + if (minimum == std.math.maxInt(u64)) { + if (device.copies != 0) return error.TelemetryInconsistentCounters; + minimum = 0; + } + const host = try counters.difference(counters.Device, resident.telemetry_device, .{ .plans = resident.stats.rounds, .hit_picks = resident.stats.hit_slots, .miss_picks = resident.stats.miss_slots, .copies = resident.stats.misses, .evictions = resident.stats.evictions, .inplace_experts = resident.telemetry_device.inplace_experts, .overflow_plans = resident.telemetry_device.overflow_plans }); + const contribution: counters.Experts = .{ + .host_plans = host.plans, + .host_hit_picks = host.hit_picks, + .host_miss_picks = host.miss_picks, + .host_mandatory_copies = host.copies, + .host_prefetch_copies = resident.stats.prefetches, + .pool_dma_copies = resident.lines.stats.copies, + .pool_dma_requested_bytes = resident.lines.stats.bytes, + .device_plans = device.plans, + .device_hit_picks = device.hit_picks, + .device_miss_picks = device.miss_picks, + .device_copies = device.copies, + .device_inplace_experts = device.inplace_experts, + .device_overflow_plans = device.overflow_plans, + .mapped_copy_payload_min_bytes = try std.math.mul(u64, device.copies, minimum), + .mapped_copy_payload_max_bytes = try std.math.mul(u64, device.copies, maximum), + }; + try counters.add(counters.Experts, &totals, contribution); + if (engine.telemetry.layers_enabled) { + var layer_sample: counters.LayerSample = .{ .owner = resident.telemetry_id, .kind = .decoder, .layer = 0, .lines = resident.opt.lines, .counts = counters.LayerCounts.fromExperts(contribution) }; + for (engine.layers[0..engine.n_layers], 0..) |layer, decoder_index| if (layer.exl3) |decoder| { + if (decoder.pool == resident) { + layer_sample.layer = decoder_index; + break; + } + }; + layers[seen_count - 1] = layer_sample; + } + } + for (sequences, 0..) |sequence, index| for (sequences[0..index]) |previous| if (previous == sequence) return error.TelemetryDuplicateOwner; + const sampled_layers = if (engine.telemetry.layers_enabled) try allocator.realloc(layers, seen_count) else layers; + return .{ .rank = rank, .scope = hash.final(), .experts = totals, .expert_pools = seen_count, .layers = sampled_layers }; +} + +test "exl3: cache telemetry pure host contracts" { + _ = counters; + _ = @import("cuda_exl3_telemetry_test.zig"); + _ = @import("../../cuda_prompt_telemetry.zig"); +} diff --git a/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig new file mode 100644 index 0000000000..69e0ded0c5 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/cuda_exl3_telemetry_test.zig @@ -0,0 +1,38 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const pool = @import("../offload/cuda_exl3_pool.zig"); +const telemetry = @import("exl3_telemetry.zig"); + +test "exl3: aggregate telemetry finds attached pools deduplicates pool aliases and separates committed device plans" { + var resident: pool.Pool = undefined; + resident.lines.pinned_bank = null; + resident.telemetry_id = 991; + resident.telemetry_valid = true; + resident.telemetry_device = .{ .plans = 5, .hit_picks = 20, .miss_picks = 10, .copies = 3, .inplace_experts = 1 }; + resident.stats = .{ .rounds = 7, .hit_slots = 25, .miss_slots = 13, .misses = 5 }; + resident.lines.stats = .{ .copies = 2, .bytes = 96 }; + const payload: [16]u8 = @splat(0); + var sources = [_][3][]const u8{.{ &payload, &payload, &payload }}; + var widths = [_][3]i32{.{ 1, 1, 1 }}; + resident.src = .{ .routed = &sources }; + resident.k2 = &widths; + var experts: @import("../pack/cuda_exl3_load.zig").Experts = undefined; + experts.pool = &resident; + var engine: @import("../../cuda_prompt.zig").Engine = undefined; + engine.telemetry = .{}; + engine.n_layers = 1; + engine.layers[0].exl3 = &experts; + const collector = @import("cuda_exl3_telemetry.zig"); + var snapshot = try collector.sampleAfterConsumers(std.testing.allocator, &engine, &.{}, .{ .rank = 0, .world = 1 }, .all_rank_consumers_joined); + defer snapshot.deinit(std.testing.allocator); + try std.testing.expectEqual(@as(usize, 1), snapshot.expert_pools); + try std.testing.expectEqual(@as(u64, 2), snapshot.experts.host_plans); + try std.testing.expectEqual(@as(u64, 5), snapshot.experts.host_hit_picks); + try std.testing.expectEqual(@as(u64, 3), snapshot.experts.host_miss_picks); + try std.testing.expectEqual(@as(u64, 2), snapshot.experts.host_mandatory_copies); + try std.testing.expectEqual(@as(u64, 5), snapshot.experts.device_plans); + try std.testing.expectEqual(@as(u64, 1), snapshot.experts.device_inplace_experts); + try std.testing.expectEqual(@as(u64, 144), snapshot.experts.mapped_copy_payload_min_bytes); + try std.testing.expectEqual(@as(u64, 144), snapshot.experts.mapped_copy_payload_max_bytes); + try std.testing.expectEqual(@as(u64, 96), snapshot.experts.pool_dma_requested_bytes); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3.zig b/zig/src/families/flashnext/exl3/model/exl3.zig new file mode 100644 index 0000000000..73e5d75ec4 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3.zig @@ -0,0 +1,33 @@ +//! Host-side EXL3 index for Qwen3.8-Flash-Next: the float64 oracles and header rules kept as tests under option C + +pub const shape = @import("../pack/exl3_shape.zig"); +pub const scan = @import("../pack/exl3_scan.zig"); +pub const cache = @import("exl3_cache.zig"); +pub const cpu = @import("exl3_cpu.zig"); +pub const place = @import("../pack/exl3_place.zig"); +pub const route = @import("exl3_route.zig"); +pub const router = @import("exl3_router.zig"); +pub const kv = @import("../kv/exl3_kv.zig"); +pub const norm = @import("exl3_norm.zig"); +pub const hc = @import("exl3_hc.zig"); +pub const moe = @import("exl3_moe.zig"); +pub const gdn_cpu = @import("exl3_gdn_cpu.zig"); +pub const attn_cpu = @import("exl3_attn_cpu.zig"); +pub const fwd_cpu = @import("exl3_fwd_cpu.zig"); + +test { + _ = shape; + _ = scan; + _ = cache; + _ = cpu; + _ = place; + _ = route; + _ = router; + _ = kv; + _ = norm; + _ = hc; + _ = moe; + _ = gdn_cpu; + _ = attn_cpu; + _ = fwd_cpu; +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_attn_cpu.zig b/zig/src/families/flashnext/exl3/model/exl3_attn_cpu.zig new file mode 100644 index 0000000000..f276d4c6f0 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_attn_cpu.zig @@ -0,0 +1,441 @@ +//! CPU oracle for one Qwen3.8-Flash-Next full-attention layer (layers 3, 7, ... 47). + +const std = @import("std"); +const shape = @import("../pack/exl3_shape.zig"); + +pub const q_heads: usize = shape.heads; +pub const kv_heads: usize = shape.kv_heads; +pub const head_dim: usize = shape.head_dim; +pub const rotary_dim: usize = 64; +pub const half: usize = rotary_dim / 2; +pub const index_heads: usize = 4; +pub const index_dim: usize = 128; +pub const ratio: usize = 4; +pub const top_blocks: usize = 512; +pub const theta: f64 = 1.0e7; +pub const eps: f64 = 1.0e-6; +/// Most keys one row ever reads: 512 selected blocks and a 3-key tail, or a dense row of up to 2051 keys. +pub const max_keys: usize = top_blocks * ratio + ratio; +pub const index_width: usize = (index_heads + 1) * index_dim; + +pub fn roundBf16(x: f64) f64 { + const f: f32 = @floatCast(x); + var u: u32 = @bitCast(f); + if (!std.math.isFinite(f)) return f; + // round to nearest even on the low 16 bits + const lsb = (u >> 16) & 1; + u +%= 0x7fff + lsb; + u &= 0xffff0000; + return @as(f32, @bitCast(u)); +} + +pub fn roundF16(x: f64) f64 { + const h: f16 = @floatCast(x); + return h; +} + +/// theta^(-2 i / rotary_dim) as the f32 table the kernels read. +pub fn invFreq() [half]f32 { + var t: [half]f32 = undefined; + for (&t, 0..) |*v, i| v.* = @floatCast(std.math.pow(f64, theta, -2.0 * @as(f64, @floatFromInt(i)) / @as(f64, rotary_dim))); + return t; +} + +pub const Norms = struct { + /// Scales with the centred offset already added: `w + offset`. + q: []const f32, + k: []const f32, + iq: []const f32, + ik: []const f32, + inv: [half]f32, +}; + +/// fp32 RMSNorm, bf16, rotate-half RoPE over the first 64 dims at `pos`, bf16. `x` and `out` are `width` long. +pub fn normRope(x: []const f64, scale: []const f32, pos: f64, inv: *const [half]f32, out: []f64) void { + const width = x.len; + var ss: f64 = 0; + for (x) |v| ss += v * v; + const rinv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(width)) + eps); + var xn: [head_dim]f64 = undefined; + for (x, 0..) |v, i| xn[i] = roundBf16(v * rinv * @as(f64, scale[i])); + const posf: f32 = @floatCast(pos); + for (0..width) |d| { + if (d < half) { + const ang: f32 = posf * inv[d]; + out[d] = roundBf16(xn[d] * @cos(@as(f64, ang)) - xn[d + half] * @sin(@as(f64, ang))); + } else if (d < 2 * half) { + const ang: f32 = posf * inv[d - half]; + out[d] = roundBf16(xn[d - half] * @sin(@as(f64, ang)) + xn[d] * @cos(@as(f64, ang))); + } else out[d] = xn[d]; + } +} + +/// How many of the 32 frequency indices each mrope axis owns with `mrope_interleaved` sections [11, 11, 10]. +pub fn mropeAxisCounts() [3]usize { + var c: [3]usize = .{ 0, 0, 0 }; + for (0..half) |i| { + if (i % 3 == 1 and i < 3 * 11) { + c[1] += 1; + } else if (i % 3 == 2 and i < 3 * 10) { + c[2] += 1; + } else c[0] += 1; + } + return c; +} + +/// One full-attention layer's state: bf16-rounded keys and values, raw index keys, pooled blocks. +pub const Cache = struct { + gpa: std.mem.Allocator, + tokens: usize, + keys: []f64, // [pos][kv head][256] + values: []f64, + raw: []f64, // [pos][128] + pooled: []f64, // [pos / 4][128] + + pub fn init(gpa: std.mem.Allocator, tokens: usize) !Cache { + const keys = try gpa.alloc(f64, tokens * kv_heads * head_dim); + errdefer gpa.free(keys); + const values = try gpa.alloc(f64, tokens * kv_heads * head_dim); + errdefer gpa.free(values); + const raw = try gpa.alloc(f64, tokens * index_dim); + errdefer gpa.free(raw); + const pooled = try gpa.alloc(f64, (tokens / ratio + 1) * index_dim); + return .{ .gpa = gpa, .tokens = tokens, .keys = keys, .values = values, .raw = raw, .pooled = pooled }; + } + + pub fn deinit(c: *Cache) void { + c.gpa.free(c.keys); + c.gpa.free(c.values); + c.gpa.free(c.raw); + c.gpa.free(c.pooled); + } +}; + +/// Projections of one token (what the EXL3 linears emit, fp16): q_proj [24 * 512], k_proj [512], v_proj [512]. +pub const Proj = struct { + qp: []const f16, + kp: []const f16, + vp: []const f16, + ip: []const f16, +}; + +/// The normed, rotated queries (bf16 values, [24 * 256]) and indexer queries ([4 * 128]) of one token. +pub fn prepare(c: *Cache, pos: usize, p: Proj, n: *const Norms, q_out: []f64, iq_out: []f64) void { + var buf: [head_dim]f64 = undefined; + for (0..q_heads) |h| { + for (0..head_dim) |d| buf[d] = p.qp[h * 2 * head_dim + d]; + normRope(buf[0..head_dim], n.q, @floatFromInt(pos), &n.inv, q_out[h * head_dim ..][0..head_dim]); + } + for (0..kv_heads) |j| { + for (0..head_dim) |d| buf[d] = p.kp[j * head_dim + d]; + normRope(buf[0..head_dim], n.k, @floatFromInt(pos), &n.inv, c.keys[(pos * kv_heads + j) * head_dim ..][0..head_dim]); + for (0..head_dim) |d| c.values[(pos * kv_heads + j) * head_dim + d] = roundBf16(p.vp[j * head_dim + d]); + } + for (0..index_heads) |h| { + for (0..index_dim) |d| buf[d] = p.ip[h * index_dim + d]; + normRope(buf[0..index_dim], n.iq, @floatFromInt(pos), &n.inv, iq_out[h * index_dim ..][0..index_dim]); + } + for (0..index_dim) |d| c.raw[pos * index_dim + d] = roundBf16(p.ip[index_heads * index_dim + d]); +} + +/// Pool block `b` (keys 4b .. 4b + 3 must be in `raw`). +pub fn poolBlock(c: *Cache, b: usize, n: *const Norms) void { + var mean: [index_dim]f64 = undefined; + for (0..index_dim) |d| { + var acc: f32 = @floatCast(c.raw[(ratio * b) * index_dim + d]); + for (1..ratio) |k| acc = acc + @as(f32, @floatCast(c.raw[(ratio * b + k) * index_dim + d])); + mean[d] = roundBf16(@as(f64, acc / @as(f32, ratio))); + } + normRope(&mean, n.ik, @floatFromInt(ratio * b), &n.inv, c.pooled[b * index_dim ..][0..index_dim]); +} + +/// Score of every complete block for one indexer query: sum_h relu(q_h . pooled_b) / sqrt(128). +pub fn blockScores(c: *const Cache, iq: []const f64, complete: usize, out: []f64) void { + for (0..complete) |b| { + var total: f64 = 0; + for (0..index_heads) |h| { + var dot: f64 = 0; + for (0..index_dim) |d| dot += iq[h * index_dim + d] * c.pooled[b * index_dim + d]; + total += @max(dot, 0); + } + out[b] = total / @sqrt(@as(f64, index_dim)); + } +} + +/// The `top_blocks` best blocks, ties to the lower id, in ascending block order. +pub fn selectBlocks(gpa: std.mem.Allocator, scores: []const f64, out: *[top_blocks]usize) !void { + const order = try gpa.alloc(usize, scores.len); + defer gpa.free(order); + for (order, 0..) |*o, i| o.* = i; + const Ctx = struct { + s: []const f64, + fn less(ctx: @This(), a: usize, b: usize) bool { + if (ctx.s[a] != ctx.s[b]) return ctx.s[a] > ctx.s[b]; + return a < b; + } + }; + std.mem.sort(usize, order, Ctx{ .s = scores }, Ctx.less); + @memcpy(out, order[0..top_blocks]); + std.mem.sort(usize, out, {}, std.sort.asc(usize)); +} + +/// Keys a row at `end = pos + 1` reads. +pub fn keyList(end: usize, blocks: ?*const [top_blocks]usize, out: *[max_keys]usize) usize { + const complete = end / ratio; + if (complete <= top_blocks) { + for (0..end) |i| out[i] = i; + return end; + } + const sel = blocks.?; + var n: usize = 0; + for (sel) |b| { + for (0..ratio) |k| { + out[n] = b * ratio + k; + n += 1; + } + } + var t = ratio * complete; + while (t < end) : (t += 1) { + out[n] = t; + n += 1; + } + return n; +} + +/// out[h * 256 + d] = fp16(bf16(bf16(o) / (1 + exp(-gate)))), o = softmax(q . K / 16) V over `keys`. +pub fn attend(c: *const Cache, q: []const f64, qp: []const f16, keys: []const usize, out: []f64) void { + var s: [max_keys]f64 = undefined; + const group = q_heads / kv_heads; + for (0..q_heads) |h| { + const kvh = h / group; + var m: f64 = -std.math.inf(f64); + for (keys, 0..) |key, j| { + var dot: f64 = 0; + for (0..head_dim) |d| dot += q[h * head_dim + d] * c.keys[(key * kv_heads + kvh) * head_dim + d]; + s[j] = dot * 0.0625; + m = @max(m, s[j]); + } + var l: f64 = 0; + for (keys, 0..) |_, j| { + s[j] = @exp(s[j] - m); + l += s[j]; + } + for (0..head_dim) |d| { + var acc: f64 = 0; + for (keys, 0..) |key, j| acc += s[j] * c.values[(key * kv_heads + kvh) * head_dim + d]; + const o = roundBf16(acc / l); + const g: f64 = qp[h * 2 * head_dim + head_dim + d]; + out[h * head_dim + d] = roundF16(roundBf16(o / (1.0 + @exp(-g)))); + } + } +} + +pub const RowInfo = struct { + sparse: bool = false, + nk: usize = 0, + /// The selected blocks of a sparse row (undefined when dense). + blocks: [top_blocks]usize = undefined, +}; + +/// The whole row: prepare, pool the blocks this token completes, select, attend. +pub fn row( + c: *Cache, + pos: usize, + p: Proj, + n: *const Norms, + scores: []f64, + q: []f64, + iq: []f64, + out: []f64, + info: *RowInfo, +) !void { + prepare(c, pos, p, n, q, iq); + const end = pos + 1; + if (end % ratio == 0) poolBlock(c, end / ratio - 1, n); + const complete = end / ratio; + info.sparse = complete > top_blocks; + if (info.sparse) { + blockScores(c, iq, complete, scores); + try selectBlocks(c.gpa, scores[0..complete], &info.blocks); + } + var keys: [max_keys]usize = undefined; + info.nk = keyList(end, if (info.sparse) &info.blocks else null, &keys); + attend(c, q, p.qp, keys[0..info.nk], out); +} + +// -- tests ----------------------------------------------------------------- + +test "mrope sections [11, 11, 10] interleave over 32 frequencies and need no per-axis table for text" { + const counts = mropeAxisCounts(); + try std.testing.expectEqual([3]usize{ 11, 11, 10 }, counts); + try std.testing.expectEqual(@as(usize, 32), counts[0] + counts[1] + counts[2]); + try std.testing.expectEqual(half, counts[0] + counts[1] + counts[2]); +} + +test "rotate-half RoPE matches an explicit rotation of the pairs (i, i + 32) and leaves dims 64.. alone" { + const inv = invFreq(); + try std.testing.expectApproxEqRel(@as(f32, 1.0), inv[0], 1e-7); + // theta^(-2*31/64) for the last frequency + try std.testing.expectApproxEqRel(@as(f32, @floatCast(std.math.pow(f64, 1e7, -62.0 / 64.0))), inv[31], 1e-6); + var x: [head_dim]f64 = undefined; + var rng = std.Random.DefaultPrng.init(7); + for (&x) |*v| v.* = rng.random().float(f64) * 2 - 1; + var scale: [head_dim]f32 = @splat(1); + scale[3] = 2.5; + var out: [head_dim]f64 = undefined; + const pos: usize = 1234; + normRope(&x, &scale, @floatFromInt(pos), &inv, &out); + + var ss: f64 = 0; + for (x) |v| ss += v * v; + const rinv = 1.0 / @sqrt(ss / head_dim + eps); + for (0..half) |i| { + const a = roundBf16(x[i] * rinv * scale[i]); + const b = roundBf16(x[i + half] * rinv * scale[i + half]); + const ang: f32 = @as(f32, @floatFromInt(pos)) * inv[i]; + try std.testing.expectEqual(roundBf16(a * @cos(@as(f64, ang)) - b * @sin(@as(f64, ang))), out[i]); + try std.testing.expectEqual(roundBf16(a * @sin(@as(f64, ang)) + b * @cos(@as(f64, ang))), out[i + half]); + } + for (rotary_dim..head_dim) |d| try std.testing.expectEqual(roundBf16(x[d] * rinv * scale[d]), out[d]); + // position 0 is the identity rotation + normRope(&x, &scale, 0, &inv, &out); + for (0..head_dim) |d| try std.testing.expectEqual(roundBf16(x[d] * rinv * scale[d]), out[d]); +} + +test "bf16 rounding is round-to-nearest-even" { + try std.testing.expectEqual(@as(f64, 1.0), roundBf16(1.0 + 1.0 / 512.0)); // below half an ulp (ulp 2^-7) + try std.testing.expectEqual(@as(f64, 1.0), roundBf16(1.0 + 1.0 / 256.0)); // tie goes to even (1.0) + try std.testing.expectEqual(@as(f64, 1.0 + 1.0 / 64.0), roundBf16(1.0 + 3.0 / 256.0)); // tie goes up to even mantissa + try std.testing.expectEqual(@as(f64, 1.0 + 1.0 / 128.0), roundBf16(1.0 + 1.0 / 128.0 + 1.0 / 1024.0)); +} + +test "selection takes the 512 best blocks, ties to the lower id, in block order" { + const gpa = std.testing.allocator; + const n: usize = 700; + const scores = try gpa.alloc(f64, n); + defer gpa.free(scores); + var rng = std.Random.DefaultPrng.init(11); + for (scores) |*v| v.* = @floor(rng.random().float(f64) * 40) / 4; // many exact ties + var sel: [top_blocks]usize = undefined; + try selectBlocks(gpa, scores, &sel); + try std.testing.expect(std.sort.isSorted(usize, &sel, {}, std.sort.asc(usize))); + // the cut: every unselected block scores no more than every selected one, and a tie at the cut keeps the lower id + var chosen: [700]bool = @splat(false); + var min_sel: f64 = std.math.inf(f64); + for (sel) |b| { + chosen[b] = true; + min_sel = @min(min_sel, scores[b]); + } + var count: usize = 0; + for (chosen) |c| count += @intFromBool(c); + try std.testing.expectEqual(top_blocks, count); + var max_out: f64 = -1; + var lowest_out_at_cut: usize = n; + for (0..n) |b| if (!chosen[b]) { + max_out = @max(max_out, scores[b]); + if (scores[b] == min_sel) lowest_out_at_cut = @min(lowest_out_at_cut, b); + }; + try std.testing.expect(max_out <= min_sel); + for (0..n) |b| if (chosen[b] and scores[b] == min_sel) try std.testing.expect(b < lowest_out_at_cut); +} + +test "key list: dense through 2051 keys, then 2048 selected plus the unfinished tail" { + var sel: [top_blocks]usize = undefined; + for (&sel, 0..) |*b, i| b.* = i + 1; // blocks 1..512 of 513 + var keys: [max_keys]usize = undefined; + try std.testing.expectEqual(@as(usize, 2051), keyList(2051, null, &keys)); + try std.testing.expectEqual(@as(usize, 2050), keys[2050]); + // end 2052: 513 complete blocks, no tail, block 0 dropped + try std.testing.expectEqual(@as(usize, 2048), keyList(2052, &sel, &keys)); + try std.testing.expectEqual(@as(usize, 4), keys[0]); + try std.testing.expectEqual(@as(usize, 2051), keys[2047]); + // end 2055: three tail keys 2052, 2053, 2054 + try std.testing.expectEqual(@as(usize, 2051), keyList(2055, &sel, &keys)); + try std.testing.expectEqual(@as(usize, 2052), keys[2048]); + try std.testing.expectEqual(@as(usize, 2054), keys[2050]); +} + +test "attend with one visible key returns that value gated" { + const gpa = std.testing.allocator; + var c = try Cache.init(gpa, 4); + defer c.deinit(); + for (0..kv_heads * head_dim) |i| { + c.keys[i] = 0.25; + c.values[i] = if (i < head_dim) 2.0 else -3.0; + } + var q: [q_heads * head_dim]f64 = @splat(0.5); + var qp: [q_heads * 2 * head_dim]f16 = @splat(0); + for (0..q_heads) |h| for (0..head_dim) |d| { + qp[h * 2 * head_dim + head_dim + d] = 0; // sigmoid(0) = 1/2 + }; + var out: [q_heads * head_dim]f64 = undefined; + attend(&c, &q, &qp, &.{0}, &out); + for (0..12) |h| try std.testing.expectEqual(@as(f64, 1.0), out[h * head_dim]); + for (12..q_heads) |h| try std.testing.expectEqual(@as(f64, -1.5), out[h * head_dim]); +} + +test "the device KV layout is exl3_kv.Table's for one layer" { + const kv = @import("../kv/exl3_kv.zig"); + const gpa = std.testing.allocator; + const tokens: usize = 8; + var t = try kv.Table.init(gpa, tokens); + defer t.deinit(gpa); + try std.testing.expectEqual(@as(usize, kv_heads), t.heads); + const slot_bytes = 2 * head_dim * 2; + for ([_]usize{ 0, 1, 2 }) |layer| for ([_]usize{ 0, 5, 7 }) |token| for ([_]usize{ 0, 1 }) |head| { + var src: [slot_bytes]u8 = undefined; + @memset(&src, @intCast(1 + layer * 40 + token * 4 + head)); + try t.write(head, layer, token, &src); + // the kernel's byte offset: layer base + (token * kv_heads + head) * slot + const off = layer * tokens * kv_heads * slot_bytes + (token * kv_heads + head) * slot_bytes; + try std.testing.expectEqualSlices(u8, &src, t.bytes[off..][0..slot_bytes]); + }; +} + +test "a full row past 2048 keys reads the selected blocks, and the dense row at 2051 keys does not" { + const gpa = std.testing.allocator; + const tokens: usize = 2100; + var c = try Cache.init(gpa, tokens); + defer c.deinit(); + var rng = std.Random.DefaultPrng.init(3); + const inv = invFreq(); + const q_scale: [head_dim]f32 = @splat(1); + const i_scale: [index_dim]f32 = @splat(1); + const norms = Norms{ .q = &q_scale, .k = &q_scale, .iq = &i_scale, .ik = &i_scale, .inv = inv }; + const qp = try gpa.alloc(f16, q_heads * 2 * head_dim); + defer gpa.free(qp); + const kp = try gpa.alloc(f16, kv_heads * head_dim); + defer gpa.free(kp); + const vp = try gpa.alloc(f16, kv_heads * head_dim); + defer gpa.free(vp); + const ip = try gpa.alloc(f16, index_width); + defer gpa.free(ip); + const scores = try gpa.alloc(f64, tokens / ratio + 1); + defer gpa.free(scores); + var q: [q_heads * head_dim]f64 = undefined; + var iq: [index_heads * index_dim]f64 = undefined; + var out: [q_heads * head_dim]f64 = undefined; + var info = RowInfo{}; + for (0..2056) |pos| { + for (qp) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + for (kp) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + for (vp) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + for (ip) |*v| v.* = @floatCast(rng.random().floatNorm(f32)); + const proj = Proj{ .qp = qp, .kp = kp, .vp = vp, .ip = ip }; + switch (pos) { + 2050, 2051, 2054 => { + try row(&c, pos, proj, &norms, scores, &q, &iq, &out, &info); + switch (pos) { + 2050 => try std.testing.expect(!info.sparse and info.nk == 2051), + 2051 => try std.testing.expect(info.sparse and info.nk == 2048), + else => try std.testing.expect(info.sparse and info.nk == 2051), + } + }, + else => { + prepare(&c, pos, proj, &norms, &q, &iq); + if ((pos + 1) % ratio == 0) poolBlock(&c, (pos + 1) / ratio - 1, &norms); + }, + } + } + for (out) |v| try std.testing.expect(std.math.isFinite(v)); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_cache.zig b/zig/src/families/flashnext/exl3/model/exl3_cache.zig new file mode 100644 index 0000000000..c193eefa6f --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_cache.zig @@ -0,0 +1,168 @@ +//! The KV bytes and the resident-set budget. + +const std = @import("std"); +const shape = @import("../pack/exl3_shape.zig"); + +/// One attention layer's K and V for `tokens` tokens, bf16. +pub fn layerBytes(tokens: usize, kv_heads: usize, head_dim: usize) u64 { + const t: u64 = @intCast(tokens); + const h: u64 = @intCast(kv_heads); + const d: u64 = @intCast(head_dim); + return t * 2 * h * d * 2; +} + +/// A kept prefix stops one token before the prompt end. An empty prompt stays at 0. +pub fn resumeAt(prompt_len: u32) u32 { + if (prompt_len == 0) return 0; + return prompt_len - 1; +} + +pub const Budget = struct { + device_bytes: u64, + /// Dense + shared + router + HC + MTP workspace, already summed by the caller. + resident_bytes: u64, + kv_bytes: u64, + routed_expert_bytes: u64, + ngram_bytes: u64, +}; + +pub const Place = enum { host, refuse }; + +/// Resident weights plus KV. Routed experts and the n-gram table are not in this sum. +pub fn deviceSum(b: Budget) u64 { + return b.resident_bytes + b.kv_bytes; +} + +pub fn fits(b: Budget) error{ExceedsDevice}!void { + if (deviceSum(b) > b.device_bytes) return error.ExceedsDevice; +} + +/// One RTX 5060 Ti. The learned expert lines are whatever this does not already use. +pub const card_bytes: u64 = 16 << 30; +/// CUDA context and the batch-1 workspace. Not a weight, and not a second copy of the bank. +pub const device_reserve: u64 = 256 << 20; + +/// Bytes left for routed experts after pinned GEMMs, KV, and the reserve. +pub fn learnedBytes(device: u64, pinned: u64, kv: u64) u64 { + if (pinned >= device) return 0; + const after_pinned = device - pinned; + if (kv >= after_pinned) return 0; + const after_kv = after_pinned - kv; + if (after_kv <= device_reserve) return 0; + return after_kv - device_reserve; +} + +fn matrixBytes(k: usize, n: usize) u64 { + const kn: u64 = @as(u64, k) * @as(u64, n); + // 3.05 bpw trellis, plus one fp16 scale along K and one along N. + return kn * 3050 / 8000 + @as(u64, k + n) * 2; +} + +/// One routed expert as stored. +pub fn expertBytes() u64 { + const k = shape.hidden; + const n = shape.expert_intermediate; + return matrixBytes(k, n) * 2 + matrixBytes(n, k); +} + +pub fn expertLines(room: u64) u64 { + return expertLinesOf(room, expertBytes()); +} + +/// Budget stride for one owned expert until a header span is larger. Not a line count. +pub const full_expert_bytes: u64 = 1_901_213; + +/// Lines that fit in `room` at `line` bytes. A zero stride is no lines, not a panic. +pub fn expertLinesOf(room: u64, line: u64) u64 { + if (line == 0) return 0; + return room / line; +} + +/// KV for every full-attention layer. +pub fn contextKv(tokens: usize) u64 { + const layers: u64 = shape.layers / shape.full_attention_interval; + return layerBytes(tokens, shape.kv_heads, shape.head_dim) * layers; +} + +/// Routed experts and the n-gram table always live on the host. +pub fn placeWeights(b: Budget) Place { + _ = b; + return .host; +} + +/// Neither the routed-expert bank nor the n-gram table may be counted as resident. +pub fn refusesResident(which: enum { routed, ngram }) bool { + _ = which; + return true; +} + +test "resumeAt stops one token before the end" { + try std.testing.expectEqual(@as(u32, 0), resumeAt(0)); + try std.testing.expectEqual(@as(u32, 7), resumeAt(8)); + try std.testing.expectEqual(@as(u32, 0), resumeAt(1)); +} + +test "layer bytes match v0.6.6 row_bytes" { + try std.testing.expectEqual(@as(u64, 2048), layerBytes(1, 2, 256)); + try std.testing.expectEqual(@as(u64, 2048), layerBytes(1, shape.kv_heads, shape.head_dim)); + try std.testing.expectEqual(@as(u64, 6144), layerBytes(3, 2, 256)); + try std.testing.expectEqual(@as(u64, 192), layerBytes(1, 2, 24)); +} + +test "device sum is resident plus kv, and a 16 GiB card refuses the overflow" { + const device: u64 = 16 << 30; + const resident: u64 = 6 << 30; + const kv: u64 = 8 << 30; + const held = Budget{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = kv, + .routed_expert_bytes = 70 << 30, + .ngram_bytes = 29 << 30, + }; + try std.testing.expectEqual(resident + kv, deviceSum(held)); + try fits(held); + try std.testing.expectEqual(Place.host, placeWeights(held)); + try std.testing.expect(refusesResident(.routed)); + try std.testing.expect(refusesResident(.ngram)); + + try std.testing.expectError(error.ExceedsDevice, fits(.{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = 11 << 30, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + })); + + const rest = device - resident; + try fits(.{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = rest, + .routed_expert_bytes = 1, + .ngram_bytes = 1, + }); + try std.testing.expectError(error.ExceedsDevice, fits(.{ + .device_bytes = device, + .resident_bytes = resident, + .kv_bytes = rest + 1, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + })); + + try std.testing.expectEqual(@as(u64, 1_893_120), expertBytes()); + try std.testing.expectEqual(@as(u64, 1_901_213), full_expert_bytes); + try std.testing.expectEqual(@as(u64, 0), expertLinesOf(1, 0)); + try std.testing.expectEqual(@as(u64, 1), expertLinesOf(full_expert_bytes, full_expert_bytes)); + + try std.testing.expectEqual(@as(u64, 0), learnedBytes(card_bytes, card_bytes, 0)); + + // Counting the expert bank as resident is what blows the same card. + try std.testing.expectError(error.ExceedsDevice, fits(.{ + .device_bytes = device, + .resident_bytes = resident + (70 << 30), + .kv_bytes = kv, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + })); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_cpu.zig b/zig/src/families/flashnext/exl3/model/exl3_cpu.zig new file mode 100644 index 0000000000..a40d6f5633 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_cpu.zig @@ -0,0 +1,380 @@ +//! CPU GEMV for one EXL3 matrix. + +const std = @import("std"); +const builtin = @import("builtin"); +const format = @import("../pack/exl3_format.zig"); + +/// AVX fp256, compiled only when this target actually has AVX2. No AVX-512. +const avx2_target = builtin.cpu.arch == .x86_64 and + builtin.zig_backend != .stage2_x86_64 and + builtin.cpu.has(.x86, .avx2); + +const Acc = @Vector(4, f64); + +/// fp32 SwiGLU from the v0.6.6 gate-up epilogue (`act_mode` 1). +pub fn swiglu(gate: []const f64, up: []const f64, out: []f64, limit: f64) error{Length}!void { + if (gate.len != up.len or gate.len != out.len) return error.Length; + for (gate, up, out) |g, u, *y| { + const gg = @min(g, limit); + const uu = @min(@max(u, -limit), limit); + y.* = (gg / (1.0 + @exp(-gg))) * uu; + } +} + +/// Weighted sum of `weights.len` expert rows. +pub fn combine(rows: []const f64, weights: []const f64, n: usize, out: []f64) error{Length}!void { + if (n == 0 or out.len != n or weights.len == 0 or rows.len != weights.len * n) return error.Length; + @memset(out, 0); + for (weights, 0..) |w, s| { + const row = rows[s * n ..][0..n]; + for (out, row) |*y, v| y.* += w * v; + } +} + +/// K = tiles_k * 16 and N = tiles_n * 16, both multiples of 128 (one Hadamard block). +pub fn gemv( + w: format.Width, + codebook: format.Codebook, + trellis: []const i16, + tiles_k: usize, + tiles_n: usize, + suh: []const f16, + svh: []const f16, + x: []const f64, + y: []f64, + scratch: []f64, +) error{Length}!void { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % format.hadamard_dim != 0 or n % format.hadamard_dim != 0) return error.Length; + if (x.len != k or y.len != n or suh.len != k or svh.len != n) return error.Length; + const cells = k * n; + if (scratch.len < 3 * cells) return error.Length; + + const weight = scratch[0..cells]; + format.dequantize(w, codebook, trellis, tiles_k, tiles_n, suh, svh, weight, scratch[cells .. 3 * cells]); + + for (0..n) |j| { + var sum: f64 = 0; + for (0..k) |i| sum += x[i] * weight[i * n + j]; + y[j] = sum; + } +} + +/// Bit-exact with `gemv`. +pub fn gemvAvx2( + w: format.Width, + codebook: format.Codebook, + trellis: []const i16, + tiles_k: usize, + tiles_n: usize, + suh: []const f16, + svh: []const f16, + x: []const f64, + y: []f64, + scratch: []f64, +) error{Length}!void { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % format.hadamard_dim != 0 or n % format.hadamard_dim != 0) return error.Length; + if (x.len != k or y.len != n or suh.len != k or svh.len != n) return error.Length; + const cells = k * n; + if (scratch.len < 3 * cells) return error.Length; + + const weight = scratch[0..cells]; + format.dequantize(w, codebook, trellis, tiles_k, tiles_n, suh, svh, weight, scratch[cells .. 3 * cells]); + + // N is a multiple of 128, so the four-column walk covers every column. + if (comptime avx2_target) { + accumulateAvx2(weight, x, y); + } else { + accumulatePortable(weight, x, y); + } +} + +/// Same four lanes as the AVX2 loop, one left-to-right sum per column. +fn accumulatePortable(weight: []const f64, x: []const f64, y: []f64) void { + @setFloatMode(.strict); + const k = x.len; + const n = y.len; + var j: usize = 0; + while (j < n) : (j += 4) { + var acc: [4]f64 = .{ 0, 0, 0, 0 }; + var i: usize = 0; + while (i < k) : (i += 1) { + const xi = x[i]; + const row = i * n + j; + inline for (0..4) |lane| acc[lane] += xi * weight[row + lane]; + } + y[j..][0..4].* = acc; + } +} + +fn avx2Load(p: [*]const f64) Acc { + return asm ( + \\vmovupd (%[p]), %[out] + : [out] "=x" (-> Acc), + : [p] "r" (p), + ); +} + +fn avx2Broadcast(p: *const f64) Acc { + return asm ( + \\vbroadcastsd (%[p]), %[out] + : [out] "=x" (-> Acc), + : [p] "r" (p), + ); +} + +fn avx2Mul(a: Acc, b: Acc) Acc { + return asm ( + \\vmulpd %[b], %[a], %[out] + : [out] "=x" (-> Acc), + : [a] "x" (a), + [b] "x" (b), + ); +} + +/// `out = sum + prod`, the same operand order as `sum += prod`. +fn avx2Add(sum: Acc, prod: Acc) Acc { + return asm ( + \\vaddpd %[prod], %[sum], %[out] + : [out] "=x" (-> Acc), + : [sum] "x" (sum), + [prod] "x" (prod), + ); +} + +fn avx2Store(p: [*]f64, v: Acc) void { + asm volatile ( + \\vmovupd %[v], (%[p]) + : + : [v] "x" (v), + [p] "r" (p), + : .{ .memory = true }); +} + +/// ymm lanes are columns j..j+3. vmovupd: +fn accumulateAvx2(weight: []const f64, x: []const f64, y: []f64) void { + comptime if (builtin.cpu.arch != .x86_64) @compileError("AVX2 GEMV is x86_64 only"); + const k = x.len; + const n = y.len; + var j: usize = 0; + while (j < n) : (j += 4) { + var acc: Acc = @splat(0); + var i: usize = 0; + while (i < k) : (i += 1) { + const prod = avx2Mul(avx2Broadcast(&x[i]), avx2Load(weight.ptr + i * n + j)); + acc = avx2Add(acc, prod); + } + avx2Store(y.ptr + j, acc); + } +} + +const sentinel: f64 = @bitCast(@as(u64, 0x7FF8_0000_0000_0001)); + +/// 128x128 mul1, width 2 halves. Trellis length is 8*8*Width.tileWords(). +const Square = struct { + allocator: std.mem.Allocator, + w: format.Width, + tiles_k: usize, + tiles_n: usize, + trellis: []i16, + suh: []f16, + svh: []f16, + x: []f64, + y: []f64, + scratch: []f64, + + fn init(allocator: std.mem.Allocator) !Square { + const w: format.Width = .{ .halves = 2 }; + const tiles_k: usize = 8; + const tiles_n: usize = 8; + const k = tiles_k * 16; + const n = tiles_n * 16; + const trellis = try allocator.alloc(i16, tiles_k * tiles_n * w.tileWords()); + errdefer allocator.free(trellis); + for (trellis, 0..) |*word, i| word.* = @bitCast(@as(u16, @truncate(i *% 4099 + 7))); + + const suh = try allocator.alloc(f16, k); + errdefer allocator.free(suh); + for (suh) |*v| v.* = 1.0; + const svh = try allocator.alloc(f16, n); + errdefer allocator.free(svh); + for (svh) |*v| v.* = 1.0; + + const x = try allocator.alloc(f64, k); + errdefer allocator.free(x); + for (x, 0..) |*v, i| v.* = @floatFromInt(i + 1); + + const y = try allocator.alloc(f64, n); + errdefer allocator.free(y); + for (y) |*v| v.* = sentinel; + + const scratch = try allocator.alloc(f64, 3 * k * n); + errdefer allocator.free(scratch); + return .{ + .allocator = allocator, + .w = w, + .tiles_k = tiles_k, + .tiles_n = tiles_n, + .trellis = trellis, + .suh = suh, + .svh = svh, + .x = x, + .y = y, + .scratch = scratch, + }; + } + + fn deinit(self: Square) void { + self.allocator.free(self.trellis); + self.allocator.free(self.suh); + self.allocator.free(self.svh); + self.allocator.free(self.x); + self.allocator.free(self.y); + self.allocator.free(self.scratch); + } +}; + +test "128x128 mul1 gemv matches dequantize then a naive matvec" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + const k = s.tiles_k * 16; + const n = s.tiles_n * 16; + try std.testing.expectEqual(@as(usize, 128), k); + try std.testing.expectEqual(@as(usize, 128), n); + try std.testing.expectEqual(@as(usize, 8 * 8 * s.w.tileWords()), s.trellis.len); + + try gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, s.scratch); + + const weight = try gpa.alloc(f64, k * n); + defer gpa.free(weight); + const dq_scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(dq_scratch); + format.dequantize(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, weight, dq_scratch); + + for (0..n) |j| { + var sum: f64 = 0; + for (0..k) |i| sum += s.x[i] * weight[i * n + j]; + try std.testing.expectEqual(@as(u64, @bitCast(sum)), @as(u64, @bitCast(s.y[j]))); + } +} + +test "scratch one f64 short returns Length and leaves y untouched" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + for (s.y) |*v| v.* = sentinel; + const short = s.scratch[0 .. s.scratch.len - 1]; + try std.testing.expectError(error.Length, gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, short)); + for (s.y) |v| try std.testing.expectEqual(@as(u64, @bitCast(sentinel)), @as(u64, @bitCast(v))); +} + +test "swiglu is silu(gate) times up, and a finite limit clamps both sides" { + var out = [_]f64{ 9, 9 }; + const gate = [_]f64{ 0, 2 }; + const up = [_]f64{ 3, 4 }; + try swiglu(&gate, &up, &out, std.math.inf(f64)); + try std.testing.expectEqual(@as(f64, 0), out[0]); + const silu2 = 2.0 / (1.0 + @exp(-2.0)); + try std.testing.expectApproxEqAbs(silu2 * 4.0, out[1], 1e-12); + + const wide_gate = [_]f64{2}; + const wide_up = [_]f64{4}; + var clamped = [_]f64{9}; + try swiglu(&wide_gate, &wide_up, &clamped, 1); + const silu1 = 1.0 / (1.0 + @exp(-1.0)); + try std.testing.expectApproxEqAbs(silu1, clamped[0], 1e-12); + + var held = [_]f64{9}; + try std.testing.expectError(error.Length, swiglu(wide_gate[0..], wide_up[0..], held[0..0], 1)); + try std.testing.expectEqual(@as(f64, 9), held[0]); +} + +test "combine is the weighted sum of expert rows" { + const rows = [_]f64{ 1, 2, 3, 4 }; + const weights = [_]f64{ 0.25, 0.75 }; + var out = [_]f64{ 9, 9 }; + try combine(&rows, &weights, 2, &out); + try std.testing.expectApproxEqAbs(@as(f64, 2.5), out[0], 1e-12); + try std.testing.expectApproxEqAbs(@as(f64, 3.5), out[1], 1e-12); + try std.testing.expectError(error.Length, combine(rows[0..3], &weights, 2, &out)); +} + +test "x shorter than K returns Length" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + for (s.y) |*v| v.* = sentinel; + const short_x = s.x[0 .. s.x.len - 1]; + try std.testing.expectError(error.Length, gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, short_x, s.y, s.scratch)); + for (s.y) |v| try std.testing.expectEqual(@as(u64, @bitCast(sentinel)), @as(u64, @bitCast(v))); +} + +fn expectGemvAvx2(allocator: std.mem.Allocator, tiles_k: usize, tiles_n: usize) !void { + const w: format.Width = .{ .halves = 2 }; + const k = tiles_k * 16; + const n = tiles_n * 16; + const trellis = try allocator.alloc(i16, tiles_k * tiles_n * w.tileWords()); + defer allocator.free(trellis); + for (trellis, 0..) |*word, i| word.* = @bitCast(@as(u16, @truncate(i *% 4099 + 7))); + + const suh = try allocator.alloc(f16, k); + defer allocator.free(suh); + for (suh) |*v| v.* = 1.0; + const svh = try allocator.alloc(f16, n); + defer allocator.free(svh); + for (svh) |*v| v.* = 1.0; + + const x = try allocator.alloc(f64, k); + defer allocator.free(x); + for (x, 0..) |*v, i| v.* = @floatFromInt(i + 1); + + const y = try allocator.alloc(f64, n); + defer allocator.free(y); + const y2 = try allocator.alloc(f64, n); + defer allocator.free(y2); + @memset(y, sentinel); + @memset(y2, sentinel); + + const scratch = try allocator.alloc(f64, 3 * k * n); + defer allocator.free(scratch); + try std.testing.expectEqual(@as(usize, 3 * k * n), scratch.len); + + try gemv(w, .mul1, trellis, tiles_k, tiles_n, suh, svh, x, y, scratch); + try gemvAvx2(w, .mul1, trellis, tiles_k, tiles_n, suh, svh, x, y2, scratch); + for (y, y2) |a, b| try std.testing.expectEqual(@as(u64, @bitCast(a)), @as(u64, @bitCast(b))); +} + +test "gemvAvx2 matches gemv on the 128x128 mul1 tile" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + const y2 = try gpa.alloc(f64, s.y.len); + defer gpa.free(y2); + @memset(s.y, sentinel); + @memset(y2, sentinel); + try gemv(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, s.scratch); + try gemvAvx2(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, y2, s.scratch); + for (s.y, y2) |a, b| try std.testing.expectEqual(@as(u64, @bitCast(a)), @as(u64, @bitCast(b))); +} + +test "gemvAvx2 matches gemv on synthetic 2560x640 and 640x2560" { + // 2560 * 640 * 3 f64. One buffer, not a second scratch shape. + try std.testing.expectEqual(@as(usize, 39_321_600), 3 * 2560 * 640 * @sizeOf(f64)); + const gpa = std.testing.allocator; + try expectGemvAvx2(gpa, 160, 40); + try expectGemvAvx2(gpa, 40, 160); +} + +test "gemvAvx2 scratch one f64 short returns Length and leaves y untouched" { + const gpa = std.testing.allocator; + const s = try Square.init(gpa); + defer s.deinit(); + for (s.y) |*v| v.* = sentinel; + const short = s.scratch[0 .. s.scratch.len - 1]; + try std.testing.expectError(error.Length, gemvAvx2(s.w, .mul1, s.trellis, s.tiles_k, s.tiles_n, s.suh, s.svh, s.x, s.y, short)); + for (s.y) |v| try std.testing.expectEqual(@as(u64, @bitCast(sentinel)), @as(u64, @bitCast(v))); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_fwd_cpu.zig b/zig/src/families/flashnext/exl3/model/exl3_fwd_cpu.zig new file mode 100644 index 0000000000..a4ebfb70a6 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_fwd_cpu.zig @@ -0,0 +1,414 @@ +//! CPU f64 oracle for the whole Qwen3.8-Flash-Next EXL3 forward, assembled from the per-op oracles (`exl3_hc` + +const std = @import("std"); +const hc = @import("exl3_hc.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); +const moe = @import("exl3_moe.zig"); +const cpu = @import("exl3_cpu.zig"); +const qmat = @import("../pack/exl3_qmat.zig"); + +pub const hidden: usize = 2560; +pub const streams: usize = 4; +pub const wide: usize = streams * hidden; +pub const topk: usize = 10; +pub const experts: usize = 512; +pub const inter: usize = 640; + +pub const QMat = qmat.QMat; +pub const parallelFor = qmat.parallelFor; +pub const fwht128 = qmat.fwht128; + +// -- small ops --------------------------------------------------------------------------------------------- + +pub fn silu(x: f64) f64 { + return x / (1.0 + @exp(-x)); +} + +pub fn sigmoid(x: f64) f64 { + return 1.0 / (1.0 + @exp(-x)); +} + +/// f64 rows [t][hidden] -> mixed hidden of the hyper-connection (pre), per row. `inject` is [t][4] or null. +pub fn hcPreRows(gpa: std.mem.Allocator, w: hc.Weights, offset: f64, h: []const f64, rows: usize, mixed: []f64, inject: ?[]f64) !void { + const low = try gpa.alloc(f64, hc.real.lowrank); + defer gpa.free(low); + const mix = try gpa.alloc(f64, wide); + defer gpa.free(mix); + const normed = try gpa.alloc(f64, wide); + defer gpa.free(normed); + for (0..rows) |t| { + const inj: ?[]f64 = if (inject) |i| i[t * streams ..][0..streams] else null; + try hc.preF64(hc.real, w, offset, 1e-6, h[t * wide ..][0..wide], mixed[t * hidden ..][0..hidden], inj, low, mix, normed); + } +} + +pub fn hcPostRows(h: []f64, branch: []const f64, inject: []const f64, rows: usize) void { + for (0..rows) |t| { + for (0..streams) |s| { + const g = inject[t * streams + s]; + for (0..hidden) |d| h[t * wide + s * hidden + d] += branch[t * hidden + d] * g; + } + } +} + +// -- layer weights ----------------------------------------------------------------------------------------- + +pub const GdnW = struct { + qkv: QMat, + z: QMat, + out: QMat, + a_w: []const f64, // [48][2560] + b_w: []const f64, + params: gdn_cpu.Params, + + pub fn deinit(w: *GdnW, gpa: std.mem.Allocator) void { + w.qkv.deinit(gpa); + w.z.deinit(gpa); + w.out.deinit(gpa); + gpa.free(w.a_w); + gpa.free(w.b_w); + gpa.free(w.params.conv_w); + gpa.free(w.params.a_log); + gpa.free(w.params.dt_bias); + gpa.free(w.params.norm_w); + } +}; + +pub const AttnW = struct { + q: QMat, + k: QMat, + v: QMat, + o: QMat, + idx: QMat, + norms: attn_cpu.Norms, + + pub fn deinit(w: *AttnW, gpa: std.mem.Allocator) void { + w.q.deinit(gpa); + w.k.deinit(gpa); + w.v.deinit(gpa); + w.o.deinit(gpa); + w.idx.deinit(gpa); + gpa.free(w.norms.q); + gpa.free(w.norms.k); + gpa.free(w.norms.iq); + gpa.free(w.norms.ik); + } +}; + +/// Where the 512 routed experts of one layer come from: decode expert `e` (gate, up, down) into `out`. +pub const ExpertSource = struct { + ctx: *anyopaque, + load: *const fn (ctx: *anyopaque, e: usize, out: *[3]QMat) anyerror!void, +}; + +pub const MoeW = struct { + gate16: []const f16, // [512][2560] + sgate16: []const f16, // [2560] + shared: [3]QMat, + source: ExpertSource, + + pub fn deinit(w: *MoeW, gpa: std.mem.Allocator) void { + for (&w.shared) |*m| m.deinit(gpa); + } +}; + +pub const Mixer = union(enum) { gdn: *GdnW, attn: *AttnW }; + +pub const LayerW = struct { + attn_hc: hc.Weights, + mlp_hc: hc.Weights, + mixer: Mixer, + moe: *MoeW, +}; + +pub const LayerState = union(enum) { gdn: gdn_cpu.State, attn: attn_cpu.Cache }; + +/// What a layer pass records for comparison with the device. +pub const Trace = struct { + attn_branch: []f64, // [rows][2560] + mlp_branch: []f64, + picks: []u16, // [rows][10] + h_mid: []f64, // [rows][10240] streams after the attention block +}; + +// -- mixers ------------------------------------------------------------------------------------------------ + +fn roundF16Slice(dst: []f16, src: []const f64) void { + for (dst, src) |*d, s| d.* = @floatCast(s); +} + +const GdnProj = struct { + w: *const GdnW, + x: []const f64, + qkv: []f64, + z: []f64, + a: []f64, + b: []f64, + gpa: std.mem.Allocator, + err: ?anyerror = null, + + fn row(j: *GdnProj, t: usize) void { + const tmp = j.gpa.alloc(f64, hidden) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(tmp); + const xr = j.x[t * hidden ..][0..hidden]; + const c = gdn_cpu.qwen.channels(); + const zw = gdn_cpu.qwen.zWidth(); + j.w.qkv.apply(xr, j.qkv[t * c ..][0..c], tmp); + j.w.z.apply(xr, j.z[t * zw ..][0..zw], tmp); + gdn_cpu.rowsDot(j.w.a_w, xr, j.a[t * 48 ..][0..48]); + gdn_cpu.rowsDot(j.w.b_w, xr, j.b[t * 48 ..][0..48]); + } +}; + +const OutJob = struct { + m: *const QMat, + inp: []const f64, + in_w: usize, + out: []f64, + gpa: std.mem.Allocator, + err: ?anyerror = null, + fn row(j: *OutJob, t: usize) void { + const tmp = j.gpa.alloc(f64, j.m.k) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(tmp); + j.m.apply(j.inp[t * j.in_w ..][0..j.in_w], j.out[t * j.m.n ..][0..j.m.n], tmp); + } +}; + +pub fn gdnBranch(gpa: std.mem.Allocator, w: *const GdnW, st: gdn_cpu.State, x: []const f64, rows: usize, out: []f64) !void { + const d = gdn_cpu.qwen; + const c = d.channels(); + const zw = d.zWidth(); + const qkv = try gpa.alloc(f64, rows * c); + defer gpa.free(qkv); + const z = try gpa.alloc(f64, rows * zw); + defer gpa.free(z); + const a = try gpa.alloc(f64, rows * 48); + defer gpa.free(a); + const b = try gpa.alloc(f64, rows * 48); + defer gpa.free(b); + var pj: GdnProj = .{ .w = w, .x = x, .qkv = qkv, .z = z, .a = a, .b = b, .gpa = gpa }; + parallelFor(rows, 48, &pj, GdnProj.row); + if (pj.err) |e| return e; + const gated = try gpa.alloc(f64, rows * zw); + defer gpa.free(gated); + try gdn_cpu.tokens(gpa, d, w.params, st, rows, qkv, z, a, b, gated); + var oj: OutJob = .{ .m = &w.out, .inp = gated, .in_w = zw, .out = out, .gpa = gpa }; + parallelFor(rows, 48, &oj, OutJob.row); + if (oj.err) |e| return e; +} + +const AttnProj = struct { + w: *const AttnW, + x: []const f64, + qp: []f16, + kp: []f16, + vp: []f16, + ip: []f16, + gpa: std.mem.Allocator, + err: ?anyerror = null, + const q_width = attn_cpu.q_heads * 2 * attn_cpu.head_dim; + const kv_width = attn_cpu.kv_heads * attn_cpu.head_dim; + + fn row(j: *AttnProj, t: usize) void { + const tmp = j.gpa.alloc(f64, hidden) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(tmp); + const big = j.gpa.alloc(f64, q_width) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(big); + const xr = j.x[t * hidden ..][0..hidden]; + j.w.q.apply(xr, big, tmp); + roundF16Slice(j.qp[t * q_width ..][0..q_width], big); + j.w.k.apply(xr, big[0..kv_width], tmp); + roundF16Slice(j.kp[t * kv_width ..][0..kv_width], big[0..kv_width]); + j.w.v.apply(xr, big[0..kv_width], tmp); + roundF16Slice(j.vp[t * kv_width ..][0..kv_width], big[0..kv_width]); + j.w.idx.apply(xr, big[0..attn_cpu.index_width], tmp); + roundF16Slice(j.ip[t * attn_cpu.index_width ..][0..attn_cpu.index_width], big[0..attn_cpu.index_width]); + } +}; + +/// Attention rows pos0 .. pos0 + rows. The projections are rounded to fp16 where the EXL3 linears emit them. +pub fn attnBranch(gpa: std.mem.Allocator, w: *const AttnW, cache: *attn_cpu.Cache, pos0: usize, x: []const f64, rows: usize, out: []f64) !void { + const q_width = AttnProj.q_width; + const kv_width = AttnProj.kv_width; + const iw = attn_cpu.index_width; + const qp = try gpa.alloc(f16, rows * q_width); + defer gpa.free(qp); + const kp = try gpa.alloc(f16, rows * kv_width); + defer gpa.free(kp); + const vp = try gpa.alloc(f16, rows * kv_width); + defer gpa.free(vp); + const ip = try gpa.alloc(f16, rows * iw); + defer gpa.free(ip); + var pj: AttnProj = .{ .w = w, .x = x, .qp = qp, .kp = kp, .vp = vp, .ip = ip, .gpa = gpa }; + parallelFor(rows, 48, &pj, AttnProj.row); + if (pj.err) |e| return e; + + const attn_width = attn_cpu.q_heads * attn_cpu.head_dim; + const o_in = try gpa.alloc(f64, rows * attn_width); + defer gpa.free(o_in); + const scores = try gpa.alloc(f64, cache.tokens / attn_cpu.ratio + 2); + defer gpa.free(scores); + var oq: [attn_width]f64 = undefined; + var oiq: [attn_cpu.index_heads * attn_cpu.index_dim]f64 = undefined; + var info = attn_cpu.RowInfo{}; + for (0..rows) |t| { + try attn_cpu.row(cache, pos0 + t, .{ + .qp = qp[t * q_width ..][0..q_width], + .kp = kp[t * kv_width ..][0..kv_width], + .vp = vp[t * kv_width ..][0..kv_width], + .ip = ip[t * iw ..][0..iw], + }, &w.norms, scores, &oq, &oiq, o_in[t * attn_width ..][0..attn_width], &info); + } + var oj: OutJob = .{ .m = &w.o, .inp = o_in, .in_w = attn_width, .out = out, .gpa = gpa }; + parallelFor(rows, 48, &oj, OutJob.row); + if (oj.err) |e| return e; +} + +// -- MoE --------------------------------------------------------------------------------------------------- + +const Loaded = struct { + id: usize, + mats: [3]QMat, +}; + +const LoadJob = struct { + src: ExpertSource, + ids: []const usize, + slots: []Loaded, + err: ?anyerror = null, + fn one(j: *LoadJob, i: usize) void { + j.slots[i].id = j.ids[i]; + j.src.load(j.src.ctx, j.ids[i], &j.slots[i].mats) catch |e| { + j.err = e; + }; + } +}; + +const MoeRow = struct { + w: *const MoeW, + x: []const f64, + out: []f64, + picks: []const u16, // [rows][10] + weights: []const f64, // [rows][10] + shared_w: []const f64, + loaded: []const Loaded, + gpa: std.mem.Allocator, + err: ?anyerror = null, + + fn find(j: *const MoeRow, id: usize) *const Loaded { + for (j.loaded) |*l| if (l.id == id) return l; + unreachable; + } + + fn row(j: *MoeRow, t: usize) void { + const buf = j.gpa.alloc(f64, hidden + 4 * inter + 3 * hidden) catch |e| { + j.err = e; + return; + }; + defer j.gpa.free(buf); + const tmp = buf[0..hidden]; + const g = buf[hidden..][0..inter]; + const u = buf[hidden + inter ..][0..inter]; + const glu = buf[hidden + 2 * inter ..][0..inter]; + const y = buf[hidden + 3 * inter ..][0..hidden]; + const acc = j.out[t * hidden ..][0..hidden]; + @memset(acc, 0); + const xr = j.x[t * hidden ..][0..hidden]; + for (0..topk) |k| { + const l = j.find(j.picks[t * topk + k]); + expertApply(&l.mats, xr, tmp, g, u, glu, y); + const wk = j.weights[t * topk + k]; + for (acc, y) |*o, v| o.* += wk * v; + } + expertApply(&j.w.shared, xr, tmp, g, u, glu, y); + const sw = j.shared_w[t]; + for (acc, y) |*o, v| o.* += sw * v; + } +}; + +fn expertApply(m: *const [3]QMat, x: []const f64, tmp: []f64, g: []f64, u: []f64, glu: []f64, y: []f64) void { + m[0].apply(x, g, tmp); + m[1].apply(x, u, tmp); + for (glu, g, u) |*o, gv, uv| o.* = silu(gv) * uv; + m[2].apply(glu, y, tmp[0..inter]); +} + +/// Routed + shared experts for `rows` rows from the f64 hidden `x` (the f64 router, not the device's). +pub fn moeBranch(gpa: std.mem.Allocator, w: *const MoeW, x: []const f64, rows: usize, out: []f64, picks_out: []u16) !void { + const picks = picks_out; + const weights = try gpa.alloc(f64, rows * topk); + defer gpa.free(weights); + const shared_w = try gpa.alloc(f64, rows); + defer gpa.free(shared_w); + var logits: [experts + 1]f64 = undefined; + var ids: std.ArrayList(usize) = .empty; + defer ids.deinit(gpa); + for (0..rows) |t| { + try moe.logitsF16(w.gate16, experts, w.sgate16, x[t * hidden ..][0..hidden], &logits); + var sw: f64 = 0; + try moe.routeFromLogits(&logits, topk, picks[t * topk ..][0..topk], weights[t * topk ..][0..topk], &sw); + shared_w[t] = sw; + for (picks[t * topk ..][0..topk]) |p| { + var seen = false; + for (ids.items) |i| seen = seen or i == p; + if (!seen) try ids.append(gpa, p); + } + } + const loaded = try gpa.alloc(Loaded, ids.items.len); + defer gpa.free(loaded); + var lj: LoadJob = .{ .src = w.source, .ids = ids.items, .slots = loaded }; + parallelFor(ids.items.len, 48, &lj, LoadJob.one); + defer if (lj.err == null) for (loaded) |*l| for (&l.mats) |*m| m.deinit(gpa); + if (lj.err) |e| return e; + var mj: MoeRow = .{ .w = w, .x = x, .out = out, .picks = picks, .weights = weights, .shared_w = shared_w, .loaded = loaded, .gpa = gpa }; + parallelFor(rows, 48, &mj, MoeRow.row); + if (mj.err) |e| return e; +} + +// -- one layer --------------------------------------------------------------------------------------------- + +/// `h` is [rows][10240] f64 streams, updated in place. `pos0` is the absolute position of row 0. +pub fn layerRows(gpa: std.mem.Allocator, w: LayerW, st: *LayerState, pos0: usize, offset: f64, h: []f64, rows: usize, trace: ?Trace) !void { + const mixed = try gpa.alloc(f64, rows * hidden); + defer gpa.free(mixed); + const inject = try gpa.alloc(f64, rows * streams); + defer gpa.free(inject); + const branch = try gpa.alloc(f64, rows * hidden); + defer gpa.free(branch); + + try hcPreRows(gpa, w.attn_hc, offset, h, rows, mixed, inject); + switch (w.mixer) { + .gdn => |g| try gdnBranch(gpa, g, st.gdn, mixed, rows, branch), + .attn => |a| try attnBranch(gpa, a, &st.attn, pos0, mixed, rows, branch), + } + if (trace) |tr| @memcpy(tr.attn_branch[0 .. rows * hidden], branch); + hcPostRows(h, branch, inject, rows); + if (trace) |tr| @memcpy(tr.h_mid[0 .. rows * wide], h[0 .. rows * wide]); + + try hcPreRows(gpa, w.mlp_hc, offset, h, rows, mixed, inject); + const picks = try gpa.alloc(u16, rows * topk); + defer gpa.free(picks); + try moeBranch(gpa, w.moe, mixed, rows, branch, picks); + if (trace) |tr| { + @memcpy(tr.mlp_branch[0 .. rows * hidden], branch); + @memcpy(tr.picks[0 .. rows * topk], picks); + } + hcPostRows(h, branch, inject, rows); +} + +comptime { + _ = qmat; +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_gdn_cpu.zig b/zig/src/families/flashnext/exl3/model/exl3_gdn_cpu.zig new file mode 100644 index 0000000000..cbf187544e --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_gdn_cpu.zig @@ -0,0 +1,399 @@ +//! CPU oracle for one Gated DeltaNet layer of Qwen3.8-Flash-Next (qwen4_exp_text), everything between the two EXL3 + +const std = @import("std"); + +pub const Dims = struct { + nk: usize, + nv: usize, + dk: usize, + dv: usize, + taps: usize, + + /// conv channels: q (nk*dk), k (nk*dk), v (nv*dv) + pub fn channels(d: Dims) usize { + return 2 * d.nk * d.dk + d.nv * d.dv; + } + pub fn zWidth(d: Dims) usize { + return d.nv * d.dv; + } + /// conv state rows kept per layer + pub fn keep(d: Dims) usize { + return d.taps - 1; + } + pub fn recElems(d: Dims) usize { + return d.nv * d.dv * d.dk; + } + pub fn convElems(d: Dims) usize { + return d.keep() * d.channels(); + } +}; + +/// Flash Next: 16 key heads, 48 value heads, 128 x 128, kernel 4. +pub const qwen: Dims = .{ .nk = 16, .nv = 48, .dk = 128, .dv = 128, .taps = 4 }; + +/// Per-sequence, per-layer state. +pub const State = struct { + conv: []f64, + rec: []f64, + + pub fn init(gpa: std.mem.Allocator, d: Dims) !State { + const conv = try gpa.alloc(f64, d.convElems()); + errdefer gpa.free(conv); + const rec = try gpa.alloc(f64, d.recElems()); + @memset(conv, 0); + @memset(rec, 0); + return .{ .conv = conv, .rec = rec }; + } + pub fn deinit(s: *State, gpa: std.mem.Allocator) void { + gpa.free(s.conv); + gpa.free(s.rec); + } + pub fn reset(s: *State) void { + @memset(s.conv, 0); + @memset(s.rec, 0); + } +}; + +/// The layer's non-EXL3 parameters, as f64. +pub const Params = struct { + conv_w: []const f64, // [channels][taps] + a_log: []const f64, // [nv] + dt_bias: []const f64, // [nv] + norm_w: []const f64, // [dv] + eps: f64 = 1e-6, +}; + +pub fn sigmoid(x: f64) f64 { + return 1.0 / (1.0 + @exp(-x)); +} + +pub fn silu(x: f64) f64 { + return x / (1.0 + @exp(-x)); +} + +/// torch.nn.functional.softplus with threshold 20. +pub fn softplus(x: f64) f64 { + return if (x > 20.0) x else std.math.log1p(@exp(x)); +} + +/// y[j] = sum_i w[j * k + i] * x[i], for a dense [rows][k] weight (in_proj_a / in_proj_b). +pub fn rowsDot(w: []const f64, x: []const f64, y: []f64) void { + const k = x.len; + for (y, 0..) |*o, j| { + var s: f64 = 0; + for (x, w[j * k ..][0..k]) |xi, wi| s += xi * wi; + o.* = s; + } +} + +/// One token of the causal depthwise conv: out = silu(sum_tap w[c][tap] * x[t - (taps-1) + tap]), x before the +pub fn convStep(d: Dims, p: Params, conv: []f64, x: []const f64, out: []f64) void { + const c = d.channels(); + const kp = d.keep(); + for (0..c) |ch| { + var acc: f64 = 0; + for (0..d.taps) |tap| { + const v = if (tap < kp) conv[tap * c + ch] else x[ch]; + acc += p.conv_w[ch * d.taps + tap] * v; + } + out[ch] = silu(acc); + } + for (0..kp) |r| { + for (0..c) |ch| conv[r * c + ch] = if (r + 1 < kp) conv[(r + 1) * c + ch] else x[ch]; + } +} + +/// Decay exp(g) and beta for each value head from the raw a and b projections. +pub fn gates(d: Dims, p: Params, a: []const f64, b: []const f64, decay: []f64, beta: []f64) void { + for (0..d.nv) |h| { + const g = -@exp(p.a_log[h]) * softplus(a[h] + p.dt_bias[h]); + decay[h] = @exp(g); + beta[h] = sigmoid(b[h]); + } +} + +/// One token of the delta rule on all value heads. `act` is the conv output [q | k | v]; y is [nv][dv]. +pub fn deltaStep(d: Dims, rec: []f64, act: []const f64, decay: []const f64, beta: []const f64, y: []f64) void { + const group = d.nv / d.nk; + const kbase = d.nk * d.dk; + const vbase = 2 * kbase; + for (0..d.nv) |hv| { + const hk = hv / group; + const q = act[hk * d.dk ..][0..d.dk]; + const k = act[kbase + hk * d.dk ..][0..d.dk]; + var qq: f64 = 0; + var kk: f64 = 0; + for (q, k) |a, b| { + qq += a * a; + kk += b * b; + } + const iq = 1.0 / @sqrt(qq + 1e-6) / @sqrt(@as(f64, @floatFromInt(d.dk))); + const ik = 1.0 / @sqrt(kk + 1e-6); + for (0..d.dv) |v| { + const row = rec[(hv * d.dv + v) * d.dk ..][0..d.dk]; + var kv: f64 = 0; + for (row, k) |*s, ki| { + s.* *= decay[hv]; + kv += s.* * ki * ik; + } + const delta = (act[vbase + hv * d.dv + v] - kv) * beta[hv]; + var o: f64 = 0; + for (row, k, q) |*s, ki, qi| { + s.* += ki * ik * delta; + o += s.* * qi * iq; + } + y[hv * d.dv + v] = o; + } + } +} + +/// Gated RMS norm per value head: w * (y * rsqrt(mean(y^2) + eps)) * sigmoid(z). +pub fn gatedNorm(d: Dims, p: Params, y: []const f64, z: []const f64, out: []f64) void { + for (0..d.nv) |hv| { + const row = y[hv * d.dv ..][0..d.dv]; + var ss: f64 = 0; + for (row) |v| ss += v * v; + const inv = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(d.dv)) + p.eps); + for (row, 0..) |v, i| out[hv * d.dv + i] = p.norm_w[i] * (v * inv) * sigmoid(z[hv * d.dv + i]); + } +} + +/// `n` tokens, in order, from the projection outputs: qkv [n][channels], z [n][nv*dv],. +pub fn tokens(gpa: std.mem.Allocator, d: Dims, p: Params, st: State, n: usize, qkv: []const f64, z: []const f64, a: []const f64, b: []const f64, out: []f64) !void { + const c = d.channels(); + const zw = d.zWidth(); + const act = try gpa.alloc(f64, c); + defer gpa.free(act); + const y = try gpa.alloc(f64, zw); + defer gpa.free(y); + const decay = try gpa.alloc(f64, d.nv); + defer gpa.free(decay); + const beta = try gpa.alloc(f64, d.nv); + defer gpa.free(beta); + for (0..n) |t| { + convStep(d, p, st.conv, qkv[t * c ..][0..c], act); + gates(d, p, a[t * d.nv ..][0..d.nv], b[t * d.nv ..][0..d.nv], decay, beta); + deltaStep(d, st.rec, act, decay, beta, y); + gatedNorm(d, p, y, z[t * zw ..][0..zw], out[t * zw ..][0..zw]); + } +} + +/// Byte sizes of the device state table: conv rows are fp16 (they are raw EXL3 projection outputs), the recurrent +pub const Table = struct { + dims: Dims, + layers: usize, + slots: usize, + + pub fn convBytesPer(t: Table) usize { + return t.dims.convElems() * 2; + } + pub fn recBytesPer(t: Table) usize { + return t.dims.recElems() * 4; + } + pub fn convBytes(t: Table) usize { + return t.convBytesPer() * t.layers * t.slots; + } + pub fn recBytes(t: Table) usize { + return t.recBytesPer() * t.layers * t.slots; + } + pub fn convOffset(t: Table, slot: usize, layer: usize) usize { + return (slot * t.layers + layer) * t.convBytesPer(); + } + pub fn recOffset(t: Table, slot: usize, layer: usize) usize { + return (slot * t.layers + layer) * t.recBytesPer(); + } +}; + +// ---- tests (small dims; the arithmetic does not depend on the Flash Next sizes) ---- + +const small: Dims = .{ .nk = 2, .nv = 4, .dk = 8, .dv = 8, .taps = 4 }; + +fn fill(rng: *std.Random.DefaultPrng, xs: []f64, scale: f64) void { + for (xs) |*x| x.* = (rng.random().float(f64) - 0.5) * 2.0 * scale; +} + +const TestParams = struct { + conv_w: []f64, + a_log: []f64, + dt_bias: []f64, + norm_w: []f64, + + fn init(gpa: std.mem.Allocator, d: Dims, rng: *std.Random.DefaultPrng) !TestParams { + const t: TestParams = .{ + .conv_w = try gpa.alloc(f64, d.channels() * d.taps), + .a_log = try gpa.alloc(f64, d.nv), + .dt_bias = try gpa.alloc(f64, d.nv), + .norm_w = try gpa.alloc(f64, d.dv), + }; + fill(rng, t.conv_w, 0.7); + fill(rng, t.a_log, 1.0); + fill(rng, t.dt_bias, 1.0); + for (t.norm_w) |*w| w.* = 0.5 + rng.random().float(f64); + return t; + } + fn deinit(t: TestParams, gpa: std.mem.Allocator) void { + gpa.free(t.conv_w); + gpa.free(t.a_log); + gpa.free(t.dt_bias); + gpa.free(t.norm_w); + } + fn view(t: TestParams) Params { + return .{ .conv_w = t.conv_w, .a_log = t.a_log, .dt_bias = t.dt_bias, .norm_w = t.norm_w }; + } +}; + +test "conv is a causal depthwise filter over the stored rows" { + const gpa = std.testing.allocator; + var rng = std.Random.DefaultPrng.init(1); + const tp = try TestParams.init(gpa, small, &rng); + defer tp.deinit(gpa); + const c = small.channels(); + const n = 7; + const xs = try gpa.alloc(f64, n * c); + defer gpa.free(xs); + fill(&rng, xs, 1.0); + var st = try State.init(gpa, small); + defer st.deinit(gpa); + const out = try gpa.alloc(f64, c); + defer gpa.free(out); + for (0..n) |t| { + convStep(small, tp.view(), st.conv, xs[t * c ..][0..c], out); + for (0..c) |ch| { + var acc: f64 = 0; + for (0..small.taps) |tap| { + // input index t - (taps - 1) + tap; before the sequence it is zero + const at: isize = @as(isize, @intCast(t)) - @as(isize, @intCast(small.taps - 1)) + @as(isize, @intCast(tap)); + if (at >= 0) acc += tp.conv_w[ch * small.taps + tap] * xs[@as(usize, @intCast(at)) * c + ch]; + } + try std.testing.expectApproxEqAbs(silu(acc), out[ch], 1e-12); + } + } + // the state ends as the last three raw rows + for (0..3) |r| for (0..c) |ch| try std.testing.expectEqual(xs[(n - 3 + r) * c + ch], st.conv[r * c + ch]); +} + +test "one scalar head follows S = g S + k (v - g S k) beta, o = S q" { + const d: Dims = .{ .nk = 1, .nv = 1, .dk = 1, .dv = 1, .taps = 4 }; + var rec = [_]f64{0.5}; + const act = [_]f64{ 2.0, 3.0, 0.75 }; // q, k, v + const decay = [_]f64{0.9}; + const beta = [_]f64{0.4}; + var y = [_]f64{0}; + deltaStep(d, &rec, &act, &decay, &beta, &y); + const ik = 1.0 / @sqrt(9.0 + 1e-6); + const iq = 1.0 / @sqrt(4.0 + 1e-6); + const k = 3.0 * ik; + const s0 = 0.5 * 0.9; + const delta = (0.75 - s0 * k) * 0.4; + const s1 = s0 + k * delta; + try std.testing.expectApproxEqAbs(s1, rec[0], 1e-14); + try std.testing.expectApproxEqAbs(s1 * 2.0 * iq, y[0], 1e-14); +} + +test "with beta 1 and no decay the state returns v for the key just written" { + const d: Dims = .{ .nk = 1, .nv = 2, .dk = 4, .dv = 3, .taps = 4 }; + var rng = std.Random.DefaultPrng.init(2); + var rec: [2 * 3 * 4]f64 = undefined; + fill(&rng, &rec, 1.0); + var act: [2 * 4 + 2 * 3]f64 = undefined; + fill(&rng, &act, 1.0); + const decay = [_]f64{ 1, 1 }; + const beta = [_]f64{ 1, 1 }; + var y: [6]f64 = undefined; + deltaStep(d, &rec, &act, &decay, &beta, &y); + var kk: f64 = 0; + for (act[4..8]) |k| kk += k * k; + const ik = 1.0 / @sqrt(kk + 1e-6); + for (0..2) |hv| for (0..3) |v| { + var read: f64 = 0; + for (0..4) |i| read += rec[(hv * 3 + v) * 4 + i] * act[4 + i] * ik; + // S k = v * (|k|^2 / (|k|^2 + 1e-6)), the l2norm eps leaves a tiny shortfall + try std.testing.expectApproxEqAbs(act[8 + hv * 3 + v], read, 1e-5); + }; +} + +test "value heads share their key head in groups (repeat_interleave)" { + const d: Dims = .{ .nk = 2, .nv = 4, .dk = 4, .dv = 2, .taps = 4 }; + var rng = std.Random.DefaultPrng.init(3); + var act: [2 * 2 * 4 + 4 * 2]f64 = undefined; + fill(&rng, &act, 1.0); + const decay = [_]f64{ 0.8, 0.8, 0.8, 0.8 }; + const beta = [_]f64{ 0.5, 0.5, 0.5, 0.5 }; + var rec: [4 * 2 * 4]f64 = @splat(0); + var y: [8]f64 = undefined; + deltaStep(d, &rec, &act, &decay, &beta, &y); + // zero state: kv = 0, so S = k (v beta)^T and o = (k . q) v beta, with q, k of key head hv / 2 + for (0..4) |hv| { + const hk = hv / 2; + const q = act[hk * 4 ..][0..4]; + const k = act[8 + hk * 4 ..][0..4]; + var qq: f64 = 0; + var kk: f64 = 0; + var qk: f64 = 0; + for (q, k) |a, b| { + qq += a * a; + kk += b * b; + qk += a * b; + } + const iq = 1.0 / @sqrt(qq + 1e-6) / @sqrt(4.0); + const ik = 1.0 / @sqrt(kk + 1e-6); + for (0..2) |v| try std.testing.expectApproxEqAbs(qk * iq * ik * act[16 + hv * 2 + v] * 0.5, y[hv * 2 + v], 1e-12); + } +} + +test "gated norm is w * rms-normalised y * sigmoid(z)" { + const d: Dims = .{ .nk = 1, .nv = 1, .dk = 1, .dv = 4, .taps = 4 }; + const p: Params = .{ .conv_w = &.{}, .a_log = &.{}, .dt_bias = &.{}, .norm_w = &.{ 1, 2, 3, 4 } }; + const y = [_]f64{ 1, -2, 3, -4 }; + const z = [_]f64{ 0, 1, -1, 2 }; + var out: [4]f64 = undefined; + gatedNorm(d, p, &y, &z, &out); + const inv = 1.0 / @sqrt((1.0 + 4 + 9 + 16) / 4.0 + 1e-6); + for (0..4) |i| try std.testing.expectApproxEqAbs(p.norm_w[i] * y[i] * inv * sigmoid(z[i]), out[i], 1e-14); +} + +test "processing tokens in two calls equals one call, state included" { + const gpa = std.testing.allocator; + var rng = std.Random.DefaultPrng.init(4); + const tp = try TestParams.init(gpa, small, &rng); + defer tp.deinit(gpa); + const n = 9; + const c = small.channels(); + const zw = small.zWidth(); + const qkv = try gpa.alloc(f64, n * c); + defer gpa.free(qkv); + const z = try gpa.alloc(f64, n * zw); + defer gpa.free(z); + const a = try gpa.alloc(f64, n * small.nv); + defer gpa.free(a); + const b = try gpa.alloc(f64, n * small.nv); + defer gpa.free(b); + fill(&rng, qkv, 1.0); + fill(&rng, z, 1.0); + fill(&rng, a, 1.0); + fill(&rng, b, 1.0); + var whole = try State.init(gpa, small); + defer whole.deinit(gpa); + var split = try State.init(gpa, small); + defer split.deinit(gpa); + const o1 = try gpa.alloc(f64, n * zw); + defer gpa.free(o1); + const o2 = try gpa.alloc(f64, n * zw); + defer gpa.free(o2); + try tokens(gpa, small, tp.view(), whole, n, qkv, z, a, b, o1); + const cut = 5; + try tokens(gpa, small, tp.view(), split, cut, qkv, z, a, b, o2); + try tokens(gpa, small, tp.view(), split, n - cut, qkv[cut * c ..], z[cut * zw ..], a[cut * small.nv ..], b[cut * small.nv ..], o2[cut * zw ..]); + try std.testing.expectEqualSlices(f64, o1, o2); + try std.testing.expectEqualSlices(f64, whole.rec, split.rec); + try std.testing.expectEqualSlices(f64, whole.conv, split.conv); +} + +test "state table addresses a (slot, layer) pair and Flash Next sizes" { + const t: Table = .{ .dims = qwen, .layers = 36, .slots = 2 }; + try std.testing.expectEqual(@as(usize, 10240), qwen.channels()); + try std.testing.expectEqual(@as(usize, 3 * 10240 * 2), t.convBytesPer()); + try std.testing.expectEqual(@as(usize, 48 * 128 * 128 * 4), t.recBytesPer()); + try std.testing.expectEqual(t.recBytesPer() * 37, t.recOffset(1, 1)); + try std.testing.expectEqual(t.convBytesPer() * 36 * 2, t.convBytes()); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_hc.zig b/zig/src/families/flashnext/exl3/model/exl3_hc.zig new file mode 100644 index 0000000000..1d4a25ef14 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_hc.zig @@ -0,0 +1,253 @@ +//! Host reference of the hyper-connection pre, post and expand steps over the unquantized hc weights. + +const std = @import("std"); +const shape = @import("../pack/exl3_shape.zig"); + +// --------------------------------------------------------------------------------------------- Hyper-connection pre /. + +pub const Dims = struct { streams: usize, dim: usize, lowrank: usize }; +pub const real: Dims = .{ .streams = shape.hc_count, .dim = shape.hidden, .lowrank = shape.hc_lowrank }; + +/// Row-major, out-major weights as f32 (exact from F16). `inject` is null for the final mixer. +pub const Weights = struct { + norm: []const f32, + down: []const f32, + up: []const f32, + inject: ?[]const f32, +}; + +pub const Error = error{Length}; + +fn silu32(x: f32) f32 { + return x / (1.0 + @exp(-x)); +} +fn sigmoid32(x: f32) f32 { + return 1.0 / (1.0 + @exp(-x)); +} + +/// f32 values `scratch` must hold: normed (wide), low (lowrank), mix (wide). +pub fn scratchLen(d: Dims) usize { + return 2 * d.streams * d.dim + d.lowrank; +} + +fn checkWeights(d: Dims, w: Weights) Error!void { + const wide = d.streams * d.dim; + if (w.norm.len != wide or w.down.len != d.lowrank * wide or w.up.len != wide * d.lowrank) return error.Length; + if (w.inject) |inj| if (inj.len != d.streams * wide) return error.Length; +} + +/// Block input `mixed` [dim] and, when the weights carry `block_inject_weight` and `inject` is non-null, the S gates. +pub fn pre(d: Dims, w: Weights, offset: f32, eps: f32, h: []const f32, mixed: []f32, inject: ?[]f32, scratch: []f32) Error!void { + const wide = d.streams * d.dim; + try checkWeights(d, w); + if (h.len != wide or mixed.len != d.dim or scratch.len < scratchLen(d)) return error.Length; + if (inject) |g| if (w.inject == null or g.len != d.streams) return error.Length; + if (!std.math.isFinite(eps) or eps < 0) return error.Length; + const normed = scratch[0..wide]; + const low = scratch[wide .. wide + d.lowrank]; + const mix = scratch[wide + d.lowrank .. 2 * wide + d.lowrank]; + + for (0..d.streams) |s| { + const row = h[s * d.dim ..][0..d.dim]; + var sum_sq: f32 = 0; + for (row) |v| sum_sq += v * v; + const rms = @sqrt(sum_sq / @as(f32, @floatFromInt(d.dim)) + eps); + for (0..d.dim) |i| normed[s * d.dim + i] = row[i] / rms * (w.norm[s * d.dim + i] + offset); + } + const inv_s: f32 = 1.0 / @as(f32, @floatFromInt(d.streams)); + for (0..d.lowrank) |r| { + var acc: f32 = 0; + const row = w.down[r * wide ..][0..wide]; + for (row, normed) |a, b| acc += a * b; + low[r] = silu32(acc * inv_s); + } + for (0..wide) |j| { + var acc: f32 = 0; + const row = w.up[j * d.lowrank ..][0..d.lowrank]; + for (row, low) |a, b| acc += a * b; + mix[j] = sigmoid32(acc); + } + for (0..d.dim) |i| { + var acc: f32 = 0; + for (0..d.streams) |s| acc += mix[s * d.dim + i] * normed[s * d.dim + i]; + mixed[i] = acc * inv_s; + } + if (inject) |g| { + const inj = w.inject.?; + for (0..d.streams) |s| { + var acc: f32 = 0; + const row = inj[s * wide ..][0..wide]; + for (row, normed) |a, b| acc += a * b; + g[s] = 2.0 * sigmoid32(acc * inv_s); + } + } +} + +/// `out[s, d] = h[s, d] + branch[d] * inject[s]`. `out` may be `h`. Lengths are checked before any store. +pub fn post(d: Dims, h: []const f32, branch: []const f32, inject: []const f32, out: []f32) Error!void { + if (h.len != d.streams * d.dim or out.len != h.len or branch.len != d.dim or inject.len != d.streams) return error.Length; + for (0..d.streams) |s| { + for (0..d.dim) |i| out[s * d.dim + i] = h[s * d.dim + i] + branch[i] * inject[s]; + } +} + +/// Model start: the embedding row copied into every stream. +pub fn expand(d: Dims, embed: []const f32, out: []f32) Error!void { + if (embed.len != d.dim or out.len != d.streams * d.dim) return error.Length; + for (0..d.streams) |s| @memcpy(out[s * d.dim ..][0..d.dim], embed); +} + +/// The same maths in f64, written as plain loops over (stream, dim) with no scratch, for the oracle. +pub fn preF64(d: Dims, w: Weights, offset: f64, eps: f64, h: []const f64, mixed: []f64, inject: ?[]f64, low: []f64, mix: []f64, normed: []f64) Error!void { + const wide = d.streams * d.dim; + try checkWeights(d, w); + if (h.len != wide or mixed.len != d.dim or low.len != d.lowrank or mix.len != wide or normed.len != wide) return error.Length; + for (0..d.streams) |s| { + var ms: f64 = 0; + for (0..d.dim) |i| ms += h[s * d.dim + i] * h[s * d.dim + i]; + ms /= @floatFromInt(d.dim); + const inv = 1.0 / @sqrt(ms + eps); + for (0..d.dim) |i| normed[s * d.dim + i] = h[s * d.dim + i] * inv * (@as(f64, w.norm[s * d.dim + i]) + offset); + } + const sf: f64 = @floatFromInt(d.streams); + for (0..d.lowrank) |r| { + var acc: f64 = 0; + for (0..wide) |c| acc += @as(f64, w.down[r * wide + c]) * normed[c]; + const z = acc / sf; + low[r] = z / (1.0 + @exp(-z)); + } + for (0..wide) |j| { + var acc: f64 = 0; + for (0..d.lowrank) |r| acc += @as(f64, w.up[j * d.lowrank + r]) * low[r]; + mix[j] = 1.0 / (1.0 + @exp(-acc)); + } + for (0..d.dim) |i| { + var acc: f64 = 0; + for (0..d.streams) |s| acc += mix[s * d.dim + i] * normed[s * d.dim + i]; + mixed[i] = acc / sf; + } + if (inject) |g| { + const inj = w.inject orelse return error.Length; + if (g.len != d.streams) return error.Length; + for (0..d.streams) |s| { + var acc: f64 = 0; + for (0..wide) |c| acc += @as(f64, inj[s * wide + c]) * normed[c]; + g[s] = 2.0 / (1.0 + @exp(-acc / sf)); + } + } +} + +pub fn postF64(d: Dims, h: []const f64, branch: []const f64, inject: []const f64, out: []f64) Error!void { + if (h.len != d.streams * d.dim or out.len != h.len or branch.len != d.dim or inject.len != d.streams) return error.Length; + for (0..d.streams) |s| for (0..d.dim) |i| { + out[s * d.dim + i] = h[s * d.dim + i] + branch[i] * inject[s]; + }; +} + +const tiny: Dims = .{ .streams = 2, .dim = 4, .lowrank = 3 }; + +fn tinyWeights(gpa: std.mem.Allocator, with_inject: bool) !struct { norm: []f32, down: []f32, up: []f32, inject: ?[]f32 } { + const wide = tiny.streams * tiny.dim; + var rng = std.Random.DefaultPrng.init(7); + const r = rng.random(); + const norm = try gpa.alloc(f32, wide); + const down = try gpa.alloc(f32, tiny.lowrank * wide); + const up = try gpa.alloc(f32, wide * tiny.lowrank); + for (norm) |*v| v.* = r.float(f32) - 0.5; + for (down) |*v| v.* = (r.float(f32) - 0.5) * 2; + for (up) |*v| v.* = (r.float(f32) - 0.5) * 2; + var inject: ?[]f32 = null; + if (with_inject) { + inject = try gpa.alloc(f32, tiny.streams * wide); + for (inject.?) |*v| v.* = (r.float(f32) - 0.5) * 2; + } + return .{ .norm = norm, .down = down, .up = up, .inject = inject }; +} + +test "pre and post match the f64 plain-loop oracle on a tiny shape" { + const gpa = std.testing.allocator; + const t = try tinyWeights(gpa, true); + defer gpa.free(t.norm); + defer gpa.free(t.down); + defer gpa.free(t.up); + defer gpa.free(t.inject.?); + const w: Weights = .{ .norm = t.norm, .down = t.down, .up = t.up, .inject = t.inject }; + const wide = tiny.streams * tiny.dim; + const h = [wide]f32{ 1, -2, 3, 0.5, 4, 1, -1, 2 }; + var mixed: [tiny.dim]f32 = undefined; + var gates: [tiny.streams]f32 = undefined; + var scratch: [2 * wide + tiny.lowrank]f32 = undefined; + try pre(tiny, w, 1.0, 1e-6, &h, &mixed, &gates, &scratch); + + var h64: [wide]f64 = undefined; + for (&h64, h) |*a, b| a.* = b; + var mixed64: [tiny.dim]f64 = undefined; + var gates64: [tiny.streams]f64 = undefined; + var low: [tiny.lowrank]f64 = undefined; + var mix: [wide]f64 = undefined; + var normed: [wide]f64 = undefined; + try preF64(tiny, w, 1.0, 1e-6, &h64, &mixed64, &gates64, &low, &mix, &normed); + for (mixed, mixed64) |a, b| try std.testing.expectApproxEqAbs(@as(f64, a), b, 1e-5); + for (gates, gates64) |a, b| try std.testing.expectApproxEqAbs(@as(f64, a), b, 1e-5); + + const branch = [tiny.dim]f32{ 0.5, -1, 2, 0 }; + var out: [wide]f32 = undefined; + try post(tiny, &h, &branch, &gates, &out); + for (0..tiny.streams) |s| for (0..tiny.dim) |i| { + try std.testing.expectApproxEqAbs(h[s * tiny.dim + i] + branch[i] * gates[s], out[s * tiny.dim + i], 1e-6); + }; + // aliasing out and h is allowed + var again = h; + try post(tiny, &again, &branch, &gates, &again); + try std.testing.expectEqualSlices(f32, &out, &again); +} + +test "zero mixing weights give mix 1/2, inject 1, and the offset-only norm" { + const gpa = std.testing.allocator; + const t = try tinyWeights(gpa, true); + defer gpa.free(t.norm); + defer gpa.free(t.down); + defer gpa.free(t.up); + defer gpa.free(t.inject.?); + @memset(t.down, 0); + @memset(t.up, 0); + @memset(t.inject.?, 0); + @memset(t.norm, 0); + const w: Weights = .{ .norm = t.norm, .down = t.down, .up = t.up, .inject = t.inject }; + const wide = tiny.streams * tiny.dim; + // stream 0 all 3s, stream 1 all -2s: each stream normalises to +-1 (offset 1, weight 0). + const h = [wide]f32{ 3, 3, 3, 3, -2, -2, -2, -2 }; + var mixed: [tiny.dim]f32 = undefined; + var gates: [tiny.streams]f32 = undefined; + var scratch: [2 * wide + tiny.lowrank]f32 = undefined; + try pre(tiny, w, 1.0, 0.0, &h, &mixed, &gates, &scratch); + // mean_s(0.5 * (+1, -1)) = 0 + for (mixed) |v| try std.testing.expectApproxEqAbs(@as(f32, 0), v, 1e-6); + for (gates) |g| try std.testing.expectApproxEqAbs(@as(f32, 1), g, 1e-6); + // gamma stored directly: offset 0 and weight 0 zeroes the normed streams + try pre(tiny, w, 0.0, 0.0, &h, &mixed, &gates, &scratch); + for (mixed) |v| try std.testing.expectEqual(@as(f32, 0), v); +} + +test "pre, post and expand refuse bad lengths before a store" { + const gpa = std.testing.allocator; + const t = try tinyWeights(gpa, false); + defer gpa.free(t.norm); + defer gpa.free(t.down); + defer gpa.free(t.up); + const w: Weights = .{ .norm = t.norm, .down = t.down, .up = t.up, .inject = null }; + const wide = tiny.streams * tiny.dim; + const h: [wide]f32 = @splat(1); + var mixed: [tiny.dim]f32 = @splat(9); + var gates: [tiny.streams]f32 = undefined; + var scratch: [2 * wide + tiny.lowrank]f32 = undefined; + try std.testing.expectError(error.Length, pre(tiny, w, 1, 1e-6, h[1..], &mixed, null, &scratch)); + try std.testing.expectError(error.Length, pre(tiny, w, 1, 1e-6, &h, &mixed, &gates, &scratch)); // no inject weights + try std.testing.expectError(error.Length, pre(tiny, w, 1, 1e-6, &h, &mixed, null, scratch[0..3])); + try std.testing.expectEqual(@as(f32, 9), mixed[0]); + var out: [wide]f32 = undefined; + try std.testing.expectError(error.Length, post(tiny, &h, mixed[0..3], &gates, &out)); + try expand(tiny, &.{ 1, 2, 3, 4 }, &out); + try std.testing.expectEqualSlices(f32, &.{ 1, 2, 3, 4, 1, 2, 3, 4 }, &out); + try std.testing.expectError(error.Length, expand(tiny, &.{ 1, 2, 3 }, &out)); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_moe.zig b/zig/src/families/flashnext/exl3/model/exl3_moe.zig new file mode 100644 index 0000000000..6220a4b9c1 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_moe.zig @@ -0,0 +1,147 @@ +//! CPU f64 oracle for one Flash Next MoE block: router, routed experts, shared expert, combine. + +const std = @import("std"); +const format = @import("../pack/exl3_format.zig"); +const cpu = @import("exl3_cpu.zig"); + +/// One decoded EXL3 matrix W[k][n] in f64. `y[j] = sum_i x[i] * w[i * n + j]`. +pub const Dense = struct { + k: usize, + n: usize, + w: []f64, + + pub fn init(gpa: std.mem.Allocator, halves: u16, trellis: []const i16, tiles_k: usize, tiles_n: usize, suh: []const f16, svh: []const f16) !Dense { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % format.hadamard_dim != 0 or n % format.hadamard_dim != 0) return error.Length; + if (suh.len != k or svh.len != n or trellis.len != tiles_k * tiles_n * @as(usize, halves) * 8) return error.Length; + const w = try gpa.alloc(f64, k * n); + errdefer gpa.free(w); + const scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(scratch); + format.dequantize(.{ .halves = halves }, .mul1, trellis, tiles_k, tiles_n, suh, svh, w, scratch); + return .{ .k = k, .n = n, .w = w }; + } + + pub fn deinit(d: *Dense, gpa: std.mem.Allocator) void { + gpa.free(d.w); + d.* = undefined; + } + + pub fn apply(d: Dense, x: []const f64, y: []f64) error{Length}!void { + if (x.len != d.k or y.len != d.n) return error.Length; + for (0..d.n) |j| { + var sum: f64 = 0; + for (0..d.k) |i| sum += x[i] * d.w[i * d.n + j]; + y[j] = sum; + } + } +}; + +/// down(silu(gate x) * up x). `gate_buf`, `up_buf`, `glu_buf` are `gate.n` long; `out` is `down.n` long. +pub fn expert(gate: Dense, up: Dense, down: Dense, x: []const f64, gate_buf: []f64, up_buf: []f64, glu_buf: []f64, out: []f64) error{Length}!void { + if (gate.n != up.n or down.k != gate.n or down.n != gate.k) return error.Length; + if (gate_buf.len != gate.n or up_buf.len != gate.n or glu_buf.len != gate.n) return error.Length; + try gate.apply(x, gate_buf); + try up.apply(x, up_buf); + try cpu.swiglu(gate_buf, up_buf, glu_buf, std.math.inf(f64)); + try down.apply(glu_buf, out); +} + +/// `logits` is `experts + 1` values, the last being the shared-gate dot. +pub fn routeFromLogits(logits: []const f64, topk: usize, picks: []u16, weights: []f64, shared_weight: *f64) error{Length}!void { + if (logits.len < 2 or topk == 0 or topk > logits.len - 1 or picks.len != topk or weights.len != topk) return error.Length; + const experts = logits.len - 1; + var top = logits[0]; + for (logits[0..experts]) |v| top = @max(top, v); + var total: f64 = 0; + for (logits[0..experts]) |v| total += @exp(v - top); + var taken: [1024]bool = @splat(false); + if (experts > taken.len) return error.Length; + var kept: f64 = 0; + for (0..topk) |k| { + var best: usize = 0; + var best_v = -std.math.inf(f64); + for (logits[0..experts], 0..) |v, i| { + if (!taken[i] and v > best_v) { + best_v = v; + best = i; + } + } + taken[best] = true; + picks[k] = @intCast(best); + weights[k] = @exp(best_v - top) / total; + kept += weights[k]; + } + for (weights) |*w| w.* /= kept; + shared_weight.* = 1.0 / (1.0 + @exp(-logits[experts])); +} + +/// Row-major `[rows][x.len]` f16 gate dotted with `x`, then the one-row shared gate. +pub fn logitsF16(gate: []const f16, rows: usize, shared_gate: []const f16, x: []const f64, logits: []f64) error{Length}!void { + if (gate.len != rows * x.len or shared_gate.len != x.len or logits.len != rows + 1) return error.Length; + for (0..rows) |r| { + var sum: f64 = 0; + for (x, 0..) |v, i| sum += v * @as(f64, gate[r * x.len + i]); + logits[r] = sum; + } + var s: f64 = 0; + for (x, 0..) |v, i| s += v * @as(f64, shared_gate[i]); + logits[rows] = s; +} + +test "full softmax then renormalise is the softmax over the selected logits" { + const route = @import("exl3_route.zig"); + var rng = std.Random.DefaultPrng.init(11); + const r = rng.random(); + for (0..20) |_| { + var logits: [65]f64 = undefined; + var logits32: [65]f32 = undefined; + for (&logits, &logits32) |*a, *b| { + const v: f32 = (r.float(f32) - 0.5) * 8; + a.* = v; + b.* = v; + } + var picks: [6]u16 = undefined; + var w: [6]f64 = undefined; + var sw: f64 = 0; + try routeFromLogits(&logits, 6, &picks, &w, &sw); + var picks32: [7]u16 = undefined; + var w32: [7]f32 = undefined; + var scratch: [64]f32 = undefined; + try route.select(&logits32, 6, &picks32, &w32, &scratch); + try std.testing.expectEqualSlices(u16, &picks, picks32[0..6]); + var sum: f64 = 0; + for (w, w32[0..6]) |a, b| { + sum += a; + // select rounds through bf16: 8 bits of mantissa, relative 2^-8 + try std.testing.expectApproxEqRel(a, @as(f64, b), 1.0 / 128.0); + } + try std.testing.expectApproxEqAbs(@as(f64, 1), sum, 1e-12); + try std.testing.expectApproxEqRel(sw, @as(f64, w32[6]), 1.0 / 128.0); + try std.testing.expectEqual(@as(u16, 64), picks32[6]); + } +} + +test "ties take the lower id and bad lengths store nothing" { + const logits = [_]f64{ 1, 3, 3, 0, 0 }; + var picks = [_]u16{ 9, 9 }; + var w = [_]f64{ 0, 0 }; + var sw: f64 = -1; + try routeFromLogits(&logits, 2, &picks, &w, &sw); + try std.testing.expectEqualSlices(u16, &.{ 1, 2 }, &picks); + try std.testing.expectEqual(@as(f64, 0.5), w[0]); + try std.testing.expectEqual(@as(f64, 0.5), sw); + var one = [_]u16{7}; + try std.testing.expectError(error.Length, routeFromLogits(&logits, 2, &one, &w, &sw)); + try std.testing.expectEqual(@as(u16, 7), one[0]); +} + +test "logitsF16 is a row dot and the shared gate is the last logit" { + const gate = [_]f16{ 1, 0, 0, 2, 0, 0 }; + const shared = [_]f16{ 0, 0, 3 }; + var logits: [3]f64 = undefined; + try logitsF16(&gate, 2, &shared, &.{ 1, 5, 7 }, &logits); + try std.testing.expectEqualSlices(f64, &.{ 1, 2, 21 }, &logits); + try std.testing.expectError(error.Length, logitsF16(&gate, 2, shared[0..2], &.{ 1, 5, 7 }, &logits)); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_norm.zig b/zig/src/families/flashnext/exl3/model/exl3_norm.zig new file mode 100644 index 0000000000..9009fa0c16 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_norm.zig @@ -0,0 +1,90 @@ +//! Centred RMSNorm apply. + +const std = @import("std"); +const scan = @import("../pack/exl3_scan.zig"); + +/// `rms = sqrt(mean(xยฒ) + eps)`, then `out[i] = x[i] / rms * (weight[i] + offset)`. +pub fn apply(x: []const f32, weight: []const f32, means: []const f32, eps: f32, out: []f32) error{ Length, NormAmbiguous }!void { + const offset = try scan.centredOffset(means); + if (!std.math.isFinite(eps) or eps < 0) return error.Length; + if (x.len == 0 or x.len != weight.len or x.len != out.len) return error.Length; + + var sum_sq: f32 = 0; + for (x) |v| sum_sq += v * v; + const rms = @sqrt(sum_sq / @as(f32, @floatFromInt(x.len)) + eps); + for (x, weight, out) |xi, wi, *yi| yi.* = xi / rms * (wi + offset); +} + +test "a unit vector changes by exactly the weight difference" { + // Same means `centredOffset` classifies as gamma-1 (1) and gamma (0). + const near_zero = [_]f32{ -0.02, 0.01, 0.0, 0.04 }; + const near_one = [_]f32{ 1.0, 0.98, 1.02, 1.01 }; + // x = [1], eps = 0: mean(xยฒ)+eps is 1, so the RMS is 1 and out = weight + offset. + const x = [_]f32{1}; + const light = [_]f32{2}; + const heavy = [_]f32{5}; + var zero_light = [_]f32{0}; + var zero_heavy = [_]f32{0}; + var one_light = [_]f32{0}; + var one_heavy = [_]f32{0}; + try apply(&x, &light, &near_zero, 0, &zero_light); + try apply(&x, &heavy, &near_zero, 0, &zero_heavy); + try apply(&x, &light, &near_one, 0, &one_light); + try apply(&x, &heavy, &near_one, 0, &one_heavy); + try std.testing.expectEqual(@as(f32, 3), zero_light[0]); + try std.testing.expectEqual(@as(f32, 6), zero_heavy[0]); + try std.testing.expectEqual(@as(f32, 2), one_light[0]); + try std.testing.expectEqual(@as(f32, 5), one_heavy[0]); + try std.testing.expectEqual(heavy[0] - light[0], zero_heavy[0] - zero_light[0]); + try std.testing.expectEqual(heavy[0] - light[0], one_heavy[0] - one_light[0]); + + // Two ones still have mean 1, so a forgotten division by length cannot pass. + const pair_x = [_]f32{ 1, 1 }; + const pair_w = [_]f32{ 2, 7 }; + var pair_out = [_]f32{ 0, 0 }; + try apply(&pair_x, &pair_w, &near_zero, 0, &pair_out); + try std.testing.expectEqual(@as(f32, 3), pair_out[0]); + try std.testing.expectEqual(@as(f32, 8), pair_out[1]); +} + +test "an ambiguous norm leaves out untouched" { + var out = [_]f32{ 9, 8 }; + const many: [65]f32 = @splat(0); + try std.testing.expectError(error.NormAmbiguous, apply(&.{ 1, 1 }, &.{ 1, 1 }, &many, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); + + // Length also disagrees. The locator runs first, so this is not Length. + try std.testing.expectError(error.NormAmbiguous, apply(&.{1}, &.{}, &.{ 0.4, 0.9 }, -1, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); +} + +test "a length mismatch leaves out untouched" { + const near_zero = [_]f32{ -0.02, 0.01, 0.0, 0.04 }; + const near_one = [_]f32{ 1.0, 0.98, 1.02, 1.01 }; + var out = [_]f32{ 9, 8 }; + try std.testing.expectError(error.Length, apply(&.{}, &.{}, &near_zero, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); + + try std.testing.expectError(error.Length, apply(&.{ 1, 1 }, &.{1}, &near_zero, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); + + try std.testing.expectError(error.Length, apply(&.{1}, &.{1}, &near_one, 0, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectEqual(@as(f32, 8), out[1]); +} + +test "a negative eps leaves out untouched" { + const near_zero = [_]f32{ -0.02, 0.01, 0.0, 0.04 }; + const near_one = [_]f32{ 1.0, 0.98, 1.02, 1.01 }; + var out = [_]f32{9}; + try std.testing.expectError(error.Length, apply(&.{1}, &.{2}, &near_zero, -1, &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectError(error.Length, apply(&.{1}, &.{2}, &near_one, std.math.nan(f32), &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); + try std.testing.expectError(error.Length, apply(&.{1}, &.{2}, &near_zero, std.math.inf(f32), &out)); + try std.testing.expectEqual(@as(f32, 9), out[0]); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_round.zig b/zig/src/families/flashnext/exl3/model/exl3_round.zig new file mode 100644 index 0000000000..9d75bb336d --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_round.zig @@ -0,0 +1,599 @@ +//! One round's miss plan for a per-layer EXL3 expert pool (the host half of `cuda_exl3_pool.Pool.plan`). + +const std = @import("std"); +const offload = @import("../offload/exl3_offload.zig"); + +const shape = offload.shape; +const Vram = offload.Vram; +const empty_id = offload.Cache.empty_id; +pub const shared_id = offload.shared_id; + +/// A copy of routed expert `id` into pool line `line`; `evicted` is the expert the line held (empty_id: none). +pub const Copy = struct { id: u16, line: u32, evicted: u16 }; + +/// Ranges into `Round.ids` (the experts the pass computes) and `Round.copies` (what it stages first). +pub const Pass = struct { id_lo: u32, id_hi: u32, copy_lo: u32, copy_hi: u32 }; + +pub const Options = struct { + swap_every: u16 = 4, + swap_cap: u16 = 96, + observe: bool = true, + /// Bytes a swap round may prefetch (the copy budget); 0 disables prefetch. + prefetch_bytes: u64 = 0, + /// Device bytes of one line (prefetch budget unit). + line_bytes: u64 = 1, + /// Double buffering: experts (of this layer) whose lines a single-pass round should not evict when it has room + protect: []const u16 = &.{}, +}; + +/// Caller-owned storage for one round's plan, sized once by `init` for the pool's line count and pick capacity. +pub const Round = struct { + ids: []u16, + /// The line each id of `ids` is read from in its pass (a later pass of the same round may reuse that line). + lines: []u32, + copies: []Copy, + passes: []Pass, + routed: []u16, // the round's routed picks as u16 (ordering scratch) + misses: []u16, + snap: []u16, // line -> id before a pass (empty_id for empty lines) + keep: []u16, // pass ids plus protected ids (stageUsedAt's keep list) + /// The round spared `protect`'s lines (double buffering held). + protected: bool = false, + n_ids: usize = 0, + n_copies: usize = 0, + n_passes: usize = 0, + n_prefetch: usize = 0, + /// Routed slots (pick in 0..511) whose expert was held / missing when the round started. + hit_slots: u32 = 0, + miss_slots: u32 = 0, + shared_slots: u32 = 0, + /// Distinct routed experts, of which `held` were already in a line (`pending` of those still in flight). + unique: u32 = 0, + held: u32 = 0, + pending: u32 = 0, + /// Copies that replaced a held expert (mandatory and prefetch). + evictions: u32 = 0, + + pub fn init(gpa: std.mem.Allocator, lines: usize, max_picks: usize) !Round { + const uniq = @min(max_picks, shape.experts); + var r: Round = undefined; + r.ids = try gpa.alloc(u16, uniq); + errdefer gpa.free(r.ids); + r.lines = try gpa.alloc(u32, uniq); + errdefer gpa.free(r.lines); + r.copies = try gpa.alloc(Copy, uniq + shape.experts); + errdefer gpa.free(r.copies); + r.passes = try gpa.alloc(Pass, uniq + 1); + errdefer gpa.free(r.passes); + r.routed = try gpa.alloc(u16, max_picks); + errdefer gpa.free(r.routed); + r.misses = try gpa.alloc(u16, uniq); + errdefer gpa.free(r.misses); + r.snap = try gpa.alloc(u16, lines); + errdefer gpa.free(r.snap); + r.keep = try gpa.alloc(u16, uniq + lines); + r.reset(); + return r; + } + + pub fn deinit(r: *Round, gpa: std.mem.Allocator) void { + gpa.free(r.ids); + gpa.free(r.lines); + gpa.free(r.copies); + gpa.free(r.passes); + gpa.free(r.routed); + gpa.free(r.misses); + gpa.free(r.snap); + gpa.free(r.keep); + r.* = undefined; + } + + fn reset(r: *Round) void { + r.n_ids = 0; + r.n_copies = 0; + r.n_passes = 0; + r.n_prefetch = 0; + r.hit_slots = 0; + r.miss_slots = 0; + r.shared_slots = 0; + r.unique = 0; + r.held = 0; + r.pending = 0; + r.evictions = 0; + r.protected = false; + } + + pub fn passIds(r: *const Round, j: usize) []const u16 { + const p = r.passes[j]; + return r.ids[p.id_lo..p.id_hi]; + } + + pub fn passLines(r: *const Round, j: usize) []const u32 { + const p = r.passes[j]; + return r.lines[p.id_lo..p.id_hi]; + } + + pub fn passCopies(r: *const Round, j: usize) []const Copy { + const p = r.passes[j]; + return r.copies[p.copy_lo..p.copy_hi]; + } + + /// Copies no pass waits for (swap-round prefetch); their lines stay in flight until reaped. + pub fn prefetch(r: *const Round) []const Copy { + return r.copies[r.n_copies - r.n_prefetch .. r.n_copies]; + } + + /// Mandatory copies of all passes. + pub fn misses_(r: *const Round) usize { + return r.n_copies - r.n_prefetch; + } + + /// More than one pass: each pass reads its own masked picks (picks of other passes' experts set to -1). + pub fn masked(r: *const Round) bool { + return r.n_passes > 1; + } + + /// True when the round needs no copy at all (the experts need no upload). + pub fn allHits(r: *const Round) bool { + return r.n_copies == 0; + } +}; + +fn contains32(xs: []const u32, x: u32) bool { + return std.mem.indexOfScalar(u32, xs, x) != null; +} + +fn contains(xs: []const u16, id: u16) bool { + return std.mem.indexOfScalar(u16, xs, id) != null; +} + +const Order = struct { + counts: []const u32, + ids: []const u16, + fn appearances(ids: []const u16, id: u16) u32 { + var n: u32 = 0; + for (ids) |x| n += @intFromBool(x == id); + return n; + } + /// Lifetime count descending, then this round's appearances, then lower id (the policy's `usedFirst`). + fn hotter(o: Order, a: u16, b: u16) bool { + if (o.counts[a] != o.counts[b]) return o.counts[a] > o.counts[b]; + const aa = appearances(o.ids, a); + const ab = appearances(o.ids, b); + if (aa != ab) return aa > ab; + return a < b; + } +}; + +/// Lines whose copy is in flight and whose expert is not in `use`: they can be neither read nor evicted this pass. +fn blocked(v: Vram, use: []const u16) usize { + var n: usize = 0; + for (v.line, v.slots) |state, id| { + if (state == .inflight and !contains(use, id)) n += 1; + } + return n; +} + +fn snapshot(v: Vram, out: []u16) void { + for (v.slots, v.line, out) |id, state, *o| o.* = if (state == .empty) empty_id else id; +} + +/// Plans one round of `layer` (see the file header). +pub fn plan( + v: *Vram, + layer: u8, + picks: []const i32, + counts: *[shape.layers][shape.experts]u32, + now: u32, + round: u32, + opt: Options, + r: *Round, +) error{ Busy, Length }!void { + r.reset(); + if (picks.len > r.routed.len or layer >= shape.layers or v.slots.len != r.snap.len) return error.Length; + // the routed picks and the distinct experts, in first-appearance order + var n_routed: usize = 0; + var n_unique: usize = 0; + for (picks) |p| { + if (p < 0) continue; + if (p > shared_id) return error.Length; + const id: u16 = @intCast(p); + if (id == shared_id) { + r.shared_slots += 1; + continue; + } + r.routed[n_routed] = id; + n_routed += 1; + if (!contains(r.ids[0..n_unique], id)) { + r.ids[n_unique] = id; + n_unique += 1; + } + } + const routed = r.routed[0..n_routed]; + const uniq = r.ids[0..n_unique]; + // capacity check before anything is written: pass 0 holds the held experts, so at least one more line. + var held_n: usize = 0; + for (uniq) |id| held_n += @intFromBool(offload.lineOf(v.*, layer, id) != null); + if (n_unique > held_n and blocked(v.*, &.{}) >= v.slots.len) return error.Busy; + + if (opt.observe) offload.observe(&counts[layer], routed) catch return error.Length; + const row: []const u32 = &counts[layer]; + + // held first (pass 0), misses hottest first; ids[] is rewritten in pass order + var n_held: usize = 0; + var n_miss: usize = 0; + for (uniq) |id| { + if (offload.lineOf(v.*, layer, id)) |line| { + r.held += 1; + if (v.line[line] == .inflight) r.pending += 1 else _ = offload.touchAt(v, layer, id, now); + } else { + r.misses[n_miss] = id; + n_miss += 1; + } + } + for (routed) |id| { + if (offload.lineOf(v.*, layer, id) != null) r.hit_slots += 1 else r.miss_slots += 1; + } + r.unique = @intCast(n_unique); + // pack held ids at the front of ids (stable), then the sorted misses + for (uniq) |id| { + if (offload.lineOf(v.*, layer, id) != null) { + r.ids[n_held] = id; + n_held += 1; + } + } + const misses = r.misses[0..n_miss]; + std.mem.sort(u16, misses, Order{ .counts = row, .ids = routed }, Order.hotter); + @memcpy(r.ids[n_held..][0..n_miss], misses); + + // passes + var at: usize = 0; // next id of ids[] to place + var mi: usize = 0; // next miss + while (at < n_unique) { + const first = r.n_passes == 0; + const lo = at; + const take_held: usize = if (first) n_held else 0; + at += take_held; + // lines usable this pass: all but in-flight lines of experts outside this pass + const usable = v.slots.len - blocked(v.*, r.ids[lo..at]); + const room = usable - take_held; + const n_new = @min(room, n_miss - mi); + std.debug.assert(n_new > 0 or take_held > 0); // the capacity check above + const hi = at + n_new; + // double buffering: also spare the protected experts' resident lines when the single pass still fits + var keep: []const u16 = r.ids[lo..hi]; + if (first and n_new == n_miss and opt.protect.len > 0) { + var spare: usize = 0; + var k: usize = hi - lo; + @memcpy(r.keep[0..k], r.ids[lo..hi]); + for (opt.protect) |id| { + if (contains(r.keep[0..k], id) or k == r.keep.len) continue; + const line = offload.lineOf(v.*, layer, id) orelse continue; + if (v.line[line] != .resident) continue; + r.keep[k] = id; + k += 1; + spare += 1; + } + if (room >= n_new + spare) { + keep = r.keep[0..k]; + r.protected = true; + } + } + snapshot(v.*, r.snap); + const copy_lo = r.n_copies; + for (r.ids[at..hi]) |id| { + const line = offload.stageUsedAt(v, layer, id, keep, counts) catch return error.Busy; + const old = r.snap[line]; + if (old != empty_id) r.evictions += 1; + r.copies[r.n_copies] = .{ .id = id, .line = @intCast(line), .evicted = old }; + r.n_copies += 1; + } + // the device orders the next pass's copies after this pass's compute: these lines are resident from here on + for (r.copies[copy_lo..r.n_copies]) |c| offload.reapAt(v, layer, c.id, now) catch unreachable; + for (r.ids[lo..hi], r.lines[lo..hi]) |id, *ln| ln.* = @intCast(offload.lineOf(v.*, layer, id).?); + r.passes[r.n_passes] = .{ .id_lo = @intCast(lo), .id_hi = @intCast(hi), .copy_lo = @intCast(copy_lo), .copy_hi = @intCast(r.n_copies) }; + r.n_passes += 1; + at = hi; + mi += n_new; + } + r.n_ids = n_unique; + if (r.n_passes == 0) { // shared and -1 picks only: one pass with no routed expert (the shared expert still runs) + r.passes[0] = .{ .id_lo = 0, .id_hi = 0, .copy_lo = 0, .copy_hi = 0 }; + r.n_passes = 1; + } + + // swap-round prefetch: hottest non-held experts into empty or strictly colder lines this round does not use + if (opt.swap_every == 0 or round == 0 or round % opt.swap_every != 0 or opt.prefetch_bytes == 0 or opt.line_bytes == 0) return; + const budget = @min(@as(u64, opt.swap_cap), opt.prefetch_bytes / opt.line_bytes); + var staged: u64 = 0; + snapshot(v.*, r.snap); + while (staged < budget) { + // hottest id not held on this layer + var best: ?u16 = null; + var id: u16 = 0; + while (id < shape.experts) : (id += 1) { + if (row[id] == 0 or offload.lineOf(v.*, layer, id) != null) continue; + if (best == null or row[id] > row[best.?]) best = id; + } + const cand = best orelse break; + // the line it would take: an empty one, else the coldest resident outside this round + var victim: ?usize = null; + for (v.line, 0..) |state, i| { + if (state == .empty) { + victim = i; + break; + } + } + if (victim == null) { + for (v.line, v.slots, 0..) |state, held_id, i| { + if (state != .resident or contains(uniq_view(r), held_id)) continue; + if (victim == null or colder(v.*, layer, counts, i, victim.?)) victim = i; + } + const vi = victim orelse break; + if (offload.countAt(counts, layer, v.slots[vi]) >= row[cand]) break; // nothing colder than the candidate + } + const line = offload.stageUsedAt(v, layer, cand, uniq_view(r), counts) catch break; + std.debug.assert(line == victim.?); + const old = r.snap[line]; + if (old != empty_id) r.evictions += 1; + r.snap[line] = cand; + r.copies[r.n_copies] = .{ .id = cand, .line = @intCast(line), .evicted = old }; + r.n_copies += 1; + r.n_prefetch += 1; + staged += 1; + } +} + +fn uniq_view(r: *const Round) []const u16 { + return r.ids[0..r.n_ids]; +} + +/// The policy's victim order (`vramCold`): lower lifetime count, then older stamp, then the higher line index. +fn colder(v: Vram, layer: u8, counts: *const [shape.layers][shape.experts]u32, i: usize, b: usize) bool { + const ci = offload.countAt(counts, layer, v.slots[i]); + const cb = offload.countAt(counts, layer, v.slots[b]); + return ci < cb or (ci == cb and v.stamp[i] < v.stamp[b]) or (ci == cb and v.stamp[i] == v.stamp[b] and i > b); +} + +/// Smallest pool that never needs a second pass for `rows` decode rows: every row's routed picks must fit at once. +pub fn minLines(top_k: usize) usize { + return top_k; +} + +/// Lines a layer gets from `free` device bytes after `reserve` and `other` (KV, scratch: the family's budget) +pub fn linesFor(free: u64, reserve: u64, other: u64, line_bytes: u64, layers: usize, min_lines: usize) error{PoolTooSmall}!usize { + if (line_bytes == 0 or layers == 0) return error.PoolTooSmall; + const spend = std.math.sub(u64, free, reserve +| other) catch return error.PoolTooSmall; + const per = spend / (line_bytes * layers); + if (per < min_lines) return error.PoolTooSmall; + return @intCast(per); +} + +test { + _ = offload; +} + +// ------------------------------------------------------------------------------------------------ tests + +const testing = std.testing; +const Line = offload.Line; + +const Fix = struct { + slots: [16]u16 = undefined, + line: [16]Line = undefined, + stamp: [16]u32 = undefined, + lay: [16]u8 = undefined, + v: Vram = undefined, + counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)), + r: Round = undefined, + + fn init(f: *Fix, lines: usize, max_picks: usize) !void { + f.v = Vram.init(f.slots[0..lines], f.line[0..lines], f.stamp[0..lines], f.lay[0..lines]); + f.counts = @splat(@splat(0)); + f.r = try Round.init(testing.allocator, lines, max_picks); + } + fn deinit(f: *Fix) void { + f.r.deinit(testing.allocator); + } + fn go(f: *Fix, picks: []const i32, now: u32, round: u32, opt: Options) !void { + return plan(&f.v, 3, picks, &f.counts, now, round, opt, &f.r); + } + fn resident(f: *const Fix, id: u16) bool { + return offload.residentAt(f.v, 3, id); + } +}; + +test "misses fill empty lines in one pass, then the same picks are all hits" { + var f: Fix = .{}; + try f.init(10, 32); + defer f.deinit(); + const picks = [_]i32{ 5, 9, 5, 512, -1, 7 }; + try f.go(&picks, 1, 1, .{}); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); + try testing.expectEqual(@as(usize, 3), f.r.misses_()); + try testing.expectEqual(@as(u32, 4), f.r.miss_slots); + try testing.expectEqual(@as(u32, 1), f.r.shared_slots); + try testing.expectEqual(@as(u32, 0), f.r.evictions); + for (f.r.passCopies(0)) |c| try testing.expectEqual(empty_id, c.evicted); + // 5 appears twice: hottest miss, copied first + try testing.expectEqual(@as(u16, 5), f.r.passCopies(0)[0].id); + try testing.expect(!f.r.masked()); + try testing.expect(f.resident(5) and f.resident(9) and f.resident(7)); + try testing.expectEqual(@as(u32, 2), f.counts[3][5]); + + try f.go(&picks, 2, 2, .{}); + try testing.expect(f.r.allHits()); + try testing.expectEqual(@as(u32, 4), f.r.hit_slots); + try testing.expectEqual(@as(u32, 3), f.r.held); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); + try testing.expectEqual(@as(usize, 3), f.r.passIds(0).len); + // a hit touched its line + try testing.expectEqual(@as(u32, 2), f.stamp[offload.lineOf(f.v, 3, 9).?]); +} + +test "synthetic prefill picks do not heat expert history" { + var f: Fix = .{}; + try f.init(10, 16); + defer f.deinit(); + try f.go(&[_]i32{ 7, 8, 9 }, 1, 1, .{ .observe = false }); + try testing.expectEqual(@as(usize, 3), f.r.misses_()); + try testing.expectEqual(@as(u32, 0), f.counts[3][7]); + try testing.expectEqual(@as(u32, 0), f.counts[3][8]); + try testing.expectEqual(@as(u32, 0), f.counts[3][9]); +} + +test "a miss evicts the coldest line outside the round, never a line the round reads" { + var f: Fix = .{}; + try f.init(3, 16); + defer f.deinit(); + try f.go(&[_]i32{ 1, 2, 3 }, 1, 1, .{}); + // heat 1 and 3, leave 2 cold + try f.go(&[_]i32{ 1, 3 }, 2, 2, .{}); + try f.go(&[_]i32{ 1, 3 }, 3, 3, .{}); + try f.go(&[_]i32{ 1, 4 }, 4, 4, .{}); + try testing.expectEqual(@as(usize, 1), f.r.misses_()); + const c = f.r.passCopies(0)[0]; + try testing.expectEqual(@as(u16, 4), c.id); + try testing.expectEqual(@as(u16, 2), c.evicted); + try testing.expect(f.resident(1) and f.resident(3) and f.resident(4) and !f.resident(2)); + // 3 is hot but not in this round, 1 is in it; the cold one now is 4 (count 1) vs 3 (count 3) -> 4 goes + try f.go(&[_]i32{ 1, 3, 6 }, 5, 5, .{}); + try testing.expectEqual(@as(u16, 4), f.r.passCopies(0)[0].evicted); +} + +test "more distinct experts than lines split into passes that each fit, every expert computed once" { + var f: Fix = .{}; + try f.init(4, 64); + defer f.deinit(); + try f.go(&[_]i32{ 20, 21 }, 1, 1, .{}); + var picks: [30]i32 = undefined; + for (&picks, 0..) |*p, i| p.* = @intCast(10 + (i % 11)); // 10..20, 11 distinct; 20 held + picks[29] = 21; // held too + try f.go(&picks, 2, 2, .{}); + try testing.expect(f.r.masked()); + try testing.expectEqual(@as(u32, 12), f.r.unique); + try testing.expectEqual(@as(u32, 2), f.r.held); + try testing.expectEqual(@as(usize, 10), f.r.misses_()); + var seen: [shape.experts]u8 = @splat(0); + for (0..f.r.n_passes) |j| { + const ids = f.r.passIds(j); + try testing.expect(ids.len >= 1 and ids.len <= 4); + // the lines a pass copies into are distinct + const cs = f.r.passCopies(j); + for (cs, 0..) |a, i| for (cs[i + 1 ..]) |b| try testing.expect(a.line != b.line); + for (ids) |id| seen[id] += 1; + } + for (10..22) |id| try testing.expectEqual(@as(u8, 1), seen[id]); + // each pass reads distinct lines; a copy of a later pass may land in a line an earlier pass read + var reused = false; + for (0..f.r.n_passes) |j| { + const ls = f.r.passLines(j); + for (ls, 0..) |a, i| for (ls[i + 1 ..]) |b| try testing.expect(a != b); + for (f.r.passCopies(j)) |c| { + try testing.expect(contains32(ls, c.line)); + if (j > 0) for (0..j) |k| { + if (contains32(f.r.passLines(k), c.line)) reused = true; + }; + } + } + try testing.expect(reused); + // pass 0 holds the two held experts first + try testing.expect(contains(f.r.passIds(0), 20) and contains(f.r.passIds(0), 21)); + try testing.expectEqual(@as(usize, 3), f.r.n_passes); // 2 held + 2 misses, then 4, then 4 + // the last pass's experts are resident after the round + for (f.r.passIds(f.r.n_passes - 1)) |id| try testing.expect(f.resident(id)); +} + +test "an in-flight prefetch line is not a victim and counts as held when the round asks for it" { + var f: Fix = .{}; + try f.init(3, 16); + defer f.deinit(); + try f.go(&[_]i32{ 1, 2, 3 }, 1, 1, .{}); + // a prefetch of 9 into line of 2 (made by hand: stage it, leave it in flight) + const line = offload.lineOf(f.v, 3, 2).?; + f.slots[line] = 9; + f.line[line] = .inflight; + try f.go(&[_]i32{ 1, 4 }, 2, 2, .{}); + // 4 cannot take 9's in-flight line; it takes 3's (not in the round) + try testing.expectEqual(@as(u16, 3), f.r.passCopies(0)[0].evicted); + try testing.expectEqual(Line.inflight, f.line[line]); + try f.go(&[_]i32{ 9, 1 }, 3, 3, .{}); + try testing.expect(f.r.allHits()); + try testing.expectEqual(@as(u32, 1), f.r.pending); + // every line in flight for another expert: busy, and nothing written + for (f.line[0..3]) |*l| l.* = .inflight; + const before = f.counts[3][7]; + try testing.expectError(error.Busy, f.go(&[_]i32{7}, 4, 4, .{})); + try testing.expectEqual(before, f.counts[3][7]); +} + +test "double buffering spares the previous round's lines when the round fits without them" { + var f: Fix = .{}; + try f.init(6, 16); + defer f.deinit(); + // previous round 1,2,3 (hot), older residents 4,5 (cold), one empty line + try f.go(&[_]i32{ 4, 5 }, 1, 1, .{}); + for (0..3) |i| try f.go(&[_]i32{ 1, 2, 3 }, @intCast(2 + i), @intCast(2 + i), .{}); + // two misses: with 1, 2, 3 spared they go to the empty line and the colder of 4, 5 + const prev = [_]u16{ 1, 2, 3 }; + try f.go(&[_]i32{ 7, 8 }, 5, 5, .{ .protect = &prev }); + try testing.expect(f.r.protected); + for (f.r.passCopies(0)) |c| try testing.expect(!contains(&prev, c.evicted)); + try testing.expect(f.resident(1) and f.resident(2) and f.resident(3)); + // a round that only fits by evicting a protected line ignores the protection + try f.go(&[_]i32{ 9, 10, 11, 12 }, 6, 6, .{ .protect = &[_]u16{ 7, 8, 1, 2, 3 } }); + try testing.expect(!f.r.protected); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); +} + +test "the shared expert and -1 are never staged; a pick above 512 is refused" { + var f: Fix = .{}; + try f.init(10, 16); + defer f.deinit(); + try f.go(&[_]i32{ 512, -1, 512 }, 1, 1, .{}); + try testing.expectEqual(@as(usize, 0), f.r.n_copies); + try testing.expectEqual(@as(u32, 2), f.r.shared_slots); + try testing.expectEqual(@as(usize, 1), f.r.n_passes); + try testing.expectEqual(@as(usize, 0), f.r.passIds(0).len); + for (f.slots[0..10]) |id| try testing.expect(id != shared_id); + try testing.expectError(error.Length, f.go(&[_]i32{513}, 2, 2, .{})); +} + +test "swap rounds prefetch the hottest non-held experts within the copy budget, only into colder lines" { + var f: Fix = .{}; + try f.init(4, 16); + defer f.deinit(); + const opt: Options = .{ .swap_every = 4, .prefetch_bytes = 200, .line_bytes = 100 }; + // lifetime counts: 30 hot (5), 31 warm (3), never in a line yet + f.counts[3][30] = 5; + f.counts[3][31] = 3; + f.counts[3][32] = 1; + try f.go(&[_]i32{ 1, 2 }, 1, 1, opt); + try testing.expectEqual(@as(usize, 0), f.r.n_prefetch); // round 1: not a swap round + try f.go(&[_]i32{ 1, 2 }, 4, 4, opt); + // two empty lines, budget 2: 30 and 31 in flight + try testing.expectEqual(@as(usize, 2), f.r.n_prefetch); + try testing.expectEqual(@as(u16, 30), f.r.prefetch()[0].id); + try testing.expectEqual(@as(u16, 31), f.r.prefetch()[1].id); + try testing.expectEqual(Line.inflight, f.line[f.r.prefetch()[0].line]); + try testing.expect(!f.resident(30)); + try offload.reapAt(&f.v, 3, 30, 5); + try offload.reapAt(&f.v, 3, 31, 5); + // round 8: 32 (count 1) is not hotter than any resident outside the round (30: 5, 31: 3) -> nothing evicted + try f.go(&[_]i32{ 1, 2 }, 8, 8, opt); + try testing.expectEqual(@as(usize, 0), f.r.n_prefetch); + // heat 32 above 31: it replaces 31, never 1 or 2 (in this round) + f.counts[3][32] = 9; + try f.go(&[_]i32{ 1, 2 }, 12, 12, opt); + try testing.expectEqual(@as(usize, 1), f.r.n_prefetch); + try testing.expectEqual(@as(u16, 31), f.r.prefetch()[0].evicted); +} + +test "line budget: free memory after reserve and the family's other bytes, at least min lines a layer" { + const line: u64 = 1_843_200; + // 12 GiB free, 1 GiB reserve, 5.5 GiB for KV and scratch: 5.5 GiB / (48 x 1,843,200 B) = 66 lines a layer + try testing.expectEqual(@as(usize, 66), try linesFor(12 << 30, 1 << 30, (11 << 30) / 2, line, 48, 10)); + try testing.expectEqual(@as(usize, 10), try linesFor(10 * 48 * line + 100, 100, 0, line, 48, 10)); + try testing.expectError(error.PoolTooSmall, linesFor(10 * 48 * line - 1, 0, 0, line, 48, 10)); + try testing.expectError(error.PoolTooSmall, linesFor(1 << 30, 2 << 30, 0, line, 48, 10)); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_route.zig b/zig/src/families/flashnext/exl3/model/exl3_route.zig new file mode 100644 index 0000000000..c6e7ab22df --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_route.zig @@ -0,0 +1,94 @@ +//! One row of the Flash Next router. v0.6.6 `src/tensorfold/cuda/moe.py` `_topk_rows`: fp32 logits. + +const std = @import("std"); + +/// `logits` is `experts + 1` fp32 values. +pub fn select( + logits: []const f32, + topk: usize, + picks: []u16, + weights: []f32, + scratch: []f32, +) error{Length}!void { + if (logits.len < 2) return error.Length; + const experts = logits.len - 1; + if (experts > std.math.maxInt(u16)) return error.Length; + if (topk == 0 or topk > experts) return error.Length; + if (picks.len != topk + 1 or weights.len != topk + 1 or scratch.len != experts) return error.Length; + + @memcpy(scratch, logits[0..experts]); + var top = scratch[0]; + for (scratch[1..]) |v| top = @max(top, v); + + var total: f32 = 0; + for (0..topk) |k| { + var best_i: usize = 0; + var best = scratch[0]; + for (scratch[1..], 1..) |v, i| { + if (v > best) { + best = v; + best_i = i; + } + } + const ex = @exp(best - top); + picks[k] = @intCast(best_i); + weights[k] = ex; + total += ex; + scratch[best_i] = -std.math.inf(f32); + } + for (weights[0..topk]) |*w| w.* = bf16(w.* / total); + + const shared: u16 = @intCast(experts); + const gate = bf16(logits[experts]); + picks[topk] = shared; + weights[topk] = bf16(1.0 / (1.0 + @exp(-gate))); +} + +/// fp32 through bf16 and back. NaN stays a NaN. Same rounding as the contract helper. +fn bf16(v: f32) f32 { + const u: u32 = @bitCast(v); + const h: u16 = if (std.math.isNan(v)) + @truncate((u >> 16) | 0x40) + else + @truncate((u +% 0x7fff +% ((u >> 16) & 1)) >> 16); + return @bitCast(@as(u32, h) << 16); +} + +test "ties take the lower id, and the shared expert is the last slot" { + const logits = [_]f32{ 1, 3, 3, 0, 0 }; + var picks = [_]u16{ 9, 9, 9 }; + var weights = [_]f32{ 9, 9, 9 }; + var scratch: [4]f32 = undefined; + try select(&logits, 2, &picks, &weights, &scratch); + + try std.testing.expectEqual(@as(u16, 1), picks[0]); + try std.testing.expectEqual(@as(u16, 2), picks[1]); + try std.testing.expectEqual(@as(u16, 4), picks[2]); + try std.testing.expectEqual(@as(f32, 0.5), weights[0]); + try std.testing.expectEqual(@as(f32, 0.5), weights[1]); + try std.testing.expectEqual(@as(f32, 0.5), weights[2]); +} + +test "one winner keeps the whole routed weight" { + const logits = [_]f32{ 5, 5, 1, 2 }; + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 9, 9 }; + var scratch: [3]f32 = undefined; + try select(&logits, 1, &picks, &weights, &scratch); + try std.testing.expectEqual(@as(u16, 0), picks[0]); + try std.testing.expectEqual(@as(u16, 3), picks[1]); + try std.testing.expectEqual(@as(f32, 1), weights[0]); + const gate = bf16(2); + const want = bf16(1.0 / (1.0 + @exp(-gate))); + try std.testing.expectEqual(want, weights[1]); +} + +test "a short buffer is refused before picks change" { + const logits = [_]f32{ 1, 2, 0 }; + var picks = [_]u16{ 7, 7 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [2]f32 = undefined; + try std.testing.expectError(error.Length, select(&logits, 2, picks[0..1], &weights, &scratch)); + try std.testing.expectEqual(@as(u16, 7), picks[0]); + try std.testing.expectEqual(@as(f32, 3), weights[0]); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_router.zig b/zig/src/families/flashnext/exl3/model/exl3_router.zig new file mode 100644 index 0000000000..48ee9f791f --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_router.zig @@ -0,0 +1,252 @@ +//! Rank 0 of the Flash Next router. + +const std = @import("std"); +const route = @import("exl3_route.zig"); + +pub const gate_rows: usize = 512; +pub const gate_cols: usize = 2560; +pub const gate_bytes: usize = gate_rows * gate_cols * @sizeOf(u16); + +pub const Error = error{ + NotRank0, + EmptyGate, + GateBytes, + MissingSharedGate, + Length, +}; + +comptime { + std.debug.assert(gate_bytes == 2_621_440); + std.debug.assert((gate_rows + 1) * gate_cols * @sizeOf(u16) == 2_626_560); +} + +/// Storage type of the two replicated gate tensors. +pub const Dtype = enum { bf16, f16 }; + +fn dot(dtype: Dtype, weight: []const u8, x: []const f32) f32 { + var sum: f32 = 0; + for (x, 0..) |v, i| { + const h = std.mem.readInt(u16, weight[i * 2 ..][0..2], .little); + const w: f32 = switch (dtype) { + .bf16 => @bitCast(@as(u32, h) << 16), + .f16 => @floatCast(@as(f16, @bitCast(h))), + }; + sum += w * v; + } + return sum; +} + +/// `gate` is row-major bf16, exactly `gate_bytes`, times `hidden` (`gate_cols`). +pub fn launch( + rank: usize, + gate: []const u8, + hidden: []const f32, + shared_gate: []const u8, + shared_x: []const f32, + topk: usize, + logits: []f32, + picks: []u16, + weights: []f32, + scratch: []f32, +) Error!void { + return launchDtype(.bf16, rank, gate, hidden, shared_gate, shared_x, topk, logits, picks, weights, scratch); +} + +/// `launch` with the gate tensors' storage type named. +pub fn launchDtype( + dtype: Dtype, + rank: usize, + gate: []const u8, + hidden: []const f32, + shared_gate: []const u8, + shared_x: []const f32, + topk: usize, + logits: []f32, + picks: []u16, + weights: []f32, + scratch: []f32, +) Error!void { + if (rank != 0) return error.NotRank0; + if (gate.len == 0) return error.EmptyGate; + if (gate.len != gate_bytes) return error.GateBytes; + if (shared_gate.len == 0) return error.MissingSharedGate; + if (shared_gate.len % 2 != 0) return error.Length; + const shared_n = shared_gate.len / 2; + if (shared_x.len != shared_n) return error.Length; + if (hidden.len != gate_cols) return error.Length; + if (logits.len != gate_rows + 1) return error.Length; + if (topk == 0 or topk > gate_rows) return error.Length; + if (picks.len != topk + 1 or weights.len != topk + 1 or scratch.len != gate_rows) return error.Length; + + const row_bytes = gate_cols * @sizeOf(u16); + for (0..gate_rows) |r| { + const row = gate[r * row_bytes ..][0..row_bytes]; + logits[r] = dot(dtype, row, hidden); + } + logits[gate_rows] = dot(dtype, shared_gate, shared_x); + try route.select(logits, topk, picks, weights, scratch); +} + +const one_bf16: u16 = 0x3f80; +const winner: usize = 42; +const shared_span: usize = 3; + +fn paint(gate: []u8) void { + @memset(gate, 0); + const at = winner * gate_cols * @sizeOf(u16); + std.mem.writeInt(u16, gate[at..][0..2], one_bf16, .little); +} + +fn sharedGate() [shared_span * 2]u8 { + var bytes: [shared_span * 2]u8 = @splat(0); + for (0..shared_span) |i| std.mem.writeInt(u16, bytes[i * 2 ..][0..2], one_bf16, .little); + return bytes; +} + +fn hiddenVec() [gate_cols]f32 { + var hidden: [gate_cols]f32 = @splat(0); + hidden[0] = 1; + return hidden; +} + +fn expectUntouched(logits: []const f32, picks: []const u16, weights: []const f32) !void { + try std.testing.expectEqual(@as(f32, 8), logits[0]); + try std.testing.expectEqual(@as(f32, 8), logits[gate_rows]); + try std.testing.expectEqualSlices(u16, &.{ 9, 9 }, picks); + try std.testing.expectEqualSlices(f32, &.{ 3, 3 }, weights); +} + +test "one nonzero gate row and a shared span of 3 select id 512" { + try std.testing.expect(shared_span != gate_cols); + try std.testing.expect(shared_span != 640); + try std.testing.expect(shared_span != 1280); + try std.testing.expect(shared_span != 61_440); + try std.testing.expectEqual(@as(usize, 2_621_440), gate_bytes); + + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + paint(gate); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = undefined; + try launch(0, gate, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch); + + try std.testing.expectEqual(@as(f32, 0), logits[0]); + try std.testing.expectEqual(@as(f32, 1), logits[winner]); + try std.testing.expectEqual(@as(f32, 0), logits[gate_rows - 1]); + try std.testing.expectEqual(@as(f32, 9), logits[gate_rows]); + try std.testing.expectEqual(@as(u16, winner), picks[0]); + try std.testing.expectEqual(@as(u16, 512), picks[1]); + try std.testing.expectEqual(@as(f32, 1), weights[0]); +} + +test "rank 1 leaves the pick buffer untouched" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + paint(gate); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.NotRank0, launch(1, gate, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); + try std.testing.expectEqual(@as(f32, 4), scratch[0]); +} + +test "a missing shared gate does not call select" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + paint(gate); + + const hidden = hiddenVec(); + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.MissingSharedGate, launch(0, gate, &hidden, &.{}, &.{}, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); + try std.testing.expectEqual(@as(f32, 4), scratch[0]); +} + +test "a 513-row gate buffer is refused" { + const gpa = std.testing.allocator; + const wide_bytes = (gate_rows + 1) * gate_cols * @sizeOf(u16); + try std.testing.expectEqual(@as(usize, 2_626_560), wide_bytes); + const gate = try gpa.alloc(u8, wide_bytes); + defer gpa.free(gate); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.GateBytes, launch(0, gate, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); +} + +test "an empty gate does not launch" { + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const shared_x = [_]f32{ 0, 5, 4 }; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.EmptyGate, launch(0, &.{}, &hidden, &shared_gate, &shared_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); + try std.testing.expectEqual(@as(f32, 4), scratch[0]); +} + +test "a shared-span mismatch stores nothing" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + @memset(gate, 0); + + const hidden = hiddenVec(); + const shared_gate = sharedGate(); + const short_x = [_]f32{1}; + var logits: [gate_rows + 1]f32 = @splat(8); + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 3, 3 }; + var scratch: [gate_rows]f32 = @splat(4); + try std.testing.expectError(error.Length, launch(0, gate, &hidden, &shared_gate, &short_x, 1, &logits, &picks, &weights, &scratch)); + try expectUntouched(&logits, &picks, &weights); +} + +test "the same bits read as F16 are not the bf16 value" { + const gpa = std.testing.allocator; + const gate = try gpa.alloc(u8, gate_bytes); + defer gpa.free(gate); + @memset(gate, 0); + const one_f16: u16 = 0x3c00; + std.mem.writeInt(u16, gate[7 * gate_cols * 2 ..][0..2], one_f16, .little); + var shared: [gate_cols * 2]u8 = @splat(0); + std.mem.writeInt(u16, shared[0..2], one_f16, .little); + var hidden: [gate_cols]f32 = @splat(0); + hidden[0] = 1; + var logits: [gate_rows + 1]f32 = undefined; + var picks = [_]u16{ 9, 9 }; + var weights = [_]f32{ 0, 0 }; + var scratch: [gate_rows]f32 = undefined; + try launchDtype(.f16, 0, gate, &hidden, &shared, &hidden, 1, &logits, &picks, &weights, &scratch); + try std.testing.expectEqual(@as(f32, 1), logits[7]); + try std.testing.expectEqual(@as(f32, 1), logits[gate_rows]); + try std.testing.expectEqual(@as(u16, 7), picks[0]); + try launchDtype(.bf16, 0, gate, &hidden, &shared, &hidden, 1, &logits, &picks, &weights, &scratch); + try std.testing.expect(logits[7] != 1); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig b/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig new file mode 100644 index 0000000000..8774d10242 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_telemetry.zig @@ -0,0 +1,226 @@ +const std = @import("std"); + +pub const Phase = enum { prefill, decode, verify, unknown }; +pub const Rank = struct { + rank: usize, + world: usize, + + pub fn validate(owner: Rank) !void { + if (owner.world == 0 or owner.rank >= owner.world) return error.TelemetryRank; + } +}; + +var next_identity: std.atomic.Value(u64) = .init(0); + +pub fn identity() !u64 { + var previous = next_identity.load(.monotonic); + while (true) { + const next = try std.math.add(u64, previous, 1); + previous = next_identity.cmpxchgWeak(previous, next, .monotonic, .monotonic) orelse return next; + } +} + +pub const Device = struct { + plans: u64 = 0, + hit_picks: u64 = 0, + miss_picks: u64 = 0, + copies: u64 = 0, + evictions: u64 = 0, + inplace_experts: u64 = 0, + overflow_plans: u64 = 0, + + pub fn pending(current: anytype, seen: @TypeOf(current)) Device { + return .{ .plans = current.rounds -% seen.rounds, .hit_picks = current.hit_slots -% seen.hit_slots, .miss_picks = current.miss_slots -% seen.miss_slots, .copies = current.copies -% seen.copies, .evictions = current.evictions -% seen.evictions, .inplace_experts = current.inplace -% seen.inplace, .overflow_plans = current.overflow -% seen.overflow }; + } + + pub fn pendingBounded(current: anytype, seen: @TypeOf(current), max_picks: u64) !Device { + const delta = pending(current, seen); + const bound = try std.math.mul(u64, delta.plans, max_picks); + if (max_picks == 0 or bound >= @as(u64, 1) << 32) return error.TelemetryExpertIntervalTooLong; + if (try std.math.add(u64, delta.hit_picks, delta.miss_picks) > bound or try std.math.add(u64, delta.copies, delta.inplace_experts) > bound or delta.overflow_plans > delta.plans) return error.TelemetryInconsistentCounters; + return delta; + } + + pub fn record(total: *Device, valid: *bool, current: anytype, seen: @TypeOf(current), max_picks: u64) void { + if (!valid.*) return; + const delta = pendingBounded(current, seen, max_picks) catch { + valid.* = false; + return; + }; + add(Device, total, delta) catch { + valid.* = false; + }; + } +}; + +pub fn add(comptime Counters: type, total: *Counters, contribution: Counters) !void { + var result = total.*; + inline for (@typeInfo(Counters).@"struct".field_names) |field| @field(result, field) = try std.math.add(u64, @field(result, field), @field(contribution, field)); + total.* = result; +} + +pub fn difference(comptime Counters: type, before: Counters, after: Counters) !Counters { + var result: Counters = undefined; + inline for (@typeInfo(Counters).@"struct".field_names) |field| { + if (@field(after, field) < @field(before, field)) return error.TelemetryCounterRegression; + @field(result, field) = @field(after, field) - @field(before, field); + } + return result; +} + +pub const Experts = struct { + host_plans: u64 = 0, + host_hit_picks: u64 = 0, + host_miss_picks: u64 = 0, + host_mandatory_copies: u64 = 0, + host_prefetch_copies: u64 = 0, + pool_dma_copies: u64 = 0, + pool_dma_requested_bytes: u64 = 0, + device_plans: u64 = 0, + device_hit_picks: u64 = 0, + device_miss_picks: u64 = 0, + device_copies: u64 = 0, + device_inplace_experts: u64 = 0, + device_overflow_plans: u64 = 0, + mapped_copy_payload_min_bytes: u64 = 0, + mapped_copy_payload_max_bytes: u64 = 0, +}; + +pub const Snapshot = struct { + rank: Rank, + scope: u64, + experts: Experts = .{}, + expert_pools: usize = 0, + layers: []LayerSample = &.{}, + + pub fn deinit(snapshot: *Snapshot, allocator: std.mem.Allocator) void { + allocator.free(snapshot.layers); + snapshot.* = undefined; + } +}; + +pub const LayerCounts = struct { + host_plans: u64 = 0, + host_hit_picks: u64 = 0, + host_miss_picks: u64 = 0, + host_mandatory_copies: u64 = 0, + host_prefetch_copies: u64 = 0, + pool_dma_copies: u64 = 0, + pool_dma_requested_bytes: u64 = 0, + device_plans: u64 = 0, + device_hit_picks: u64 = 0, + device_miss_picks: u64 = 0, + device_copies: u64 = 0, + device_inplace_experts: u64 = 0, + device_overflow_plans: u64 = 0, + mapped_copy_payload_min_bytes: u64 = 0, + mapped_copy_payload_max_bytes: u64 = 0, + + pub fn fromExperts(experts: Experts) LayerCounts { + var result: LayerCounts = undefined; + inline for (@typeInfo(LayerCounts).@"struct".field_names) |field| @field(result, field) = @field(experts, field); + return result; + } +}; + +pub const LayerSample = struct { + owner: u64, + kind: enum { decoder }, + layer: usize, + lines: usize, + counts: LayerCounts = .{}, + + fn samePool(first: LayerSample, last: LayerSample) bool { + return first.owner != 0 and first.owner == last.owner and first.kind == last.kind and first.layer == last.layer and first.lines == last.lines; + } +}; + +pub fn LayerCursor(comptime capacity: usize) type { + return struct { + lock: std.atomic.Mutex = .unlocked, + scope: u64 = 0, + rank: Rank = .{ .rank = 0, .world = 1 }, + samples: [capacity]LayerSample = undefined, + len: usize = 0, + + pub fn claim(cursor: *@This(), allocator: std.mem.Allocator, before: Snapshot, after: Snapshot) ![]LayerSample { + while (!cursor.lock.tryLock()) std.atomic.spinLoopHint(); + defer cursor.lock.unlock(); + if (before.layers.len != after.layers.len or after.layers.len > capacity) return error.TelemetryOwnerChanged; + if (!std.meta.eql(before.rank, after.rank) or before.scope != after.scope) return error.TelemetryOwnerChanged; + if (cursor.len != 0 and (cursor.scope != after.scope or !std.meta.eql(cursor.rank, after.rank) or cursor.len != after.layers.len)) return error.TelemetryOwnerChanged; + const result = try allocator.alloc(LayerSample, after.layers.len); + errdefer allocator.free(result); + for (before.layers, after.layers, result, 0..) |first, last, *delta, index| { + if (!first.samePool(last)) return error.TelemetryOwnerChanged; + for (after.layers[0..index]) |previous| if (previous.owner == last.owner) return error.TelemetryDuplicateOwner; + if (cursor.len != 0 and !cursor.samples[index].samePool(last)) return error.TelemetryOwnerChanged; + delta.* = last; + delta.counts = try difference(LayerCounts, first.counts, last.counts); + } + for (after.layers, 0..) |last, index| { + cursor.samples[index] = last; + } + cursor.scope = after.scope; + cursor.rank = after.rank; + cursor.len = after.layers.len; + return result; + } + }; +} + +pub const Interval = struct { + phase: Phase, + model_rows: u64, + steps: u64, +}; + +pub fn fraction(misses: u64, hits: u64) !?f64 { + const denominator = try std.math.add(u64, misses, hits); + return if (denominator == 0) null else @as(f64, @floatFromInt(misses)) / @as(f64, @floatFromInt(denominator)); +} + +pub const Report = struct { + rank: Rank, + interval: Interval, + experts: Experts, + expert_pools: usize = 0, + layers: []const LayerSample = &.{}, + + pub fn write(report: Report, writer: *std.Io.Writer) !void { + try report.header(writer, "expert"); + inline for (@typeInfo(Experts).@"struct".field_names) |field| try writer.print(" {s}={d}", .{ field, @field(report.experts, field) }); + try writer.writeByte('\n'); + for (report.layers) |layer| { + try report.header(writer, "expert-layer"); + try writer.print(" pool={d} kind={s} layer={d} lines={d}", .{ layer.owner, @tagName(layer.kind), layer.layer, layer.lines }); + inline for (@typeInfo(LayerCounts).@"struct".field_names) |field| try writer.print(" {s}={d}", .{ field, @field(layer.counts, field) }); + try writer.writeByte('\n'); + } + } + + fn header(report: Report, writer: *std.Io.Writer, cache: []const u8) !void { + try writer.print("EXL3-CACHE v=1 rank={d} world={d} phase={s} rows={d} steps={d} expert_pools={d} cache={s}", .{ report.rank.rank, report.rank.world, @tagName(report.interval.phase), report.interval.model_rows, report.interval.steps, report.expert_pools, cache }); + } + + pub fn formatAlloc(report: Report, allocator: std.mem.Allocator) ![]u8 { + var writer: std.Io.Writer.Allocating = .init(allocator); + errdefer writer.deinit(); + try report.write(&writer.writer); + return writer.toOwnedSlice(); + } + + pub fn between(before: Snapshot, after: Snapshot, interval: Interval) !Report { + try before.rank.validate(); + try after.rank.validate(); + if (!std.meta.eql(before.rank, after.rank) or before.scope != after.scope or before.expert_pools != after.expert_pools) return error.TelemetryOwnerChanged; + if (interval.model_rows == 0 or interval.steps == 0) return error.TelemetryDenominator; + const experts = try difference(Experts, before.experts, after.experts); + if (experts.device_overflow_plans != 0) return error.TelemetryExpertOverflow; + return .{ .rank = after.rank, .interval = interval, .experts = experts, .expert_pools = after.expert_pools }; + } +}; + +test "exl3: telemetry host contracts" { + _ = @import("exl3_telemetry_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3/model/exl3_telemetry_test.zig b/zig/src/families/flashnext/exl3/model/exl3_telemetry_test.zig new file mode 100644 index 0000000000..459bace311 --- /dev/null +++ b/zig/src/families/flashnext/exl3/model/exl3_telemetry_test.zig @@ -0,0 +1,96 @@ +const std = @import("std"); +const telemetry = @import("exl3_telemetry.zig"); +const expectEqual = std.testing.expectEqual; +const expectError = std.testing.expectError; + +test "exl3: checked telemetry aggregation exceeds u32 without overflow or partial mutation" { + var totals: telemetry.Experts = .{ .device_hit_picks = std.math.maxInt(u32) }; + try telemetry.add(telemetry.Experts, &totals, .{ .device_hit_picks = 10 }); + try expectEqual(@as(u64, std.math.maxInt(u32)) + 10, totals.device_hit_picks); + totals.device_hit_picks = std.math.maxInt(u64); + const before = totals; + try expectError(error.Overflow, telemetry.add(telemetry.Experts, &totals, .{ .host_plans = 5, .device_hit_picks = 1 })); + try std.testing.expectEqualDeep(before, totals); + try expectError(error.TelemetryCounterRegression, telemetry.difference(telemetry.Experts, totals, .{})); + try expectEqual(@as(?f64, null), try telemetry.fraction(0, 0)); + try expectEqual(@as(?f64, 0.25), try telemetry.fraction(1, 3)); + try expectError(error.Overflow, telemetry.fraction(std.math.maxInt(u64), 1)); +} + +test "exl3: invalid device telemetry cannot fail normal ownership transfer or publish partial totals" { + const Misc = struct { now: u32 = 0, rounds: u32 = 0, hit_slots: u32 = 0, miss_slots: u32 = 0, copies: u32 = 0, evictions: u32 = 0, overflow: u32 = 0, inplace: u32 = 0 }; + var total: telemetry.Device = .{ .copies = 10 }; + var valid = true; + telemetry.Device.record(&total, &valid, Misc{ .rounds = 1, .hit_slots = 11 }, Misc{}, 10); + try std.testing.expect(!valid); + try expectEqual(@as(u64, 10), total.copies); + telemetry.Device.record(&total, &valid, Misc{ .rounds = 1, .copies = 1 }, Misc{}, 10); + try expectEqual(@as(u64, 10), total.copies); + try expectError(error.TelemetryExpertIntervalTooLong, telemetry.Device.pendingBounded(Misc{ .rounds = 1 << 30 }, Misc{}, 10)); +} + +test "exl3: phase report stable schema rank and model-row denominators" { + const before: telemetry.Snapshot = .{ .rank = .{ .rank = 1, .world = 2 }, .scope = 7 }; + var after = before; + after.experts.device_hit_picks = 40; + after.experts.device_miss_picks = 10; + const report = try telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 16, .steps = 8 }); + const text = try report.formatAlloc(std.testing.allocator); + defer std.testing.allocator.free(text); + try std.testing.expect(std.mem.startsWith(u8, text, "EXL3-CACHE v=1 rank=1 world=2 phase=decode rows=16 steps=8 expert_pools=0 cache=expert ")); + try std.testing.expect(std.mem.indexOf(u8, text, "device_hit_picks=40 device_miss_picks=10") != null); + try expectError(error.TelemetryDenominator, telemetry.Report.between(before, after, .{ .phase = .prefill, .model_rows = 0, .steps = 1 })); + after.rank.rank = 0; + try expectError(error.TelemetryOwnerChanged, telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 })); + after = before; + after.scope += 1; + try expectError(error.TelemetryOwnerChanged, telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 })); +} + +test "exl3: per-layer deltas sum to aggregate counters and overlapping reports claim counters once" { + const rank: telemetry.Rank = .{ .rank = 0, .world = 2 }; + var before_layers = [_]telemetry.LayerSample{.{ .owner = 11, .kind = .decoder, .layer = 4, .lines = 20 }}; + var after_layers = before_layers; + after_layers[0].counts.host_miss_picks = 5; + after_layers[0].counts.host_mandatory_copies = 2; + after_layers[0].counts.device_miss_picks = 3; + after_layers[0].counts.device_copies = 1; + var before: telemetry.Snapshot = .{ .rank = rank, .scope = 42, .layers = &before_layers, .expert_pools = 1 }; + var after = before; + after.layers = &after_layers; + after.experts.host_miss_picks = 5; + after.experts.host_mandatory_copies = 2; + after.experts.device_miss_picks = 3; + after.experts.device_copies = 1; + var cursor: telemetry.LayerCursor(2) = .{}; + const first = try cursor.claim(std.testing.allocator, before, after); + defer std.testing.allocator.free(first); + const report = try telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 }); + try std.testing.expectEqual(report.experts.host_miss_picks, first[0].counts.host_miss_picks); + try std.testing.expectEqual(report.experts.host_mandatory_copies, first[0].counts.host_mandatory_copies); + try std.testing.expectEqual(report.experts.device_miss_picks, first[0].counts.device_miss_picks); + try std.testing.expectEqual(report.experts.device_copies, first[0].counts.device_copies); + + var concurrent_before_layers = after_layers; + var concurrent_after_layers = after_layers; + concurrent_before_layers[0].counts.host_miss_picks = 3; + concurrent_before_layers[0].counts.host_mandatory_copies = 1; + concurrent_before_layers[0].counts.device_miss_picks = 2; + concurrent_before_layers[0].counts.device_copies = 0; + concurrent_after_layers[0].counts.host_miss_picks = 8; + concurrent_after_layers[0].counts.host_mandatory_copies = 4; + concurrent_after_layers[0].counts.device_miss_picks = 6; + concurrent_after_layers[0].counts.device_copies = 2; + before.layers = &concurrent_before_layers; + after.layers = &concurrent_after_layers; + before.experts = .{ .host_miss_picks = 5, .host_mandatory_copies = 2, .device_miss_picks = 3, .device_copies = 1 }; + after.experts = .{ .host_miss_picks = 10, .host_mandatory_copies = 5, .device_miss_picks = 7, .device_copies = 3 }; + const second = try cursor.claim(std.testing.allocator, before, after); + defer std.testing.allocator.free(second); + const second_report = try telemetry.Report.between(before, after, .{ .phase = .decode, .model_rows = 1, .steps = 1 }); + try std.testing.expectEqual(second_report.experts.host_miss_picks, second[0].counts.host_miss_picks); + try std.testing.expectEqual(second_report.experts.host_mandatory_copies, second[0].counts.host_mandatory_copies); + try std.testing.expectEqual(second_report.experts.device_miss_picks, second[0].counts.device_miss_picks); + try std.testing.expectEqual(second_report.experts.device_copies, second[0].counts.device_copies); + try std.testing.expectEqual(@as(u64, 10), first[0].counts.host_miss_picks + second[0].counts.host_miss_picks); +} diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency.zig new file mode 100644 index 0000000000..4395f177df --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency.zig @@ -0,0 +1,122 @@ +const std = @import("std"); +const Io = std.Io; +const Sha256 = std.crypto.hash.sha2.Sha256; + +pub const Identity = struct { + model: [32]u8, + world: u32, + rank: u32, +}; + +pub const decoder_layers = 48; +pub const experts = 512; +pub const format_version: u32 = 1; + +pub const Counts = struct { + decoder: [decoder_layers][experts]u32 = @splat(@splat(0)), +}; + +const magic = "TFEXL3FQ"; +const header_bytes = 64; +const decoder_bytes = decoder_layers * experts * @sizeOf(u32); +pub const max_file_bytes = header_bytes + decoder_bytes + Sha256.digest_length; + +pub fn read(gpa: std.mem.Allocator, io: Io, path: []const u8, identity: Identity) !Counts { + try validateIdentity(identity); + const file: Io.File = if (comptime @import("builtin").os.tag == .windows) return error.UnsupportedPlatform else .{ + .handle = try std.posix.openat(Io.Dir.cwd().handle, path, .{ + .ACCMODE = .RDONLY, + .NONBLOCK = true, + .CLOEXEC = true, + .NOCTTY = true, + }, 0), + .flags = .{ .nonblocking = true }, + }; + defer file.close(io); + const stat = try file.stat(io); + if (stat.kind != .file) return error.ProfileNotRegular; + if (stat.size > max_file_bytes) return error.ProfileTooLarge; + const buffer = try gpa.alloc(u8, max_file_bytes + 1); + defer gpa.free(buffer); + var reader = file.reader(io, &.{}); + const length = reader.interface.readSliceShort(buffer) catch return reader.err.?; + if (length > max_file_bytes) return error.ProfileTooLarge; + return decode(buffer[0..length], identity); +} + +pub fn write(gpa: std.mem.Allocator, io: Io, path: []const u8, identity: Identity, counts: *const Counts) !void { + try validateIdentity(identity); + if (comptime @import("builtin").os.tag == .windows) return error.UnsupportedPlatform; + const basename = std.fs.path.basename(path); + if (path.len == 0 or std.mem.endsWith(u8, path, "/") or + std.mem.eql(u8, basename, ".") or std.mem.eql(u8, basename, "..")) return error.InvalidPath; + const bytes = try gpa.alloc(u8, max_file_bytes); + defer gpa.free(bytes); + encode(bytes, identity, counts); + var atomic = try Io.Dir.cwd().createFileAtomic(io, path, .{ + .permissions = .fromMode(0o600), + .make_path = false, + .replace = true, + }); + defer atomic.deinit(io); + try atomic.file.writePositionalAll(io, bytes, 0); + try atomic.file.sync(io); + try atomic.replace(io); +} + +fn validateIdentity(identity: Identity) !void { + if (identity.world == 0 or identity.rank >= identity.world) return error.InvalidIdentity; +} + +fn encode(bytes: []u8, identity: Identity, counts: *const Counts) void { + @memcpy(bytes[0..8], magic); + putWord(bytes, 8, format_version); + putWord(bytes, 12, 0); + @memcpy(bytes[16..48], &identity.model); + putWord(bytes, 48, identity.world); + putWord(bytes, 52, identity.rank); + putWord(bytes, 56, decoder_layers); + putWord(bytes, 60, experts); + var offset: usize = header_bytes; + for (&counts.decoder) |*layer| for (layer) |count| { + putWord(bytes, offset, count); + offset += 4; + }; + var checksum: [Sha256.digest_length]u8 = undefined; + Sha256.hash(bytes[0..offset], &checksum, .{}); + @memcpy(bytes[offset..], &checksum); +} + +fn decode(bytes: []const u8, identity: Identity) !Counts { + if (bytes.len < header_bytes + Sha256.digest_length) return error.TruncatedProfile; + if (!std.mem.eql(u8, bytes[0..8], magic)) return error.MalformedProfile; + if (getWord(bytes, 8) != format_version) return error.UnsupportedVersion; + if (getWord(bytes, 12) != 0 or getWord(bytes, 56) != decoder_layers or getWord(bytes, 60) != experts) return error.MalformedProfile; + const payload_end = header_bytes + decoder_bytes; + if (bytes.len < max_file_bytes) return error.TruncatedProfile; + if (bytes.len != max_file_bytes) return error.MalformedProfile; + var checksum: [Sha256.digest_length]u8 = undefined; + Sha256.hash(bytes[0..payload_end], &checksum, .{}); + if (!std.mem.eql(u8, bytes[payload_end..], &checksum)) return error.ChecksumMismatch; + if (!std.mem.eql(u8, bytes[16..48], &identity.model) or + getWord(bytes, 48) != identity.world or getWord(bytes, 52) != identity.rank) return error.IdentityMismatch; + var counts: Counts = .{}; + var offset: usize = header_bytes; + for (&counts.decoder) |*layer| for (layer) |*count| { + count.* = getWord(bytes, offset); + offset += 4; + }; + return counts; +} + +fn putWord(bytes: []u8, offset: usize, value: u32) void { + std.mem.writeInt(u32, bytes[offset..][0..4], value, .little); +} + +fn getWord(bytes: []const u8, offset: usize) u32 { + return std.mem.readInt(u32, bytes[offset..][0..4], .little); +} + +test { + _ = @import("cuda_exl3_frequency_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_state.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_state.zig new file mode 100644 index 0000000000..18578046d8 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_state.zig @@ -0,0 +1,444 @@ +//! Lifetime placement priors: routed/metadata fingerprint; startup warmth may be discarded by prefill. +const std = @import("std"); +const codec = @import("cuda_exl3_frequency.zig"); +const Sha256 = std.crypto.hash.sha2.Sha256; +const namespace = "flashnext-exl3-placement-v1"; + +pub const State = struct { + gpa: std.mem.Allocator, + io: std.Io, + path: []u8, + identity: codec.Identity, + lock: std.Io.File, + + pub fn deinit(state: *State) void { + state.lock.close(state.io); + state.gpa.free(state.path); + state.* = undefined; + } + + pub fn save(state: *State, own: anytype) void { + const counts = snapshot(own) catch |err| { + std.log.err("EXL3 placement profile rank {d}: unsafe snapshot skipped ({s})", .{ state.identity.rank, @errorName(err) }); + return; + }; + codec.write(state.gpa, state.io, state.path, state.identity, &counts) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: save failed ({s})", .{ state.identity.rank, @errorName(err) }); + return; + }; + std.log.info("EXL3 placement profile rank {d}: saved decoder lifetime placement snapshot to {s}", .{ state.identity.rank, state.path }); + } +}; + +pub const Prepared = struct { + state: State, + counts: ?codec.Counts, +}; + +pub fn load(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, own: anytype, world: u32, rank: u32) !?State { + const prepared = try prepare(gpa, io, dir, own.engine.layers[0..own.engine.n_layers], world, rank) orelse return null; + var state = prepared.state; + errdefer state.deinit(); + if (prepared.counts) |*counts| try finish(own, counts); + return state; +} + +pub fn prepare(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anytype, world: u32, rank: u32) !?Prepared { + const root = std.c.getenv("TF_FLASHNEXT_EXL3_PROFILE_DIR") orelse return null; + if (std.mem.span(root).len == 0) return null; + if (comptime @import("builtin").os.tag == .windows) { + std.log.warn("EXL3 placement profiles require POSIX storage; persistence disabled", .{}); + return null; + } + if (!available(layers)) { + std.log.warn("EXL3 placement profile rank {d}: unavailable decoder pools; persistence disabled", .{rank}); + return null; + } + const before = std.Io.Clock.awake.now(io).toNanoseconds(); + const model = fingerprint(gpa, io, dir, layers) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: fingerprint unavailable ({s}); persistence disabled", .{ rank, @errorName(err) }); + return null; + }; + std.log.info("EXL3 placement profile rank {d}: routed-bank/metadata fingerprint took {d:.3} s (excludes router/dense weights)", .{ rank, @as(f64, @floatFromInt(std.Io.Clock.awake.now(io).toNanoseconds() - before)) / 1e9 }); + var state = open(gpa, io, std.mem.span(root), .{ .model = model, .world = world, .rank = rank }) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: lock/storage unavailable ({s}); persistence disabled", .{ rank, @errorName(err) }); + return null; + }; + errdefer state.deinit(); + var counts: ?codec.Counts = codec.read(gpa, io, state.path, state.identity) catch |err| blk: { + std.log.warn("EXL3 placement profile rank {d}: cold start ({s})", .{ rank, @errorName(err) }); + break :blk null; + }; + if (counts) |*value| { + decay(value); + std.log.info("EXL3 placement profile rank {d}: loaded and decayed rank-local placement counts before pool allocation", .{rank}); + } + return .{ .state = state, .counts = counts }; +} + +pub fn finish(own: anytype, counts: *const codec.Counts) !void { + const layers = own.engine.layers[0..own.engine.n_layers]; + if (!available(layers)) return error.ProfileUnavailable; + for (layers, 0..) |layer, index| { + const pool = layer.exl3.?.pool orelse return error.ProfileUnavailable; + pool.counts[pool.opt.layer] = counts.decoder[index]; + } + var warmed: usize = 0; + for (layers) |layer| warmed += try layer.exl3.?.pool.?.warmProfile(); + std.log.info("EXL3 placement profile: imported and warmed {d} expert lines after pool allocation", .{warmed}); +} + +fn open(gpa: std.mem.Allocator, io: std.Io, root: []const u8, identity: codec.Identity) !State { + const private = try std.fs.path.join(gpa, &.{ root, namespace }); + defer gpa.free(private); + var directory = try std.Io.Dir.cwd().createDirPathOpen(io, private, .{ + .permissions = .fromMode(0o700), + .open_options = .{ .iterate = true }, + }); + defer directory.close(io); + try directory.setPermissions(io, .fromMode(0o700)); + const name = try std.fmt.allocPrint(gpa, "{s}-w{d}-r{d}.bin", .{ std.fmt.bytesToHex(&identity.model, .lower), identity.world, identity.rank }); + defer gpa.free(name); + const path = try std.fs.path.join(gpa, &.{ private, name }); + errdefer gpa.free(path); + const lock_path = try std.fmt.allocPrint(gpa, "{s}.lock", .{path}); + defer gpa.free(lock_path); + const lock = try openLock(io, lock_path); + errdefer lock.close(io); + try lock.setPermissions(io, .fromMode(0o600)); + return .{ .gpa = gpa, .io = io, .path = path, .identity = identity, .lock = lock }; +} + +fn openLock(io: std.Io, path: []const u8) !std.Io.File { + const file: std.Io.File = if (comptime @import("builtin").os.tag == .windows) return error.UnsupportedPlatform else .{ + .handle = try std.posix.openat(std.Io.Dir.cwd().handle, path, .{ + .ACCMODE = .RDWR, + .CREAT = true, + .NONBLOCK = true, + .CLOEXEC = true, + .NOCTTY = true, + .NOFOLLOW = true, + }, 0o600), + .flags = .{ .nonblocking = true }, + }; + errdefer file.close(io); + if ((try file.stat(io)).kind != .file) return error.ProfileLockNotRegular; + if (!try file.tryLock(io, .exclusive)) return error.WouldBlock; + return file; +} + +fn restore(state: *State, own: anytype) !bool { + var counts = codec.read(state.gpa, state.io, state.path, state.identity) catch |err| { + std.log.warn("EXL3 placement profile rank {d}: cold start ({s})", .{ state.identity.rank, @errorName(err) }); + return false; + }; + decay(&counts); + const layers = own.engine.layers[0..own.engine.n_layers]; + for (layers, 0..) |layer, index| { + const pool = layer.exl3.?.pool.?; + pool.counts[pool.opt.layer] = counts.decoder[index]; + } + std.log.info("EXL3 placement profile rank {d}: imported decoder lifetime counts, decayed once by 2 (nonzero minimum 1)", .{state.identity.rank}); + var warmed: usize = 0; + for (layers) |layer| warmed += try layer.exl3.?.pool.?.warmProfile(); + std.log.info("EXL3 placement profile rank {d}: warmed {d} expert lines", .{ state.identity.rank, warmed }); + return true; +} + +fn available(layers: anytype) bool { + if (layers.len != 48) return false; + for (layers) |layer| { + const experts = layer.exl3 orelse return false; + if (experts.bank == null or experts.routed.len != codec.experts) return false; + if (experts.pool) |pool| if (pool.opt.layer >= 48) return false; + } + return true; +} + +fn decay(counts: *codec.Counts) void { + for (&counts.decoder) |*layer| decayRow(layer); +} + +fn decayRow(row: *[512]u32) void { + for (row) |*count| { + if (count.* != 0) count.* = @max(1, count.* / 2); + } +} + +fn snapshot(own: anytype) !codec.Counts { + const layers = own.engine.layers[0..own.engine.n_layers]; + if (!available(layers)) return error.ProfileUnavailable; + try own.stream.synchronize(); + var counts: codec.Counts = .{}; + for (layers, 0..) |layer, index| { + const pool = layer.exl3.?.pool.?; + try pool.lines.drain(); + counts.decoder[index] = pool.counts[pool.opt.layer]; + } + return counts; +} + +fn hashBytes(hash: *Sha256, bytes: []const u8) void { + var length: [8]u8 = undefined; + std.mem.writeInt(u64, &length, bytes.len, .little); + hash.update(&length); + hash.update(bytes); +} + +fn hashRouted(hash: *Sha256, routed: []const [3][]const u8) void { + for (routed) |expert| for (expert) |projection| hashBytes(hash, projection); +} + +fn fingerprint(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, layers: anytype) ![32]u8 { + var hash: Sha256 = .init(.{}); + hashBytes(&hash, namespace); + for ([_][]const u8{ "config.json", "model.safetensors.index.json" }) |name| { + const path = try std.fs.path.join(gpa, &.{ dir, name }); + defer gpa.free(path); + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); + defer gpa.free(bytes); + hashBytes(&hash, name); + hashBytes(&hash, bytes); + } + hashBytes(&hash, "decoder-routed-trellis"); + for (layers) |layer| hashRouted(&hash, expertRouted(layer.exl3.?)); + return hash.finalResult(); +} + +fn expertRouted(experts: anytype) []const [3][]const u8 { + if (experts.pool) |pool| return pool.src.routed; + return experts.routed; +} + +const MockPool = struct { + counts: *[48][512]u32, + opt: struct { layer: u8, lines: usize = 10 }, + src: struct { routed: []const [3][]const u8 = &.{.{ "gate", "up", "down" }} } = .{}, + lines: struct { + drained: bool = false, + pub fn drain(copies: *@This()) !void { + copies.drained = true; + } + } = .{}, + warm_calls: usize = 0, + fail_warm: bool = false, + + pub fn warmProfile(pool: *MockPool) !usize { + pool.warm_calls += 1; + if (pool.fail_warm) return error.MockCudaWarmFailure; + return 1; + } +}; + +const mock_routed: [512][3][]const u8 = @splat(.{ "gate", "up", "down" }); +const MockExperts = struct { pool: ?*MockPool, bank: ?[]const u8 = "bank", routed: []const [3][]const u8 = &mock_routed }; +const MockLayer = struct { exl3: ?*MockExperts }; +const Mock = struct { + decoder: [48][512]u32 = @splat(@splat(0)), + pools: [48]MockPool = undefined, + experts: [48]MockExperts = undefined, + engine: struct { layers: [48]MockLayer, n_layers: usize = 48 } = undefined, + + fn init(mock: *Mock) void { + mock.* = .{}; + mock.engine.n_layers = 48; + for (&mock.pools, &mock.experts, 0..) |*pool, *experts, index| { + pool.* = .{ .counts = &mock.decoder, .opt = .{ .layer = @intCast(index) }, .src = .{ .routed = &mock_routed } }; + experts.* = .{ .pool = pool }; + mock.engine.layers[index] = .{ .exl3 = experts }; + } + } + + const Stream = struct { + syncs: usize = 0, + fail: bool = false, + pub fn synchronize(stream: *@This()) !void { + if (stream.fail) return error.MockCudaSyncFailure; + stream.syncs += 1; + } + }; + fn owner(mock: *Mock) struct { engine: *@TypeOf(mock.engine), stream: Stream } { + return .{ .engine = &mock.engine, .stream = .{} }; + } +}; + +const test_identity: codec.Identity = .{ .model = @splat(42), .world = 2, .rank = 0 }; + +test "exl3: placement state lifetime lock conflicts, rank separation and private storage" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + var state = try open(allocator, io, root, test_identity); + var held = true; + defer if (held) state.deinit(); + try std.testing.expectError(error.WouldBlock, open(allocator, io, root, test_identity)); + const empty: codec.Counts = .{}; + try codec.write(allocator, io, state.path, state.identity, &empty); + try std.testing.expectError(error.WouldBlock, open(allocator, io, root, test_identity)); + var rank_identity = test_identity; + rank_identity.rank = 1; + var other_rank = try open(allocator, io, root, rank_identity); + defer other_rank.deinit(); + try std.testing.expect(!std.mem.eql(u8, state.path, other_rank.path)); + const lock_stat = try state.lock.stat(io); + try std.testing.expectEqual(@as(u32, 0o600), lock_stat.permissions.toMode() & 0o777); + const directory = try std.Io.Dir.cwd().openDir(io, std.fs.path.dirname(state.path).?, .{}); + defer directory.close(io); + const directory_stat = try directory.stat(io); + try std.testing.expectEqual(@as(u32, 0o700), directory_stat.permissions.toMode() & 0o777); + state.deinit(); + held = false; + var reopened = try open(allocator, io, root, test_identity); + defer reopened.deinit(); +} + +test "exl3: placement state cold IO never imports or warms; CUDA warm errors propagate" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + var state = try open(allocator, io, root, test_identity); + defer state.deinit(); + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(); + try std.testing.expect(!try restore(&state, &own)); + try std.Io.Dir.cwd().writeFile(io, .{ .sub_path = state.path, .data = "corrupt" }); + try std.testing.expect(!try restore(&state, &own)); + var counts: codec.Counts = .{}; + counts.decoder[0][1] = 123; + try codec.write(allocator, io, state.path, state.identity, &counts); + mock.pools[0].fail_warm = true; + try std.testing.expectError(error.MockCudaWarmFailure, restore(&state, &own)); + try std.testing.expectEqual(@as(u32, 61), mock.decoder[0][1]); + mock.engine.n_layers = 47; + try std.testing.expect(!available(mock.engine.layers[0..mock.engine.n_layers])); +} + +test "exl3: placement state decay and replacement snapshot do not add imported prior twice" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + var state = try open(allocator, io, root, test_identity); + defer state.deinit(); + var counts: codec.Counts = .{}; + counts.decoder[0][0..5].* = .{ 0, 1, 2, 3, std.math.maxInt(u32) }; + try codec.write(allocator, io, state.path, state.identity, &counts); + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(); + try std.testing.expect(try restore(&state, &own)); + try std.testing.expectEqualSlices(u32, &.{ 0, 1, 1, 1, 2147483647 }, mock.decoder[0][0..5]); + for (mock.pools) |pool| try std.testing.expectEqual(@as(usize, 1), pool.warm_calls); + mock.decoder[0][1] += 8; + const saved = try snapshot(&own); + try codec.write(allocator, io, state.path, state.identity, &saved); + const disk = try codec.read(allocator, io, state.path, state.identity); + try std.testing.expectEqual(@as(u32, 9), disk.decoder[0][1]); + mock.init(); + own = mock.owner(); + try std.testing.expect(try restore(&state, &own)); + try std.testing.expectEqual(@as(u32, 4), mock.decoder[0][1]); + const missing = try std.fs.path.join(allocator, &.{ root, "missing", "profile" }); + defer allocator.free(missing); + try std.testing.expectError(error.FileNotFound, codec.write(allocator, io, missing, state.identity, &saved)); +} + +test "exl3: placement state snapshot synchronizes once and drains decoder counts" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(); + for (&mock.pools, 0..) |*pool, index| pool.counts[pool.opt.layer][7] = @intCast(100 + index); + const counts = try snapshot(&own); + try std.testing.expectEqual(@as(usize, 1), own.stream.syncs); + try std.testing.expectEqual(@as(u32, 100), counts.decoder[0][7]); + try std.testing.expectEqual(@as(u32, 147), counts.decoder[47][7]); + for (mock.pools) |pool| try std.testing.expect(pool.lines.drained); + own.stream.fail = true; + try std.testing.expectError(error.MockCudaSyncFailure, snapshot(&own)); +} + +test "exl3: prepared placement counts import and warm only after pools attach" { + var mock: Mock = undefined; + mock.init(); + var own = mock.owner(); + var counts: codec.Counts = .{}; + for (&counts.decoder, 0..) |*row, layer| row[7] = @intCast(layer + 1); + + try finish(&own, &counts); + + for (mock.pools[0..48], 0..) |pool, layer| { + try std.testing.expectEqual(@as(u32, @intCast(layer + 1)), mock.decoder[layer][7]); + try std.testing.expectEqual(@as(usize, 1), pool.warm_calls); + } +} + +test "exl3: placement fingerprint reads routed and metadata bytes and ignores pool size" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + try temporary.dir.writeFile(io, .{ .sub_path = "config.json", .data = "config" }); + try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "index" }); + var mock: Mock = undefined; + mock.init(); + const original = try fingerprint(allocator, io, root, &mock.engine.layers); + for (&mock.experts) |*experts| experts.pool = null; + try std.testing.expectEqual(original, try fingerprint(allocator, io, root, &mock.engine.layers)); + for (&mock.experts, &mock.pools) |*experts, *pool| experts.pool = pool; + mock.pools[0].opt.lines = 23; + try std.testing.expectEqual(original, try fingerprint(allocator, io, root, &mock.engine.layers)); + mock.pools[47].src.routed = &.{.{ "changed gate", "up", "down" }}; + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers)))); + mock.pools[47].src.routed = &.{.{ "gate", "up", "down" }}; + try temporary.dir.writeFile(io, .{ .sub_path = "config.json", .data = "changed config" }); + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers)))); + try temporary.dir.writeFile(io, .{ .sub_path = "config.json", .data = "config" }); + try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "changed index" }); + try std.testing.expect(!std.mem.eql(u8, &original, &(try fingerprint(allocator, io, root, &mock.engine.layers)))); + try temporary.dir.writeFile(io, .{ .sub_path = "model.safetensors.index.json", .data = "index" }); +} + +test "exl3: placement lock rejects nonregular storage without waiting for FIFO readers" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + const allocator = std.testing.allocator; + const io = std.testing.io; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(path); + try std.testing.expectError(error.ProfileLockNotRegular, openLock(io, "/dev/null")); + const fifo = try std.fs.path.join(allocator, &.{ path, "profile.lock" }); + defer allocator.free(fifo); + const terminated = try allocator.dupeSentinel(u8, fifo, 0); + defer allocator.free(terminated); + switch (@import("builtin").os.tag) { + .linux => { + const result = std.os.linux.mknodat(temporary.dir.handle, "profile.lock", std.posix.S.IFIFO | 0o600, 0); + if (std.posix.errno(result) != .SUCCESS) return error.FifoCreateFailed; + }, + .macos => { + const Native = struct { + extern "c" fn mkfifo(pathname: [*:0]const u8, mode: std.posix.mode_t) c_int; + }; + if (Native.mkfifo(terminated, 0o600) != 0) return error.FifoCreateFailed; + }, + else => return error.SkipZigTest, + } + try std.testing.expectError(error.ProfileLockNotRegular, openLock(io, fifo)); +} diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_test.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_test.zig new file mode 100644 index 0000000000..f39617e4a8 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_frequency_test.zig @@ -0,0 +1,299 @@ +const std = @import("std"); +const profile = @import("cuda_exl3_frequency.zig"); +const Io = std.Io; +const allocator = std.testing.allocator; +const io = std.testing.io; +const identity: profile.Identity = .{ .model = @splat(0xa7), .world = 2, .rank = 1 }; +const header_bytes = 64; +const checksum_bytes = 32; + +fn profilePath(directory: Io.Dir) ![]u8 { + const root = try directory.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + return std.fs.path.join(allocator, &.{ root, "frequency.bin" }); +} + +fn sample(counts: *profile.Counts) void { + counts.* = .{}; + for (&counts.decoder, 0..) |*layer, layer_index| for (layer, 0..) |*count, expert_index| { + count.* = @intCast((layer_index + 1) * 100_000 + expert_index); + }; + counts.decoder[0][0] = 0; + counts.decoder[47][511] = std.math.maxInt(u32); +} + +fn expectCounts(expected: *const profile.Counts, actual: *const profile.Counts) !void { + for (&expected.decoder, &actual.decoder) |*expected_layer, *actual_layer| { + try std.testing.expectEqualSlices(u32, expected_layer, actual_layer); + } +} + +fn expectOnlyProfile(directory: Io.Dir) !void { + var iterator = directory.iterate(); + var entries: usize = 0; + while (try iterator.next(io)) |entry| { + entries += 1; + try std.testing.expectEqualStrings("frequency.bin", entry.name); + } + try std.testing.expectEqual(@as(usize, 1), entries); +} + +fn bytesAt(path: []const u8) ![]u8 { + return Io.Dir.cwd().readFileAlloc(io, path, allocator, .limited(profile.max_file_bytes + 2)); +} + +fn expectPreserved(path: []const u8, before: []const u8) !void { + const after = try bytesAt(path); + defer allocator.free(after); + try std.testing.expectEqualSlices(u8, before, after); +} + +test "exl3: frequency full decoder round trip with private permissions" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + sample(&counts); + try profile.write(allocator, io, path, identity, &counts); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); + const stat = try temporary.dir.statFile(io, "frequency.bin", .{}); + try std.testing.expectEqual(@as(u64, profile.max_file_bytes), stat.size); + try std.testing.expectEqual(@as(std.posix.mode_t, 0o600), stat.permissions.toMode() & 0o777); + const bytes = try bytesAt(path); + defer allocator.free(bytes); + try std.testing.expectEqualSlices(u8, "TFEXL3FQ", bytes[0..8]); + try std.testing.expectEqualSlices(u8, &identity.model, bytes[16..48]); + try std.testing.expectEqualSlices(u8, &.{ 1, 0, 0, 0 }, bytes[8..12]); + try std.testing.expectEqualSlices(u8, &.{ 0, 0, 0, 0 }, bytes[64..68]); + try std.testing.expectEqualSlices(u8, &.{ 0xa1, 0x86, 0x01, 0x00 }, bytes[68..72]); + try std.testing.expectEqualSlices(u8, &.{ 0xff, 0xff, 0xff, 0xff }, bytes[64 + 48 * 512 * 4 - 4 ..][0..4]); + try expectOnlyProfile(temporary.dir); +} + +test "exl3: frequency missing profile and every exact identity mismatch are recoverable" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + try std.testing.expectError(error.FileNotFound, profile.read(allocator, io, path, identity)); + const counts: profile.Counts = .{}; + try profile.write(allocator, io, path, identity, &counts); + for (0..32) |model_index| { + var different = identity; + different.model[model_index] ^= 1; + try std.testing.expectError(error.IdentityMismatch, profile.read(allocator, io, path, different)); + } + var different = identity; + different.world = 3; + try std.testing.expectError(error.IdentityMismatch, profile.read(allocator, io, path, different)); + different = identity; + different.rank = 0; + try std.testing.expectError(error.IdentityMismatch, profile.read(allocator, io, path, different)); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); +} + +test "exl3: frequency rejects truncated oversized malformed and corrupt profiles" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + sample(&counts); + try profile.write(allocator, io, path, identity, &counts); + const original = try bytesAt(path); + defer allocator.free(original); + for ([_]usize{ 0, 7, 63, 64, 95, 96, profile.max_file_bytes - 1 }) |length| { + try temporary.dir.writeFile(io, .{ .sub_path = "frequency.bin", .data = original[0..length] }); + try std.testing.expectError(error.TruncatedProfile, profile.read(allocator, io, path, identity)); + } + const changed = try allocator.dupe(u8, original); + defer allocator.free(changed); + const Mutation = struct { offset: usize, value: u8, expected: anyerror }; + for ([_]Mutation{ + .{ .offset = 0, .value = 0, .expected = error.MalformedProfile }, + .{ .offset = 8, .value = 2, .expected = error.UnsupportedVersion }, + .{ .offset = 12, .value = 2, .expected = error.MalformedProfile }, + .{ .offset = 56, .value = 47, .expected = error.MalformedProfile }, + .{ .offset = 60, .value = 1, .expected = error.MalformedProfile }, + .{ .offset = 16, .value = 0, .expected = error.ChecksumMismatch }, + .{ .offset = 48, .value = 3, .expected = error.ChecksumMismatch }, + .{ .offset = 64, .value = 1, .expected = error.ChecksumMismatch }, + .{ .offset = 64 + 48 * 512 * 4, .value = 1, .expected = error.ChecksumMismatch }, + .{ .offset = profile.max_file_bytes - 1, .value = original[profile.max_file_bytes - 1] ^ 1, .expected = error.ChecksumMismatch }, + }) |mutation| { + @memcpy(changed, original); + changed[mutation.offset] = mutation.value; + try temporary.dir.writeFile(io, .{ .sub_path = "frequency.bin", .data = changed }); + try std.testing.expectError(mutation.expected, profile.read(allocator, io, path, identity)); + } + const oversized = try allocator.alloc(u8, profile.max_file_bytes + 1); + defer allocator.free(oversized); + @memcpy(oversized[0..original.len], original); + oversized[original.len] = 0; + try temporary.dir.writeFile(io, .{ .sub_path = "frequency.bin", .data = oversized }); + try std.testing.expectError(error.ProfileTooLarge, profile.read(allocator, io, path, identity)); +} + +fn failWrite(userdata: ?*anyopaque, file: Io.File, header: []const u8, data: []const []const u8, splat: usize, offset: u64) Io.File.WritePositionalError!usize { + if (offset != 0) return error.NoSpaceLeft; + const prefix = data[0][0..13]; + return io.vtable.fileWritePositional(userdata, file, header, &.{prefix}, splat, offset); +} + +fn failSync(_: ?*anyopaque, file: Io.File) Io.File.SyncError!void { + const stat = file.stat(io) catch return error.InputOutput; + if (stat.permissions.toMode() & 0o777 != 0o600 or stat.size != profile.max_file_bytes) return error.AccessDenied; + return error.InputOutput; +} + +fn failRename(_: ?*anyopaque, _: Io.Dir, _: []const u8, _: Io.Dir, _: []const u8) Io.Dir.RenameError!void { + return error.AccessDenied; +} + +test "exl3: frequency partial write sync and rename failures preserve previous bytes and remove temporaries" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + try profile.write(allocator, io, path, identity, &counts); + const before = try bytesAt(path); + defer allocator.free(before); + sample(&counts); + for (0..3) |failure_index| { + var failing_vtable = io.vtable.*; + const expected: anyerror = switch (failure_index) { + 0 => failure: { + failing_vtable.fileWritePositional = failWrite; + break :failure error.NoSpaceLeft; + }, + 1 => failure: { + failing_vtable.fileSync = failSync; + break :failure error.InputOutput; + }, + else => failure: { + failing_vtable.dirRename = failRename; + break :failure error.AccessDenied; + }, + }; + var failing_io = io; + failing_io.vtable = &failing_vtable; + try std.testing.expectError(expected, profile.write(allocator, failing_io, path, identity, &counts)); + try expectPreserved(path, before); + try expectOnlyProfile(temporary.dir); + const restored = try profile.read(allocator, io, path, identity); + try std.testing.expectEqual(@as(u32, 0), restored.decoder[47][511]); + } + try profile.write(allocator, io, path, identity, &counts); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); +} + +test "exl3: frequency allocation and invalid identity failures preserve previous profile" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + const counts: profile.Counts = .{}; + try profile.write(allocator, io, path, identity, &counts); + const before = try bytesAt(path); + defer allocator.free(before); + var failing = std.testing.FailingAllocator.init(allocator, .{ .fail_index = 0 }); + try std.testing.expectError(error.OutOfMemory, profile.write(failing.allocator(), io, path, identity, &counts)); + try std.testing.expectError(error.OutOfMemory, profile.read(failing.allocator(), io, path, identity)); + for ([_]profile.Identity{ + .{ .model = identity.model, .world = 0, .rank = 0 }, + .{ .model = identity.model, .world = 2, .rank = 2 }, + }) |invalid| { + try std.testing.expectError(error.InvalidIdentity, profile.write(allocator, io, path, invalid, &counts)); + try std.testing.expectError(error.InvalidIdentity, profile.read(allocator, io, path, invalid)); + } + try std.testing.expectError(error.InvalidPath, profile.write(allocator, io, "", identity, &counts)); + try std.testing.expectError(error.InvalidPath, profile.write(allocator, io, ".", identity, &counts)); + try std.testing.expectError(error.InvalidPath, profile.write(allocator, io, "..", identity, &counts)); + try expectPreserved(path, before); + try expectOnlyProfile(temporary.dir); +} + +test "exl3: frequency read allocation stays fixed and oversized input needs no allocation" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + const single: profile.Identity = .{ .model = @splat(0), .world = 1, .rank = 0 }; + var counts: profile.Counts = .{}; + sample(&counts); + try profile.write(allocator, io, path, single, &counts); + const backing = try allocator.alloc(u8, profile.max_file_bytes + 1); + defer allocator.free(backing); + var bounded = std.heap.FixedBufferAllocator.init(backing); + const restored = try profile.read(bounded.allocator(), io, path, single); + try expectCounts(&counts, &restored); + const file = try temporary.dir.openFile(io, "frequency.bin", .{ .mode = .read_write }); + defer file.close(io); + try file.setLength(io, profile.max_file_bytes + 1); + var failing = std.testing.FailingAllocator.init(allocator, .{ .fail_index = 0 }); + try std.testing.expectError(error.ProfileTooLarge, profile.read(failing.allocator(), io, path, single)); + try std.testing.expect(!failing.has_induced_failure); +} + +test "exl3: frequency failed first save leaves no destination or temporary file" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{ .iterate = true }); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + var counts: profile.Counts = .{}; + sample(&counts); + var failing_vtable = io.vtable.*; + failing_vtable.fileSync = failSync; + var failing_io = io; + failing_io.vtable = &failing_vtable; + try std.testing.expectError(error.InputOutput, profile.write(allocator, failing_io, path, identity, &counts)); + try std.testing.expectError(error.FileNotFound, profile.read(allocator, io, path, identity)); + var iterator = temporary.dir.iterate(); + try std.testing.expectEqual(@as(?Io.Dir.Entry, null), try iterator.next(io)); + try profile.write(allocator, io, path, identity, &counts); + const restored = try profile.read(allocator, io, path, identity); + try expectCounts(&counts, &restored); + try expectOnlyProfile(temporary.dir); +} + +test "exl3: frequency rejects a FIFO without a writer and other nonregular input" { + if (comptime @import("builtin").os.tag == .windows) return error.SkipZigTest; + var temporary = std.testing.tmpDir(.{}); + defer temporary.cleanup(); + const path = try profilePath(temporary.dir); + defer allocator.free(path); + const terminated_path = try allocator.dupeSentinel(u8, path, 0); + defer allocator.free(terminated_path); + switch (@import("builtin").os.tag) { + .linux => { + const result = std.os.linux.mknodat(temporary.dir.handle, "frequency.bin", std.posix.S.IFIFO | 0o600, 0); + if (std.posix.errno(result) != .SUCCESS) return error.FifoCreateFailed; + }, + .macos => { + const Native = struct { + extern "c" fn mkfifo(pathname: [*:0]const u8, mode: std.posix.mode_t) c_int; + }; + if (Native.mkfifo(terminated_path, 0o600) != 0) return error.FifoCreateFailed; + }, + else => return error.SkipZigTest, + } + var failing = std.testing.FailingAllocator.init(allocator, .{ .fail_index = 0 }); + try std.testing.expectError(error.ProfileNotRegular, profile.read(failing.allocator(), io, path, identity)); + try std.testing.expectError(error.ProfileNotRegular, profile.read(allocator, io, "/dev/null", identity)); + const root = try temporary.dir.realPathFileAlloc(io, ".", allocator); + defer allocator.free(root); + try std.testing.expectError(error.ProfileNotRegular, profile.read(allocator, io, root, identity)); + try std.testing.expect(!failing.has_induced_failure); +} diff --git a/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig b/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig new file mode 100644 index 0000000000..da076ee7e6 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/cuda_exl3_pool.zig @@ -0,0 +1,345 @@ +//! One layer's EXL3 routed-expert offload pool on upstream's `Experts` (plan step 6 of + +const std = @import("std"); +const cuda = @import("cuda"); +const round_ = @import("../model/exl3_round.zig"); +const exl3e = cuda.exl3_experts; +const Lines = cuda.exl3_pool.Lines; +const telemetry = @import("../model/exl3_telemetry.zig"); + +pub const offload = @import("exl3_offload.zig"); +pub const round = round_; +pub const Round = round_.Round; +pub const linesFor = round_.linesFor; +pub const shape = offload.shape; +pub const shared_id = offload.shared_id; +/// The fewest lines a layer may have: one row's ten routed picks always fit one pass. +pub const min_lines: usize = 10; + +test { + _ = @import("exl3_frequency_seed.zig"); +} + +/// The layer's experts in host memory: routed [512] gate/up/down trellis bytes (stored layout), the shared expert's. +pub const Source = struct { + routed: []const [3][]const u8, + shared: ?[3][]const u8 = null, + pinned_bank: ?[]const u8 = null, +}; + +/// Device addresses of the resident, stacked fp16 scales ([E, D] or [E, I], E = 513). +pub const Scales = struct { suh_g: u64, suh_u: u64, svh_g: u64, svh_u: u64, suh_d: u64, svh_d: u64 }; + +pub const Options = struct { + layer: u8, + lines: usize, + dims: usize = 2560, + width: usize = 640, + /// Rows and slots a round may hold (sizes the masked-pick buffers). + max_rows: usize = 16, + slots: usize = 11, + staging: usize = 8, + /// Copy from the source bytes directly (they are pinned already); otherwise through pinned staging lines. + direct: bool = false, + codebook: cuda.exl3.Codebook = .mul1, + act_mode: i32 = 1, + swap_every: u16 = 4, + swap_cap: u16 = 96, + /// Bytes a swap round may prefetch (the copy budget); 0: no prefetch. + prefetch_bytes: u64 = 0, + /// Double buffering: a round that fits without them does not evict the previous round's lines. + double_buffer: bool = true, + /// Masked-pick buffers for this many passes of a split round (0: the round's own bound, uniq + 1). + max_passes: usize = 0, +}; + +/// Passes a round of a `lines`-line pool can split into when nothing prefetches (`prefetch_bytes` 0): pass 0 and one +pub fn passesFor(lines: usize) usize { + return 2 + (shape.experts + lines - 1) / lines; +} + +pub const Stats = struct { + rounds: u64 = 0, + passes: u64 = 0, + hit_slots: u64 = 0, + miss_slots: u64 = 0, + unique: u64 = 0, + misses: u64 = 0, // mandatory copies + prefetches: u64 = 0, + evictions: u64 = 0, + busy: u64 = 0, // plans that had to drain in-flight copies first + protected: u64 = 0, // rounds that spared the previous round's lines +}; + +pub const Pool = struct { + gpa: std.mem.Allocator, + io: std.Io, + opt: Options, + src: Source, + scales: Scales, + lines: Lines, + k2: [][3]i32, // per routed expert, from its trellis sizes + slots_tab: []u16, + line_tab: []offload.Line, + stamp_tab: []u32, + layer_tab: []u8, + v: offload.Vram, + counts: *[shape.layers][shape.experts]u32, + own_counts: bool, + round: Round, + now: u32 = 0, + round_no: u32 = 0, + picks_host: cuda.HostBuffer, + picks_dev: cuda.DeviceBuffer, + picks_ev: cuda.Event, + picks_pending: bool = false, + picks_n: usize = 0, // the planned round's picks (rows * slots) + prev: []u16, // the previous round's routed experts (double buffering) + prev_n: usize = 0, + stats: Stats = .{}, + telemetry_id: u64 = 0, + telemetry_valid: bool = true, + telemetry_device: telemetry.Device = .{}, + + /// `counts`: the family's lifetime use table shared by every layer's pool (null: the pool owns one). + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, src: Source, scales: Scales, opt: Options, counts: ?*[shape.layers][shape.experts]u32) !Pool { + if (opt.lines < min_lines) return error.PoolTooSmall; + if (src.routed.len != shape.experts) return error.Invalid; + const tiles = (opt.dims / 16) * (opt.width / 16); + var bytes: [3]usize = .{ 0, 0, 0 }; + const k2 = try gpa.alloc([3]i32, shape.experts); + errdefer gpa.free(k2); + var held: usize = 0; + for (src.routed, k2) |e, *w| { + // an expert with no source bytes: never picked here + if (e[0].len == 0 and e[1].len == 0 and e[2].len == 0) { + w.* = .{ 0, 0, 0 }; + continue; + } + held += 1; + for (0..3) |p| { + // trellis bytes = tiles * 16 * k2 (k2 half-bits a value, 8 * k2 int16 words a 16x16 tile) + if (e[p].len == 0 or e[p].len % (tiles * 16) != 0) return error.Invalid; + w[p] = @intCast(e[p].len / (tiles * 16)); + bytes[p] = @max(bytes[p], e[p].len); + } + } + if (held == 0) return error.Invalid; + const table_e: usize = if (src.shared != null) shape.experts + 1 else shape.experts; + var lines = try Lines.init(gpa, d, .{ .experts = table_e, .bytes = bytes, .lines = opt.lines, .staging = if (opt.direct) 0 else opt.staging }); + errdefer lines.deinit(); + lines.pinned_bank = if (opt.direct) src.pinned_bank else null; + if (src.shared) |sh| { + var w: [3]i32 = undefined; + for (0..3) |p| { + if (sh[p].len == 0 or sh[p].len % (tiles * 16) != 0) return error.Invalid; + w[p] = @intCast(sh[p].len / (tiles * 16)); + } + // the shared expert gets its own resident region (it may be wider than the routed experts) + try lines.setResident(shared_id, sh, w); + } + const n = opt.lines; + const slots_tab = try gpa.alloc(u16, n); + errdefer gpa.free(slots_tab); + const line_tab = try gpa.alloc(offload.Line, n); + errdefer gpa.free(line_tab); + const stamp_tab = try gpa.alloc(u32, n); + errdefer gpa.free(stamp_tab); + const layer_tab = try gpa.alloc(u8, n); + errdefer gpa.free(layer_tab); + const own = counts == null; + const c = counts orelse blk: { + const t = try gpa.create([shape.layers][shape.experts]u32); + t.* = @splat(@splat(0)); + break :blk t; + }; + errdefer if (own) gpa.destroy(c); + var rnd = try Round.init(gpa, n, opt.max_rows * opt.slots); + errdefer rnd.deinit(gpa); + const passes_n = if (opt.max_passes == 0) rnd.passes.len else @min(opt.max_passes, rnd.passes.len); + if (opt.max_passes != 0 and opt.prefetch_bytes != 0) return error.Invalid; // in-flight prefetches lengthen rounds + const pick_bytes = passes_n * opt.max_rows * opt.slots * 4; + var picks_host = try cuda.HostBuffer.alloc(d, pick_bytes); + errdefer picks_host.free(); + var picks_dev = try cuda.DeviceBuffer.alloc(d, pick_bytes); + errdefer picks_dev.free(); + var picks_ev = try cuda.Event.init(d, false); + errdefer picks_ev.deinit(); + const prev = try gpa.alloc(u16, shape.experts); + errdefer gpa.free(prev); + return .{ + .telemetry_id = try telemetry.identity(), + .prev = prev, + .gpa = gpa, + .io = io, + .opt = opt, + .src = src, + .scales = scales, + .lines = lines, + .k2 = k2, + .slots_tab = slots_tab, + .line_tab = line_tab, + .stamp_tab = stamp_tab, + .layer_tab = layer_tab, + .v = offload.Vram.init(slots_tab, line_tab, stamp_tab, layer_tab), + .counts = c, + .own_counts = own, + .round = rnd, + .picks_host = picks_host, + .picks_dev = picks_dev, + .picks_ev = picks_ev, + }; + } + + pub fn deinit(p: *Pool) void { + p.lines.deinit(); + p.picks_ev.deinit(); + p.picks_dev.free(); + p.picks_host.free(); + p.round.deinit(p.gpa); + if (p.own_counts) p.gpa.destroy(p.counts); + p.gpa.free(p.prev); + p.gpa.free(p.layer_tab); + p.gpa.free(p.stamp_tab); + p.gpa.free(p.line_tab); + p.gpa.free(p.slots_tab); + p.gpa.free(p.k2); + p.* = undefined; + } + + /// The `Experts` layer: the pool's device tables (fixed addresses) and the resident scales. + pub fn tables(p: *const Pool) exl3e.Layer { + return .{ + .gate = p.lines.table(0), + .up = p.lines.table(1), + .down = p.lines.table(2), + .gate_k2 = p.lines.widths(0), + .up_k2 = p.lines.widths(1), + .down_k2 = p.lines.widths(2), + .suh_g = p.scales.suh_g, + .suh_u = p.scales.suh_u, + .svh_g = p.scales.svh_g, + .svh_u = p.scales.svh_u, + .suh_d = p.scales.suh_d, + .svh_d = p.scales.svh_d, + .experts = p.lines.spec.experts, + .dims = p.opt.dims, + .width = p.opt.width, + .codebook = p.opt.codebook, + .act_mode = p.opt.act_mode, + }; + } + + /// Promotes prefetch lines whose copy landed (the policy's reap); never blocks. + fn reap(p: *Pool) !void { + for (p.line_tab, p.slots_tab, 0..) |state, id, i| { + if (state != .inflight) continue; + if (try p.lines.landed(i)) offload.reapAt(&p.v, p.opt.layer, id, p.now) catch unreachable; + } + } + + /// Host step of a round: `picks` [rows * slots] (int32, -1 skipped, 512 the shared expert), as the device holds + pub fn plan(p: *Pool, picks: []const i32, rows: usize) !*const Round { + return p.planInternal(picks, rows, true); + } + + fn planInternal(p: *Pool, picks: []const i32, rows: usize, observe: bool) !*const Round { + if (rows == 0 or rows > p.opt.max_rows or picks.len != rows * p.opt.slots) return error.Invalid; + p.now += 1; + p.round_no += 1; + try p.reap(); + const ropt: round_.Options = .{ + .swap_every = p.opt.swap_every, + .swap_cap = p.opt.swap_cap, + .prefetch_bytes = p.opt.prefetch_bytes, + .line_bytes = p.lines.layout.line_bytes, + .protect = if (p.opt.double_buffer) p.prev[0..p.prev_n] else &.{}, + .observe = observe, + }; + round_.plan(&p.v, p.opt.layer, picks, p.counts, p.now, p.round_no, ropt, &p.round) catch |e| switch (e) { + error.Busy => { + // every line is held by a copy in flight: let them land, then plan again (nothing was written) + p.stats.busy += 1; + try p.lines.drain(); + try p.reap(); + try round_.plan(&p.v, p.opt.layer, picks, p.counts, p.now, p.round_no, ropt, &p.round); + }, + else => return e, + }; + const r = &p.round; + if (r.n_passes * p.opt.max_rows * p.opt.slots * 4 > p.picks_dev.len) return error.PoolPasses; + p.picks_n = picks.len; + p.prev_n = r.n_ids; + @memcpy(p.prev[0..r.n_ids], r.ids[0..r.n_ids]); + if (r.protected) p.stats.protected += 1; + if (r.masked()) { + if (p.picks_pending) try p.picks_ev.synchronize(); // the previous round's uploads read these bytes + p.picks_pending = false; + const stride = p.opt.max_rows * p.opt.slots; + const host = p.picks_host.slice(i32); + var member: [shape.experts + 1]bool = undefined; + for (0..r.n_passes) |j| { + @memset(&member, false); + for (r.passIds(j)) |id| member[id] = true; + member[shared_id] = j == 0; + for (picks, host[j * stride ..][0..picks.len]) |pk, *o| { + o.* = if (pk >= 0 and pk <= shared_id and member[@intCast(pk)]) pk else -1; + } + } + } + p.stats.rounds += 1; + p.stats.passes += r.n_passes; + p.stats.hit_slots += r.hit_slots; + p.stats.miss_slots += r.miss_slots; + p.stats.unique += r.unique; + p.stats.misses += r.misses_(); + p.stats.evictions += r.evictions; + return r; + } + + fn loadCopy(p: *Pool, c: round_.Copy) !void { + if (p.k2[c.id][0] == 0) return error.Invalid; // an expert without source bytes reached the planner + const ev: ?usize = if (c.evicted == offload.Cache.empty_id) null else c.evicted; + try p.lines.load(c.line, c.id, p.src.routed[c.id], p.k2[c.id], ev, p.opt.direct, p.io); + } + + /// Device step of pass `j`: its copies and table rewrites on the copy stream, its masked picks (if any) on `s` + pub fn stage(p: *Pool, r: *const Round, j: usize, s: cuda.Stream) !void { + for (r.passCopies(j)) |c| try p.loadCopy(c); + if (r.masked()) { + const stride = p.opt.max_rows * p.opt.slots * 4; + const n = p.picks_n * 4; + try p.picks_dev.d.check(p.picks_dev.d.api.cuMemcpyHtoDAsync_v2(p.picks_dev.ptr + j * stride, p.picks_host.bytes[j * stride ..].ptr, n, s.handle), "cuMemcpyHtoDAsync"); + try p.picks_ev.record(s); + p.picks_pending = true; + } + for (r.passLines(j)) |line| try p.lines.gate(s, line); + } + + /// The picks pass `j` reads: the round's own device picks, or the pool's masked copy for a split round. + pub fn picksFor(p: *const Pool, r: *const Round, j: usize, own: u64) u64 { + if (!r.masked()) return own; + return p.picks_dev.ptr + j * p.opt.max_rows * p.opt.slots * 4; + } + + /// After pass `j` was enqueued on `s`: its lines are protected until it finishes. + pub fn release(p: *Pool, r: *const Round, j: usize, s: cuda.Stream) !void { + try p.lines.release(s, r.passLines(j)); + if (j + 1 == r.n_passes) { + for (r.prefetch()) |c| try p.loadCopy(c); + p.stats.prefetches += r.n_prefetch; + } + } + + pub fn warmProfile(pool: *Pool) !usize { + return @import("exl3_frequency_seed.zig").warm(pool); + } + + pub fn deviceBytes(p: *const Pool) usize { + return p.lines.deviceBytes() + p.picks_dev.len; + } + + pub fn lineBytes(p: *const Pool) usize { + return p.lines.layout.line_bytes; + } +}; diff --git a/zig/src/families/flashnext/exl3/offload/exl3_frequency_seed.zig b/zig/src/families/flashnext/exl3/offload/exl3_frequency_seed.zig new file mode 100644 index 0000000000..274e8aa982 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_frequency_seed.zig @@ -0,0 +1,139 @@ +const std = @import("std"); + +pub fn select(counts: []const u32, widths: []const [3]i32, output: []u16) !usize { + if (counts.len != 512 or widths.len != counts.len or output.len > counts.len) return error.InvalidProfileShape; + var order: [512]u16 = undefined; + var used: usize = 0; + for (counts, widths, 0..) |count, width, expert| { + if (count == 0 or width[0] <= 0 or width[1] <= 0 or width[2] <= 0) continue; + order[used] = @intCast(expert); + used += 1; + } + const Ranking = struct { + fn hotter(histogram: []const u32, left: u16, right: u16) bool { + if (histogram[left] != histogram[right]) return histogram[left] > histogram[right]; + return left < right; + } + }; + std.mem.sort(u16, order[0..used], counts, Ranking.hotter); + const selected = @min(output.len, used); + @memcpy(output[0..selected], order[0..selected]); + return selected; +} + +pub fn warm(pool: anytype) !usize { + const policy = @import("exl3_offload.zig"); + for (pool.line_tab, pool.slots_tab) |state, expert| { + if (state != .empty or expert != policy.Cache.empty_id) return error.ProfilePoolNotEmpty; + } + if (pool.round_no != 0) return error.ProfilePoolNotEmpty; + var selected: [512]u16 = undefined; + const count = try select(&pool.counts[pool.opt.layer], pool.k2, selected[0..@min(pool.opt.lines, selected.len)]); + for (selected[0..count], 0..) |expert, held| { + const line = try policy.stageUsedAt(&pool.v, pool.opt.layer, expert, selected[0..held], pool.counts); + try pool.lines.load(line, expert, pool.src.routed[expert], pool.k2[expert], null, pool.opt.direct, pool.io); + } + try pool.lines.drain(); + for (selected[0..count]) |expert| try policy.reapAt(&pool.v, pool.opt.layer, expert, pool.now); + return count; +} + +test "exl3: frequency seed ranks local positive experts without heating observations" { + var histogram: [512]u32 = @splat(0); + var widths: [512][3]i32 = @splat(.{ 0, 0, 0 }); + histogram[2] = 19; + histogram[3] = 19; + histogram[4] = 20; + histogram[511] = std.math.maxInt(u32); + widths[2] = .{ 1, 1, 1 }; + widths[3] = .{ 1, 1, 1 }; + widths[4] = .{ 1, 0, 1 }; + var output: [2]u16 = undefined; + const before = histogram; + try std.testing.expectEqual(@as(usize, 2), try select(&histogram, &widths, &output)); + try std.testing.expectEqualSlices(u16, &.{ 2, 3 }, &output); + try std.testing.expectEqualSlices(u32, &before, &histogram); +} + +test "exl3: frequency seed leaves spare lines for unknown and unobserved experts" { + var histogram: [512]u32 = @splat(0); + const widths: [512][3]i32 = @splat(.{ 1, 1, 1 }); + var output: [512]u16 = undefined; + try std.testing.expectEqual(@as(usize, 0), try select(&histogram, &widths, &output)); + histogram[511] = 1; + try std.testing.expectEqual(@as(usize, 1), try select(&histogram, &widths, &output)); + try std.testing.expectEqual(@as(u16, 511), output[0]); + try std.testing.expectEqual(@as(usize, 0), try select(&histogram, &widths, output[0..0])); + try std.testing.expectError(error.InvalidProfileShape, select(histogram[0..511], &widths, &output)); +} + +test "exl3: frequency warm publishes residency only after copies drain and never heats counts" { + const policy = @import("exl3_offload.zig"); + const Copies = struct { + ids: [2]u16 = undefined, + used: usize = 0, + drained: bool = false, + states: []policy.Line, + fail: bool = false, + + pub fn load(copies: *@This(), line: usize, expert: usize, _: [3][]const u8, _: [3]i32, evicted: ?usize, _: bool, _: void) !void { + try std.testing.expect(evicted == null); + try std.testing.expectEqual(policy.Line.inflight, copies.states[line]); + if (copies.fail) return error.CopyFailed; + copies.ids[copies.used] = @intCast(expert); + copies.used += 1; + } + + pub fn drain(copies: *@This()) !void { + for (copies.states) |state| try std.testing.expectEqual(policy.Line.inflight, state); + copies.drained = true; + } + }; + const Fake = struct { + line_tab: []policy.Line, + slots_tab: []u16, + round_no: u32 = 0, + counts: *[48][512]u32, + opt: struct { layer: u8 = 7, lines: usize = 2, direct: bool = true } = .{}, + k2: []const [3]i32, + v: policy.Vram, + lines: Copies, + src: struct { routed: []const [3][]const u8 }, + io: void = {}, + now: u32 = 0, + }; + const allocator = std.testing.allocator; + const counts = try allocator.create([48][512]u32); + defer allocator.destroy(counts); + counts.* = std.mem.zeroes([48][512]u32); + counts[7][2] = 100; + counts[7][4] = 200; + const before = counts[7]; + const widths: [512][3]i32 = @splat(.{ 1, 1, 1 }); + const sources: [512][3][]const u8 = @splat(.{ &.{}, &.{}, &.{} }); + var ids: [2]u16 = undefined; + var states: [2]policy.Line = undefined; + var stamps: [2]u32 = undefined; + var layers: [2]u8 = undefined; + var fake: Fake = .{ + .line_tab = &states, + .slots_tab = &ids, + .counts = counts, + .k2 = &widths, + .v = policy.Vram.init(&ids, &states, &stamps, &layers), + .lines = .{ .states = &states }, + .src = .{ .routed = &sources }, + }; + try std.testing.expectEqual(@as(usize, 2), try warm(&fake)); + try std.testing.expect(fake.lines.drained); + try std.testing.expectEqualSlices(u16, &.{ 4, 2 }, &ids); + try std.testing.expectEqualSlices(u32, &before, &counts[7]); + for (states) |state| try std.testing.expectEqual(policy.Line.resident, state); + try std.testing.expectError(error.ProfilePoolNotEmpty, warm(&fake)); + fake.v = policy.Vram.init(&ids, &states, &stamps, &layers); + fake.lines = .{ .states = &states, .fail = true }; + try std.testing.expectError(error.CopyFailed, warm(&fake)); + try std.testing.expect(!fake.lines.drained); + try std.testing.expectEqual(policy.Line.inflight, states[0]); + try std.testing.expectEqualSlices(u32, &before, &counts[7]); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_line_deltas.zig b/zig/src/families/flashnext/exl3/offload/exl3_line_deltas.zig new file mode 100644 index 0000000000..50e9eca705 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_line_deltas.zig @@ -0,0 +1,120 @@ +const std = @import("std"); + +pub const Plan = struct { + pool_bytes: u64, +}; + +pub fn apply(text: []const u8, uniform_lines: usize, decoder_pools: usize, line_bytes: []const usize, min_lines: usize, fixed_bytes: u64, startup_reservation: u64, counts: []usize) !Plan { + if (counts.len != line_bytes.len or decoder_pools > counts.len or min_lines == 0) return error.InvalidLineDeltaIndex; + if (uniform_lines < min_lines) return error.PoolTooSmall; + @memset(counts, 0); + var added: u64 = 0; + var removed: u64 = 0; + var entries = std.mem.splitScalar(u8, text, ','); + while (entries.next()) |entry| { + const colon = std.mem.indexOfScalar(u8, entry, ':') orelse return error.InvalidLineDeltas; + const index = try decimal(entry[0..colon]); + if (index >= decoder_pools) return error.InvalidLineDeltaIndex; + if (counts[index] != 0) return error.DuplicateLineDelta; + const delta = entry[colon + 1 ..]; + if (delta.len < 2 or (delta[0] != '+' and delta[0] != '-')) return error.InvalidLineDeltas; + const magnitude = try decimal(delta[1..]); + const result = if (delta[0] == '+') result: { + added = try std.math.add(u64, added, magnitude); + break :result try std.math.add(usize, uniform_lines, magnitude); + } else result: { + removed = try std.math.add(u64, removed, magnitude); + if (magnitude > uniform_lines) return error.PoolTooSmall; + break :result uniform_lines - magnitude; + }; + if (result < min_lines) return error.PoolTooSmall; + counts[index] = result; + } + if (added != removed) return error.NonZeroLineDeltas; + var before: u64 = 0; + var after: u64 = 0; + var expert_bytes: u64 = 0; + for (counts, line_bytes) |*count, bytes| { + if (count.* == 0) count.* = uniform_lines; + before = try std.math.add(u64, before, try std.math.mul(u64, uniform_lines, bytes)); + after = try std.math.add(u64, after, try std.math.mul(u64, count.*, bytes)); + expert_bytes = try std.math.add(u64, expert_bytes, try std.math.mul(u64, count.*, bytes)); + } + if (before != after) return error.LineDeltaByteBudgetChanged; + const pool_bytes = try std.math.add(u64, expert_bytes, fixed_bytes); + if (pool_bytes > startup_reservation) return error.InsufficientStartupMemory; + return .{ .pool_bytes = pool_bytes }; +} + +fn decimal(text: []const u8) !usize { + if (text.len == 0) return error.InvalidLineDeltas; + for (text) |digit| if (digit < '0' or digit > '9') return error.InvalidLineDeltas; + return std.fmt.parseInt(usize, text, 10); +} + +test "exl3: line deltas parse signed decoder ordinals and leave MTP unchanged" { + var counts: [22]usize = undefined; + const bytes: [22]usize = @splat(256); + const plan = try apply("20:+1,15:-1", 32, 21, &bytes, 10, 100, std.math.maxInt(u64), &counts); + for (counts, 0..) |count, index| try std.testing.expectEqual(@as(usize, if (index == 20) 33 else if (index == 15) 31 else 32), count); + try std.testing.expectEqual(22 * 32 * 256 + 100, plan.pool_bytes); + try std.testing.expectEqual(@as(usize, 32), counts[21]); +} + +test "exl3: line deltas reject malformed empty duplicate and out of range maps" { + var counts: [3]usize = undefined; + const bytes = [_]usize{ 256, 256, 512 }; + for ([_][]const u8{ "", "0", "0:1", "0:+", ":+1", "0:++1", "0:+1,", ",0:+1", "0:+1,,1:-1", " 0:+1,1:-1", "0:+1, 1:-1", "0:+1 ,1:-1", "-1:+1", "+0:+1", "0x0:+1", "0:+1_0", "0:+1:0" }) |text| { + try std.testing.expectError(error.InvalidLineDeltas, apply(text, 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + for ([_][]const u8{ "0:+1,0:-1", "0:+0,0:-0", "00:+0,0:+0" }) |text| { + try std.testing.expectError(error.DuplicateLineDelta, apply(text, 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + for ([_][]const u8{ "2:+0", "3:+0", "99:+0" }) |text| { + try std.testing.expectError(error.InvalidLineDeltaIndex, apply(text, 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + try std.testing.expectError(error.InvalidLineDeltaIndex, apply("0:+0", 32, 0, &bytes, 10, 0, std.math.maxInt(u64), &counts)); +} + +test "exl3: line deltas conserve lines and actual bytes independently" { + var counts: [4]usize = undefined; + const bytes = [_]usize{ 256, 512, 256, 1024 }; + for ([_][]const u8{ "0:+1", "0:-1", "0:+2,2:-1" }) |text| { + try std.testing.expectError(error.NonZeroLineDeltas, apply(text, 32, 3, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + for ([_][]const u8{ "0:+1,1:-1", "0:-1,1:+1" }) |text| { + try std.testing.expectError(error.LineDeltaByteBudgetChanged, apply(text, 32, 3, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + const expected_counts = [_]usize{ 33, 32, 31, 32 }; + const expected_bytes = 32 * (256 + 512 + 256 + 1024) + 100; + const plan = try apply("0:+1,2:-1", 32, 3, &bytes, 10, 100, expected_bytes, &counts); + try std.testing.expectEqualSlices(usize, &expected_counts, &counts); + try std.testing.expectEqual(expected_bytes, plan.pool_bytes); + try std.testing.expectError(error.InsufficientStartupMemory, apply("0:+1,2:-1", 32, 3, &bytes, 10, 100, expected_bytes - 1, &counts)); +} + +test "exl3: line deltas enforce floor before allocation and reject arithmetic overflow" { + var counts: [2]usize = undefined; + const bytes = [_]usize{ 256, 256 }; + _ = try apply("0:+22,1:-22", 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts); + try std.testing.expectEqualSlices(usize, &.{ 54, 10 }, &counts); + for ([_][]const u8{ "0:+23,1:-23", "0:+32,1:-32", "0:+33,1:-33" }) |text| { + try std.testing.expectError(error.PoolTooSmall, apply(text, 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + } + try std.testing.expectError(error.PoolTooSmall, apply("0:+0", 9, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + try std.testing.expectError(error.Overflow, apply("0:+18446744073709551615,1:-18446744073709551615", 32, 2, &bytes, 10, 0, std.math.maxInt(u64), &counts)); + try std.testing.expectError(error.Overflow, apply("0:+1,1:-1", 32, 2, &.{ std.math.maxInt(usize), std.math.maxInt(usize) }, 10, 0, std.math.maxInt(u64), &counts)); +} + +test "exl3: line deltas allow byte neutral multiway transfers within conservative startup sizing" { + var counts: [4]usize = undefined; + const bytes = [_]usize{ 256, 512, 768, 1024 }; + const fixed: u64 = 100; + const reservation: u64 = 32 * (256 + 512 + 768 + 1024) + fixed; + const plan = try apply("0:+1,1:-2,2:+1", 32, 3, &bytes, 10, fixed, reservation, &counts); + try std.testing.expectEqualSlices(usize, &.{ 33, 30, 33, 32 }, &counts); + try std.testing.expect(plan.pool_bytes <= reservation); + try std.testing.expectEqual(reservation, plan.pool_bytes); + _ = try apply("0:+0", 32, 3, &bytes, 10, fixed, reservation, &counts); + try std.testing.expectEqualSlices(usize, &.{ 32, 32, 32, 32 }, &counts); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload.zig new file mode 100644 index 0000000000..29376ad764 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload.zig @@ -0,0 +1,52 @@ +//! Flash Next EXL3 offload hub: re-exports the policy, VRAM, count, fill and scheduling modules under one name. + +const policy = @import("exl3_offload_policy.zig"); +const vram = @import("exl3_offload_vram.zig"); +const counts = @import("exl3_offload_counts.zig"); +const fill = @import("exl3_offload_fill.zig"); +const used = @import("exl3_offload_used.zig"); + +pub const shape = policy.shape; +pub const shared_id = policy.shared_id; +pub const Decision = policy.Decision; +pub const Policy = policy.Policy; +pub const Cache = policy.Cache; +pub const init = policy.init; +pub const find = policy.find; +pub const planRound = policy.planRound; +pub const swap = policy.swap; +pub const installProfile = policy.installProfile; +pub const Line = vram.Line; +pub const Vram = vram.Vram; +pub const onDevice = vram.onDevice; +pub const residentAt = vram.residentAt; +pub const touch = vram.touch; +pub const stage = vram.stage; +pub const stageKeeping = vram.stageKeeping; +pub const schedule = vram.schedule; +pub const reapAt = vram.reapAt; +pub const reap = vram.reap; +pub const vramIndexAt = vram.vramIndexAt; +pub const touchAt = vram.touchAt; +pub const lineOf = vram.lineOf; +pub const Copy = vram.Copy; +pub const lineIds = vram.lineIds; +pub const stagedCopies = vram.stagedCopies; +pub const countAt = counts.countAt; +pub const observe = counts.observe; +pub const Owned = counts.Owned; +pub const allOwned = counts.allOwned; +pub const seedHottest = fill.seedHottest; +pub const fillOwned = fill.fillOwned; +pub const fillOwnedLayers = fill.fillOwnedLayers; +pub const stageUsed = used.stageUsed; +pub const stageUsedAt = used.stageUsedAt; +pub const scheduleUsed = used.scheduleUsed; +pub const scheduleUsedAt = used.scheduleUsedAt; + +comptime { + _ = @import("exl3_offload_policy_test.zig"); + _ = @import("exl3_offload_vram_test.zig"); + _ = @import("exl3_offload_used_test.zig"); + _ = @import("exl3_offload_fill_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_counts.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_counts.zig new file mode 100644 index 0000000000..4b212ab68c --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_counts.zig @@ -0,0 +1,53 @@ +//! Flash Next EXL3 offload: lifetime routing counts and the expert-id span a rank owns. + +const std = @import("std"); +const pol = @import("exl3_offload_policy.zig"); + +const shape = pol.shape; +const appearances = pol.appearances; + +pub fn countOf(counts: []const u32, id: u16) u32 { + if (id >= counts.len) return 0; + return counts[id]; +} + +pub fn countAt(counts: *const [shape.layers][shape.experts]u32, layer: u8, id: u16) u32 { + if (layer >= shape.layers or id >= shape.experts) return 0; + return counts[layer][id]; +} + +pub const UsedOrder = struct { + counts: []const u32, + ids: []const u16, +}; + +/// Lifetime count descending, then this round's appearances, then lower id. +pub fn usedFirst(ctx: UsedOrder, a: u16, b: u16) bool { + const ca = countOf(ctx.counts, a); + const cb = countOf(ctx.counts, b); + if (ca != cb) return ca > cb; + const aa = appearances(ctx.ids, a); + const ab = appearances(ctx.ids, b); + if (aa != ab) return aa > ab; + return a < b; +} + +/// Add one to each routed id in `ids`. +pub fn observe(counts: []u32, ids: []const u16) error{Length}!void { + if (counts.len != shape.experts) return error.Length; + for (ids) |id| { + if (id >= shape.experts) continue; + counts[id] +|= 1; + } +} + +/// Expert ids this rank stores. Rank 0 is `0..256`, rank 1 is `256..512`. +pub const Owned = struct { lo: u16, hi: u16 }; + +pub fn allOwned() Owned { + return .{ .lo = 0, .hi = @intCast(shape.experts) }; +} + +pub fn inOwned(owned: Owned, id: u16) bool { + return id >= owned.lo and id < owned.hi; +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_fill.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_fill.zig new file mode 100644 index 0000000000..c81bd6df4b --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_fill.zig @@ -0,0 +1,145 @@ +//! Flash Next EXL3 offload: first fill of the VRAM lines, hottest-first and striped across layers. + +const std = @import("std"); +const pol = @import("exl3_offload_policy.zig"); +const vram_ = @import("exl3_offload_vram.zig"); +const counts_ = @import("exl3_offload_counts.zig"); + +const shape = pol.shape; +const shared_id = pol.shared_id; +const Cache = pol.Cache; +const Vram = vram_.Vram; +const stage = vram_.stage; +const vramEmpty = vram_.vramEmpty; +const vramIndex = vram_.vramIndex; +const vramIndexAt = vram_.vramIndexAt; +const Owned = counts_.Owned; +const allOwned = counts_.allOwned; + +/// Fill empty VRAM lines with the hottest counted ids. +pub fn seedHottest(v: *Vram, counts: []const u32, limit: u16) error{Length}!u16 { + return seedSpan(v, counts, limit, allOwned()); +} + +fn seedSpan(v: *Vram, counts: []const u32, limit: u16, owned: Owned) error{Length}!u16 { + if (counts.len != shape.experts) return error.Length; + var seeded: u16 = 0; + while (seeded < limit) { + if (vramEmpty(v.*) == null) break; + var best: ?u16 = null; + var best_n: u32 = 0; + var id: u16 = owned.lo; + while (id < owned.hi) : (id += 1) { + const count = counts[id]; + if (count == 0) continue; + if (vramIndex(v.*, id) != null) continue; + if (best == null or count > best_n) { + best = id; + best_n = count; + } + } + if (best == null) break; + _ = stage(v, best.?) catch break; + seeded += 1; + } + return seeded; +} + +/// Use every empty line. Hottest owned ids go in first. Lines that are still empty take the remaining owned ids. +pub fn fillOwned(v: *Vram, counts: []const u32, owned: Owned) error{Length}!u16 { + if (owned.hi > shape.experts or owned.lo > owned.hi) return error.Length; + var n: u16 = try seedSpan(v, counts, std.math.maxInt(u16), owned); + var id: u16 = owned.lo; + while (id < owned.hi) : (id += 1) { + if (vramEmpty(v.*) == null) break; + if (vramIndex(v.*, id) != null) continue; + _ = stage(v, id) catch break; + n += 1; + } + return n; +} + +fn quotaAt(lines: u16, layer: usize) u16 { + const n_layers: u16 = @intCast(shape.layers); + const base = lines / n_layers; + if (layer < @as(usize, lines % n_layers)) return base + 1; + return base; +} + +fn placeOwned(v: *Vram, layer: u8, id: u16) bool { + // Shared 512 is not a routed line. Callers only pass an owned id. + std.debug.assert(id != shared_id and id != Cache.empty_id); + const slot = vramEmpty(v.*) orelse return false; + v.slots[slot] = id; + v.layer[slot] = layer; + v.line[slot] = .inflight; + v.stamp[slot] = 0; + return true; +} + +/// Hottest owned `(layer, id)` that still has stripe room, one line at a time. +fn placeHottest( + v: *Vram, + lines: u16, + counts: []const [shape.experts]u32, + owned: Owned, + used: *[shape.layers]u16, +) u16 { + var placed: u16 = 0; + while (placed < lines and vramEmpty(v.*) != null) { + var found = false; + var best_layer: u8 = 0; + var best_id: u16 = 0; + var best_c: u32 = 0; + for (counts, 0..) |row, layer_i| { + if (used[layer_i] >= quotaAt(lines, layer_i)) continue; + var id = owned.lo; + while (id < owned.hi) : (id += 1) { + const c = row[id]; + if (c == 0) continue; + if (vramIndexAt(v.*, @intCast(layer_i), id) != null) continue; + const layer: u8 = @intCast(layer_i); + const better = !found or c > best_c or + (c == best_c and id < best_id) or + (c == best_c and id == best_id and layer < best_layer); + if (!better) continue; + found = true; + best_layer = layer; + best_id = id; + best_c = c; + } + } + if (!found) break; + if (!placeOwned(v, best_layer, best_id)) break; + used[best_layer] += 1; + placed += 1; + } + return placed; +} + +/// Id, then layer. An uncapped layer-major walk would park every line on layer 0. +fn fillStripe(v: *Vram, lines: u16, owned: Owned, used: *[shape.layers]u16, placed: u16) u16 { + var n = placed; + var id = owned.lo; + while (id < owned.hi and n < lines) : (id += 1) { + var layer: usize = 0; + while (layer < shape.layers and n < lines) : (layer += 1) { + if (used[layer] >= quotaAt(lines, layer)) continue; + if (vramIndexAt(v.*, @intCast(layer), id) != null) continue; + if (!placeOwned(v, @intCast(layer), id)) return n; + used[layer] += 1; + n += 1; + } + } + return n; +} + +/// Stage up to `lines` owned experts. +pub fn fillOwnedLayers(v: *Vram, lines: u16, counts: []const [shape.experts]u32, owned: Owned) error{Length}!u16 { + if (counts.len != shape.layers) return error.Length; + if (owned.hi > shape.experts or owned.lo > owned.hi) return error.Length; + std.debug.assert(v.layer.len == v.slots.len and v.line.len == v.slots.len); + var used: [shape.layers]u16 = @splat(0); + const hot = placeHottest(v, lines, counts, owned, &used); + return fillStripe(v, lines, owned, &used, hot); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_fill_test.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_fill_test.zig new file mode 100644 index 0000000000..d886f29332 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_fill_test.zig @@ -0,0 +1,207 @@ +//! Flash Next EXL3 offload tests: the first fill of the VRAM lines. + +const std = @import("std"); +const hub = @import("exl3_offload.zig"); +const vram_ = @import("exl3_offload_vram.zig"); + +const shape = hub.shape; +const shared_id = hub.shared_id; +const Cache = hub.Cache; +const Decision = hub.Decision; +const Policy = hub.Policy; +const Line = hub.Line; +const Vram = hub.Vram; +const Owned = hub.Owned; +const Copy = hub.Copy; +const init = hub.init; +const find = hub.find; +const planRound = hub.planRound; +const swap = hub.swap; +const installProfile = hub.installProfile; +const onDevice = hub.onDevice; +const residentAt = hub.residentAt; +const touch = hub.touch; +const stage = hub.stage; +const stageKeeping = hub.stageKeeping; +const schedule = hub.schedule; +const reap = hub.reap; +const reapAt = hub.reapAt; +const vramIndexAt = hub.vramIndexAt; +const vramIndex = vram_.vramIndex; +const observe = hub.observe; +const allOwned = hub.allOwned; +const seedHottest = hub.seedHottest; +const fillOwned = hub.fillOwned; +const fillOwnedLayers = hub.fillOwnedLayers; +const scheduleUsed = hub.scheduleUsed; +const scheduleUsedAt = hub.scheduleUsedAt; +const lineOf = hub.lineOf; +const lineIds = hub.lineIds; +const stagedCopies = hub.stagedCopies; + +fn benchPolicy(slot_n: u16, expert_bytes: u64, copy_bytes: u64) Policy { + return .{ + .slots = slot_n, + .expert_bytes = expert_bytes, + .copy_bytes = copy_bytes, + }; +} + +test "observed counts seed empty lines and a later round keeps the hot id" { + var counts: [shape.experts]u32 = @splat(0); + try observe(&counts, &.{ 1, shared_id, 1, 3 }); + try std.testing.expectEqual(@as(u32, 2), counts[1]); + try std.testing.expectEqual(@as(u32, 1), counts[3]); + try observe(&counts, &.{1}); + try std.testing.expectEqual(@as(u32, 3), counts[1]); + var short: [1]u32 = .{9}; + try std.testing.expectError(error.Length, observe(&short, &.{1})); + try std.testing.expectEqual(@as(u32, 9), short[0]); + + var seeded_counts: [shape.experts]u32 = @splat(0); + seeded_counts[3] = 10; + seeded_counts[1] = 10; + seeded_counts[8] = 2; + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + const n = try seedHottest(&vram, &seeded_counts, 2); + try std.testing.expectEqual(@as(u16, 2), n); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u16, 3), slots[1]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expect(!onDevice(vram, 1)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 8)); + try reap(&vram, 1, 1); + try std.testing.expect(onDevice(vram, 1)); + + var wide: [3]u16 = .{ Cache.empty_id, Cache.empty_id, Cache.empty_id }; + var wide_line: [3]Line = .{ .empty, .empty, .empty }; + var wide_stamp: [3]u32 = .{ 0, 0, 0 }; + var wide_layer: [3]u8 = .{ 0, 0, 0 }; + var wide_vram = Vram.init(&wide, &wide_line, &wide_stamp, &wide_layer); + try std.testing.expectError(error.Length, seedHottest(&wide_vram, counts[0..3], 1)); + try std.testing.expectEqual(Line.empty, wide_line[0]); +} + +test "a rank fills its own lines and does not copy the other rank's expert" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.experts]u32 = @splat(0); + counts[400] = 100; + counts[7] = 3; + const rank0 = Owned{ .lo = 0, .hi = 256 }; + const n = try fillOwned(&vram, &counts, rank0); + try std.testing.expectEqual(@as(u16, 2), n); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + try std.testing.expectEqual(@as(u16, 0), slots[1]); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 400)); + try std.testing.expectEqual(Line.inflight, line[0]); + try reap(&vram, 7, 1); + try reap(&vram, 0, 1); + + var decisions: [2]Decision = undefined; + var scratch: [2]u16 = undefined; + const asked = [_]u16{ 400, 9 }; + try scheduleUsed(&vram, &asked, &counts, benchPolicy(2, 100, 10_000), 4, &decisions, &scratch, rank0); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[1]); + try std.testing.expect(onDevice(vram, 7)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 400)); + try std.testing.expect(vramIndex(vram, 9) != null); +} + +test "fillOwnedLayers stripes an explicit line count and skips the other rank" { + const wide: u16 = @intCast(shape.layers + 2); + var slots: [shape.layers + 2]u16 = @splat(Cache.empty_id); + var line: [shape.layers + 2]Line = @splat(.empty); + var stamp: [shape.layers + 2]u32 = @splat(0); + var layer: [shape.layers + 2]u8 = @splat(0); + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + const rank0 = Owned{ .lo = 0, .hi = 256 }; + var short: [2][shape.experts]u32 = @splat(@splat(0)); + try std.testing.expectError(error.Length, fillOwnedLayers(&vram, wide, &short, rank0)); + try std.testing.expectError(error.Length, fillOwnedLayers(&vram, wide, &counts, .{ .lo = 4, .hi = 2 })); + try std.testing.expectEqual(Line.empty, line[0]); + + const n = try fillOwnedLayers(&vram, wide, &counts, rank0); + try std.testing.expectEqual(wide, n); + // Id, then layer: slot 1 is layer 1's lowest id, not layer 0's second id. + try std.testing.expectEqual(@as(u16, 0), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(@as(u16, 0), slots[1]); + try std.testing.expectEqual(@as(u8, 1), layer[1]); + try std.testing.expectEqual(@as(u16, 0), slots[shape.layers - 1]); + try std.testing.expectEqual(@as(u8, @intCast(shape.layers - 1)), layer[shape.layers - 1]); + try std.testing.expectEqual(@as(u16, 1), slots[shape.layers]); + try std.testing.expectEqual(@as(u8, 0), layer[shape.layers]); + try std.testing.expectEqual(@as(u16, 1), slots[shape.layers + 1]); + try std.testing.expectEqual(@as(u8, 1), layer[shape.layers + 1]); + var on_layer0: usize = 0; + for (slots, layer) |id, lay| { + try std.testing.expect(id < 256); + try std.testing.expect(id != shared_id); + if (lay == 0) on_layer0 += 1; + } + try std.testing.expectEqual(@as(usize, 2), on_layer0); + + // Fewer lines than layers: the remainder is layers 0..3, one lowest id each. + var few_slots: [4]u16 = @splat(Cache.empty_id); + var few_line: [4]Line = @splat(.empty); + var few_stamp: [4]u32 = @splat(0); + var few_layer: [4]u8 = @splat(0); + var few = Vram.init(&few_slots, &few_line, &few_stamp, &few_layer); + try std.testing.expectEqual(@as(u16, 4), try fillOwnedLayers(&few, 4, &counts, rank0)); + try std.testing.expectEqual([_]u16{ 0, 0, 0, 0 }, few_slots); + try std.testing.expectEqual([_]u8{ 0, 1, 2, 3 }, few_layer); + + // Rank 1. A hotter id outside Owned is not staged. Layer 0 keeps its hottest owned id. + var rank_slots: [4]u16 = @splat(Cache.empty_id); + var rank_line: [4]Line = @splat(.empty); + var rank_stamp: [4]u32 = @splat(0); + var rank_layer: [4]u8 = @splat(0); + var rank = Vram.init(&rank_slots, &rank_line, &rank_stamp, &rank_layer); + counts[0][1] = 999; + counts[0][300] = 5; + const rank1 = Owned{ .lo = 256, .hi = 512 }; + try std.testing.expectEqual(@as(u16, 4), try fillOwnedLayers(&rank, 4, &counts, rank1)); + try std.testing.expect(vramIndexAt(rank, 0, 300) != null); + try std.testing.expect(vramIndexAt(rank, 1, 256) != null); + try std.testing.expect(vramIndexAt(rank, 2, 256) != null); + try std.testing.expect(vramIndexAt(rank, 3, 256) != null); + try std.testing.expectEqual(@as(?usize, null), vramIndex(rank, 1)); + for (rank_slots) |id| { + try std.testing.expect(id >= 256 and id < 512); + try std.testing.expect(id != shared_id); + } +} + +test "fillOwnedLayers keeps a hot id inside the stripe and does not give its layer a second line" { + const lines: u16 = @intCast(shape.layers); + var slots: [shape.layers]u16 = @splat(Cache.empty_id); + var line: [shape.layers]Line = @splat(.empty); + var stamp: [shape.layers]u32 = @splat(0); + var layer: [shape.layers]u8 = @splat(0); + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[0][9] = 10; + counts[0][8] = 9; + const n = try fillOwnedLayers(&vram, lines, &counts, .{ .lo = 0, .hi = 256 }); + try std.testing.expectEqual(lines, n); + try std.testing.expect(vramIndexAt(vram, 0, 9) != null); + try std.testing.expectEqual(@as(?usize, null), vramIndexAt(vram, 0, 8)); + try std.testing.expect(vramIndexAt(vram, 1, 0) != null); + try std.testing.expect(vramIndexAt(vram, 40, 0) != null); + var on_layer0: usize = 0; + for (layer) |lay| { + if (lay == 0) on_layer0 += 1; + } + try std.testing.expectEqual(@as(usize, 1), on_layer0); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_policy.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_policy.zig new file mode 100644 index 0000000000..5c8dd9dc79 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_policy.zig @@ -0,0 +1,191 @@ +//! Flash Next EXL3 offload: the host slot cache and its per-round copy/cpu/hit planner. + +const std = @import("std"); +/// Flash Next's routed geometry (was exl3_shape.zig): 48 decoder layers, 512 routed experts. +pub const shape = struct { + pub const layers: usize = 48; + pub const experts: usize = 512; + pub const shared_expert: usize = 512; +}; + +pub const shared_id: u16 = @intCast(shape.shared_expert); + +pub const Decision = enum { hit, copy, cpu }; + +pub const Policy = struct { + slots: u16, + expert_bytes: u64, + /// Bytes the copy engine may move this round. copy_count = expert_bytes==0 ? 0 : copy_bytes / expert_bytes. + copy_bytes: u64, + swap_every: u16 = 4, + swap_cap: u16 = 96, +}; + +pub const Cache = struct { + slots: []u16, // expert id, or empty_id + freq: []u32, // frequency of the resident id, parallel to slots + pub const empty_id: u16 = 0xffff; +}; + +/// Clear every slot to `empty_id` and every frequency to 0. Lengths must match. +pub fn init(slots: []u16, freq: []u32) Cache { + std.debug.assert(slots.len == freq.len); + @memset(slots, Cache.empty_id); + @memset(freq, 0); + return .{ .slots = slots, .freq = freq }; +} + +/// Slot index of a resident expert, or null. `empty_id` is a vacant line, not an expert. +pub fn find(c: Cache, id: u16) ?usize { + if (id == Cache.empty_id) return null; + return indexOf(c.slots, id); +} + +/// One decision per id. Does not write the slot table. Unique misses sort by count descending. +pub fn planRound( + c: Cache, + ids: []const u16, + policy: Policy, + decisions: []Decision, + scratch: []u16, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + std.debug.assert(c.slots.len == c.freq.len); + + var n: usize = 0; + for (ids) |id| { + if (id == shared_id) continue; + if (find(c, id) != null) continue; + try pushUnique(scratch, &n, id); + } + std.mem.sort(u16, scratch[0..n], ids, hotterFirst); + + const copies = copyCount(policy, n); + for (ids, decisions) |id, *decision| { + if (id == shared_id or find(c, id) != null) { + decision.* = .hit; + continue; + } + const rank = indexOf(scratch[0..n], id).?; + decision.* = if (rank < copies) .copy else .cpu; + } +} + +/// Install on swap rounds only (`round > 0` and `round % swap_every == 0`). +pub fn swap( + c: *Cache, + requested: []const u16, + round: u32, + policy: Policy, + scratch: []u16, +) error{Scratch}!u16 { + if (policy.swap_every == 0 or round == 0 or round % @as(u32, policy.swap_every) != 0) return 0; + std.debug.assert(c.slots.len == c.freq.len); + + var n: usize = 0; + for (requested) |id| { + if (id == shared_id) continue; + if (find(c.*, id) != null) continue; + try pushUnique(scratch, &n, id); + } + std.mem.sort(u16, scratch[0..n], requested, hotterFirst); + heatResidents(c, requested); + + const cap = @min(n, @as(usize, policy.swap_cap)); + var installed: usize = 0; + for (scratch[0..cap]) |id| { + const slot = indexOf(c.slots, Cache.empty_id) orelse victim(c.*, scratch[0..installed]) orelse break; + c.slots[slot] = id; + c.freq[slot] = appearances(requested, id); + installed += 1; + } + return @intCast(installed); +} + +/// Fill slots from a routed histogram of length 512. +pub fn installProfile(c: *Cache, histogram: []const u32) void { + std.debug.assert(histogram.len == shape.experts); + std.debug.assert(c.slots.len == c.freq.len); + @memset(c.slots, Cache.empty_id); + @memset(c.freq, 0); + + const n = @min(c.slots.len, histogram.len); + var filled: usize = 0; + while (filled < n) : (filled += 1) { + var best: u16 = 0; + var best_n: u32 = 0; + var any = false; + for (histogram, 0..) |count, i| { + const id: u16 = @intCast(i); + if (indexOf(c.slots[0..filled], id) != null) continue; + // Scan is id-ascending, so the first id that reaches the max count wins ties. + if (!any or count > best_n) { + best = id; + best_n = count; + any = true; + } + } + if (!any) break; + c.slots[filled] = best; + c.freq[filled] = best_n; + } +} + +pub fn copyCount(policy: Policy, unique_misses: usize) usize { + if (policy.expert_bytes == 0 or unique_misses == 0) return 0; + const raw = policy.copy_bytes / policy.expert_bytes; + if (raw >= unique_misses) return unique_misses; + return @intCast(raw); +} + +pub fn appearances(ids: []const u16, id: u16) u32 { + var n: u32 = 0; + for (ids) |x| { + if (x == id) n +|= 1; + } + return n; +} + +/// Count descending, id ascending. `ctx` is the multiset the counts are taken from. +pub fn hotterFirst(ctx: []const u16, a: u16, b: u16) bool { + const ca = appearances(ctx, a); + const cb = appearances(ctx, b); + if (ca != cb) return ca > cb; + return a < b; +} + +pub fn indexOf(ids: []const u16, id: u16) ?usize { + for (ids, 0..) |x, i| { + if (x == id) return i; + } + return null; +} + +pub fn pushUnique(dst: []u16, n: *usize, id: u16) error{Scratch}!void { + if (indexOf(dst[0..n.*], id) != null) return; + if (n.* >= dst.len) return error.Scratch; + dst[n.*] = id; + n.* += 1; +} + +fn heatResidents(c: *Cache, requested: []const u16) void { + for (c.slots, 0..) |id, i| { + if (id == Cache.empty_id) continue; + const n = appearances(requested, id); + if (n == 0) continue; + c.freq[i] = c.freq[i] +| n; + } +} + +/// Lowest frequency, skipping lines already placed in this call. +fn victim(c: Cache, skip: []const u16) ?usize { + var best: ?usize = null; + for (c.slots, 0..) |id, i| { + if (id == Cache.empty_id) continue; + if (indexOf(skip, id) != null) continue; + if (best) |b| { + if (c.freq[i] < c.freq[b] or (c.freq[i] == c.freq[b] and i > b)) best = i; + } else best = i; + } + return best; +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_policy_test.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_policy_test.zig new file mode 100644 index 0000000000..a12152217e --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_policy_test.zig @@ -0,0 +1,359 @@ +//! Flash Next EXL3 offload tests: the slot cache planner, swap and profile install. + +const std = @import("std"); +const hub = @import("exl3_offload.zig"); +const vram_ = @import("exl3_offload_vram.zig"); + +const shape = hub.shape; +const shared_id = hub.shared_id; +const Cache = hub.Cache; +const Decision = hub.Decision; +const Policy = hub.Policy; +const Line = hub.Line; +const Vram = hub.Vram; +const Owned = hub.Owned; +const Copy = hub.Copy; +const init = hub.init; +const find = hub.find; +const planRound = hub.planRound; +const swap = hub.swap; +const installProfile = hub.installProfile; +const onDevice = hub.onDevice; +const residentAt = hub.residentAt; +const touch = hub.touch; +const stage = hub.stage; +const stageKeeping = hub.stageKeeping; +const schedule = hub.schedule; +const reap = hub.reap; +const reapAt = hub.reapAt; +const vramIndexAt = hub.vramIndexAt; +const vramIndex = vram_.vramIndex; +const observe = hub.observe; +const allOwned = hub.allOwned; +const seedHottest = hub.seedHottest; +const fillOwned = hub.fillOwned; +const fillOwnedLayers = hub.fillOwnedLayers; +const scheduleUsed = hub.scheduleUsed; +const scheduleUsedAt = hub.scheduleUsedAt; +const lineOf = hub.lineOf; +const lineIds = hub.lineIds; +const stagedCopies = hub.stagedCopies; + +fn benchPolicy(slot_n: u16, expert_bytes: u64, copy_bytes: u64) Policy { + return .{ + .slots = slot_n, + .expert_bytes = expert_bytes, + .copy_bytes = copy_bytes, + }; +} + +test "resident ids are hits and planRound does not change slots" { + var slots = [_]u16{ 4, 1, 9 }; + var freq = [_]u32{ 3, 8, 1 }; + const cache = init(&slots, &freq); + slots = .{ 4, 1, 9 }; + freq = .{ 3, 8, 1 }; + const before_slots = slots; + const before_freq = freq; + + const ids = [_]u16{ 4, shared_id, 1, 9, 4 }; + var decisions: [5]Decision = undefined; + try planRound(cache, &ids, benchPolicy(3, 100, 10_000), &decisions, &[_]u16{}); + + for (decisions) |d| try std.testing.expectEqual(Decision.hit, d); + try std.testing.expectEqual(before_slots, slots); + try std.testing.expectEqual(before_freq, freq); + try std.testing.expectEqual(@as(?usize, 1), find(cache, 1)); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); +} + +test "misses past the copy budget are cpu and slots stay put" { + var slots = [_]u16{42}; + var freq = [_]u32{6}; + const cache = init(&slots, &freq); + slots[0] = 42; + freq[0] = 6; + const before_slots = slots; + const before_freq = freq; + + // Counts: 5 x4, 1 x3, 8 x1, 3 x1. Budget 250/100 = 2, so 5 and 1 are copied. + const ids = [_]u16{ 8, 5, 1, 5, 3, 5, 1, 42, 5, 1 }; + var decisions: [ids.len]Decision = undefined; + var scratch: [8]u16 = undefined; + try planRound(cache, &ids, benchPolicy(1, 100, 250), &decisions, &scratch); + + const want = [_]Decision{ .cpu, .copy, .copy, .copy, .cpu, .copy, .copy, .hit, .copy, .copy }; + try std.testing.expectEqualSlices(Decision, &want, &decisions); + try std.testing.expectEqual(before_slots, slots); + try std.testing.expectEqual(before_freq, freq); +} + +test "duplicate miss ids share one decision" { + var slots: [1]u16 = undefined; + var freq: [1]u32 = undefined; + const cache = init(&slots, &freq); + const ids = [_]u16{ 6, 2, 6 }; + var decisions: [3]Decision = undefined; + var scratch: [4]u16 = undefined; + // One copy. 6 appears twice, 2 once, so both 6s are copy and 2 is cpu. + try planRound(cache, &ids, benchPolicy(1, 100, 100), &decisions, &scratch); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(Decision.cpu, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[2]); +} + +test "equal miss counts break ties toward the lower id" { + var slots: [0]u16 = .{}; + var freq: [0]u32 = .{}; + const cache = init(&slots, &freq); + // 10 x5, 6 x4, 2 x4. One copy would be 10; two copies are 10 and 2, not 6. + const ids = [_]u16{ 6, 10, 2, 6, 10, 2, 10, 6, 2, 10, 6, 10, 2 }; + var decisions: [ids.len]Decision = undefined; + var scratch: [4]u16 = undefined; + try planRound(cache, &ids, benchPolicy(0, 50, 100), &decisions, &scratch); + for (ids, decisions) |id, d| { + const want: Decision = if (id == 10 or id == 2) .copy else .cpu; + try std.testing.expectEqual(want, d); + } +} + +test "shared id is a hit and is never stored" { + var slots: [4]u16 = @splat(Cache.empty_id); + var freq: [4]u32 = @splat(0); + var cache = init(&slots, &freq); + slots[0] = 7; + freq[0] = 2; + + const ids = [_]u16{ shared_id, 4, shared_id }; + var decisions: [3]Decision = undefined; + var scratch: [4]u16 = undefined; + try planRound(cache, &ids, benchPolicy(4, 100, 0), &decisions, &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.cpu, decisions[1]); + try std.testing.expectEqual(Decision.hit, decisions[2]); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + + const requested = [_]u16{ shared_id, 4, shared_id, 9 }; + const installed = try swap(&cache, &requested, 4, benchPolicy(4, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 2), installed); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); + for (slots) |id| try std.testing.expect(id != shared_id); + + var hist: [shape.experts]u32 = @splat(0); + hist[4] = 3; + installProfile(&cache, &hist); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); + for (slots) |id| try std.testing.expect(id != shared_id); +} + +test "round 2 installs nothing and round 4 installs min of cap and missing" { + var slots: [8]u16 = @splat(Cache.empty_id); + var freq: [8]u32 = @splat(0); + var cache = init(&slots, &freq); + const requested = [_]u16{ 3, 1, 3, 8 }; + var scratch: [8]u16 = undefined; + var policy = benchPolicy(8, 1, 0); + policy.swap_every = 4; + + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &requested, 0, policy, &scratch)); + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &requested, 2, policy, &scratch)); + try std.testing.expectEqual(@as(u16, Cache.empty_id), slots[0]); + try std.testing.expectEqual(@as(u32, 0), freq[0]); + + // Three unique misses, under the cap. + const n = try swap(&cache, &requested, 4, policy, &scratch); + try std.testing.expectEqual(@as(u16, 3), n); + try std.testing.expectEqual(@as(u16, 3), slots[0]); + try std.testing.expectEqual(@as(u32, 2), freq[0]); + try std.testing.expectEqual(@as(u16, 1), slots[1]); + try std.testing.expectEqual(@as(u16, 8), slots[2]); + try std.testing.expectEqual(@as(u16, Cache.empty_id), slots[3]); +} + +test "a cold line is replaced by the missing id" { + var slots = [_]u16{41}; + var freq = [_]u32{0}; + var cache = init(&slots, &freq); + slots[0] = 41; + freq[0] = 0; + var scratch: [2]u16 = undefined; + const requested = [_]u16{9}; + const n = try swap(&cache, &requested, 4, benchPolicy(1, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u32, 1), freq[0]); + try std.testing.expectEqual(@as(?usize, null), find(cache, 41)); +} + +test "eviction ties take the later slot" { + var slots = [_]u16{ 30, 10, 20 }; + var freq = [_]u32{ 4, 4, 4 }; + var cache = init(&slots, &freq); + slots = .{ 30, 10, 20 }; + freq = .{ 4, 4, 4 }; + var scratch: [2]u16 = undefined; + const n = try swap(&cache, &[_]u16{7}, 4, benchPolicy(3, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + // Higher id would have dropped 30. Higher slot index drops 20. + try std.testing.expectEqual([_]u16{ 30, 10, 7 }, slots); + try std.testing.expectEqual(@as(u32, 1), freq[2]); +} + +test "a resident named this round is heated before victims are chosen" { + var slots = [_]u16{ 2, 1 }; + var freq = [_]u32{ 0, 0 }; + var cache = init(&slots, &freq); + slots = .{ 2, 1 }; + freq = .{ 0, 0 }; + var scratch: [2]u16 = undefined; + // Without the heat, both freqs stay 0 and the later slot (id 1) would be the victim. + const n = try swap(&cache, &[_]u16{ 1, 99 }, 4, benchPolicy(2, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + try std.testing.expectEqual([_]u16{ 99, 1 }, slots); + try std.testing.expectEqual(@as(u32, 1), freq[0]); + try std.testing.expectEqual(@as(u32, 1), freq[1]); +} + +test "resident frequency adds this round's request count" { + var slots = [_]u16{ 5, Cache.empty_id }; + var freq = [_]u32{ 4, 0 }; + var cache = init(&slots, &freq); + slots = .{ 5, Cache.empty_id }; + freq = .{ 4, 0 }; + var scratch: [2]u16 = undefined; + const requested = [_]u16{ 5, 5, 5, 4 }; + const n = try swap(&cache, &requested, 4, benchPolicy(2, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 1), n); + try std.testing.expectEqual(@as(u16, 5), slots[0]); + try std.testing.expectEqual(@as(u32, 7), freq[0]); + try std.testing.expectEqual(@as(u16, 4), slots[1]); + try std.testing.expectEqual(@as(u32, 1), freq[1]); +} + +test "empty lines are filled before a cold resident is evicted" { + var slots = [_]u16{ 30, Cache.empty_id, 40 }; + var freq = [_]u32{ 5, 0, 1 }; + var cache = init(&slots, &freq); + slots = .{ 30, Cache.empty_id, 40 }; + freq = .{ 5, 0, 1 }; + var scratch: [4]u16 = undefined; + const n = try swap(&cache, &[_]u16{ 7, 8 }, 4, benchPolicy(3, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 2), n); + try std.testing.expectEqual([_]u16{ 30, 7, 8 }, slots); + try std.testing.expectEqual(@as(u32, 5), freq[0]); +} + +test "a candidate placed this call is not evicted by a later one" { + var slots = [_]u16{ 1, 2, 3, 4 }; + var freq = [_]u32{ 1, 50, 2, 40 }; + var cache = init(&slots, &freq); + slots = .{ 1, 2, 3, 4 }; + freq = .{ 1, 50, 2, 40 }; + var scratch: [4]u16 = undefined; + const n = try swap(&cache, &[_]u16{ 10, 11, 12 }, 4, benchPolicy(4, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 3), n); + try std.testing.expectEqual([_]u16{ 10, 2, 11, 12 }, slots); + try std.testing.expectEqual(@as(u32, 50), freq[1]); +} + +test "installProfile keeps the hottest routed ids" { + var slots: [3]u16 = @splat(shared_id); + var freq = [_]u32{ 9, 9, 9 }; + var cache = init(&slots, &freq); + var hist: [shape.experts]u32 = @splat(0); + hist[10] = 1; + hist[4] = 7; + hist[8] = 7; + hist[2] = 3; + installProfile(&cache, &hist); + try std.testing.expectEqual([_]u16{ 4, 8, 2 }, slots); + try std.testing.expectEqual([_]u32{ 7, 7, 3 }, freq); + try std.testing.expectEqual(@as(?usize, null), find(cache, shared_id)); + + var wide_slots: [513]u16 = undefined; + var wide_freq: [513]u32 = undefined; + var wide = init(&wide_slots, &wide_freq); + var flat: [shape.experts]u32 = @splat(0); + flat[0] = 1; + installProfile(&wide, &flat); + for (wide_slots[0..shape.experts], 0..) |id, i| { + try std.testing.expectEqual(@as(u16, @intCast(i)), id); + } + try std.testing.expectEqual(Cache.empty_id, wide_slots[512]); + try std.testing.expectEqual(@as(u32, 1), wide_freq[0]); + try std.testing.expectEqual(@as(u32, 0), wide_freq[1]); + try std.testing.expectEqual(@as(?usize, null), find(wide, shared_id)); +} + +test "swap installs at most the cap even when slots and misses are larger" { + var slots: [200]u16 = @splat(Cache.empty_id); + var freq: [200]u32 = @splat(0); + var cache = init(&slots, &freq); + var requested: [200]u16 = undefined; + for (&requested, 0..) |*id, i| id.* = @intCast(i); + var scratch: [200]u16 = undefined; + const n = try swap(&cache, &requested, 4, benchPolicy(200, 1, 0), &scratch); + try std.testing.expectEqual(@as(u16, 96), n); + for (slots[0..96], 0..) |id, i| try std.testing.expectEqual(@as(u16, @intCast(i)), id); + try std.testing.expectEqual(Cache.empty_id, slots[96]); + try std.testing.expectEqual(@as(u32, 1), freq[0]); + try std.testing.expectEqual(@as(u32, 0), freq[96]); +} + +test "an empty slot table misses every routed id and installs nothing" { + var slots: [0]u16 = .{}; + var freq: [0]u32 = .{}; + var cache = init(&slots, &freq); + const ids = [_]u16{ shared_id, 4, 4, 8 }; + var decisions: [4]Decision = undefined; + var scratch: [2]u16 = undefined; + const policy = benchPolicy(0, 10, 10); + try planRound(cache, &ids, policy, &decisions, &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[2]); + try std.testing.expectEqual(Decision.cpu, decisions[3]); + + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &ids, 4, policy, &scratch)); + try std.testing.expectEqual(@as(usize, 0), slots.len); +} + +test "zero expert bytes copies nothing" { + var slots: [0]u16 = .{}; + var freq: [0]u32 = .{}; + const cache = init(&slots, &freq); + const ids = [_]u16{ 1, 2 }; + var decisions: [2]Decision = undefined; + var scratch: [2]u16 = undefined; + try planRound(cache, &ids, benchPolicy(0, 0, 1000), &decisions, &scratch); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.cpu, decisions[1]); +} + +test "a short scratch or a short decision buffer is refused without writes" { + var slots = [_]u16{1}; + var freq = [_]u32{5}; + var cache = init(&slots, &freq); + slots[0] = 1; + freq[0] = 5; + + var decisions: [1]Decision = undefined; + try std.testing.expectError(error.Length, planRound(cache, &[_]u16{ 1, 2 }, benchPolicy(1, 1, 1), &decisions, &[_]u16{})); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u32, 5), freq[0]); + + var one: [2]Decision = undefined; + try std.testing.expectError(error.Scratch, planRound(cache, &[_]u16{ 2, 3 }, benchPolicy(1, 1, 1), &one, &[_]u16{})); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + + var tiny: [1]u16 = undefined; + try std.testing.expectError(error.Scratch, swap(&cache, &[_]u16{ 2, 3 }, 4, benchPolicy(1, 1, 0), &tiny)); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u32, 5), freq[0]); + + var off = benchPolicy(1, 1, 0); + off.swap_every = 0; + var scratch: [2]u16 = undefined; + try std.testing.expectEqual(@as(u16, 0), try swap(&cache, &[_]u16{2}, 4, off, &scratch)); + try std.testing.expectEqual(@as(u32, 5), freq[0]); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_used.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_used.zig new file mode 100644 index 0000000000..602935cd22 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_used.zig @@ -0,0 +1,146 @@ +//! Flash Next EXL3 offload: VRAM scheduling ordered by lifetime counts, per layer and per rank. + +const std = @import("std"); +const pol = @import("exl3_offload_policy.zig"); +const vram_ = @import("exl3_offload_vram.zig"); +const counts_ = @import("exl3_offload_counts.zig"); + +const shape = pol.shape; +const shared_id = pol.shared_id; +const Cache = pol.Cache; +const Decision = pol.Decision; +const Policy = pol.Policy; +const indexOf = pol.indexOf; +const pushUnique = pol.pushUnique; +const copyCount = pol.copyCount; +const Vram = vram_.Vram; +const vramEmpty = vram_.vramEmpty; +const vramIndexAt = vram_.vramIndexAt; +const residentAt = vram_.residentAt; +const touchAt = vram_.touchAt; +const Owned = counts_.Owned; +const inOwned = counts_.inOwned; +const countAt = counts_.countAt; +const UsedOrder = counts_.UsedOrder; +const usedFirst = counts_.usedFirst; + +/// Lowest `counts[layer][id]`, then oldest stamp, then the higher slot index. +fn vramCold( + v: Vram, + counts: *const [shape.layers][shape.experts]u32, + keep_layer: u8, + keep_ids: []const u16, +) ?usize { + std.debug.assert(v.layer.len == v.slots.len); + var best: ?usize = null; + for (v.line, 0..) |state, i| { + if (state != .resident) continue; + if (v.layer[i] == keep_layer and indexOf(keep_ids, v.slots[i]) != null) continue; + if (best) |b| { + const ci = countAt(counts, v.layer[i], v.slots[i]); + const cb = countAt(counts, v.layer[b], v.slots[b]); + const colder = ci < cb or + (ci == cb and v.stamp[i] < v.stamp[b]) or + (ci == cb and v.stamp[i] == v.stamp[b] and i > b); + if (colder) best = i; + } else best = i; + } + return best; +} + +/// `stageKeeping`, but the victim is the coldest lifetime count. +pub fn stageUsed(v: *Vram, id: u16, keep: []const u16, counts: []const u32) error{ Busy, Shared, Length }!usize { + if (counts.len != shape.experts) return error.Length; + var table: [shape.layers][shape.experts]u32 = std.mem.zeroes([shape.layers][shape.experts]u32); + @memcpy(table[0][0..], counts); + return stageUsedAt(v, 0, id, keep, &table); +} + +pub fn stageUsedAt( + v: *Vram, + layer: u8, + id: u16, + keep_ids: []const u16, + counts: *const [shape.layers][shape.experts]u32, +) error{ Busy, Shared, Length }!usize { + if (layer >= shape.layers) return error.Length; + if (id == shared_id or id == Cache.empty_id) return error.Shared; + if (vramIndexAt(v.*, layer, id)) |i| { + if (v.line[i] == .inflight) return i; + v.line[i] = .inflight; + return i; + } + const slot = vramEmpty(v.*) orelse vramCold(v.*, counts, layer, keep_ids) orelse return error.Busy; + v.slots[slot] = id; + v.layer[slot] = layer; + v.line[slot] = .inflight; + v.stamp[slot] = 0; + return slot; +} + +/// `schedule`, ordered by lifetime counts from `observe`. +pub fn scheduleUsed( + v: *Vram, + ids: []const u16, + counts: []const u32, + policy: Policy, + now: u32, + decisions: []Decision, + scratch: []u16, + owned: Owned, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + if (counts.len != shape.experts) return error.Length; + var table: [shape.layers][shape.experts]u32 = std.mem.zeroes([shape.layers][shape.experts]u32); + @memcpy(table[0][0..], counts); + return scheduleUsedAt(v, 0, ids, &table, policy, now, decisions, scratch, owned); +} + +/// Generate path. Misses sort by `counts[layer]`. A victim scores `counts[victim.layer][victim.id]`. Keep is `(layer). +pub fn scheduleUsedAt( + v: *Vram, + layer: u8, + ids: []const u16, + counts: *const [shape.layers][shape.experts]u32, + policy: Policy, + now: u32, + decisions: []Decision, + scratch: []u16, + owned: Owned, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + if (layer >= shape.layers) return error.Length; + + var n: usize = 0; + for (ids, decisions) |id, *decision| { + if (id == shared_id) { + decision.* = .hit; + continue; + } + if (residentAt(v.*, layer, id)) { + _ = touchAt(v, layer, id, now); + decision.* = .hit; + continue; + } + decision.* = .cpu; + if (!inOwned(owned, id)) continue; + if (vramIndexAt(v.*, layer, id) != null) continue; + try pushUnique(scratch, &n, id); + } + const row: []const u32 = &counts[layer]; + std.mem.sort(u16, scratch[0..n], UsedOrder{ .counts = row, .ids = ids }, usedFirst); + + const copies = copyCount(policy, n); + var staged: usize = 0; + for (scratch[0..n]) |id| { + if (staged >= copies) break; + _ = stageUsedAt(v, layer, id, ids, counts) catch |err| switch (err) { + error.Length => return error.Length, + error.Busy, error.Shared => break, + }; + staged += 1; + for (ids, decisions) |asked, *decision| { + if (asked == id) decision.* = .copy; + } + } +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_used_test.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_used_test.zig new file mode 100644 index 0000000000..0256922ca4 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_used_test.zig @@ -0,0 +1,276 @@ +//! Flash Next EXL3 offload tests: count-ordered scheduling and line addressing. + +const std = @import("std"); +const hub = @import("exl3_offload.zig"); +const vram_ = @import("exl3_offload_vram.zig"); + +const shape = hub.shape; +const shared_id = hub.shared_id; +const Cache = hub.Cache; +const Decision = hub.Decision; +const Policy = hub.Policy; +const Line = hub.Line; +const Vram = hub.Vram; +const Owned = hub.Owned; +const Copy = hub.Copy; +const init = hub.init; +const find = hub.find; +const planRound = hub.planRound; +const swap = hub.swap; +const installProfile = hub.installProfile; +const onDevice = hub.onDevice; +const residentAt = hub.residentAt; +const touch = hub.touch; +const stage = hub.stage; +const stageKeeping = hub.stageKeeping; +const schedule = hub.schedule; +const reap = hub.reap; +const reapAt = hub.reapAt; +const vramIndexAt = hub.vramIndexAt; +const vramIndex = vram_.vramIndex; +const observe = hub.observe; +const allOwned = hub.allOwned; +const seedHottest = hub.seedHottest; +const fillOwned = hub.fillOwned; +const fillOwnedLayers = hub.fillOwnedLayers; +const scheduleUsed = hub.scheduleUsed; +const scheduleUsedAt = hub.scheduleUsedAt; +const lineOf = hub.lineOf; +const lineIds = hub.lineIds; +const stagedCopies = hub.stagedCopies; + +fn benchPolicy(slot_n: u16, expert_bytes: u64, copy_bytes: u64) Policy { + return .{ + .slots = slot_n, + .expert_bytes = expert_bytes, + .copy_bytes = copy_bytes, + }; +} + +test "a high count keeps an old stamp and a one-off is the victim" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 100 }; + var counts: [shape.experts]u32 = @splat(0); + counts[1] = 50; + counts[4] = 1; + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + + const miss = [_]u16{9}; + try scheduleUsed(&vram, &miss, &counts, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(Line.resident, line[0]); + try std.testing.expectEqual(@as(u16, 9), slots[1]); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expect(onDevice(vram, 1)); + try std.testing.expect(!onDevice(vram, 9)); +} + +test "zero lifetime counts fall back to stamp LRU" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 4, 7 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 5 }; + var counts: [shape.experts]u32 = @splat(0); + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + + const miss = [_]u16{9}; + try scheduleUsed(&vram, &miss, &counts, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(onDevice(vram, 7)); +} + +test "lifetime count beats this round's appearances" { + var slots = [_]u16{Cache.empty_id}; + var line = [_]Line{.empty}; + var stamp = [_]u32{0}; + var layer = [_]u8{0}; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.experts]u32 = @splat(0); + counts[8] = 5; + const ids = [_]u16{ 3, 3, 3, 3, 8 }; + var decisions: [ids.len]Decision = undefined; + var scratch: [2]u16 = undefined; + try scheduleUsed(&vram, &ids, &counts, benchPolicy(1, 100, 100), 1, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[4]); + try std.testing.expectEqual(@as(u16, 8), slots[0]); + try std.testing.expectEqual(Line.inflight, line[0]); +} + +test "a resident named this step is kept even when its count is low" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 100 }; + var counts: [shape.experts]u32 = @splat(0); + counts[4] = 1; + const ids = [_]u16{ 1, 4, 9 }; + var decisions: [3]Decision = undefined; + var scratch: [2]u16 = undefined; + try scheduleUsed(&vram, &ids, &counts, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.cpu, decisions[2]); + try std.testing.expect(onDevice(vram, 1)); + try std.testing.expect(onDevice(vram, 4)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 9)); +} + +test "a keep on layer 0 does not protect that id on layer 40" { + var slots = [_]u16{ 7, 7, 3 }; + var line = [_]Line{ .resident, .resident, .resident }; + var stamp = [_]u32{ 4, 4, 4 }; + var layer = [_]u8{ 0, 40, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 7, 7, 3 }; + line = .{ .resident, .resident, .resident }; + stamp = .{ 4, 4, 4 }; + layer = .{ 0, 40, 0 }; + + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[0][7] = 10; + counts[40][7] = 1; + counts[0][3] = 8; + const ids = [_]u16{ 7, 9 }; + var decisions: [2]Decision = undefined; + var scratch: [2]u16 = undefined; + try scheduleUsedAt(&vram, 0, &ids, &counts, benchPolicy(3, 100, 100), 8, &decisions, &scratch, allOwned()); + + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[1]); + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expect(residentAt(vram, 0, 3)); + try std.testing.expect(!residentAt(vram, 40, 7)); + try std.testing.expect(vramIndexAt(vram, 0, 9) != null); + try std.testing.expectEqual(@as(u32, 8), stamp[0]); + try std.testing.expectEqual(@as(u32, 4), stamp[2]); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + try std.testing.expectEqual(@as(u16, 3), slots[2]); +} + +test "scheduleUsed is the layer-0 wrapper and is not the multi-layer path" { + var slots = [_]u16{ 1, 4 }; + var line = [_]Line{ .resident, .resident }; + var stamp = [_]u32{ 1, 50 }; + var layer = [_]u8{ 0, 40 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 50 }; + layer = .{ 0, 40 }; + + // The flat vector says id 4 is hot. + var flat: [shape.experts]u32 = @splat(0); + flat[1] = 1; + flat[4] = 100; + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + const miss = [_]u16{9}; + try scheduleUsed(&vram, &miss, &flat, benchPolicy(2, 100, 100), 6, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(@as(u16, 9), slots[1]); + try std.testing.expectEqual(@as(u8, 0), layer[1]); + try std.testing.expect(!residentAt(vram, 40, 4)); + + slots = .{ 1, 4 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 50 }; + layer = .{ 0, 40 }; + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[0][1] = 1; + counts[40][4] = 100; + try scheduleUsedAt(&vram, 0, &miss, &counts, benchPolicy(2, 100, 100), 7, &decisions, &scratch, allOwned()); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(@as(u16, 4), slots[1]); + try std.testing.expectEqual(@as(u8, 40), layer[1]); + try std.testing.expect(residentAt(vram, 40, 4)); + try std.testing.expect(!residentAt(vram, 0, 1)); +} + +test "scheduleUsedAt does not stage another rank or the shared expert" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + counts[3][400] = 50; + counts[3][9] = 1; + const ids = [_]u16{ 400, shared_id, 9 }; + var decisions: [3]Decision = undefined; + var scratch: [4]u16 = undefined; + const rank0 = Owned{ .lo = 0, .hi = 256 }; + const policy = benchPolicy(2, 100, 10_000); + try std.testing.expectError(error.Length, scheduleUsedAt(&vram, @intCast(shape.layers), &ids, &counts, policy, 4, &decisions, &scratch, rank0)); + try std.testing.expectEqual(Line.empty, line[0]); + + try scheduleUsedAt(&vram, 3, &ids, &counts, policy, 4, &decisions, &scratch, rank0); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[2]); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 400)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, shared_id)); + try std.testing.expect(!residentAt(vram, 3, 9)); + try std.testing.expect(vramIndexAt(vram, 3, 9) != null); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u8, 3), layer[0]); + try std.testing.expectEqual(Line.empty, line[1]); + try std.testing.expect(slots[0] != shared_id and slots[1] != shared_id); +} + +test "lineIds maps hits to their line and everything else to -1, stagedCopies lists each staged line once" { + var slots: [3]u16 = undefined; + var line: [3]Line = undefined; + var stamp: [3]u32 = undefined; + var layer: [3]u8 = undefined; + var v = Vram.init(&slots, &line, &stamp, &layer); + var counts: [shape.layers][shape.experts]u32 = @splat(@splat(0)); + const pol = Policy{ .slots = 3, .expert_bytes = 10, .copy_bytes = 20 }; + var ids = [_]u16{ 7, 9, 9, 11, 512 }; + var dec: [5]Decision = undefined; + var scratch: [8]u16 = undefined; + try observe(&counts[2], &ids); + try scheduleUsedAt(&v, 2, &ids, &counts, pol, 1, &dec, &scratch, allOwned()); + // 3 unique misses (9 twice is hottest), budget for two copies: 9 and 7 copy, 11 cpu, 512 hit. + try std.testing.expectEqual(Decision.copy, dec[0]); + try std.testing.expectEqual(Decision.copy, dec[1]); + try std.testing.expectEqual(Decision.copy, dec[2]); + try std.testing.expectEqual(Decision.cpu, dec[3]); + try std.testing.expectEqual(Decision.hit, dec[4]); + var copies: [4]Copy = undefined; + const n = try stagedCopies(v, 2, &ids, &dec, &copies); + try std.testing.expectEqual(@as(usize, 2), n); + var ln: [5]i32 = undefined; + try lineIds(v, 2, &ids, &dec, &ln); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, -1, -1, -1 }, &ln); + try std.testing.expectError(error.Length, stagedCopies(v, 2, &ids, &dec, copies[0..1])); + for (copies[0..n]) |c| try reapAt(&v, 2, c.id, 2); + try scheduleUsedAt(&v, 2, &ids, &counts, pol, 3, &dec, &scratch, allOwned()); + try lineIds(v, 2, &ids, &dec, &ln); + try std.testing.expectEqual(@as(i32, @intCast(lineOf(v, 2, 7).?)), ln[0]); + try std.testing.expectEqual(@as(i32, @intCast(lineOf(v, 2, 9).?)), ln[1]); + try std.testing.expectEqual(ln[1], ln[2]); + try std.testing.expect(ln[3] == -1 and ln[4] == -1); + try std.testing.expectEqual(@as(?usize, null), lineOf(v, 3, 7)); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_vram.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_vram.zig new file mode 100644 index 0000000000..e20bd2c4a0 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_vram.zig @@ -0,0 +1,216 @@ +//! Flash Next EXL3 offload: the VRAM line table, its LRU scheduling, and the line addressing the device pool reads. + +const std = @import("std"); +const pol = @import("exl3_offload_policy.zig"); + +const Cache = pol.Cache; +const Decision = pol.Decision; +const Policy = pol.Policy; +const shared_id = pol.shared_id; +const indexOf = pol.indexOf; +const pushUnique = pol.pushUnique; +const hotterFirst = pol.hotterFirst; +const copyCount = pol.copyCount; + +pub const Line = enum { empty, inflight, resident }; + +/// VRAM working set. Host RAM still holds every routed expert. `stamp` is the last reap or touch time. A higher stamp. +pub const Vram = struct { + slots: []u16, + line: []Line, + stamp: []u32, + /// Layer of `slots[i]`. Parallel to the other three slices. + layer: []u8, + + pub fn init(slots: []u16, line: []Line, stamp: []u32, layer: []u8) Vram { + std.debug.assert(slots.len == line.len and slots.len == stamp.len and slots.len == layer.len); + @memset(slots, Cache.empty_id); + @memset(line, .empty); + @memset(stamp, 0); + @memset(layer, 0); + return .{ .slots = slots, .line = line, .stamp = stamp, .layer = layer }; + } +}; + +/// True once `id` has been reaped into VRAM. Id-only: the first matching line wins. +pub fn onDevice(v: Vram, id: u16) bool { + const i = vramIndex(v, id) orelse return false; + return v.line[i] == .resident; +} + +/// True when this layer's copy of `id` is resident. Another layer's copy does not count. +pub fn residentAt(v: Vram, layer: u8, id: u16) bool { + const i = vramIndexAt(v, layer, id) orelse return false; + return v.line[i] == .resident; +} + +/// A resident hit. Moves that line to `now` so LRU keeps it. +pub fn touch(v: *Vram, id: u16, now: u32) bool { + const i = vramIndex(v.*, id) orelse return false; + if (v.line[i] != .resident) return false; + v.stamp[i] = now; + return true; +} + +/// Start a host-to-VRAM copy of `id`. +pub fn stage(v: *Vram, id: u16) error{ Busy, Shared }!usize { + return stageKeeping(v, id, &.{}); +} + +/// Same as `stage`, except a resident whose id is in `keep` is not a victim. +pub fn stageKeeping(v: *Vram, id: u16, keep: []const u16) error{ Busy, Shared }!usize { + if (id == shared_id or id == Cache.empty_id) return error.Shared; + if (vramIndex(v.*, id)) |i| { + if (v.line[i] == .inflight) return i; + v.line[i] = .inflight; + return i; + } + const slot = vramEmpty(v.*) orelse vramLruKeep(v.*, keep) orelse return error.Busy; + v.slots[slot] = id; + v.layer[slot] = 0; + v.line[slot] = .inflight; + v.stamp[slot] = 0; + return slot; +} + +/// One step of the VRAM cache. +pub fn schedule( + v: *Vram, + ids: []const u16, + policy: Policy, + now: u32, + decisions: []Decision, + scratch: []u16, +) error{ Length, Scratch }!void { + if (decisions.len != ids.len) return error.Length; + + var n: usize = 0; + for (ids, decisions) |id, *decision| { + if (id == shared_id) { + decision.* = .hit; + continue; + } + if (onDevice(v.*, id)) { + _ = touch(v, id, now); + decision.* = .hit; + continue; + } + decision.* = .cpu; + if (vramIndex(v.*, id) != null) continue; + try pushUnique(scratch, &n, id); + } + std.mem.sort(u16, scratch[0..n], ids, hotterFirst); + + const copies = copyCount(policy, n); + var staged: usize = 0; + for (scratch[0..n]) |id| { + if (staged >= copies) break; + _ = stageKeeping(v, id, ids) catch break; + staged += 1; + for (ids, decisions) |asked, *decision| { + if (asked == id) decision.* = .copy; + } + } +} + +/// The copy of `(layer, id)` has landed. +pub fn reapAt(v: *Vram, layer: u8, id: u16, now: u32) error{NotInflight}!void { + const i = vramIndexAt(v.*, layer, id) orelse return error.NotInflight; + if (v.line[i] != .inflight) return error.NotInflight; + v.line[i] = .resident; + v.stamp[i] = now; +} + +/// Layer-0 wrapper. Does not promote another layer's inflight copy. Old tests are layer 0. Generate calls `reapAt`. +pub fn reap(v: *Vram, id: u16, now: u32) error{NotInflight}!void { + return reapAt(v, 0, id, now); +} + +/// First occupied slot with this id, on any layer. +pub fn vramIndex(v: Vram, id: u16) ?usize { + if (id == Cache.empty_id) return null; + for (v.slots, 0..) |slot, i| { + if (slot == id and v.line[i] != .empty) return i; + } + return null; +} + +/// Slot whose layer and id both match. The id-only search aliases two layers. +pub fn vramIndexAt(v: Vram, layer: u8, id: u16) ?usize { + if (id == Cache.empty_id) return null; + std.debug.assert(v.layer.len == v.slots.len and v.line.len == v.slots.len); + for (v.slots, v.layer, 0..) |slot, lay, i| { + if (slot == id and lay == layer and v.line[i] != .empty) return i; + } + return null; +} + +pub fn touchAt(v: *Vram, layer: u8, id: u16, now: u32) bool { + const i = vramIndexAt(v.*, layer, id) orelse return false; + if (v.line[i] != .resident) return false; + v.stamp[i] = now; + return true; +} + +pub fn vramEmpty(v: Vram) ?usize { + for (v.line, 0..) |state, i| if (state == .empty) return i; + return null; +} + +/// Oldest resident stamp. Equal stamps take the higher slot index. +fn vramLru(v: Vram) ?usize { + return vramLruKeep(v, &.{}); +} + +/// Oldest resident whose id is not in `keep`. Equal stamps take the higher slot index. +fn vramLruKeep(v: Vram, keep: []const u16) ?usize { + var best: ?usize = null; + for (v.line, 0..) |state, i| { + if (state != .resident) continue; + if (indexOf(keep, v.slots[i]) != null) continue; + if (best) |b| { + if (v.stamp[i] < v.stamp[b] or (v.stamp[i] == v.stamp[b] and i > b)) best = i; + } else best = i; + } + return best; +} + +/// Line index of `(layer, id)` in any state, or null. The device pool addresses its buffers by this index. +pub fn lineOf(v: Vram, layer: u8, id: u16) ?usize { + return vramIndexAt(v, layer, id); +} + +/// A host-to-device copy the policy asked for: expert `id` into pool line `line`. +pub const Copy = struct { id: u16, line: u32 }; + +/// Per pick, the pool line the hit kernel reads, or -1 when the kernel must skip the slot (a `.copy` or `.cpu` +pub fn lineIds(v: Vram, layer: u8, ids: []const u16, decisions: []const Decision, out: []i32) error{Length}!void { + if (decisions.len != ids.len or out.len != ids.len) return error.Length; + for (ids, decisions, out) |id, d, *o| { + o.* = -1; + if (d != .hit or id == shared_id) continue; + const i = vramIndexAt(v, layer, id) orelse continue; + if (v.line[i] != .resident) continue; + o.* = @intCast(i); + } +} + +/// Unique `(id, line)` pairs the round staged: a `.copy` decision whose line is in flight. +pub fn stagedCopies(v: Vram, layer: u8, ids: []const u16, decisions: []const Decision, out: []Copy) error{Length}!usize { + if (decisions.len != ids.len) return error.Length; + var n: usize = 0; + for (ids, decisions) |id, d| { + if (d != .copy) continue; + const i = vramIndexAt(v, layer, id) orelse continue; + if (v.line[i] != .inflight) continue; + var seen = false; + for (out[0..n]) |c| { + if (c.id == id) seen = true; + } + if (seen) continue; + if (n >= out.len) return error.Length; + out[n] = .{ .id = id, .line = @intCast(i) }; + n += 1; + } + return n; +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_offload_vram_test.zig b/zig/src/families/flashnext/exl3/offload/exl3_offload_vram_test.zig new file mode 100644 index 0000000000..8a1d4f8f87 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_offload_vram_test.zig @@ -0,0 +1,246 @@ +//! Flash Next EXL3 offload tests: the VRAM line states, LRU and layer tagging. + +const std = @import("std"); +const hub = @import("exl3_offload.zig"); +const vram_ = @import("exl3_offload_vram.zig"); + +const shape = hub.shape; +const shared_id = hub.shared_id; +const Cache = hub.Cache; +const Decision = hub.Decision; +const Policy = hub.Policy; +const Line = hub.Line; +const Vram = hub.Vram; +const Owned = hub.Owned; +const Copy = hub.Copy; +const init = hub.init; +const find = hub.find; +const planRound = hub.planRound; +const swap = hub.swap; +const installProfile = hub.installProfile; +const onDevice = hub.onDevice; +const residentAt = hub.residentAt; +const touch = hub.touch; +const stage = hub.stage; +const stageKeeping = hub.stageKeeping; +const schedule = hub.schedule; +const reap = hub.reap; +const reapAt = hub.reapAt; +const vramIndexAt = hub.vramIndexAt; +const vramIndex = vram_.vramIndex; +const observe = hub.observe; +const allOwned = hub.allOwned; +const seedHottest = hub.seedHottest; +const fillOwned = hub.fillOwned; +const fillOwnedLayers = hub.fillOwnedLayers; +const scheduleUsed = hub.scheduleUsed; +const scheduleUsedAt = hub.scheduleUsedAt; +const lineOf = hub.lineOf; +const lineIds = hub.lineIds; +const stagedCopies = hub.stagedCopies; + +fn benchPolicy(slot_n: u16, expert_bytes: u64, copy_bytes: u64) Policy { + return .{ + .slots = slot_n, + .expert_bytes = expert_bytes, + .copy_bytes = copy_bytes, + }; +} + +test "a staged copy is not on device until reap, and LRU drops the older resident" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + + try std.testing.expectEqual(@as(usize, 0), try stage(&vram, 4)); + try std.testing.expectEqual(@as(usize, 1), try stage(&vram, 7)); + try std.testing.expect(!onDevice(vram, 4)); + try std.testing.expectError(error.NotInflight, reap(&vram, 9, 1)); + try reap(&vram, 4, 1); + try reap(&vram, 7, 2); + try std.testing.expect(onDevice(vram, 4)); + try std.testing.expect(touch(&vram, 7, 5)); + try std.testing.expect(onDevice(vram, 4)); + + // 4 was reaped at 1 and not touched again. 7 was touched at 5. 4 is the victim. + try std.testing.expectEqual(@as(usize, 0), try stage(&vram, 8)); + try std.testing.expect(!onDevice(vram, 8)); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expectEqual(@as(u16, 7), slots[1]); + try reap(&vram, 8, 6); + try std.testing.expect(onDevice(vram, 8)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 4)); +} + +test "an in-flight line is not reused, and a full in-flight set is busy" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + _ = try stage(&vram, 1); + _ = try stage(&vram, 2); + try std.testing.expectError(error.Busy, stage(&vram, 3)); + try std.testing.expectEqual(@as(u16, 1), slots[0]); + try std.testing.expectEqual(@as(u16, 2), slots[1]); + try std.testing.expectError(error.Shared, stage(&vram, shared_id)); +} + +test "schedule touches residents, stages the hottest miss, and will not evict a hit" { + var slots = [_]u16{ Cache.empty_id, Cache.empty_id }; + var line = [_]Line{ .empty, .empty }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + var decisions: [4]Decision = undefined; + var scratch: [4]u16 = undefined; + const policy = benchPolicy(2, 100, 100); + + const first = [_]u16{ 7, shared_id, 4, 7 }; + try schedule(&vram, &first, policy, 3, &decisions, &scratch); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(Decision.copy, decisions[3]); + try std.testing.expectEqual(Decision.cpu, decisions[2]); + try std.testing.expect(!onDevice(vram, 7)); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expectEqual(@as(u16, 7), slots[0]); + try std.testing.expectEqual(Line.empty, line[1]); + + try reap(&vram, 7, 4); + const held = [_]u16{7}; + try schedule(&vram, &held, policy, 9, decisions[0..1], &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(@as(u32, 9), stamp[0]); + + // Both lines are residents this step asked for, so the miss cannot take either. + _ = try stage(&vram, 3); + try reap(&vram, 3, 10); + const protected = [_]u16{ 7, 3, 8 }; + try schedule(&vram, &protected, benchPolicy(2, 100, 10_000), 11, decisions[0..3], &scratch); + try std.testing.expectEqual(Decision.hit, decisions[0]); + try std.testing.expectEqual(Decision.hit, decisions[1]); + try std.testing.expectEqual(Decision.cpu, decisions[2]); + try std.testing.expect(onDevice(vram, 7)); + try std.testing.expect(onDevice(vram, 3)); + try std.testing.expectEqual(@as(?usize, null), vramIndex(vram, 8)); +} + +test "a zero copy budget stages nothing, and an older resident is the one that moves" { + var slots = [_]u16{ 4, 7 }; + var line = [_]Line{ .resident, .resident }; + var stamp = [_]u32{ 1, 5 }; + var layer = [_]u8{ 0, 0 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 4, 7 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 5 }; + var decisions: [1]Decision = undefined; + var scratch: [2]u16 = undefined; + + const miss = [_]u16{9}; + try schedule(&vram, &miss, benchPolicy(2, 0, 10_000), 6, &decisions, &scratch); + try std.testing.expectEqual(Decision.cpu, decisions[0]); + try std.testing.expectEqual(@as(u16, 4), slots[0]); + try std.testing.expectEqual(Line.resident, line[0]); + + try schedule(&vram, &miss, benchPolicy(2, 100, 100), 6, &decisions, &scratch); + try std.testing.expectEqual(Decision.copy, decisions[0]); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(onDevice(vram, 7)); + try std.testing.expect(!onDevice(vram, 9)); +} + +test "residentAt distinguishes the same expert id on two layers" { + var slots = [_]u16{ 7, 7 }; + var line = [_]Line{ .resident, .resident }; + var stamp = [_]u32{ 1, 2 }; + var layer = [_]u8{ 0, 40 }; + const vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 7, 7 }; + line = .{ .resident, .resident }; + stamp = .{ 1, 2 }; + layer = .{ 0, 40 }; + + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expect(!residentAt(vram, 1, 7)); + try std.testing.expect(!residentAt(vram, 0, 8)); + // Id-only lookup still aliases. Generate does not call it on this table. + try std.testing.expect(onDevice(vram, 7)); + + line[1] = .inflight; + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expect(!residentAt(vram, 40, 7)); +} + +test "reapAt promotes one inflight layer and leaves the other inflight" { + var slots = [_]u16{ 7, 7 }; + var line = [_]Line{ .inflight, .inflight }; + var stamp = [_]u32{ 0, 0 }; + var layer = [_]u8{ 0, 40 }; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{ 7, 7 }; + line = .{ .inflight, .inflight }; + stamp = .{ 0, 0 }; + layer = .{ 0, 40 }; + + try reapAt(&vram, 0, 7, 3); + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expect(!residentAt(vram, 40, 7)); + try std.testing.expectEqual(@as(u32, 3), stamp[0]); + + // The other layer's copy is still inflight, then promotable on its own reap. + try reapAt(&vram, 40, 7, 4); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expect(residentAt(vram, 0, 7)); + try std.testing.expectEqual(@as(u32, 4), stamp[1]); + try std.testing.expectEqual(@as(u32, 3), stamp[0]); + + // Layer 0 is already resident. + line[1] = .inflight; + stamp[1] = 0; + try std.testing.expectError(error.NotInflight, reap(&vram, 7, 9)); + try std.testing.expectEqual(Line.resident, line[0]); + try std.testing.expectEqual(Line.inflight, line[1]); + try std.testing.expectEqual(@as(u32, 3), stamp[0]); + try reapAt(&vram, 40, 7, 9); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expectEqual(@as(u32, 9), stamp[1]); + + // An inflight copy that exists only on layer 40 is not promoted by `reap`. + slots = .{ 7, Cache.empty_id }; + line = .{ .inflight, .empty }; + stamp = .{ 0, 0 }; + layer = .{ 40, 0 }; + try std.testing.expectError(error.NotInflight, reap(&vram, 7, 1)); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(!residentAt(vram, 40, 7)); + try reapAt(&vram, 40, 7, 2); + try std.testing.expect(residentAt(vram, 40, 7)); + try std.testing.expectEqual(@as(u8, 40), layer[0]); +} + +test "stageKeeping tags a newly assigned line as layer 0" { + var slots = [_]u16{4}; + var line = [_]Line{.resident}; + var stamp = [_]u32{1}; + var layer = [_]u8{40}; + var vram = Vram.init(&slots, &line, &stamp, &layer); + slots = .{4}; + line = .{.resident}; + stamp = .{1}; + layer = .{40}; + + try std.testing.expectEqual(@as(usize, 0), try stageKeeping(&vram, 9, &.{})); + try std.testing.expectEqual(@as(u16, 9), slots[0]); + try std.testing.expectEqual(@as(u8, 0), layer[0]); + try std.testing.expectEqual(Line.inflight, line[0]); + try std.testing.expect(vramIndexAt(vram, 0, 9) != null); + try std.testing.expectEqual(@as(?usize, null), vramIndexAt(vram, 40, 9)); + try std.testing.expect(slots[0] != shared_id); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_profile_line_budget.zig b/zig/src/families/flashnext/exl3/offload/exl3_profile_line_budget.zig new file mode 100644 index 0000000000..16ba9055e9 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_profile_line_budget.zig @@ -0,0 +1,468 @@ +const std = @import("std"); + +pub const heuristic = "Greedy marginal histogram coverage per line byte, not real miss-cost optimization. Equal ratios prefer fewer lines, then the lower decoder ordinal. Exact line completion uses bounded search, potentially exponential with many distinct costs. Candidates must also preserve exact device picks bytes; otherwise sizing stays uniform."; +pub const experts: usize = 512; + +pub const Request = struct { + counts: []const []const u32, + rank: usize, + world: usize, + line_bytes: []const usize, + uniform_lines: usize, + min_lines: usize, + max_rows: usize, + slots: usize, +}; + +pub const Budget = struct { + baseline_bytes: u64, + allocated_bytes: u64, + uniform_fallback: bool = false, +}; + +pub const Error = std.mem.Allocator.Error || error{ InvalidDimensions, InvalidRank, InvalidLineCost, InvalidBounds, Overflow, UnfillableBudget }; + +fn poolBytes(line_bytes: usize, lines: usize, max_rows: usize, slots: usize, min_lines: usize) Error!u64 { + if (lines == 0) return error.InvalidBounds; + if (min_lines == 0) return error.InvalidBounds; + if (max_rows == 0 or slots == 0) return error.InvalidDimensions; + const picks = try std.math.mul(usize, max_rows, slots); + const max_passes = 2 + experts / min_lines + @intFromBool(experts % min_lines != 0); + const passes = @min(max_passes, @min(picks, experts) + 1); + const picks_bytes = try std.math.mul(usize, try std.math.mul(usize, passes, picks), @sizeOf(i32)); + return std.math.add(u64, try std.math.mul(u64, lines, line_bytes), picks_bytes); +} + +const Group = struct { + bytes: u64, + available: usize, + + fn descending(_: void, left: Group, right: Group) bool { + return left.bytes > right.bytes; + } +}; + +pub fn allocate(gpa: std.mem.Allocator, request: Request, output: []usize) Error!Budget { + if ((request.world != 1 and request.world != 2) or request.rank >= request.world) return error.InvalidRank; + const local_experts = experts / request.world; + if (request.min_lines == 0 or request.min_lines > request.uniform_lines or request.uniform_lines > local_experts) return error.InvalidBounds; + if (request.counts.len == 0 or request.counts.len != request.line_bytes.len or output.len != request.counts.len) return error.InvalidDimensions; + var baseline_bytes: u64 = 0; + var line_budget_bytes: u64 = 0; + var floor_bytes: u64 = 0; + for (request.counts, request.line_bytes) |histogram, bytes| { + if (histogram.len != experts) return error.InvalidDimensions; + if (bytes == 0) return error.InvalidLineCost; + baseline_bytes = try std.math.add(u64, baseline_bytes, try poolBytes(bytes, request.uniform_lines, request.max_rows, request.slots, request.min_lines)); + line_budget_bytes = try std.math.add(u64, line_budget_bytes, try std.math.mul(u64, request.uniform_lines, bytes)); + floor_bytes = try std.math.add(u64, floor_bytes, try std.math.mul(u64, request.min_lines, bytes)); + } + const ranked = try gpa.alloc([experts]u32, request.counts.len); + defer gpa.free(ranked); + const lines = try gpa.alloc(usize, request.counts.len); + defer gpa.free(lines); + const storage = try gpa.alloc(Group, request.counts.len); + defer gpa.free(storage); + const layer_groups = try gpa.alloc(usize, request.counts.len); + defer gpa.free(layer_groups); + const attempted = try gpa.alloc(bool, request.counts.len); + defer gpa.free(attempted); + const suffix_length = try std.math.add(usize, request.counts.len, 1); + const suffix_bytes = try gpa.alloc(u64, suffix_length); + defer gpa.free(suffix_bytes); + const suffix_gcd = try gpa.alloc(u64, suffix_length); + defer gpa.free(suffix_gcd); + const first = request.rank * local_experts; + for (request.counts, ranked) |histogram, *ordered| { + @memcpy(ordered[0..local_experts], histogram[first..][0..local_experts]); + std.mem.sort(u32, ordered[0..local_experts], {}, std.sort.desc(u32)); + } + @memset(lines, request.min_lines); + var group_count: usize = 0; + for (request.line_bytes) |bytes| { + var found: ?usize = null; + for (storage[0..group_count], 0..) |group, index| { + if (group.bytes == bytes) { + found = index; + break; + } + } + const group_index = found orelse new: { + storage[group_count] = .{ .bytes = bytes, .available = 0 }; + group_count += 1; + break :new group_count - 1; + }; + storage[group_index].available = try std.math.add(usize, storage[group_index].available, local_experts - request.min_lines); + } + const groups = storage[0..group_count]; + std.mem.sort(Group, groups, {}, Group.descending); + for (request.line_bytes, layer_groups) |bytes, *group_index| { + for (groups, 0..) |group, index| { + if (group.bytes == bytes) { + group_index.* = index; + break; + } + } + } + var remaining = line_budget_bytes - floor_bytes; + while (remaining != 0) { + @memset(attempted, false); + var selected: ?usize = null; + for (0..lines.len) |_| { + var best: ?usize = null; + for (lines, request.line_bytes, 0..) |count, bytes, layer| { + if (attempted[layer] or count == local_experts or bytes > remaining) continue; + if (best == null or better(layer, best.?, ranked, lines, request.line_bytes)) best = layer; + } + const candidate = best orelse break; + attempted[candidate] = true; + const group_index = layer_groups[candidate]; + groups[group_index].available -= 1; + const possible = canFill(groups, remaining - request.line_bytes[candidate], suffix_bytes, suffix_gcd); + groups[group_index].available += 1; + if (possible) { + selected = candidate; + break; + } + } + const layer = selected orelse return error.UnfillableBudget; + lines[layer] += 1; + groups[layer_groups[layer]].available -= 1; + remaining -= request.line_bytes[layer]; + } + var allocated_bytes: u64 = 0; + for (lines, request.line_bytes) |count, bytes| allocated_bytes = try std.math.add(u64, allocated_bytes, try poolBytes(bytes, count, request.max_rows, request.slots, request.min_lines)); + const uniform_fallback = allocated_bytes != baseline_bytes; + if (uniform_fallback) { + @memset(lines, request.uniform_lines); + allocated_bytes = baseline_bytes; + } + @memcpy(output, lines); + return .{ .baseline_bytes = baseline_bytes, .allocated_bytes = allocated_bytes, .uniform_fallback = uniform_fallback }; +} + +fn better(left: usize, right: usize, ranked: []const [experts]u32, lines: []const usize, bytes: []const usize) bool { + const left_score = @as(u128, ranked[left][lines[left]]) * bytes[right]; + const right_score = @as(u128, ranked[right][lines[right]]) * bytes[left]; + if (left_score != right_score) return left_score > right_score; + if (lines[left] != lines[right]) return lines[left] < lines[right]; + return left < right; +} + +fn gcd(left: u64, right: u64) u64 { + var numerator = left; + var denominator = right; + while (denominator != 0) { + const remainder = numerator % denominator; + numerator = denominator; + denominator = remainder; + } + return numerator; +} + +fn canFill(groups: []const Group, target: u64, suffix_bytes: []u64, suffix_gcd: []u64) bool { + suffix_bytes[groups.len] = 0; + suffix_gcd[groups.len] = 0; + var index = groups.len; + while (index != 0) { + index -= 1; + const group = groups[index]; + const capacity = @as(u128, group.bytes) * group.available + suffix_bytes[index + 1]; + suffix_bytes[index] = @intCast(@min(capacity, target)); + suffix_gcd[index] = if (group.available == 0) suffix_gcd[index + 1] else gcd(group.bytes, suffix_gcd[index + 1]); + } + return fill(groups, 0, target, suffix_bytes, suffix_gcd); +} + +fn fill(groups: []const Group, index: usize, target: u64, suffix_bytes: []const u64, suffix_gcd: []const u64) bool { + if (target == 0) return true; + if (index == groups.len or target > suffix_bytes[index] or suffix_gcd[index] == 0 or target % suffix_gcd[index] != 0) return false; + const group = groups[index]; + if (group.available == 0) return fill(groups, index + 1, target, suffix_bytes, suffix_gcd); + const max_take = @min(group.available, target / group.bytes); + const required = target - @min(target, suffix_bytes[index + 1]); + const min_take = required / group.bytes + @intFromBool(required % group.bytes != 0); + if (min_take > max_take) return false; + var take = max_take; + while (true) { + if (fill(groups, index + 1, target - take * group.bytes, suffix_bytes, suffix_gcd)) return true; + if (take == min_take) break; + take -= 1; + } + return false; +} + +fn testRequest(counts: []const []const u32, bytes: []const usize, uniform: usize, minimum: usize) Request { + return .{ .counts = counts, .rank = 0, .world = 2, .line_bytes = bytes, .uniform_lines = uniform, .min_lines = minimum, .max_rows = 1, .slots = 1 }; +} + +test "profile budget: sorted marginal counts favor hot layers with exact bytes" { + var hot: [experts]u32 = @splat(0); + const cold: [experts]u32 = @splat(1); + @memset(hot[0..6], 100); + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 4, 6 }, 3, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); + try std.testing.expectEqual(@as(u64, 30 + 2 * 8), plan.baseline_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); + hot[0] = 0; + hot[255] = 100; + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 4, 6 }, 3, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); +} + +test "profile budget: rank filtering handles asymmetric counts and TP1" { + var first: [experts]u32 = @splat(0); + var second: [experts]u32 = @splat(0); + @memset(first[0..4], 10); + @memset(second[257..260], 20); + second[511] = 20; + var output: [2]usize = undefined; + var request = testRequest(&.{ &first, &second }, &.{ 8, 8 }, 2, 1); + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 3, 1 }, &output); + request.rank = 1; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 1, 3 }, &output); + request.rank = 0; + request.world = 1; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 1, 3 }, &output); +} + +test "profile budget: zero and flat histograms retain uniform sizing deterministically" { + for ([_]u32{ 0, 7, std.math.maxInt(u32) }) |value| { + const histogram: [experts]u32 = @splat(value); + var output: [3]usize = undefined; + for (0..3) |_| { + _ = try allocate(std.testing.allocator, testRequest(&.{ &histogram, &histogram, &histogram }, &.{ 16, 16, 16 }, 17, 10), &output); + try std.testing.expectEqualSlices(usize, &.{ 17, 17, 17 }, &output); + } + } +} + +test "profile budget: tied ratios prefer fewer lines then lower decoder ordinal" { + var first: [experts]u32 = @splat(0); + var second: [experts]u32 = @splat(0); + const third: [experts]u32 = @splat(0); + @memset(first[0..3], 10); + @memset(second[0..3], 10); + var output: [3]usize = undefined; + _ = try allocate(std.testing.allocator, testRequest(&.{ &first, &second, &third }, &.{ 8, 8, 8 }, 2, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 3, 2, 1 }, &output); +} + +test "profile budget: ratios use wide products and avoid greedy dead ends" { + const hot: [experts]u32 = @splat(std.math.maxInt(u32) - 1); + const cold: [experts]u32 = @splat(std.math.maxInt(u32)); + const unit: usize = @as(usize, 1) << 32; + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 2 * unit, 3 * unit }, 3, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); + try std.testing.expectEqual(@as(u64, 15) * unit + 2 * 8, plan.allocated_bytes); + const low: [experts]u32 = @splat(1); + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &low }, &.{ 4, 6 }, 2, 1), &output); + try std.testing.expectEqualSlices(usize, &.{ 2, 2 }, &output); +} + +test "profile budget: floors caps and single layer budgets" { + const hot: [experts]u32 = @splat(100); + const cold: [experts]u32 = @splat(0); + var output: [2]usize = undefined; + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 8, 8 }, 200, 10), &output); + try std.testing.expectEqualSlices(usize, &.{ 256, 144 }, &output); + var request = testRequest(&.{ &hot, &cold }, &.{ 8, 8 }, 400, 10); + request.world = 1; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 512, 288 }, &output); + _ = try allocate(std.testing.allocator, testRequest(&.{ &hot, &cold }, &.{ 7, 11 }, 10, 10), &output); + try std.testing.expectEqualSlices(usize, &.{ 10, 10 }, &output); + _ = try allocate(std.testing.allocator, testRequest(&.{&hot}, &.{13}, 47, 10), output[0..1]); + try std.testing.expectEqual(@as(usize, 47), output[0]); +} + +test "profile budget: malformed inputs and overflow leave output unchanged" { + const histogram: [experts]u32 = @splat(0); + const shared: [513]u32 = @splat(0); + var output = [_]usize{ 999, 999 }; + const valid = testRequest(&.{&histogram}, &.{8}, 20, 10); + for ([_][]const []const u32{ &.{}, &.{histogram[0..511]}, &.{&shared} }) |counts| { + var request = valid; + request.counts = counts; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, output[0..1])); + } + var request = valid; + request.line_bytes = &.{}; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, output[0..1])); + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, valid, &output)); + request = valid; + request.line_bytes = &.{0}; + try std.testing.expectError(error.InvalidLineCost, allocate(std.testing.allocator, request, output[0..1])); + for ([_][2]usize{ .{ 0, 0 }, .{ 3, 0 }, .{ 2, 2 }, .{ 1, 1 }, .{ 2, std.math.maxInt(usize) } }) |identity| { + request = valid; + request.world = identity[0]; + request.rank = identity[1]; + try std.testing.expectError(error.InvalidRank, allocate(std.testing.allocator, request, output[0..1])); + } + for ([_][2]usize{ .{ 20, 0 }, .{ 9, 10 }, .{ 257, 10 }, .{ 0, 0 } }) |bounds| { + request = valid; + request.uniform_lines = bounds[0]; + request.min_lines = bounds[1]; + try std.testing.expectError(error.InvalidBounds, allocate(std.testing.allocator, request, output[0..1])); + } + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, testRequest(&.{&histogram}, &.{std.math.maxInt(usize)}, 2, 1), output[0..1])); + const bytes = std.math.maxInt(usize) / 2 + 1; + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, testRequest(&.{ &histogram, &histogram }, &.{ bytes, bytes }, 1, 1), &output)); + try std.testing.expectEqualSlices(usize, &.{ 999, 999 }, &output); +} + +fn allocationFailureProbe(gpa: std.mem.Allocator) !void { + const histogram: [experts]u32 = @splat(7); + var output = [_]usize{ 999, 999 }; + const result = allocate(gpa, testRequest(&.{ &histogram, &histogram }, &.{ 4, 6 }, 3, 1), &output) catch |err| { + try std.testing.expectEqualSlices(usize, &.{ 999, 999 }, &output); + return err; + }; + try std.testing.expectEqual(result.baseline_bytes, result.allocated_bytes); +} + +test "profile budget: allocation failures release scratch and preserve output" { + try std.testing.checkAllAllocationFailures(std.testing.allocator, allocationFailureProbe, .{}); +} + +test "profile budget: feasibility matches exhaustive bounded enumeration" { + const groups = [_]Group{ .{ .bytes = 11, .available = 1 }, .{ .bytes = 7, .available = 2 }, .{ .bytes = 5, .available = 3 }, .{ .bytes = 3, .available = 4 } }; + var suffix_bytes: [5]u64 = undefined; + var suffix_gcd: [5]u64 = undefined; + for (0..65) |target| { + var expected = false; + for (0..2) |first| { + for (0..3) |second| { + for (0..4) |third| { + for (0..5) |fourth| { + if (11 * first + 7 * second + 5 * third + 3 * fourth == target) expected = true; + } + } + } + } + try std.testing.expectEqual(expected, canFill(&groups, target, &suffix_bytes, &suffix_gcd)); + } + try std.testing.expect(!canFill(&.{ .{ .bytes = 6, .available = 10 }, .{ .bytes = 4, .available = 10 } }, 2, &suffix_bytes, &suffix_gcd)); +} + +test "profile budget: heterogeneous allocations preserve the feasible uniform witness" { + const first: [experts]u32 = @splat(17); + const second: [experts]u32 = @splat(11); + const third: [experts]u32 = @splat(5); + var output: [3]usize = undefined; + for (1..6) |first_cost| { + for (1..6) |second_cost| { + for (1..6) |third_cost| { + const bytes = [_]usize{ first_cost, second_cost, third_cost }; + const plan = try allocate(std.testing.allocator, testRequest(&.{ &first, &second, &third }, &bytes, 2, 1), &output); + try std.testing.expectEqual(@as(u64, 2 * (first_cost + second_cost + third_cost) + 3 * 8), plan.allocated_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); + var observed: u64 = 0; + for (output, bytes) |count, cost| { + try std.testing.expect(count >= 1 and count <= 256); + observed += try poolBytes(cost, count, 1, 1, 1); + } + try std.testing.expectEqual(plan.allocated_bytes, observed); + } + } + } +} + +test "profile budget: device picks capacity stays uniform across pool sizes" { + const line_bytes = (1 << 20) + 8; + for ([_]usize{ 10, 11, 31, 32, 33, 64, 256, 512 }) |lines| { + const actual = try poolBytes(line_bytes, lines, 64, 11, 10); + try std.testing.expectEqual(@as(u64, lines * line_bytes + 54 * 64 * 11 * 4), actual); + try std.testing.expectEqual(@as(usize, 54), @divExact(@as(usize, @intCast(actual - lines * line_bytes)), 64 * 11 * 4)); + } + try std.testing.expectEqual(@as(u64, 10 * line_bytes + 12 * 11 * 4), try poolBytes(line_bytes, 10, 1, 11, 10)); + try std.testing.expectEqual(@as(u64, 10 * line_bytes + 2 * 4), try poolBytes(line_bytes, 10, 1, 1, 10)); + try std.testing.expectEqual(@as(u64, 256 * line_bytes + 12 * 11 * 4), try poolBytes(line_bytes, 256, 1, 11, 10)); +} + +test "profile budget: redistribution preserves expert planner and picks bytes exactly" { + var hot: [experts]u32 = @splat(0); + @memset(hot[0..55], 100); + const cold: [experts]u32 = @splat(1); + const line_bytes = (1 << 20) + 8; + var request = testRequest(&.{ &hot, &cold }, &.{ line_bytes, line_bytes }, 54, 53); + request.max_rows = 64; + request.slots = 11; + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 55, 53 }, &output); + try std.testing.expect(!plan.uniform_fallback); + try std.testing.expectEqual(@as(u64, 108 * line_bytes + 2 * 12 * 64 * 11 * 4), plan.baseline_bytes); + var observed: u64 = 0; + for (output) |lines| { + try std.testing.expect(lines >= request.min_lines); + observed += try poolBytes(line_bytes, lines, request.max_rows, request.slots, request.min_lines); + } + try std.testing.expectEqual(plan.baseline_bytes, observed); + try std.testing.expectEqual(observed, plan.allocated_bytes); +} + +test "profile budget: fixed picks capacity preserves exact bytes across pass thresholds" { + var hot: [experts]u32 = @splat(0); + @memset(hot[0..33], 100); + const cold: [experts]u32 = @splat(1); + const line_bytes = (1 << 20) + 8; + var request = testRequest(&.{ &hot, &cold }, &.{ line_bytes, line_bytes }, 32, 31); + var output: [2]usize = undefined; + _ = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 33, 31 }, &output); + request.max_rows = 64; + request.slots = 11; + const plan = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 33, 31 }, &output); + try std.testing.expect(!plan.uniform_fallback); + try std.testing.expectEqual(@as(u64, 64 * line_bytes + 2 * 19 * 64 * 11 * 4), plan.baseline_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); +} + +test "profile budget: fewer picks bytes also retain the exact uniform budget" { + const hot: [experts]u32 = @splat(100); + const cold: [experts]u32 = @splat(1); + var request = testRequest(&.{ &hot, &cold }, &.{ 4, 6 }, 3, 1); + request.max_rows = 16; + request.slots = 11; + var output: [2]usize = undefined; + const plan = try allocate(std.testing.allocator, request, &output); + try std.testing.expectEqualSlices(usize, &.{ 6, 1 }, &output); + try std.testing.expect(!plan.uniform_fallback); + try std.testing.expectEqual(@as(u64, 30 + 2 * 177 * 16 * 11 * 4), plan.baseline_bytes); + try std.testing.expectEqual(plan.baseline_bytes, plan.allocated_bytes); +} + +test "profile budget: TP2 line count above the local expert bank is rejected for caller fallback" { + const histogram: [experts]u32 = @splat(1); + var request = testRequest(&.{&histogram}, &.{8}, 257, 10); + var output: [1]usize = .{257}; + try std.testing.expectError(error.InvalidBounds, allocate(std.testing.allocator, request, &output)); + try std.testing.expectEqual(@as(usize, 257), output[0]); + request.uniform_lines = 256; + _ = try allocate(std.testing.allocator, request, &output); +} + +test "profile budget: picks dimensions and byte overflow preserve output" { + const histogram: [experts]u32 = @splat(0); + var request = testRequest(&.{&histogram}, &.{8}, 20, 10); + var output = [_]usize{999}; + request.max_rows = 0; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, &output)); + request.max_rows = 16; + request.slots = 0; + try std.testing.expectError(error.InvalidDimensions, allocate(std.testing.allocator, request, &output)); + request.slots = 11; + request.max_rows = std.math.maxInt(usize); + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, request, &output)); + request.slots = 1; + request.max_rows = std.math.maxInt(usize) / 4; + try std.testing.expectError(error.Overflow, allocate(std.testing.allocator, request, &output)); + try std.testing.expectEqual(@as(usize, 999), output[0]); +} diff --git a/zig/src/families/flashnext/exl3/offload/exl3_recent_demand.zig b/zig/src/families/flashnext/exl3/offload/exl3_recent_demand.zig new file mode 100644 index 0000000000..a836408b20 --- /dev/null +++ b/zig/src/families/flashnext/exl3/offload/exl3_recent_demand.zig @@ -0,0 +1,212 @@ +const std = @import("std"); + +pub const expert_count: usize = 512; +pub const Count = u32; +const Row = [expert_count / @bitSizeOf(u64)]u64; + +pub const RecentDemand = struct { + history: []Row, + counts: [expert_count]Count = @splat(0), + next_row: usize = 0, + retained_rows: usize = 0, + + pub fn init(allocator: std.mem.Allocator, window_rows: usize) !RecentDemand { + if (window_rows == 0 or window_rows > std.math.maxInt(Count)) return error.InvalidWindowSize; + const history = try allocator.alloc(Row, window_rows); + @memset(history, @splat(0)); + return .{ .history = history }; + } + + pub fn deinit(self: *RecentDemand, allocator: std.mem.Allocator) void { + allocator.free(self.history); + self.* = undefined; + } + + pub fn reset(self: *RecentDemand) void { + @memset(self.history, @splat(0)); + self.counts = @splat(0); + self.next_row = 0; + self.retained_rows = 0; + } + + pub fn windowRows(self: *const RecentDemand) usize { + return self.history.len; + } + + pub fn rowCount(self: *const RecentDemand) usize { + return self.retained_rows; + } + + pub fn scores(self: *const RecentDemand) *const [expert_count]Count { + return &self.counts; + } + + pub fn score(self: *const RecentDemand, expert_id: i32) Count { + if (expert_id < 0 or expert_id >= expert_count) return 0; + return self.counts[@intCast(expert_id)]; + } + + pub fn observeRow(self: *RecentDemand, picks: []const i32) void { + var incoming: Row = @splat(0); + for (picks) |expert_id| { + if (expert_id < 0 or expert_id >= expert_count) continue; + const expert: usize = @intCast(expert_id); + incoming[expert / @bitSizeOf(u64)] |= @as(u64, 1) << @as(u6, @intCast(expert % @bitSizeOf(u64))); + } + + const outgoing = self.history[self.next_row]; + for (outgoing, incoming, 0..) |expired_word, incoming_word, word_index| { + var removed = expired_word & ~incoming_word; + while (removed != 0) { + const expert = word_index * @bitSizeOf(u64) + @ctz(removed); + self.counts[expert] -= 1; + removed &= removed - 1; + } + var added = incoming_word & ~expired_word; + while (added != 0) { + const expert = word_index * @bitSizeOf(u64) + @ctz(added); + self.counts[expert] += 1; + added &= added - 1; + } + } + self.history[self.next_row] = incoming; + self.next_row += 1; + if (self.next_row == self.history.len) self.next_row = 0; + if (self.retained_rows < self.history.len) self.retained_rows += 1; + } +}; + +test "demand expires after exactly the configured number of rows" { + var demand = try RecentDemand.init(std.testing.allocator, 3); + defer demand.deinit(std.testing.allocator); + + demand.observeRow(&.{0}); + demand.observeRow(&.{ 1, 511 }); + demand.observeRow(&.{0}); + try std.testing.expectEqual(@as(Count, 2), demand.score(0)); + try std.testing.expectEqual(@as(Count, 1), demand.score(511)); + try std.testing.expectEqual(@as(usize, 3), demand.rowCount()); + + demand.observeRow(&.{}); + try std.testing.expectEqual(@as(Count, 1), demand.score(0)); + try std.testing.expectEqual(@as(Count, 1), demand.score(1)); + demand.observeRow(&.{ -1, 512 }); + try std.testing.expectEqual(@as(Count, 0), demand.score(1)); + try std.testing.expectEqual(@as(Count, 0), demand.score(511)); + demand.observeRow(&.{}); + try std.testing.expectEqual(@as([expert_count]Count, @splat(0)), demand.scores().*); + try std.testing.expectEqual(@as(usize, 3), demand.rowCount()); +} + +test "duplicates contribute once and shared masked and invalid picks are ignored" { + var demand = try RecentDemand.init(std.testing.allocator, 2); + defer demand.deinit(std.testing.allocator); + + demand.observeRow(&.{ -1, 0, 0, 63, 64, 64, 511, 511, 512, 512, -2, 513, std.math.maxInt(i32), std.math.minInt(i32) }); + for (demand.scores(), 0..) |count, expert| { + const expected: Count = if (expert == 0 or expert == 63 or expert == 64 or expert == 511) 1 else 0; + try std.testing.expectEqual(expected, count); + } + try std.testing.expectEqual(@as(Count, 0), demand.score(-1)); + try std.testing.expectEqual(@as(Count, 0), demand.score(512)); + try std.testing.expectEqual(@as(Count, 0), demand.score(std.math.maxInt(i32))); +} + +test "pick permutations produce identical state through wraparound" { + var ordered = try RecentDemand.init(std.testing.allocator, 3); + defer ordered.deinit(std.testing.allocator); + var permuted = try RecentDemand.init(std.testing.allocator, 3); + defer permuted.deinit(std.testing.allocator); + + for (0..10) |_| { + ordered.observeRow(&.{ 0, 63, 64, 511 }); + permuted.observeRow(&.{ 511, 64, 0, 63, 0, 512, -1 }); + ordered.observeRow(&.{ 7, 8 }); + permuted.observeRow(&.{ 8, 7, 8 }); + ordered.observeRow(&.{}); + permuted.observeRow(&.{ -1, 512 }); + try std.testing.expectEqual(ordered.counts, permuted.counts); + try std.testing.expectEqualSlices(Row, ordered.history, permuted.history); + try std.testing.expectEqual(ordered.next_row, permuted.next_row); + } +} + +test "expiration follows row order" { + var forward = try RecentDemand.init(std.testing.allocator, 2); + defer forward.deinit(std.testing.allocator); + var reverse = try RecentDemand.init(std.testing.allocator, 2); + defer reverse.deinit(std.testing.allocator); + + forward.observeRow(&.{1}); + forward.observeRow(&.{2}); + reverse.observeRow(&.{2}); + reverse.observeRow(&.{1}); + try std.testing.expectEqual(forward.counts, reverse.counts); + forward.observeRow(&.{}); + reverse.observeRow(&.{}); + try std.testing.expectEqual(@as(Count, 0), forward.score(1)); + try std.testing.expectEqual(@as(Count, 1), forward.score(2)); + try std.testing.expectEqual(@as(Count, 1), reverse.score(1)); + try std.testing.expectEqual(@as(Count, 0), reverse.score(2)); +} + +test "reset and replay reproduce every state transition" { + const rows = [_][]const i32{ &.{ 1, 2, 2 }, &.{511}, &.{}, &.{ -1, 512 }, &.{ 2, 511 }, &.{7}, &.{ 1, 7 } }; + var replayed = try RecentDemand.init(std.testing.allocator, 3); + defer replayed.deinit(std.testing.allocator); + var fresh = try RecentDemand.init(std.testing.allocator, 3); + defer fresh.deinit(std.testing.allocator); + + for (rows) |row| replayed.observeRow(row); + replayed.reset(); + try std.testing.expectEqual(fresh.counts, replayed.counts); + try std.testing.expectEqualSlices(Row, fresh.history, replayed.history); + try std.testing.expectEqual(@as(usize, 0), replayed.rowCount()); + try std.testing.expectEqual(@as(usize, 0), replayed.next_row); + try std.testing.expectEqual(@as(usize, 3), replayed.windowRows()); + + for (rows) |row| { + replayed.observeRow(row); + fresh.observeRow(row); + try std.testing.expectEqual(fresh.counts, replayed.counts); + try std.testing.expectEqualSlices(Row, fresh.history, replayed.history); + try std.testing.expectEqual(fresh.next_row, replayed.next_row); + try std.testing.expectEqual(fresh.rowCount(), replayed.rowCount()); + } + replayed.reset(); + replayed.reset(); + try std.testing.expectEqual(@as([expert_count]Count, @splat(0)), replayed.scores().*); +} + +test "all expert scores stay exact at capacity over repeated replacements" { + var demand = try RecentDemand.init(std.testing.allocator, 1); + defer demand.deinit(std.testing.allocator); + var picks: [expert_count]i32 = undefined; + for (&picks, 0..) |*expert, index| expert.* = @intCast(index); + for (0..1024) |_| { + demand.observeRow(&picks); + try std.testing.expectEqual(@as([expert_count]Count, @splat(1)), demand.scores().*); + } + demand.observeRow(&.{}); + try std.testing.expectEqual(@as([expert_count]Count, @splat(0)), demand.scores().*); + + var larger = try RecentDemand.init(std.testing.allocator, 4); + defer larger.deinit(std.testing.allocator); + for (0..16) |index| { + larger.observeRow(&picks); + const expected: Count = @intCast(@min(index + 1, larger.windowRows())); + for (larger.scores()) |count| try std.testing.expectEqual(expected, count); + } + for (0..4) |index| { + larger.observeRow(&.{}); + const expected: Count = @intCast(3 - index); + for (larger.scores()) |count| try std.testing.expectEqual(expected, count); + } +} + +test "invalid capacities are rejected before allocation" { + try std.testing.expectError(error.InvalidWindowSize, RecentDemand.init(std.testing.allocator, 0)); + if (@bitSizeOf(usize) > @bitSizeOf(Count)) { + try std.testing.expectError(error.InvalidWindowSize, RecentDemand.init(std.testing.allocator, @as(usize, std.math.maxInt(Count)) + 1)); + } +} diff --git a/zig/src/families/flashnext/exl3/pack/cuda_exl3_config.zig b/zig/src/families/flashnext/exl3/pack/cuda_exl3_config.zig new file mode 100644 index 0000000000..b4db09d186 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_config.zig @@ -0,0 +1,113 @@ +//! EXL3 admission for the Flash Next CUDA family: config.json's quantization_config names quant_method "exl3" and the + +const std = @import("std"); + +/// The one codebook the family reads (the kernels' `mul1`, cuda.exl3.Codebook.mul1). +pub const codebook = "mul1"; +/// The weight format the family lists for an EXL3 checkpoint. +pub const format = "exl3-mul1"; + +fn object(o: std.json.ObjectMap, name: []const u8) ?std.json.ObjectMap { + const v = o.get(name) orelse return null; + return if (v == .object) v.object else null; +} + +fn word(o: std.json.ObjectMap, name: []const u8) ?[]const u8 { + const v = o.get(name) orelse return null; + return if (v == .string) v.string else null; +} + +/// config.json's quantization_config: top level, else the text config's. +pub fn block(root: std.json.ObjectMap) ?std.json.ObjectMap { + if (object(root, "quantization_config")) |q| return q; + if (object(root, "text_config")) |t| if (object(t, "quantization_config")) |q| return q; + return null; +} + +/// Whether config.json declares an EXL3 checkpoint (quant_method "exl3"), whatever its codebook. +pub fn named(root: std.json.ObjectMap) bool { + const q = block(root) orelse return false; + const m = word(q, "quant_method") orelse return false; + return std.mem.eql(u8, m, "exl3"); +} + +/// The first field an EXL3 checkpoint differs in from what the family serves, else null. +pub fn mismatch(root: std.json.ObjectMap, text: std.json.ObjectMap, comptime geometry: fn (std.json.ObjectMap) ?[]const u8) ?[]const u8 { + const q = block(root) orelse return "quantization_config"; + const m = word(q, "quant_method") orelse return "quantization_config.quant_method"; + if (!std.mem.eql(u8, m, "exl3")) return "quantization_config.quant_method"; + const book = word(q, "codebook") orelse return "quantization_config.codebook"; + if (!std.mem.eql(u8, book, codebook)) return "quantization_config.codebook"; + return geometry(text); +} + +const fixture = @embedFile("fixtures/exl3_config.json"); + +test "exl3: the Qwen3.8-Flash-Next EXL3 pack's config.json is admitted as Mode.exl3" { + const config = @import("../../cuda_config.zig"); + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = fixture }); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + defer a.free(dir); + try std.testing.expectEqual(config.Mode.exl3, try config.mode(a, std.testing.io, dir)); +} + +test "exl3: other codebooks and other methods are refused by name; bits are never read" { + const config = @import("../../cuda_config.zig"); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const Case = struct { field: []const u8, value: std.json.Value, refused: ?[]const u8 }; + const cases = [_]Case{ + .{ .field = "codebook", .value = .{ .string = "mcg" }, .refused = "quantization_config.codebook" }, + .{ .field = "codebook", .value = .{ .string = "inst3" }, .refused = "quantization_config.codebook" }, + .{ .field = "codebook", .value = .null, .refused = "quantization_config.codebook" }, + .{ .field = "quant_method", .value = .{ .string = "gptq" }, .refused = "quantization_config.quant_method" }, + // the average width and the head's are not a contract: the trellis says each matrix's own width + .{ .field = "bits", .value = .{ .float = 4.25 }, .refused = null }, + .{ .field = "bits", .value = .{ .integer = 2 }, .refused = null }, + .{ .field = "head_bits", .value = .{ .integer = 8 }, .refused = null }, + }; + for (cases) |c| { + const parsed = try std.json.parseFromSlice(std.json.Value, a, fixture, .{}); + const root = parsed.value.object; + const text = root.get("text_config").?.object; + try std.testing.expect(named(root)); + try std.testing.expectEqual(@as(?[]const u8, null), mismatch(root, text, config.geometry)); + try root.getPtr("quantization_config").?.object.put(a, c.field, c.value); + const got = mismatch(parsed.value.object, text, config.geometry); + if (c.refused) |want| try std.testing.expectEqualStrings(want, got orelse "accepted") else try std.testing.expectEqual(@as(?[]const u8, null), got); + } +} + +test "exl3: the shapes and semantics checks still apply to an EXL3 checkpoint" { + const config = @import("../../cuda_config.zig"); + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const Case = struct { field: []const u8, value: std.json.Value, refused: []const u8 }; + const cases = [_]Case{ + .{ .field = "hidden_size", .value = .{ .integer = 4096 }, .refused = "hidden_size" }, + .{ .field = "num_experts_per_tok", .value = .{ .integer = 8 }, .refused = "num_experts_per_tok" }, + .{ .field = "rms_norm_eps", .value = .{ .float = 1e-5 }, .refused = "rms_norm_eps" }, + .{ .field = "output_gate_type", .value = .{ .string = "silu" }, .refused = "output_gate_type" }, + }; + for (cases) |c| { + const parsed = try std.json.parseFromSlice(std.json.Value, a, fixture, .{}); + const root = parsed.value.object; + const text = &root.getPtr("text_config").?.object; + try text.put(a, c.field, c.value); + try std.testing.expectEqualStrings(c.refused, mismatch(root, text.*, config.geometry) orelse "accepted"); + } +} + +test "exl3: a non-EXL3 config is not named EXL3" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const parsed = try std.json.parseFromSlice(std.json.Value, arena.allocator(), "{\"quantization_config\":{\"quant_method\":\"modelopt\"}}", .{}); + try std.testing.expect(!named(parsed.value.object)); + const bare = try std.json.parseFromSlice(std.json.Value, arena.allocator(), "{\"text_config\":{}}", .{}); + try std.testing.expect(!named(bare.value.object)); +} diff --git a/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig b/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig new file mode 100644 index 0000000000..cec7e8da81 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_load.zig @@ -0,0 +1,527 @@ +//! EXL3 weights for the Flash Next CUDA family's own loaders (cuda_projection_load.zig, cuda_prompt_load.zig) + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const wt = @import("../../cuda_prompt_wt.zig"); +const projection = @import("../../cuda_projection.zig"); +const pool_ = @import("../offload/cuda_exl3_pool.zig"); + +const Tensor = core.safetensors.Tensor; +const Store = wt.Store; + +/// format.py's mul1 marker value (core.exl3_load.markers). +const mul1_magic: u32 = blk: { + for (core.exl3_load.markers) |m| if (m.codebook == .mul1) break :blk m.value; + unreachable; +}; + +/// The EXL3 kernels and scratch one engine's calls share, on the engine's one compute stream. +pub const Ops = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + refs: usize = 1, + fork: ?*@import("../tp/cuda_exl3_fork.zig").Fork = null, + dense: ?cuda.exl3.Dense = null, + experts: ?cuda.exl3_experts.Experts = null, + /// A part's rows before its columns land in the stitched row (bf16 [rows, part n]). + tmp: ?cuda.DeviceBuffer = null, + /// The routed experts' scratch for `moe_rows` rows of 11 slots. + moe: ?cuda.DeviceBuffer = null, + moe_rows: usize = 0, + /// fp16 rows split into bf16 hi and lo [2, rows, k], and their fp32 products [2, rows, n] (an fp16-input Proj). + halves: ?cuda.DeviceBuffer = null, + sums: ?cuda.DeviceBuffer = null, + glue: ?cuda.Module = null, + split_fn: cuda.Function = undefined, + /// Sizing maxima over what was loaded (`reserve` allocates every scratch for them once, before any round): the + max_k: usize = 0, + max_z: usize = 0, + max_counters: usize = 0, + max_part_n: usize = 0, + max_f16_k: usize = 0, + max_f16_n: usize = 0, + moe_layer: ?cuda.exl3_experts.Layer = null, + /// Pinned [moe_rows, slots] int32: one block's router picks on the host (the offload pool's A/B cut). + picks: ?cuda.HostBuffer = null, + + pub fn create(gpa: std.mem.Allocator, d: *const cuda.Driver) !*Ops { + const o = try gpa.create(Ops); + o.* = .{ .gpa = gpa, .d = d }; + return o; + } + + pub fn retain(o: *Ops) *Ops { + o.refs += 1; + return o; + } + + pub fn release(o: *Ops) void { + o.refs -= 1; + if (o.refs != 0) return; + @import("../tp/cuda_exl3_fork.zig").release(o); + if (o.dense) |*x| x.unload(); + if (o.experts) |*x| x.unload(); + if (o.tmp) |*b| b.free(); + if (o.moe) |*b| b.free(); + if (o.halves) |*b| b.free(); + if (o.sums) |*b| b.free(); + if (o.picks) |*b| b.free(); + if (o.glue) |*m| m.unload(); + o.gpa.destroy(o); + } + + /// The dense EXL3 linear, loaded on first use (no GPU work at checkpoint load). + pub fn linear(o: *Ops) !*cuda.exl3.Dense { + if (o.dense == null) o.dense = try cuda.exl3.Dense.load(o.d); + return &o.dense.?; + } + + /// fn_exl3_split_f16 (flashnext_exl3.cu), loaded on first use. + pub fn split(o: *Ops) !cuda.Function { + if (o.glue == null) { + if (!cuda.kernels.available) return error.BuiltWithoutKernels; + var m = try cuda.Module.load(o.d, cuda.kernels.flashnext_exl3); + errdefer m.unload(); + o.split_fn = try m.function("fn_exl3_split_f16"); + o.glue = m; + } + return o.split_fn; + } + + pub fn grouped(o: *Ops) !*const cuda.exl3_experts.Experts { + if (o.experts == null) o.experts = try cuda.exl3_experts.Experts.load(o.d); + return &o.experts.?; + } + + pub fn noteWeight(o: *Ops, w: cuda.exl3.Weight, row_n: usize) void { + o.max_k = @max(o.max_k, w.k); + o.max_z = @max(o.max_z, cuda.exl3.Linear.zBytes(w, 128)); + o.max_counters = @max(o.max_counters, cuda.exl3.Linear.counterBytes(w)); + if (w.n != row_n) o.max_part_n = @max(o.max_part_n, w.n); + } + + /// Every scratch the EXL3 calls of a pass of up to `rows` rows read, allocated now (after `s` drains), so no call + pub fn reserve(o: *Ops, s: cuda.Stream, rows: usize, moe_rows: usize, slots: usize) !void { + try s.synchronize(); + if (o.max_k > 0) { + const de = try o.linear(); + try grow(o.d, &de.xh, 128 * o.max_k * 2); + if (o.max_z > 0) try grow(o.d, &de.z, o.max_z); + try grow(o.d, &de.counters, o.max_counters); + try @import("../tp/cuda_exl3_fork.zig").reserve(o); // parallel parts' branches (decode projections) + } + if (o.max_part_n > 0) try grow(o.d, &o.tmp, rows * o.max_part_n * 2); + if (o.max_f16_k > 0) { + _ = try o.split(); + try grow(o.d, &o.halves, 2 * rows * o.max_f16_k * 2); + try grow(o.d, &o.sums, 2 * rows * o.max_f16_n * 4); + } + if (o.moe_layer) |l| { + _ = try o.grouped(); + try grow(o.d, &o.moe, try cuda.exl3_experts.scratchBytes(l, moe_rows, slots)); + if (o.picks == null) o.picks = try cuda.HostBuffer.alloc(o.d, moe_rows * slots * 4); + } + } + + fn grow(d: *const cuda.Driver, slot: *?cuda.DeviceBuffer, bytes: usize) !void { + if (bytes == 0) return; + if (slot.*) |b| if (b.len >= bytes) return; + if (slot.*) |*b| b.free(); + slot.* = null; + slot.* = try cuda.DeviceBuffer.alloc(d, bytes); + } + + /// Every scratch buffer's address and size: unchanged across a pass means nothing grew (nor synchronised) in it. + pub fn fingerprint(o: *const Ops) [8][2]u64 { + var f: [8][2]u64 = @splat(.{ 0, 0 }); + const de: ?cuda.exl3.Dense = o.dense; + const bufs = [8]?cuda.DeviceBuffer{ o.tmp, o.moe, o.halves, o.sums, if (de) |x| x.xh else null, if (de) |x| x.z else null, if (de) |x| x.counters else null, null }; + for (bufs, &f) |b, *e| if (b) |x| { + e.* = .{ x.ptr, x.len }; + }; + if (o.picks) |h| f[7] = .{ @intFromPtr(h.bytes.ptr), h.bytes.len }; + return f; + } + + /// Device bytes `reserve` holds (for the pool's budget). + pub fn reservedBytes(o: *const Ops) usize { + var n: usize = 0; + for ([_]?cuda.DeviceBuffer{ o.tmp, o.moe, o.halves, o.sums }) |b| if (b) |x| { + n += x.len; + }; + if (o.dense) |de| for ([_]?cuda.DeviceBuffer{ de.xh, de.z, de.counters }) |b| if (b) |x| { + n += x.len; + }; + return n; + } + + /// At least `bytes` of part scratch; a growth waits for `s` first (an older call may still read the old buffer). + pub fn scratch(o: *Ops, s: cuda.Stream, slot: *?cuda.DeviceBuffer, bytes: usize) !u64 { + if (slot.*) |b| if (b.len >= bytes) return b.ptr; + try s.synchronize(); + if (slot.*) |*b| b.free(); + slot.* = null; + slot.* = try cuda.DeviceBuffer.alloc(o.d, bytes); + return slot.*.?.ptr; + } +}; + +/// The store's shared Ops, made on first use; every holder retains it. +pub fn opsOf(store: *Store, d: *const cuda.Driver) !*Ops { + if (store.exl3_ops == null) store.exl3_ops = try Ops.create(store.gpa, d); + return store.exl3_ops.?.retain(); +} + +/// One EXL3 matrix read from the store: words as stored, its width and shape. +pub const Mat = struct { + trellis: []const u8, + suh: []const u8, + svh: []const u8, + k: usize, + n: usize, + halves: u32, +}; + +/// Whether `prefix` names an EXL3 matrix (its trellis is indexed). +pub fn has(store: *const Store, prefix: []const u8) bool { + var buf: [256]u8 = undefined; + return store.has(std.fmt.bufPrint(&buf, "{s}.trellis", .{prefix}) catch return false); +} + +/// `prefix`'s trellis, scales and mul1 marker, checked: int16 [K/16, N/16, 16 * bits], fp16 suh [K] and svh [N]. +pub fn mat(store: *Store, prefix: []const u8) !Mat { + var buf: [256]u8 = undefined; + const t = try store.hold(try std.fmt.bufPrint(&buf, "{s}.trellis", .{prefix})); + if (t.dtype != .i16 or t.rank != 3 or t.dim(2) % 8 != 0) return error.UnexpectedTensor; + const halves = t.dim(2) / 8; + // the mul1 half steps are odd halves; 2..16 halves (1 to 8 bits) are what the kernels instantiate + if (halves < 2 or halves > 16) return error.UnsupportedExl3Width; + const k = t.dim(0) * 16; + const n = t.dim(1) * 16; + if (k % 128 != 0 or n % 128 != 0) return error.UnexpectedTensor; + const marker = store.hold(try std.fmt.bufPrint(&buf, "{s}.mul1", .{prefix})) catch return error.NotExl3Mul1; + if (marker.dtype != .i32 or marker.bytes.len != 4 or std.mem.readInt(u32, marker.bytes[0..4], .little) != mul1_magic) return error.NotExl3Mul1; + const suh = try store.hold(try std.fmt.bufPrint(&buf, "{s}.suh", .{prefix})); + const svh = try store.hold(try std.fmt.bufPrint(&buf, "{s}.svh", .{prefix})); + if (!suh.is(.f16, &.{k}) or !svh.is(.f16, &.{n})) return error.UnexpectedTensor; + if (store.has(try std.fmt.bufPrint(&buf, "{s}.bias", .{prefix}))) return error.UnexpectedTensor; // none in the text model + return .{ .trellis = t.bytes, .suh = suh.bytes, .svh = svh.bytes, .k = k, .n = n, .halves = @intCast(halves) }; +} + +/// One column range of a projection's output row. +pub const Part = struct { + col: usize, + n: usize, + kind: union(enum) { exl3: cuda.exl3.Weight, native: cuda.qmmf.Weight }, +}; + +/// A projection's parts and the device buffers they own. +pub const Proj = struct { + gpa: std.mem.Allocator, + ops: *Ops, + k: usize, + n: usize, + parts: []Part, + /// The rows the pass hands in: bf16, or fp16 bit patterns (an EXL3 pack's n-gram rows, cuda_exl3_ple.zig). + x: enum { bf16, f16 } = .bf16, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + natives: std.ArrayList(projection.Projection) = .empty, + + pub fn destroy(p: *Proj) void { + for (p.bufs.items) |*b| b.free(); + p.bufs.deinit(p.gpa); + for (p.natives.items) |*q| q.free(); + p.natives.deinit(p.gpa); + p.gpa.free(p.parts); + p.ops.release(); + p.gpa.destroy(p); + } + + fn up(p: *Proj, d: *const cuda.Driver, bytes: []const u8) !u64 { + try p.bufs.ensureUnusedCapacity(p.gpa, 1); + const b = try cuda.DeviceBuffer.fromHost(d, bytes); + p.bufs.appendAssumeCapacity(b); + return b.ptr; + } + + /// An EXL3 matrix as the stored-layout weight (trellis [K/16][N/16][W]: strides(k, n, halves, false)). + fn exl3Part(p: *Proj, d: *const cuda.Driver, m: Mat) !cuda.exl3.Weight { + const st = cuda.exl3.strides(m.k, m.n, m.halves, false); + return .{ .trellis = try p.up(d, m.trellis), .suh = try p.up(d, m.suh), .svh = try p.up(d, m.svh), .k = @intCast(m.k), .n = @intCast(m.n), .halves = m.halves, .codebook = .mul1, .stride_k = st.k, .stride_nb = st.nb }; + } +}; + +/// A Projection whose rows come from `parts` (EXL3 matrices or dense BF16 rows), `n` columns of `k` inputs. +fn wrap(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, prefix: []const u8, names: []const []const u8, n: usize, k: usize) !projection.Projection { + const p = try gpa.create(Proj); + p.* = .{ .gpa = gpa, .ops = try opsOf(store, d), .k = k, .n = n, .parts = &.{} }; + errdefer p.destroy(); + var parts: std.ArrayList(Part) = .empty; + errdefer parts.deinit(gpa); + var buf: [256]u8 = undefined; + var col: usize = 0; + var i: usize = 0; + while (i < names.len) { + const name = try std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, names[i] }); + if (has(store, name)) { + const m = try mat(store, name); + if (m.k != k) return error.UnexpectedTensor; + const w = try p.exl3Part(d, m); + p.ops.noteWeight(w, n); + try parts.append(gpa, .{ .col = col, .n = m.n, .kind = .{ .exl3 = w } }); + col += m.n; + i += 1; + continue; + } + // a run of dense BF16 parts (after the Store's exact widening) packs as one lane weight + var rows: usize = 0; + var j = i; + var bytes: std.ArrayList(u8) = .empty; + defer bytes.deinit(gpa); + while (j < names.len) : (j += 1) { + var nb: [256]u8 = undefined; + const dense_name = try std.fmt.bufPrint(&nb, "{s}{s}", .{ prefix, names[j] }); + if (has(store, dense_name)) break; + const t = try store.tensor(try std.fmt.bufPrint(&nb, "{s}{s}.weight", .{ prefix, names[j] })); + if (t.dtype != .bf16 or t.rank != 2 or t.dim(1) != k or t.bytes.len != t.dim(0) * k * 2) return error.UnexpectedTensor; + try bytes.appendSlice(gpa, t.bytes); + rows += t.dim(0); + } + try p.natives.ensureUnusedCapacity(gpa, 1); + const q = try projection.dense(gpa, d, .{ .dtype = .bf16, .rank = 2, .shape = .{ rows, k, 1, 1, 1 }, .bytes = bytes.items }); + p.natives.appendAssumeCapacity(q); + try parts.append(gpa, .{ .col = col, .n = rows, .kind = .{ .native = q.native.? } }); + col += rows; + i = j; + } + if (col != n) return error.UnexpectedTensor; + p.parts = try parts.toOwnedSlice(gpa); + var out = projection.Projection.empty(d); + out.n = n; + out.k = k; + out.exl3 = p; + return out; +} + +/// cuda_projection_load.load's EXL3 branch: `prefix` one EXL3 matrix of n outputs and k inputs. +pub fn load(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, prefix: []const u8, n: usize, k: usize) !projection.Projection { + return wrap(gpa, d, store, prefix, &.{""}, n, k); +} + +/// Engine.initRank's head for an EXL3 pack: one rank holds the whole vocabulary. +pub fn head(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, world: usize, rank: usize) !?projection.Projection { + if (!has(store, "language_model.lm_head")) return null; + if (world != 1 or rank != 0) return error.UnsupportedQuantization; + if (try @import("../../cuda_head.zig").enabled()) return error.BadHeadSharding; + return try load(gpa, d, store, "language_model.lm_head", wt.vocab, wt.dims); +} + +/// Whether any of `prefix`'s parts is an EXL3 matrix (cuda_projection_load.loadStitched's branch). +pub fn hasAny(store: *const Store, prefix: []const u8, names: []const []const u8) bool { + var buf: [256]u8 = undefined; + for (names) |name| if (has(store, std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, name }) catch return false)) return true; + return false; +} + +/// cuda_projection_load.loadStitched's EXL3 branch: parts side by side in one output row, EXL3 or dense BF16. +pub fn stitched(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, prefix: []const u8, names: []const []const u8, n: usize, k: usize) !projection.Projection { + return wrap(gpa, d, store, prefix, names, n, k); +} + +/// cuda_prompt_load's PLE branch: an EXL3 pack's n-gram rows reach key_proj and value_proj as fp16 bit patterns +pub fn pleRows(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, key: *projection.Projection, value: *projection.Projection) !void { + if (store.exl3_ops == null) return; // not an EXL3 pack: the n-gram rows are bf16 + for ([_]*projection.Projection{ key, value }) |q| { + const native = q.native orelse return error.UnexpectedTensor; + const p = try gpa.create(Proj); + p.* = .{ .gpa = gpa, .ops = try opsOf(store, d), .k = q.k, .n = q.n, .parts = &.{}, .x = .f16 }; + errdefer p.destroy(); + p.ops.max_f16_k = @max(p.ops.max_f16_k, q.k); + p.ops.max_f16_n = @max(p.ops.max_f16_n, q.n); + p.parts = try gpa.alloc(Part, 1); + p.parts[0] = .{ .col = 0, .n = q.n, .kind = .{ .native = native } }; + try p.natives.append(gpa, q.*); + var out = projection.Projection.empty(d); + out.n = q.n; + out.k = q.k; + out.exl3 = p; + q.* = out; + } +} + +/// One layer's routed experts and the shared expert (id `wt.experts_n`) for exl3_experts.Experts. +pub const Experts = struct { + gpa: std.mem.Allocator, + ops: *Ops, + layer: cuda.exl3_experts.Layer, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + /// Each expert's and projection's trellis words on the device, for a pool to rewrite (gate, up, down; 513 each). + trellis: [3][]u64 = .{ &.{}, &.{}, &.{} }, + halves: [3][]i32 = .{ &.{}, &.{}, &.{} }, + bank: ?cuda.HostBuffer = null, + /// The pool's Source: every routed expert's gate/up/down bytes inside `bank`, and the shared expert's. + routed: [][3][]const u8 = &.{}, + shared: [3][]const u8 = .{ &.{}, &.{}, &.{} }, + pool: ?*pool_.Pool = null, + /// Host time the bank's copies took at load (ns). + bank_ns: u64 = 0, + + pub fn destroy(x: *Experts) void { + if (x.pool) |p| { + p.deinit(); + x.gpa.destroy(p); + } + for (x.bufs.items) |*b| b.free(); + x.bufs.deinit(x.gpa); + for (x.trellis) |t| x.gpa.free(t); + for (x.halves) |h| x.gpa.free(h); + x.gpa.free(x.routed); + if (x.bank) |*b| b.free(); + x.ops.release(); + x.gpa.destroy(x); + } + + fn up(x: *Experts, d: *const cuda.Driver, bytes: []const u8) !u64 { + try x.bufs.ensureUnusedCapacity(x.gpa, 1); + const b = try cuda.DeviceBuffer.fromHost(d, bytes); + x.bufs.appendAssumeCapacity(b); + return b.ptr; + } + + /// Whether the routed experts wait for a pool (offloaded, none attached yet). + pub fn needsPool(x: *const Experts) bool { + return x.bank != null and x.pool == null; + } + + /// Bytes of one pool line of this layer (gate, up and down of its widest expert, 256-aligned). + pub fn lineBytes(x: *const Experts) usize { + var bytes: [3]usize = .{ 0, 0, 0 }; + for (x.routed) |e| for (0..3) |j| { + bytes[j] = @max(bytes[j], e[j].len); + }; + return cuda.exl3_pool.Layout.init(bytes).line_bytes; + } + + /// Makes the pool: `lines` device lines, `counts` the family's lifetime use table (null: the pool's own), `layer` + pub fn attach(x: *Experts, io: std.Io, d: *const cuda.Driver, which: u8, lines: usize, max_rows: usize, counts: ?*[pool_.shape.layers][pool_.shape.experts]u32) !void { + if (x.bank == null or x.pool != null) return error.Invalid; + const p = try x.gpa.create(pool_.Pool); + errdefer x.gpa.destroy(p); + const l = x.layer; + p.* = try pool_.Pool.init(x.gpa, io, d, .{ .routed = x.routed, .shared = if (x.shared[0].len == 0) null else x.shared, .pinned_bank = x.bank.?.bytes }, .{ .suh_g = l.suh_g, .suh_u = l.suh_u, .svh_g = l.svh_g, .svh_u = l.svh_u, .suh_d = l.suh_d, .svh_d = l.svh_d }, .{ + .layer = which, + .lines = lines, + .dims = wt.dims, + .width = wt.moe_w, + .max_rows = max_rows, + .slots = wt.slots, + .direct = true, // the bank is pinned: one H2D a projection on the copy engine, no staging memcpy + .codebook = .mul1, + .act_mode = l.act_mode, + .max_passes = pool_.passesFor(pool_.min_lines), + }, counts); + x.pool = p; + x.layer = p.tables(); + } +}; + +/// Whether `mlp` (".mlp" of a layer) holds EXL3 experts. +pub fn hasExperts(store: *const Store, mlp: []const u8) bool { + var buf: [256]u8 = undefined; + return store.has(std.fmt.bufPrint(&buf, "{s}.shared_expert.gate_proj.trellis", .{mlp}) catch return false); +} + +/// TF_FLASHNEXT_EXL3_ROUTED=device keeps every routed expert on the device (cuda_exl3_memory.zig reads the same). +pub fn offloaded() bool { + return @import("cuda_exl3_memory.zig").options().routed_on_host; +} + +fn align256(n: usize) usize { + return (n + 255) / 256 * 256; +} + +/// Every routed expert of `mlp`, the shared expert as id 512 (the same 2560 x 640 x 2560 shape): resident, or with +pub fn experts(gpa: std.mem.Allocator, d: *cuda.Driver, store: *Store, mlp: []const u8, world: usize, rank: usize) !*Experts { + if (world != 1 or rank != 0) return error.UnsupportedQuantization; // one rank holds all 513 + const count = wt.experts_n + 1; + const offload = offloaded(); + const x = try gpa.create(Experts); + x.* = .{ .gpa = gpa, .ops = try opsOf(store, d), .layer = undefined }; + errdefer x.destroy(); + for (&x.trellis, &x.halves) |*t, *h| { + t.* = try gpa.alloc(u64, count); + @memset(t.*, 0); + h.* = try gpa.alloc(i32, count); + @memset(h.*, 0); + } + const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + const ks = [3]usize{ wt.dims, wt.dims, wt.moe_w }; + const ns = [3]usize{ wt.moe_w, wt.moe_w, wt.dims }; + var scales: [6]std.ArrayList(u8) = @splat(.empty); + defer for (&scales) |*l| l.deinit(gpa); + var buf: [256]u8 = undefined; + const nameOf = struct { + fn f(b: []u8, m: []const u8, e: usize, proj: []const u8) ![]const u8 { + if (e == wt.experts_n) return std.fmt.bufPrint(b, "{s}.shared_expert.{s}", .{ m, proj }); + return std.fmt.bufPrint(b, "{s}.experts.{d}.{s}", .{ m, e, proj }); + } + }.f; + var at: usize = 0; + if (offload) { + // the bank's size first (headers only), then one pinned region for the layer + var total: usize = 0; + for (0..count) |e| { + for (projs) |proj| total += align256((try mat(store, try nameOf(&buf, mlp, e, proj))).trellis.len); + store.release(); + } + x.bank = try cuda.HostBuffer.alloc(d, total); + x.routed = try gpa.alloc([3][]const u8, wt.experts_n); + } + const t0 = std.Io.Clock.awake.now(store.io).toNanoseconds(); + for (0..count) |e| { + for (projs, 0..) |proj, j| { + const m = try mat(store, try nameOf(&buf, mlp, e, proj)); + if (m.k != ks[j] or m.n != ns[j]) return error.UnexpectedTensor; + if (x.bank) |bank| { + const dst = bank.bytes[at..][0..m.trellis.len]; + @memcpy(dst, m.trellis); + at += align256(m.trellis.len); + if (e == wt.experts_n) x.shared[j] = dst else x.routed[e][j] = dst; + } else x.trellis[j][e] = try x.up(d, m.trellis); + x.halves[j][e] = @intCast(m.halves); + try scales[2 * j].appendSlice(gpa, m.suh); + try scales[2 * j + 1].appendSlice(gpa, m.svh); + } + // each expert's mapped pages are read once: release them so the 513 experts never pin a shard's worth + store.release(); + } + x.bank_ns = @intCast(std.Io.Clock.awake.now(store.io).toNanoseconds() - t0); + x.layer = .{ + // offloaded: the pool's tables replace these at attach (0 until then; a pass refuses an unattached layer) + .gate = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.trellis[0])), + .up = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.trellis[1])), + .down = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.trellis[2])), + .gate_k2 = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.halves[0])), + .up_k2 = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.halves[1])), + .down_k2 = if (offload) 0 else try x.up(d, std.mem.sliceAsBytes(x.halves[2])), + .suh_g = try x.up(d, scales[0].items), + .svh_g = try x.up(d, scales[1].items), + .suh_u = try x.up(d, scales[2].items), + .svh_u = try x.up(d, scales[3].items), + .suh_d = try x.up(d, scales[4].items), + .svh_d = try x.up(d, scales[5].items), + .experts = count, + .dims = wt.dims, + .width = wt.moe_w, + .codebook = .mul1, + // fp32 SwiGLU between the grouped GEMMs (GLM's mode 0 rounds to bf16 there) + .act_mode = 1, + }; + x.ops.moe_layer = x.layer; + return x; +} + +test "exl3: the mul1 marker value is format.py's" { + try std.testing.expectEqual(@as(u32, 0x83DCD12D), mul1_magic); +} diff --git a/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig b/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig new file mode 100644 index 0000000000..ae55c588b9 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_memory.zig @@ -0,0 +1,134 @@ +//! Device bytes of an EXL3 Flash Next checkpoint as cuda_exl3_load.zig places it: + +const std = @import("std"); +const core = @import("core"); +const cuda = @import("cuda"); +const wt = @import("../../cuda_prompt_wt.zig"); +const config = @import("../../cuda_config.zig"); +const Entry = core.safetensors.Entry; + +pub const Options = struct { + /// Routed experts' trellis words live in pinned host memory (the offload pool's), not counted here. + routed_on_host: bool = true, +}; + +/// TF_FLASHNEXT_EXL3_ROUTED=device counts every routed expert as resident; the default is the offload pool's layout. +pub fn options() Options { + const v = std.c.getenv("TF_FLASHNEXT_EXL3_ROUTED") orelse return .{}; + return .{ .routed_on_host = !std.mem.eql(u8, std.mem.span(v), "device") }; +} + +/// Whether `dir`'s config.json admits as Mode.exl3 (any other mode, or a refusal: false). +pub fn isExl3(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) bool { + return (config.mode(gpa, io, dir) catch return false) == .exl3; +} + +/// One tensor's device bytes (0: kept on the host). +pub fn tensorBytes(name: []const u8, e: Entry, o: Options) !u64 { + const bytes: u64 = e.end - e.begin; + if (wt.hostOnly(name)) return 0; + // the BF16 embedding table is pinned host memory the gather reads in place (cuda_embed.Table.uploadDenseHost) + if (e.dtype == .bf16 and std.mem.indexOf(u8, name, "embed_tokens") != null) return 0; + // the codebook marker is checked at load and never uploaded + if (std.mem.endsWith(u8, name, ".mul1") or std.mem.endsWith(u8, name, ".mcg")) return 0; + if (std.mem.indexOf(u8, name, ".mlp.experts.") != null) { + if (std.mem.endsWith(u8, name, ".trellis")) return if (o.routed_on_host) 0 else bytes; + return bytes; // stacked suh / svh, resident + } + if (std.mem.endsWith(u8, name, ".trellis") or std.mem.endsWith(u8, name, ".suh") or std.mem.endsWith(u8, name, ".svh")) return bytes; + if (!std.mem.endsWith(u8, name, ".weight") and e.rank == 1) return try std.math.mul(u64, e.dim(0), 4); // A_log, dt_bias as fp32 + if ((e.dtype == .f16 or e.dtype == .bf16) and e.rank == 2 and std.mem.endsWith(u8, name, ".weight")) { + // the router's rows and the embedding table go up as stored (BF16). + if (std.mem.indexOf(u8, name, ".mlp.gate.") != null or std.mem.indexOf(u8, name, ".shared_expert_gate.") != null or std.mem.indexOf(u8, name, "embed_tokens") != null) return bytes; + return try std.math.mul(u64, try std.math.mul(u64, cuda.nvfp4.padded(e.dim(0)), e.dim(1)), 2); + } + // norms the loaders scale go up as fp32 (cuda_rank_memory.zig's list), the GDN gated norm and convs as stored + if (e.rank == 1 and centered(name)) return try std.math.mul(u64, e.dim(0), 4); + return bytes; +} + +fn centered(name: []const u8) bool { + for ([_][]const u8{ ".hc_norm.weight", ".q_norm.weight", ".k_norm.weight", ".q_layernorm.weight", ".k_layernorm.weight", ".norm_key.weight", ".norm_query.weight", ".norm_conv.weight" }) |suffix| { + if (std.mem.endsWith(u8, name, suffix)) return true; + } + return false; +} + +/// Per layer: three 513-entry pointer tables (int64) and three width tables (int32) for the experts' grouped GEMMs. +const tables_per_layer: u64 = 3 * (wt.experts_n + 1) * 8 + 3 * (wt.experts_n + 1) * 4; + +/// Every shard the index names (and the mixer patch, which Store.open adds), each read for its header alone. +pub fn deviceBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, o: Options) !u64 { + var store = try wt.Store.open(gpa, io, dir); + defer store.deinit(); + const wm = store.parsed.value.object.get("weight_map") orelse return error.MissingTensor; + var files: std.StringHashMapUnmanaged(void) = .empty; + defer files.deinit(gpa); + var it = wm.object.iterator(); + while (it.next()) |kv| try files.put(gpa, kv.value_ptr.string, {}); + var total: u64 = tables_per_layer * wt.layers_n; + var fi = files.keyIterator(); + while (fi.next()) |rel| { + const path = try std.fs.path.join(gpa, &.{ dir, rel.* }); + defer gpa.free(path); + var file = try core.safetensors.File.open(gpa, io, path); + defer file.close(io); + var ni = file.names.iterator(); + while (ni.next()) |e| total = try std.math.add(u64, total, try tensorBytes(e.key_ptr.*, e.value_ptr.*, o)); + } + return total; +} + +/// The family's deviceWeightBytes for an EXL3 checkpoint; 0 if unread. +pub fn weightBytes(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) u64 { + return deviceBytes(gpa, io, dir, options()) catch 0; +} + +/// A rank's placement bytes: one rank holds the whole model (no MTP head: an EXL3 pack decodes plain). +pub fn placement(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, world: usize, rank: usize) !u64 { + if (world != 1 or rank != 0) return error.BadRank; + return deviceBytes(gpa, io, dir, options()); +} + +fn entry(dtype: core.safetensors.DType, shape: []const usize, bytes: u64) Entry { + var e: Entry = .{ .dtype = dtype, .rank = @intCast(shape.len), .shape = @splat(1), .begin = 0, .end = bytes }; + for (shape, 0..) |d, i| e.shape[i] = d; + return e; +} + +test "exl3: routed trellis stays on the host under offload, its scales and every other EXL3 matrix are resident" { + const tr = entry(.i16, &.{ 160, 40, 24 }, 160 * 40 * 24 * 2); + const routed = "model.language_model.layers.0.mlp.experts.7.gate_proj.trellis"; + try std.testing.expectEqual(@as(u64, 0), try tensorBytes(routed, tr, .{})); + try std.testing.expectEqual(tr.end, try tensorBytes(routed, tr, .{ .routed_on_host = false })); + const suh = entry(.f16, &.{2560}, 5120); + try std.testing.expectEqual(@as(u64, 5120), try tensorBytes("model.language_model.layers.0.mlp.experts.7.gate_proj.suh", suh, .{})); + try std.testing.expectEqual(tr.end, try tensorBytes("model.language_model.layers.0.mlp.shared_expert.gate_proj.trellis", tr, .{})); + try std.testing.expectEqual(tr.end, try tensorBytes("lm_head.trellis", tr, .{})); + try std.testing.expectEqual(@as(u64, 0), try tensorBytes("lm_head.mul1", entry(.i32, &.{}, 4), .{})); +} + +test "exl3: dense F16 weights count as their BF16 device packing, norms as fp32" { + const router = entry(.f16, &.{ 512, 2560 }, 512 * 2560 * 2); + try std.testing.expectEqual(router.end, try tensorBytes("model.language_model.layers.3.mlp.gate.weight", router, .{})); + const down = entry(.f16, &.{ 320, 10240 }, 320 * 10240 * 2); + try std.testing.expectEqual(@as(u64, cuda.nvfp4.padded(320) * 10240 * 2), try tensorBytes("model.language_model.layers.3.attn_hyper_connection.input_mix_weight_down.weight", down, .{})); + const norm = entry(.f16, &.{10240}, 20480); + try std.testing.expectEqual(@as(u64, 40960), try tensorBytes("model.language_model.layers.3.attn_hyper_connection.hc_norm.weight", norm, .{})); + const gnorm = entry(.bf16, &.{128}, 256); + try std.testing.expectEqual(@as(u64, 256), try tensorBytes("model.language_model.layers.0.linear_attn.norm.weight", gnorm, .{})); + try std.testing.expectEqual(@as(u64, 48 * 4), try tensorBytes("model.language_model.layers.0.linear_attn.A_log", entry(.bf16, &.{48}, 96), .{})); + try std.testing.expectEqual(@as(u64, 0), try tensorBytes("model.visual.blocks.0.attn.qkv.weight", entry(.bf16, &.{ 3456, 1152 }, 3456 * 1152 * 2), .{})); +} + +test "exl3: placement names an EXL3 checkpoint and refuses any world but one" { + const a = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(std.testing.io, .{ .sub_path = "config.json", .data = @embedFile("fixtures/exl3_config.json") }); + const dir = try tmp.dir.realPathFileAlloc(std.testing.io, ".", a); + defer a.free(dir); + try std.testing.expect(isExl3(a, std.testing.io, dir)); + try std.testing.expectError(error.BadRank, placement(a, std.testing.io, dir, 2, 0)); + try std.testing.expectError(error.BadRank, placement(a, std.testing.io, dir, 1, 1)); +} diff --git a/zig/src/families/flashnext/exl3/pack/cuda_exl3_store.zig b/zig/src/families/flashnext/exl3/pack/cuda_exl3_store.zig new file mode 100644 index 0000000000..147d2dbffb --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/cuda_exl3_store.zig @@ -0,0 +1,121 @@ +//! The EXL3 pack's tensors as the Flash Next loaders read them (cuda_prompt_wt.zig `Store`). + +const std = @import("std"); +const core = @import("core"); + +/// The mixer file the index does not name (the MTP head's `mtp.hyper_connection_mixer.*`). +pub const patch_file = "mtp_hyper_connection_mixer_patch.safetensors"; + +/// Adds every tensor of `dir`'s mixer patch to the index's weight map (the file absent: nothing). +pub fn addPatch(io: std.Io, dir: []const u8, parsed: *std.json.Parsed(std.json.Value)) !void { + const a = parsed.arena.allocator(); + const wm = (parsed.value.object.getPtr("weight_map") orelse return); + if (wm.* != .object) return; + const path = try std.fs.path.join(a, &.{ dir, patch_file }); + var file = core.safetensors.File.open(a, io, path) catch |err| switch (err) { + error.FileNotFound => return, + else => return err, + }; + defer file.close(io); + var it = file.names.iterator(); + while (it.next()) |e| { + if (wm.object.contains(e.key_ptr.*)) continue; + try wm.object.put(a, try a.dupe(u8, e.key_ptr.*), .{ .string = patch_file }); + } +} + +/// Whether the loaders read `name` widened: a dense F16 `.weight` (an EXL3 part never ends in `.weight`). +pub fn widens(name: []const u8, t: core.safetensors.Tensor) bool { + return t.dtype == .f16 and std.mem.endsWith(u8, name, ".weight"); +} + +/// RNE to BF16 of an F32 (cuda_prompt_wt.toBf16's rounding, without its NaN quieting: a NaN is inexact here). +fn bf16Of(v: f32) u16 { + const bits: u32 = @bitCast(v); + return @intCast((bits + 0x7FFF + ((bits >> 16) & 1)) >> 16); +} + +/// The F16 words as BF16, exactly, or error.F16NotBf16Exact at the first value BF16 cannot hold (NaN included). +pub fn exactBf16(out: []u16, f16_bytes: []const u8) !void { + if (f16_bytes.len != out.len * 2) return error.UnexpectedTensor; + for (out, 0..) |*o, i| { + const h: f16 = @bitCast(std.mem.readInt(u16, f16_bytes[2 * i ..][0..2], .little)); + const v: f32 = h; + const b = bf16Of(v); + const back: f32 = @bitCast(@as(u32, b) << 16); + if (back != v or std.math.isNan(v)) return error.F16NotBf16Exact; + o.* = b; + } +} + +/// A widened copy of `t` as a BF16 tensor of the same shape, in `owned` (the caller frees it). +pub fn widen(gpa: std.mem.Allocator, name: []const u8, t: core.safetensors.Tensor, owned: *std.ArrayList([]u16)) !core.safetensors.Tensor { + const count = t.bytes.len / 2; + const words = try gpa.alloc(u16, count); + errdefer gpa.free(words); + exactBf16(words, t.bytes) catch |err| { + if (err == error.F16NotBf16Exact) std.log.err("flashnext: {s} is F16 with a value BF16 cannot hold exactly", .{name}); + return err; + }; + try owned.append(gpa, words); + var out = t; + out.dtype = .bf16; + out.bytes = std.mem.sliceAsBytes(words); + return out; +} + +test "exl3: an F16 that came from BF16 widens bit for bit and one extra mantissa bit is refused" { + const values = [_]f32{ 0, -0.0, 1, -2.5, 0.0078125, 65280, -3.0517578125e-05, 5.9604644775390625e-08, std.math.inf(f32) }; + var bytes: [values.len * 2]u8 = undefined; + for (values, 0..) |v, i| { + const b: f32 = @bitCast(@as(u32, bf16Of(v)) << 16); + try std.testing.expectEqual(v, b); // every case is a BF16 value + std.mem.writeInt(u16, bytes[2 * i ..][0..2], @bitCast(@as(f16, @floatCast(v))), .little); + } + var out: [values.len]u16 = undefined; + try exactBf16(&out, &bytes); + for (values, out) |v, o| try std.testing.expectEqual(v, @as(f32, @bitCast(@as(u32, o) << 16))); + // 1 + 2^-10 needs 11 significant bits: F16 holds it, BF16 does not + std.mem.writeInt(u16, bytes[0..2], @bitCast(@as(f16, 1.0009765625)), .little); + try std.testing.expectError(error.F16NotBf16Exact, exactBf16(&out, &bytes)); + std.mem.writeInt(u16, bytes[0..2], @bitCast(std.math.nan(f16)), .little); + try std.testing.expectError(error.F16NotBf16Exact, exactBf16(&out, &bytes)); +} + +test "exl3: only dense F16 weights widen, EXL3 parts keep their dtype" { + const t: core.safetensors.Tensor = .{ .dtype = .f16, .rank = 1, .shape = .{ 4, 1, 1, 1, 1 }, .bytes = &.{} }; + try std.testing.expect(widens("model.language_model.layers.0.mlp.gate.weight", t)); + try std.testing.expect(widens("mtp.hyper_connection_mixer.hc_norm.weight", t)); + try std.testing.expect(!widens("model.language_model.layers.0.self_attn.q_proj.suh", t)); + try std.testing.expect(!widens("lm_head.svh", t)); + var b = t; + b.dtype = .bf16; + try std.testing.expect(!widens("model.language_model.layers.0.linear_attn.norm.weight", b)); +} + +test "exl3: the mixer patch's names join the weight map, an indexed name keeps its shard" { + const a = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const header = "{\"mtp.hyper_connection_mixer.hc_norm.weight\":{\"dtype\":\"F16\",\"shape\":[1],\"data_offsets\":[0,2]},\"x\":{\"dtype\":\"U8\",\"shape\":[1],\"data_offsets\":[2,3]}}"; + var data: [8 + header.len + 3]u8 = undefined; + std.mem.writeInt(u64, data[0..8], header.len, .little); + @memcpy(data[8..][0..header.len], header); + @memset(data[8 + header.len ..], 0); + try tmp.dir.writeFile(io, .{ .sub_path = patch_file, .data = &data }); + const dir = try tmp.dir.realPathFileAlloc(io, ".", a); + defer a.free(dir); + var parsed = try std.json.parseFromSlice(std.json.Value, a, "{\"weight_map\":{\"x\":\"model-00001-of-00007.safetensors\"}}", .{}); + defer parsed.deinit(); + try addPatch(io, dir, &parsed); + const wm = parsed.value.object.get("weight_map").?.object; + try std.testing.expectEqualStrings(patch_file, wm.get("mtp.hyper_connection_mixer.hc_norm.weight").?.string); + try std.testing.expectEqualStrings("model-00001-of-00007.safetensors", wm.get("x").?.string); + // a folder without the patch file opens as before + var other = std.testing.tmpDir(.{}); + defer other.cleanup(); + const empty = try other.dir.realPathFileAlloc(io, ".", a); + defer a.free(empty); + try addPatch(io, empty, &parsed); +} diff --git a/zig/src/families/flashnext/exl3/pack/exl3_format.zig b/zig/src/families/flashnext/exl3/pack/exl3_format.zig new file mode 120000 index 0000000000..66d9fdff5e --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_format.zig @@ -0,0 +1 @@ +../../../../core/exl3_format.zig \ No newline at end of file diff --git a/zig/src/families/flashnext/exl3/pack/exl3_place.zig b/zig/src/families/flashnext/exl3/pack/exl3_place.zig new file mode 100644 index 0000000000..554ccf6881 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_place.zig @@ -0,0 +1,309 @@ +//! Header-only home for one tensor name. + +const std = @import("std"); +const scan = @import("exl3_scan.zig"); +const shape = @import("exl3_shape.zig"); +const cache = @import("../model/exl3_cache.zig"); + +pub const Home = enum { + ngram_host, // PLE rows, page cache or SSD, never a VRAM alloc + routed_host, // routed expert bank, host RAM, cache policy decides GPU copies + shared_resident, // shared expert, replicated, stays on device + resident, // dense, router, HC, MTP head, unquantized fp16 linears + centred_norm, // add centredOffset in fp32 at load + mtp_patch, // the mixer file kept outside the index + vision_skip, // vision encoder, not this milestone + plain, +}; + +pub const Tensor = struct { + name: []const u8, + rank: usize, + /// Filename of the shard, used only for the MTP mixer patch. + file: []const u8 = "", +}; + +const centred_names = [_][]const u8{ + "hc_norm", + "q_norm", + "k_norm", + "q_layernorm", + "k_layernorm", + "norm_key", + "norm_query", + "norm_conv", + "pre_fc_norm_embedding", + "pre_fc_norm_hidden", +}; + +const resident_names = [_][]const u8{ + "lm_head", + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "out_proj", + "router", + "hyper_connection", + "in_proj", + "conv1d", + "A_log", + "dt_bias", + "indexer", + "mtp", + // Router logits and the PLE pair. `k_proj` is not a substring of `key_proj`. + "mlp.gate", + "key_proj", + "value_proj", +}; + +comptime { + std.debug.assert(shape.experts == 512); + std.debug.assert(shape.shared_expert == shape.experts); +} + +pub fn home(t: Tensor) Home { + if (scan.isMtpMixerPatch(t.file)) return .mtp_patch; + + const kind = scan.classifyName(t.name, t.rank); + if (kind == .ngram_row) return .ngram_host; + if (isVision(t.name)) return .vision_skip; + if (contains(t.name, "shared_expert")) return .shared_resident; + if (isRouted(t.name)) return .routed_host; + + // Gamma is a rank-1 `.plain` weight. + if (isCentred(t.name)) { + if (kind == .exl3_part) return .plain; + return .centred_norm; + } + + // `mtp` pins the MTP head. + if (contains(t.name, "mtp.fc_embedding")) return .plain; + + if ((kind == .exl3_part or kind == .plain) and isResident(t.name)) return .resident; + return .plain; +} + +/// True when `which` is the cache's host place and both host banks refuse a resident copy. +pub fn hostBank(which: cache.Place) bool { + const budget = cache.Budget{ + .device_bytes = 0, + .resident_bytes = 0, + .kv_bytes = 0, + .routed_expert_bytes = 0, + .ngram_bytes = 0, + }; + const host = cache.refusesResident(.routed) and cache.refusesResident(.ngram) and cache.placeWeights(budget) == .host; + return host and which == .host; +} + +fn contains(hay: []const u8, needle: []const u8) bool { + return std.mem.indexOf(u8, hay, needle) != null; +} + +fn isVision(name: []const u8) bool { + return contains(name, "visual") or std.mem.startsWith(u8, name, "vision") or contains(name, ".vision."); +} + +/// `.experts.` / `experts.` gate, up, and down for ids 0..511. Any `shared` substring stays out. +fn isRouted(name: []const u8) bool { + if (contains(name, "shared")) return false; + return contains(name, ".experts.") or contains(name, "experts."); +} + +/// Token is a full dot-component, or the name's suffix. +pub fn isCentred(name: []const u8) bool { + for (centred_names) |token| { + if (std.mem.endsWith(u8, name, token)) { + if (name.len == token.len or name[name.len - token.len - 1] == '.') return true; + } + var it = std.mem.splitScalar(u8, name, '.'); + while (it.next()) |part| if (std.mem.eql(u8, part, token)) return true; + } + return false; +} + +fn isResident(name: []const u8) bool { + for (resident_names) |token| { + if (!contains(name, token)) continue; + // `mlp.switch_mlp.gate_proj` contains `mlp.gate` and is not the router logits. + if (std.mem.eql(u8, token, "mlp.gate") and contains(name, "gate_proj")) continue; + return true; + } + return false; +} + +test "a tensor name picks one home and the host banks stay off device" { + const ngram = "model.language_model.layers.2.ple.ngram_embedding.shard_0.trellis"; + try std.testing.expectEqual(Home.ngram_host, home(.{ .name = ngram, .rank = 2 })); + + const wide = home(.{ .name = ngram, .rank = 3 }); + try std.testing.expect(wide != .ngram_host); + try std.testing.expectEqual(Home.plain, wide); + + try std.testing.expectEqual(Home.routed_host, home(.{ + .name = "model.layers.0.mlp.experts.0.up_proj.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.shared_resident, home(.{ + .name = "model.layers.0.mlp.shared_expert.up_proj.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.centred_norm, home(.{ + .name = "model.layers.0.self_attn.q_norm.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.centred_norm, home(.{ + .name = "model.layers.0.attn_hyper_connection.hc_norm.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.centred_norm, home(.{ + .name = "language_model.mtp.pre_fc_norm_embedding.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.plain, home(.{ + .name = "model.layers.0.self_attn.q_norm.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.mtp_patch, home(.{ + .name = "weight", + .rank = 1, + .file = "mtp_hyper_connection_mixer_patch.safetensors", + })); + try std.testing.expectEqual(Home.vision_skip, home(.{ + .name = "model.visual.blocks.0.weight", + .rank = 1, + })); + try std.testing.expectEqual(Home.plain, home(.{ + .name = "model.embed_tokens.weight", + .rank = 2, + })); + try std.testing.expectEqual(Home.resident, home(.{ + .name = "model.layers.0.linear_attn.out_proj.weight", + .rank = 2, + })); + try std.testing.expectEqual(Home.plain, home(.{ .name = "readme", .rank = 0 })); + try std.testing.expect(hostBank(.host)); +} + +/// Where the bytes live. +pub const Stay = enum { vram_pinned, vram_learned, host, disk, skip }; + +/// Integer after `.experts.`, or null when the name is not a routed expert. +pub fn expertId(name: []const u8) ?u16 { + const key = ".experts."; + const at = std.mem.indexOf(u8, name, key) orelse return null; + var i = at + key.len; + if (i >= name.len or name[i] < '0' or name[i] > '9') return null; + var n: u32 = 0; + while (i < name.len and name[i] >= '0' and name[i] <= '9') : (i += 1) { + n = n * 10 + (name[i] - '0'); + if (n > std.math.maxInt(u16)) return null; + } + if (i >= name.len or name[i] != '.') return null; + return @intCast(n); +} + +pub fn stay(which: Home, id: ?u16) error{IllegalCut}!Stay { + switch (which) { + .ngram_host => return .disk, + .vision_skip => return .skip, + .routed_host => { + _ = id orelse return error.IllegalCut; + return .vram_learned; + }, + .plain => return .host, + .shared_resident, .resident, .centred_norm, .mtp_patch => return .vram_pinned, + } +} + +pub fn stayOf(t: Tensor) error{IllegalCut}!Stay { + return stay(home(t), expertId(t.name)); +} + +test "the decode path stays pinned, the n-gram on disk, and the routed experts are learned" { + const ngram = "model.language_model.layers.2.ple.ngram_embedding.shard_0.trellis"; + try std.testing.expectEqual(Stay.disk, try stayOf(.{ .name = ngram, .rank = 2 })); + + const pinned = [_][]const u8{ + "model.layers.0.mlp.shared_expert.up_proj.trellis", + "lm_head.weight", + "model.layers.0.mlp.router.weight", + "model.layers.0.mlp_hyper_connection.weight", + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + "model.layers.0.self_attn.o_proj.weight", + "model.layers.0.linear_attn.in_proj_qkv.weight", + "model.layers.0.linear_attn.out_proj.weight", + "model.layers.0.linear_attn.conv1d.weight", + "model.layers.0.linear_attn.A_log", + "model.layers.0.linear_attn.dt_bias", + "model.language_model.layers.0.mlp.gate.weight", + "model.language_model.layers.1.ple.key_proj.weight", + "model.language_model.layers.1.ple.value_proj.weight", + }; + try std.testing.expectEqual(Stay.host, try stayOf(.{ + .name = "model.embed_tokens.weight", + .rank = 2, + })); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = "readme", .rank = 0 })); + for (pinned) |name| { + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 })); + } + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ + .name = "model.layers.0.self_attn.q_norm.weight", + .rank = 1, + })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ + .name = "weight", + .rank = 1, + .file = "mtp_hyper_connection_mixer_patch.safetensors", + })); + try std.testing.expectEqual(Stay.skip, try stayOf(.{ + .name = "model.visual.blocks.0.weight", + .rank = 1, + })); + + try std.testing.expectEqual(@as(?u16, 0), expertId("model.layers.0.mlp.experts.0.up_proj.trellis")); + try std.testing.expectEqual(@as(?u16, 255), expertId("model.layers.0.mlp.experts.255.down_proj.trellis")); + try std.testing.expectEqual(@as(?u16, 511), expertId("model.layers.0.mlp.experts.511.gate_proj.trellis")); + try std.testing.expectEqual(@as(?u16, null), expertId("model.layers.0.mlp.shared_expert.up_proj.trellis")); + + const low = "model.layers.0.mlp.experts.0.up_proj.trellis"; + const high = "model.layers.0.mlp.experts.400.up_proj.trellis"; + try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = low, .rank = 3 })); + try std.testing.expectEqual(Stay.vram_learned, try stayOf(.{ .name = high, .rank = 3 })); + + try std.testing.expectError(error.IllegalCut, stay(.routed_host, null)); + try std.testing.expect(std.mem.indexOf(u8, "key_proj", "k_proj") == null); + try std.testing.expect(std.mem.indexOf(u8, "value_proj", "v_proj") == null); + try std.testing.expect(std.mem.indexOf(u8, "mlp.shared_expert_gate", "mlp.gate") == null); +} + +test "mlp.gate, key_proj, and value_proj stay pinned, and so does mlp.router.weight" { + const gate = "model.language_model.layers.0.mlp.gate.weight"; + const key = "model.language_model.layers.1.ple.key_proj.weight"; + const value = "model.language_model.layers.1.ple.value_proj.weight"; + const router = "model.layers.0.mlp.router.weight"; + const shared_gate = "model.layers.0.mlp.shared_expert_gate.weight"; + for ([_][]const u8{ gate, key, value, router }) |name| { + try std.testing.expectEqual(Home.resident, home(.{ .name = name, .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = name, .rank = 2 })); + } + // `shared_expert` matches before the resident tokens. This is not the 2560ร—640 expert. + try std.testing.expectEqual(Home.shared_resident, home(.{ .name = shared_gate, .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = shared_gate, .rank = 1 })); + + const switch_gate = "model.layers.0.mlp.switch_mlp.gate_proj.weight"; + try std.testing.expectEqual(Home.plain, home(.{ .name = switch_gate, .rank = 2 })); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = switch_gate, .rank = 2 })); + try std.testing.expectEqual(Home.routed_host, home(.{ + .name = "model.layers.0.mlp.experts.0.gate_proj.trellis", + .rank = 3, + })); + try std.testing.expectEqual(Home.plain, home(.{ .name = "mtp.fc_embedding.weight", .rank = 2 })); + try std.testing.expectEqual(Stay.host, try stayOf(.{ .name = "mtp.fc_embedding.weight", .rank = 2 })); + try std.testing.expectEqual(Home.resident, home(.{ .name = "language_model.mtp.fc_hidden.weight", .rank = 2 })); + try std.testing.expectEqual(Stay.vram_pinned, try stayOf(.{ .name = "language_model.mtp.fc_hidden.weight", .rank = 2 })); +} diff --git a/zig/src/families/flashnext/exl3/pack/exl3_qmat.zig b/zig/src/families/flashnext/exl3/pack/exl3_qmat.zig new file mode 100644 index 0000000000..34e4fc3b5b --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_qmat.zig @@ -0,0 +1,229 @@ +//! Flash Next EXL3 CPU oracle: the unpacked quantized matrix, its Hadamard transform and the row-parallel helper. + +const std = @import("std"); +const format = @import("exl3_format.zig"); + +// -- parallel for ------------------------------------------------------------------------------------------ + +/// Runs `f(ctx, i)` for i in 0..n over up to `max_threads` threads (the caller included). +pub fn parallelFor(n: usize, max_threads: usize, ctx: anytype, comptime f: anytype) void { + const Ctx = @TypeOf(ctx); + const W = struct { + fn run(next: *std.atomic.Value(usize), count: usize, c: Ctx) void { + while (true) { + const i = next.fetchAdd(1, .monotonic); + if (i >= count) return; + f(c, i); + } + } + }; + var next = std.atomic.Value(usize).init(0); + const cpus = std.Thread.getCpuCount() catch 4; + const want = @min(@min(cpus, max_threads), n); + var handles: [128]std.Thread = undefined; + var spawned: usize = 0; + while (spawned + 1 < want and spawned < handles.len) : (spawned += 1) { + handles[spawned] = std.Thread.spawn(.{}, W.run, .{ &next, n, ctx }) catch break; + } + W.run(&next, n, ctx); + for (handles[0..spawned]) |h| h.join(); +} + +// -- QMat -------------------------------------------------------------------------------------------------- + +/// In-place orthonormal Walsh-Hadamard transform of each 128 block of `x` (Sylvester order, H[i][j] =). +pub fn fwht128(x: []f64) void { + std.debug.assert(x.len % 128 == 0); + const inv = 1.0 / @sqrt(128.0); + var base: usize = 0; + while (base < x.len) : (base += 128) { + const b = x[base..][0..128]; + var h: usize = 1; + while (h < 128) : (h *= 2) { + var i: usize = 0; + while (i < 128) : (i += 2 * h) { + for (i..i + h) |j| { + const lo = b[j]; + const hi = b[j + h]; + b[j] = lo + hi; + b[j + h] = lo - hi; + } + } + } + for (b) |*v| v.* *= inv; + } +} + +pub const QMat = struct { + k: usize, + n: usize, + wq: []f32, + suh: []f32, + svh: []f32, + + /// `threads` parallelises the trellis unpack over tile rows. + pub fn init(gpa: std.mem.Allocator, halves: u16, trellis: []const i16, tiles_k: usize, tiles_n: usize, suh: []const f16, svh: []const f16, threads: usize) !QMat { + const k = tiles_k * 16; + const n = tiles_n * 16; + if (k % 128 != 0 or n % 128 != 0 or suh.len != k or svh.len != n) return error.Length; + if (trellis.len != tiles_k * tiles_n * @as(usize, halves) * 8) return error.Length; + const wq = try gpa.alloc(f32, k * n); + errdefer gpa.free(wq); + const su = try gpa.alloc(f32, k); + errdefer gpa.free(su); + const sv = try gpa.alloc(f32, n); + errdefer gpa.free(sv); + for (su, suh) |*d, s| d.* = @floatCast(s); + for (sv, svh) |*d, s| d.* = @floatCast(s); + var job: UnpackJob = .{ .halves = halves, .trellis = trellis, .tiles_n = tiles_n, .wq = wq, .n = n }; + parallelFor(tiles_k, threads, &job, UnpackJob.row); + return .{ .k = k, .n = n, .wq = wq, .suh = su, .svh = sv }; + } + + pub fn deinit(m: *QMat, gpa: std.mem.Allocator) void { + gpa.free(m.wq); + gpa.free(m.suh); + gpa.free(m.svh); + m.* = undefined; + } + + /// y[0..n] = x[0..k] W. `tmp` is k long. + pub fn apply(m: QMat, x: []const f64, y: []f64, tmp: []f64) void { + std.debug.assert(x.len == m.k and y.len == m.n and tmp.len == m.k); + for (tmp, x, m.suh) |*d, v, s| d.* = v * s; + fwht128(tmp); + @memset(y, 0); + for (0..m.k) |i| { + const u = tmp[i]; + if (u == 0) continue; + const row = m.wq[i * m.n ..][0..m.n]; + for (y, row) |*o, w| o.* += u * @as(f64, w); + } + fwht128(y); + for (y, m.svh) |*o, s| o.* *= s; + } + + /// Same as `apply` with the output columns [lo, hi) only (the Hadamard blocks must be whole: lo. + pub fn applyCols(m: QMat, tmp_h: []const f64, y: []f64, lo: usize, hi: usize) void { + std.debug.assert(lo % 128 == 0 and hi % 128 == 0 and y.len == hi - lo); + @memset(y, 0); + for (0..m.k) |i| { + const u = tmp_h[i]; + if (u == 0) continue; + const row = m.wq[i * m.n + lo ..][0 .. hi - lo]; + for (y, row) |*o, w| o.* += u * @as(f64, w); + } + fwht128(y); + for (y, m.svh[lo..hi]) |*o, s| o.* *= s; + } + + /// The transformed input (x * suh) H_K that `applyCols` consumes. + pub fn prepare(m: QMat, x: []const f64, out: []f64) void { + for (out, x, m.suh) |*d, v, s| d.* = v * s; + fwht128(out); + } +}; + +const UnpackJob = struct { + halves: u16, + trellis: []const i16, + tiles_n: usize, + wq: []f32, + n: usize, + fn row(j: *UnpackJob, tk: usize) void { + var placement: [256]format.Placement = undefined; + format.tilePlacement(&placement); + const w: format.Width = .{ .halves = j.halves }; + const words = w.tileWords(); + for (0..j.tiles_n) |tn| { + const tile = j.trellis[(tk * j.tiles_n + tn) * words ..][0..words]; + var states: [256]u32 = undefined; + format.tileStates(w, tile, &states); + for (0..256) |p| { + const v: f16 = format.codebookValue(.mul1, @intCast(states[p])); + j.wq[(tk * 16 + placement[p].row) * j.n + tn * 16 + placement[p].col] = @floatCast(v); + } + } + } +}; + +test "fwht128 is orthonormal and equals the dense Sylvester matrix" { + var rng = std.Random.DefaultPrng.init(3); + var x: [256]f64 = undefined; + for (&x) |*v| v.* = rng.random().float(f64) * 2 - 1; + var y = x; + fwht128(&y); + for (0..2) |blk| { + for (0..128) |i| { + var s: f64 = 0; + for (0..128) |j| { + const sign: f64 = if (@popCount(i & j) & 1 == 1) -1 else 1; + s += sign * x[blk * 128 + j]; + } + try std.testing.expectApproxEqAbs(s / @sqrt(128.0), y[blk * 128 + i], 1e-12); + } + } + fwht128(&y); + for (x, y) |a, b| try std.testing.expectApproxEqAbs(a, b, 1e-12); +} + +test "QMat.apply equals the dense decoder's dequantize then matvec" { + const gpa = std.testing.allocator; + var rng = std.Random.DefaultPrng.init(17); + const r = rng.random(); + for ([_]u16{ 4, 6, 10 }) |halves| { + const tiles_k: usize = 8; + const tiles_n: usize = 16; + const k = tiles_k * 16; + const n = tiles_n * 16; + const words = @as(usize, halves) * 8; + const trellis = try gpa.alloc(i16, tiles_k * tiles_n * words); + defer gpa.free(trellis); + for (trellis) |*v| v.* = @bitCast(r.int(u16)); + const suh = try gpa.alloc(f16, k); + defer gpa.free(suh); + const svh = try gpa.alloc(f16, n); + defer gpa.free(svh); + for (suh) |*v| v.* = if (r.boolean()) 1.0 else -1.0; + for (svh) |*v| v.* = @floatCast(0.5 + r.float(f32)); + const dense = try gpa.alloc(f64, k * n); + defer gpa.free(dense); + const scratch = try gpa.alloc(f64, 2 * k * n); + defer gpa.free(scratch); + format.dequantize(.{ .halves = halves }, .mul1, trellis, tiles_k, tiles_n, suh, svh, dense, scratch); + var q = try QMat.init(gpa, halves, trellis, tiles_k, tiles_n, suh, svh, 4); + defer q.deinit(gpa); + var x: [k]f64 = undefined; + for (&x) |*v| v.* = r.floatNorm(f64); + var want: [n]f64 = undefined; + for (0..n) |j| { + var s: f64 = 0; + for (0..k) |i| s += x[i] * dense[i * n + j]; + want[j] = s; + } + var got: [n]f64 = undefined; + var tmp: [k]f64 = undefined; + q.apply(&x, &got, &tmp); + for (want, got) |a, b| try std.testing.expectApproxEqAbs(a, b, 1e-9 * (1 + @abs(a))); + // column slices equal the whole + var prep: [k]f64 = undefined; + q.prepare(&x, &prep); + var part: [128]f64 = undefined; + q.applyCols(&prep, &part, 128, 256); + for (part, got[128..256]) |a, b| try std.testing.expectApproxEqAbs(a, b, 1e-12); + } +} + +test "parallelFor visits every index once" { + const Ctx = struct { + hits: []std.atomic.Value(u8), + fn f(c: *@This(), i: usize) void { + _ = c.hits[i].fetchAdd(1, .monotonic); + } + }; + var hits: [1000]std.atomic.Value(u8) = undefined; + for (&hits) |*h| h.* = std.atomic.Value(u8).init(0); + var c: Ctx = .{ .hits = &hits }; + parallelFor(hits.len, 8, &c, Ctx.f); + for (&hits) |*h| try std.testing.expectEqual(@as(u8, 1), h.load(.monotonic)); +} diff --git a/zig/src/families/flashnext/exl3/pack/exl3_scan.zig b/zig/src/families/flashnext/exl3/pack/exl3_scan.zig new file mode 100644 index 0000000000..06b53571d8 --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_scan.zig @@ -0,0 +1,206 @@ +//! EXL3 safetensors header scan. + +const std = @import("std"); + +/// Same tags, in the same order, as `exl3_format.Codebook`. +pub const format = struct { + pub const Codebook = enum { inst3, mcg, mul1 }; +}; + +pub const Scales = enum { suh, su, svh, sv }; + +pub const Group = struct { + halves: u16, + codebook: format.Codebook, + k: usize, + n: usize, + in_scales: Scales, + out_scales: Scales, + bias: bool, +}; + +pub const Parts = struct { + trellis_words: usize, + k_tiles: usize, + n_tiles: usize, + suh_len: ?usize = null, + su_len: ?usize = null, + svh_len: ?usize = null, + sv_len: ?usize = null, + bias: bool = false, + mcg: bool = false, + mul1: bool = false, +}; + +pub const Kind = enum { exl3_part, ngram_row, plain }; + +const suffixes = [_][]const u8{ "trellis", "suh", "su", "svh", "sv", "bias", "mcg", "mul1" }; + +/// One EXL3 linear group from its header parts. Half-bit widths need the mul1 marker. +pub fn parseGroup(parts: Parts) error{NotExl3}!Group { + if (parts.trellis_words == 0 or parts.trellis_words % 8 != 0) return error.NotExl3; + const halves_usize = parts.trellis_words / 8; + if (halves_usize < 2 or halves_usize > 16) return error.NotExl3; + const halves: u16 = @intCast(halves_usize); + if (parts.mcg and parts.mul1) return error.NotExl3; + const codebook: format.Codebook = if (parts.mul1) .mul1 else if (parts.mcg) .mcg else .inst3; + if (halves % 2 == 1 and codebook != .mul1) return error.NotExl3; + + const k = parts.k_tiles * 16; + const n = parts.n_tiles * 16; + if (k == 0 or n == 0 or k % 128 != 0 or n % 128 != 0) return error.NotExl3; + + const in_scales: Scales = if (parts.suh_len != null and parts.su_len == null) + .suh + else if (parts.su_len != null and parts.suh_len == null) + .su + else + return error.NotExl3; + if (in_scales == .suh and parts.suh_len.? != k) return error.NotExl3; + if (in_scales == .su and parts.su_len.? != k / 16) return error.NotExl3; + + const out_scales: Scales = if (parts.svh_len != null and parts.sv_len == null) + .svh + else if (parts.sv_len != null and parts.svh_len == null) + .sv + else + return error.NotExl3; + if (out_scales == .svh and parts.svh_len.? != n) return error.NotExl3; + if (out_scales == .sv and parts.sv_len.? != n / 16) return error.NotExl3; + + return .{ + .halves = halves, + .codebook = codebook, + .k = k, + .n = n, + .in_scales = in_scales, + .out_scales = out_scales, + .bias = parts.bias, + }; +} + +/// An n-gram shard is a rank-2 `ngram_embedding` `.trellis`, not an EXL3 tile tensor. +pub fn classifyName(name: []const u8, rank: usize) Kind { + if (rank == 2 and std.mem.indexOf(u8, name, "ngram_embedding") != null and std.mem.endsWith(u8, name, ".trellis")) { + return .ngram_row; + } + const dot = std.mem.lastIndexOfScalar(u8, name, '.') orelse return .plain; + const suffix = name[dot + 1 ..]; + for (suffixes) |part| if (std.mem.eql(u8, suffix, part)) return .exl3_part; + return .plain; +} + +/// The MTP mixer's extra file, the one turboderp keeps outside the index. +pub fn isMtpMixerPatch(filename: []const u8) bool { + return std.mem.eql(u8, filename, "mtp_hyper_connection_mixer_patch.safetensors"); +} + +fn codebookName(name: []const u8) error{QuantRefuse}!void { + if (std.mem.eql(u8, name, "mul1") or std.mem.eql(u8, name, "mcg") or std.mem.eql(u8, name, "3inst") or std.mem.eql(u8, name, "inst3")) return; + return error.QuantRefuse; +} + +fn integerBits(bits: f32) error{QuantRefuse}!void { + if (!std.math.isFinite(bits) or bits < 1 or bits > 8) return error.QuantRefuse; + if (bits != @round(bits)) return error.QuantRefuse; +} + +/// Admit a pack from config fields alone. Average bits may be fractional. Head and MTP bits may not. +pub fn admit(codebook: []const u8, average_bits: f32, head_bits: ?f32, mtp_bits: ?f32) error{QuantRefuse}!void { + try codebookName(codebook); + if (!std.math.isFinite(average_bits) or average_bits < 1 or average_bits > 8) return error.QuantRefuse; + if (head_bits) |bits| try integerBits(bits); + if (mtp_bits) |bits| try integerBits(bits); +} + +/// 1 when centred norms are stored as gamma-1, 0 when stored as gamma. +pub fn centredOffset(means: []const f32) error{NormAmbiguous}!f32 { + if (means.len == 0) return 1; + if (means.len > 64) return error.NormAmbiguous; + var copy: [64]f32 = undefined; + @memcpy(copy[0..means.len], means); + var i: usize = 1; + while (i < means.len) : (i += 1) { + const v = copy[i]; + var j = i; + while (j > 0 and copy[j - 1] > v) : (j -= 1) copy[j] = copy[j - 1]; + copy[j] = v; + } + const median = copy[means.len / 2]; + var above: usize = 0; + for (means) |m| if (m > 0.5) { + above += 1; + }; + const frac = @as(f32, @floatFromInt(above)) / @as(f32, @floatFromInt(means.len)); + const around_zero = frac <= 0.1 and median >= -0.5 and median <= 0.25; + const around_one = frac >= 0.9 and median >= 0.75 and median <= 1.5; + if (around_zero == around_one) return error.NormAmbiguous; + return if (around_zero) 1 else 0; +} + +test "a 3-bit mul1 gate group parses and a half-bit group needs mul1" { + const gate = try parseGroup(.{ + .trellis_words = 48, + .k_tiles = 160, + .n_tiles = 40, + .suh_len = 2560, + .svh_len = 640, + .mul1 = true, + }); + try std.testing.expectEqual(@as(u16, 6), gate.halves); + try std.testing.expectEqual(format.Codebook.mul1, gate.codebook); + try std.testing.expectEqual(@as(usize, 2560), gate.k); + try std.testing.expectEqual(@as(usize, 640), gate.n); + try std.testing.expectEqual(Scales.suh, gate.in_scales); + try std.testing.expectEqual(Scales.svh, gate.out_scales); + + const half = Parts{ + .trellis_words = 40, + .k_tiles = 160, + .n_tiles = 40, + .suh_len = 2560, + .svh_len = 640, + .mul1 = true, + }; + try std.testing.expectEqual(@as(u16, 5), (try parseGroup(half)).halves); + var no_marker = half; + no_marker.mul1 = false; + try std.testing.expectError(error.NotExl3, parseGroup(no_marker)); + var both = half; + both.mcg = true; + try std.testing.expectError(error.NotExl3, parseGroup(both)); +} + +test "K must be a multiple of 128" { + try std.testing.expectError(error.NotExl3, parseGroup(.{ + .trellis_words = 48, + .k_tiles = 20, + .n_tiles = 40, + .suh_len = 320, + .svh_len = 640, + .mul1 = true, + })); +} + +test "names sort into n-gram rows, EXL3 parts and plain tensors" { + try std.testing.expectEqual(Kind.exl3_part, classifyName("model.language_model.layers.0.mlp.experts.0.gate_proj.trellis", 3)); + try std.testing.expectEqual(Kind.ngram_row, classifyName("ple.ngram_embedding.shard_3.trellis", 2)); + try std.testing.expectEqual(Kind.exl3_part, classifyName("ple.ngram_embedding.shard_3.trellis", 3)); + try std.testing.expectEqual(Kind.plain, classifyName("model.embed_tokens.weight", 2)); + try std.testing.expect(isMtpMixerPatch("mtp_hyper_connection_mixer_patch.safetensors")); + try std.testing.expect(!isMtpMixerPatch("model.safetensors")); +} + +test "the 3.05 config is admitted and a fractional head width is not" { + try admit("mul1", 3.05, 5, 3); + try admit("3inst", 4, null, null); + try std.testing.expectError(error.QuantRefuse, admit("mul1", 3.05, 3.05, null)); + try std.testing.expectError(error.QuantRefuse, admit("exl2", 3, null, null)); +} + +test "centred norms are gamma-1 when the means sit near zero" { + try std.testing.expectEqual(@as(f32, 1), try centredOffset(&.{ -0.02, 0.01, 0.0, 0.04 })); + try std.testing.expectEqual(@as(f32, 0), try centredOffset(&.{ 1.0, 0.98, 1.02, 1.01 })); + try std.testing.expectError(error.NormAmbiguous, centredOffset(&.{ 0.4, 0.9 })); + try std.testing.expectEqual(@as(f32, 1), try centredOffset(&.{})); +} diff --git a/zig/src/families/flashnext/exl3/pack/exl3_shape.zig b/zig/src/families/flashnext/exl3/pack/exl3_shape.zig new file mode 100644 index 0000000000..ec97e5d1df --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_shape.zig @@ -0,0 +1,66 @@ +//! Published text shape of turboderp/Qwen3.8-Flash-Next-exl3, revision 3.05bpw_h5_ng5. + +const std = @import("std"); + +pub const model = "turboderp/Qwen3.8-Flash-Next-exl3"; +pub const revision = "3.05bpw_h5_ng5"; + +pub const hidden: usize = 2560; +pub const layers: usize = 48; +pub const vocab: usize = 248320; +pub const heads: usize = 24; +pub const kv_heads: usize = 2; +pub const head_dim: usize = 256; +pub const experts: usize = 512; +pub const topk: usize = 10; +pub const expert_intermediate: usize = 640; +pub const shared_intermediate: usize = 640; +/// Routed ids are 0..511. The shared expert is the next id in the same table. +pub const shared_expert: usize = experts; +pub const hc_count: usize = 4; +pub const hc_lowrank: usize = 320; + +pub const ngram: usize = 3; +pub const ngram_heads_per: usize = 8; +pub const ngram_dh: usize = 160; +/// Branch name `ng5`: the n-gram row codec is 5 bits. Confirmed by words-per-row once a shard header is read. +pub const ngram_bits: usize = 5; +pub const ngram_shards: usize = 128; +pub const ngram_vocab_base: usize = 20_000_000; + +pub const quant_method = "exl3"; +pub const quant_version = "1.4.4"; +pub const average_bits: f32 = 3.05; +pub const head_bits: usize = 5; +pub const mtp_bits: usize = 3; +pub const codebook = "mul1"; + +pub const tile: usize = 16; +pub const hadamard: usize = 128; + +pub const linear_layers_per_cycle: usize = 3; +pub const full_attention_interval: usize = 4; + +/// Heads in one n-gram lookup: (ngram - 1) * heads_per_ngram. +pub fn ngramHeadCount() usize { + return (ngram - 1) * ngram_heads_per; +} + +/// q_proj outputs. 24 * 256 = 6144. +pub fn qOut() usize { + return heads * head_dim; +} + +/// One attention layer's kv width, keys or values. 2 * 256 = 512. +pub fn kvOut() usize { + return kv_heads * head_dim; +} + +test "the published 3.05 shape is the one the budget is written against" { + try std.testing.expectEqual(@as(usize, 16), ngramHeadCount()); + try std.testing.expectEqual(@as(usize, 6144), qOut()); + try std.testing.expectEqual(@as(usize, 512), kvOut()); + try std.testing.expect(hidden % hadamard == 0); + try std.testing.expect(expert_intermediate % hadamard == 0); + try std.testing.expect(qOut() % hadamard == 0); +} diff --git a/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig b/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig new file mode 100644 index 0000000000..d924759dec --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/exl3_startup_memory.zig @@ -0,0 +1,72 @@ +const std = @import("std"); + +pub const Reservation = struct { + stream_bytes: u64, + streams: u64, + + pub fn total(self: Reservation) !u64 { + return std.math.mul(u64, self.stream_bytes, self.streams); + } +}; + +pub const Ledger = struct { + persistent: u64, + scratch: u64, + pool_scratch: u64, + expert_pools: u64, + reservation: Reservation, + + pub fn total(self: Ledger) !u64 { + var bytes = try self.reservation.total(); + inline for (.{ self.persistent, self.scratch, self.pool_scratch, self.expert_pools }) |part| bytes = try std.math.add(u64, bytes, part); + return bytes; + } +}; + +pub fn mebibytes(text: []const u8) !u64 { + return std.math.mul(u64, try std.fmt.parseInt(u64, text, 10), 1 << 20); +} + +pub fn available(free: u64, reserve: u64, future: u64) !u64 { + const held = try std.math.add(u64, reserve, future); + if (held > free) return error.InsufficientStartupMemory; + return free - held; +} + +pub fn poolBytes(lines: u64, line_bytes: u64, layers: u64, fixed: u64) !u64 { + return std.math.add(u64, try std.math.mul(u64, try std.math.mul(u64, lines, line_bytes), layers), fixed); +} + +test "BF16 c1 c8 preserve the requested stream bytes" { + const bf16_stream: u64 = 262144 * 256 * 2; + for ([_]u64{ 1, 8 }) |streams| { + const reservation: Reservation = .{ .stream_bytes = bf16_stream, .streams = streams }; + try std.testing.expectEqual(bf16_stream * streams, try reservation.total()); + const free = try available(16 << 30, 1 << 30, try reservation.total()); + try std.testing.expectEqual(@as(u64, 16 << 30), free + (1 << 30) + try reservation.total()); + } +} + +test "startup accounting fails closed on overflow and underfunded explicit pools" { + try std.testing.expectError(error.Overflow, (Reservation{ .stream_bytes = std.math.maxInt(u64), .streams = 8 }).total()); + try std.testing.expectError(error.Overflow, mebibytes("18446744073709551615")); + try std.testing.expectError(error.Overflow, poolBytes(std.math.maxInt(u64), 2, 49, 0)); + try std.testing.expectError(error.InsufficientStartupMemory, available(100, 20, try poolBytes(10, 1, 8, 1))); + try std.testing.expectError(error.Overflow, available(100, std.math.maxInt(u64), 1)); +} + +test "ledger includes all allocations after the pool checkpoint" { + const ledger: Ledger = .{ .persistent = 100, .scratch = 20, .pool_scratch = 40, .expert_pools = 50, .reservation = .{ .stream_bytes = 10, .streams = 8 } }; + try std.testing.expectEqual(@as(u64, 290), try ledger.total()); +} + +test "pool setup may consume its allocator margin but not lane or server reservations" { + const server: u64 = 100; + const margin: u64 = 25; + const future: u64 = 80; + const planner: u64 = 10; + const planned_free = server + margin + future + planner; + try std.testing.expectEqual(margin, try available(planned_free, server, future + planner)); + try std.testing.expectEqual(@as(u64, 0), try available(planned_free - margin, server, future + planner)); + try std.testing.expectError(error.InsufficientStartupMemory, available(planned_free - margin - 1, server, future + planner)); +} diff --git a/zig/src/families/flashnext/exl3/pack/fixtures/exl3_config.json b/zig/src/families/flashnext/exl3/pack/fixtures/exl3_config.json new file mode 100644 index 0000000000..ea4b025ebf --- /dev/null +++ b/zig/src/families/flashnext/exl3/pack/fixtures/exl3_config.json @@ -0,0 +1,168 @@ +{ + "architectures": [ + "Qwen4ExpForConditionalGeneration" + ], + "image_token_id": 248056, + "language_model_only": false, + "model_type": "qwen4_exp", + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 248044, + "dtype": "bfloat16", + "eos_token_id": 248044, + "full_attention_interval": 4, + "hc_count": 4, + "hc_lowrank": 320, + "head_dim": 256, + "heads_per_ngram": 8, + "hidden_act": "silu", + "hidden_size": 2560, + "indexer_budget": 2048, + "indexer_compress_ratio": 4, + "indexer_head_dim": 128, + "indexer_kv_heads": 1, + "indexer_n_heads": 4, + "initializer_range": 0.02, + "layer_types": [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention" + ], + "linear_conv_kernel_dim": 4, + "linear_key_head_dim": 128, + "linear_num_key_heads": 16, + "linear_num_value_heads": 48, + "linear_value_head_dim": 128, + "make_ngram_vocab_size_divisible_by": 128, + "mamba_ssm_dtype": "float32", + "max_position_embeddings": 262144, + "model_type": "qwen4_exp_text", + "moe_intermediate_size": 640, + "mtp": { + "hybrid": true, + "layer_types": [ + "full_attention" + ], + "mtp_use_hidden_state_from_layer": null, + "num_hidden_layers": 1, + "rope_theta": 10000000 + }, + "mtp_num_hidden_layers": 1, + "mtp_use_dedicated_embeddings": false, + "ngram_size": 3, + "ngram_vocab_size_base": 20000000, + "num_attention_heads": 24, + "num_experts": 512, + "num_experts_per_tok": 10, + "num_hidden_layers": 48, + "num_key_value_heads": 2, + "output_gate_type": "sigmoid", + "output_router_logits": false, + "pad_token_id": null, + "partial_rotary_factor": 0.25, + "ple_conv_kernel_size": 4, + "ple_embed_dim": 2560, + "ple_layer_ids": [ + 2 + ], + "rms_norm_eps": 1e-06, + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 11, + 11, + 10 + ], + "partial_rotary_factor": 0.25, + "rope_theta": 10000000, + "rope_type": "default" + }, + "router_aux_loss_coef": 0.001, + "shared_expert_intermediate_size": 640, + "split_ngram_parts": 128, + "tie_word_embeddings": false, + "use_cache": true, + "vocab_size": 248320 + }, + "tie_word_embeddings": false, + "transformers_version": "5.8.0.dev0", + "video_token_id": 248057, + "vision_config": { + "deepstack_visual_indexes": [], + "depth": 27, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 1152, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 4304, + "model_type": "qwen4_exp", + "num_heads": 16, + "num_position_embeddings": 2304, + "out_hidden_size": 2560, + "patch_size": 16, + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "vision_end_token_id": 248054, + "vision_start_token_id": 248053, + "quantization_config": { + "quant_method": "exl3", + "version": "1.4.4", + "bits": 3.05, + "head_bits": 5, + "calibration": { + "rows": 250, + "cols": 2048 + }, + "out_scales": "always", + "codebook": "mul1", + "vision_bits": 5, + "mtp_bits": 3 + } +} \ No newline at end of file diff --git a/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork.zig b/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork.zig new file mode 100644 index 0000000000..4f296e551d --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork.zig @@ -0,0 +1,126 @@ +//! A multi-part EXL3 projection's parts on parallel streams (a DeltaNet layer's q|k|v, z, b, a; attention's q, k, v). + +const std = @import("std"); +const cuda = @import("cuda"); +const load = @import("../pack/cuda_exl3_load.zig"); + +pub const branches = 3; +/// Rows a forked projection takes at most (decode rounds and verify windows; prompts keep one stream). +pub const max_rows: usize = 16; + +pub const Fork = struct { + side: [branches]cuda.Stream, + dense: [branches]cuda.exl3.Dense, + tmp: [branches]cuda.DeviceBuffer, + ev_fork: cuda.Event, + ev_join: [branches]cuda.Event, + tmp_bytes: usize, +}; + +pub fn enabled() bool { + return force orelse true; +} + +/// A test's override of the fork. +pub var force: ?bool = null; + +/// Ops.reserve's hook: the branches' streams, kernels and scratch for `o`'s widest projection, before any pass. +pub fn reserve(o: *load.Ops) !void { + if (o.max_k == 0 or o.fork != null) return; + const f = try o.gpa.create(Fork); + errdefer o.gpa.destroy(f); + f.tmp_bytes = max_rows * @max(o.max_part_n, 128) * 2; + var made: usize = 0; + errdefer for (0..made) |i| { + f.dense[i].unload(); + f.side[i].deinit(); + f.tmp[i].free(); + f.ev_join[i].deinit(); + }; + f.ev_fork = try cuda.Event.init(o.d, false); + errdefer f.ev_fork.deinit(); + for (0..branches) |i| { + var side = try cuda.Stream.init(o.d, true); + errdefer side.deinit(); + var dense = try cuda.exl3.Dense.load(o.d); + errdefer dense.unload(); + var tmp = try cuda.DeviceBuffer.alloc(o.d, f.tmp_bytes); + errdefer tmp.free(); + var join = try cuda.Event.init(o.d, false); + errdefer join.deinit(); + dense.xh = try cuda.DeviceBuffer.alloc(o.d, 128 * o.max_k * 2); + if (o.max_z > 0) dense.z = try cuda.DeviceBuffer.alloc(o.d, o.max_z); + dense.counters = try cuda.DeviceBuffer.alloc(o.d, o.max_counters); + f.side[i] = side; + f.dense[i] = dense; + f.tmp[i] = tmp; + f.ev_join[i] = join; + made += 1; + } + o.fork = f; +} + +pub fn release(o: *load.Ops) void { + const f = o.fork orelse return; + o.fork = null; + for (0..branches) |i| { + f.side[i].synchronize() catch {}; + f.dense[i].unload(); + f.side[i].deinit(); + f.tmp[i].free(); + f.ev_join[i].deinit(); + } + f.ev_fork.deinit(); + o.gpa.destroy(f); +} + +fn copy2d(s: cuda.Stream, dst: u64, dpitch: usize, src: u64, spitch: usize, width: usize, height: usize) !void { + const m: cuda.abi.Memcpy2D = .{ .src = src, .src_pitch = spitch, .dst = dst, .dst_pitch = dpitch, .width = width, .height = height }; + try s.d.check(s.d.api.cuMemcpy2DAsync_v2(&m, s.handle), "cuMemcpy2DAsync"); +} + +/// Pass.lane()'s parts loop (bf16 out, sk 1) with the EXL3 parts after the first on side streams; false when the +pub fn lane(p: anytype, x: u64, e: *const load.Proj, out: u64, m: usize, comptime runMain: anytype) !bool { + if (!enabled() or m > max_rows) return false; + var n_exl3: usize = 0; + for (e.parts) |pt| n_exl3 += @intFromBool(pt.kind == .exl3); + if (n_exl3 < 2) return false; + const f = e.ops.fork orelse return false; + for (e.parts) |pt| if (pt.kind == .exl3 and m * pt.n * 2 > f.tmp_bytes and pt.n != e.n and m != 1) return false; + const s = p.stream; + try f.ev_fork.record(s); + var b: usize = 0; + var first = true; + for (e.parts) |pt| { + const w = switch (pt.kind) { + .exl3 => |w| w, + .native => { + try runMain(p, e, pt, x, m, out); + continue; + }, + }; + if (first or b == branches) { + first = false; + try runMain(p, e, pt, x, m, out); + continue; + } + const st = f.side[b]; + try st.wait(f.ev_fork); + const de = &f.dense[b]; + if (pt.n == e.n or m == 1) { + try de.decode(st, w, x, .bf16, m, out + pt.col * 2, .bf16); + } else { + const tmp = f.tmp[b].ptr; + try de.decode(st, w, x, .bf16, m, tmp, .bf16); + try copy2d(st, out + pt.col * 2, e.n * 2, tmp, pt.n * 2, pt.n * 2, m); + } + try f.ev_join[b].record(st); + b += 1; + } + for (f.ev_join[0..b]) |ev| try s.wait(ev); + return true; +} + +test { + _ = @import("cuda_exl3_fork_test.zig"); +} diff --git a/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork_test.zig b/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork_test.zig new file mode 100644 index 0000000000..e690f9e32e --- /dev/null +++ b/zig/src/families/flashnext/exl3/tp/cuda_exl3_fork_test.zig @@ -0,0 +1,105 @@ +const std = @import("std"); +const cuda = @import("cuda"); +const load = @import("../pack/cuda_exl3_load.zig"); +const fork = @import("cuda_exl3_fork.zig"); + +const Fake = struct { + var streams: usize = 0; + var modules: usize = 0; + var events: usize = 0; + var buffers: usize = 0; + fn sync(_: cuda.abi.Stream) callconv(.c) c_int { + return 0; + } + fn stream(_: cuda.abi.Stream) callconv(.c) c_int { + streams += 1; + return 0; + } + fn module(_: cuda.abi.Module) callconv(.c) c_int { + modules += 1; + return 0; + } + fn event(_: cuda.abi.Event) callconv(.c) c_int { + events += 1; + return 0; + } + fn alloc(out: *u64, _: usize) callconv(.c) c_int { + out.* = 0x1000; + return 0; + } + fn buffer(_: u64) callconv(.c) c_int { + buffers += 1; + return 0; + } + fn createStream(out: *cuda.abi.Stream, _: c_uint) callconv(.c) c_int { + out.* = null; + return 0; + } + fn createEvent(out: *cuda.abi.Event, _: c_uint) callconv(.c) c_int { + out.* = null; + return 0; + } + fn driver() cuda.Driver { + streams = 0; + modules = 0; + events = 0; + buffers = 0; + var d: cuda.Driver = undefined; + d.api.cuStreamCreate = createStream; + d.api.cuEventCreate = createEvent; + d.api.cuStreamSynchronize = sync; + d.api.cuStreamDestroy_v2 = stream; + d.api.cuModuleUnload = module; + d.api.cuEventDestroy_v2 = event; + d.api.cuMemAlloc_v2 = alloc; + d.api.cuMemFree_v2 = buffer; + return d; + } +}; + +fn attach(o: *load.Ops) !*fork.Fork { + const f = try o.gpa.create(fork.Fork); + f.tmp_bytes = 64; + f.ev_fork = .{ .d = o.d, .handle = null }; + for (0..fork.branches) |i| { + f.side[i] = .{ .d = o.d, .handle = null }; + f.dense[i] = .{ .d = o.d, .lin = .{ .mod = .{ .d = o.d, .handle = null }, .fns = undefined, .rot_in = undefined } }; + f.tmp[i] = try cuda.DeviceBuffer.alloc(o.d, 64); + f.ev_join[i] = .{ .d = o.d, .handle = null }; + } + o.fork = f; + return f; +} + +test "exl3: projection forks retire with their Ops owner and preserve other owners" { + var d = Fake.driver(); + const left = try load.Ops.create(std.testing.allocator, &d); + const right = try load.Ops.create(std.testing.allocator, &d); + const l = try attach(left); + const r = try attach(right); + _ = left.retain(); + left.release(); + try std.testing.expect(left.fork == l and right.fork == r); + try std.testing.expectEqual(@as(usize, 0), Fake.modules); + left.release(); + try std.testing.expect(right.fork == r); + try std.testing.expectEqual(fork.branches, Fake.modules); + right.release(); + try std.testing.expectEqual(2 * fork.branches, Fake.streams); + try std.testing.expectEqual(2 * fork.branches, Fake.modules); + try std.testing.expectEqual(2 * (fork.branches + 1), Fake.events); + try std.testing.expectEqual(2 * fork.branches, Fake.buffers); +} + +test "exl3: fork reservation failure releases the current branch before publication" { + if (cuda.kernels.available) return error.SkipZigTest; + var d = Fake.driver(); + const o = try load.Ops.create(std.testing.allocator, &d); + defer o.release(); + o.max_k = 128; + try std.testing.expectError(error.BuiltWithoutKernels, fork.reserve(o)); + try std.testing.expect(o.fork == null); + try std.testing.expectEqual(@as(usize, 1), Fake.streams); + try std.testing.expectEqual(@as(usize, 1), Fake.events); + try std.testing.expectEqual(@as(usize, 0), Fake.modules); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index ab6a6e9443..e6b5d9c796 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -358,7 +358,8 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O low = try std.fmt.allocPrint(a, " on rank {d}, the smallest room of the ranks", .{pr.rank}); }; const free = loaded.free_streams; - const grow = budget.growRoom(room, !@hasField(F.Options, "lanes")); + const preallocates_lanes = if (@hasDecl(F, "preallocates_lanes")) F.preallocates_lanes else @hasField(F.Options, "lanes"); + const grow = budget.growRoom(room, !preallocates_lanes); const wrapped = kv_pages.wrapped(F, loaded, free, served.stream_bytes, room); const layout = wrapped.layout; var paged: ?kv_pages.Plan = null; @@ -412,7 +413,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O // rules and budget from the family's snap_* declarations (cuda_snaps.zig) const image: ?*const pc.rows.Image = if (F == nemotron.native) &cuda_disk.image else if (comptime cuda_disk.keptDisk(F)) |K| &K.image else null; const copies: pc.Snapshots = .{ .ptr = loaded.target, .vtable = &snaps.Snaps(F, image).vtable }; - store.* = pc.Store.init(gpa, copies, snaps.passRules(F, @import("cuda_startup.zig").rows(F, loaded)), if (@hasDecl(F, "snap_slots")) cache else try snaps.cacheBudget(F, loaded.target, cache)); + store.* = pc.Store.init(gpa, copies, snaps.passRules(F, @import("cuda_startup.zig").rows(F, loaded), snaps.loadedGrid(loaded)), if (@hasDecl(F, "snap_slots")) cache else try snaps.cacheBudget(F, loaded.target, cache)); errdefer { const disk = store.disk; store.deinit(); diff --git a/zig/src/native/cuda_snaps.zig b/zig/src/native/cuda_snaps.zig index 35a5785f1a..eda5ba048f 100644 --- a/zig/src/native/cuda_snaps.zig +++ b/zig/src/native/cuda_snaps.zig @@ -15,15 +15,24 @@ pub fn cacheRules(comptime F: type) api.prompt_cache.Rules { return rules; } -/// Native CUDA families own their grid, independently of the server's template chunk starts. -pub fn passRules(comptime F: type, prompt_rows: u32) api.prompt_cache.Rules { +/// Native CUDA families own their grid; a loaded `snap_grid` (the prefill step) keeps states on chunk starts only. +pub fn passRules(comptime F: type, prompt_rows: u32, snap_grid: u32) api.prompt_cache.Rules { var rules = cacheRules(F); rules.chunk = prompt_rows; if (@hasDecl(F, "snap_repeat_min")) rules.repeat_min = if (std.c.getenv("TF_FLASHNEXT_REPEAT_CACHE")) |v| (if (std.mem.eql(u8, std.mem.span(v), "0")) 0 else F.snap_repeat_min) else F.snap_repeat_min; rules.chunk_starts = false; + if (snap_grid > 0) { + rules.planned = true; + rules.grid = snap_grid; + } return rules; } +/// The grid a loaded family declares (0: any row). +pub fn loadedGrid(loaded: anytype) u32 { + return if (@hasField(@TypeOf(loaded), "snap_grid")) loaded.snap_grid else 0; +} + /// The cache's bytes less what the family sets aside for its states (`snap_ready`, as Nemotron's spare buffers). pub fn cacheBudget(comptime F: type, target: anytype, cache: u64) !u64 { return if (@hasDecl(F, "snap_ready")) cache - try F.snap_ready(target, cache) else cache; @@ -91,3 +100,34 @@ test "a family with snap_grid keeps on its chunk starts, any other at any row as test { _ = @import("cuda_snaps_test.zig"); } + +test "a family that names its prefill step keeps marks on chunk starts, so a resumed pass cuts as a fresh one" { + const Exl3 = struct {}; + const Stub = struct { + fn bytes(_: *anyopaque, _: u32) u64 { + return 0; + } + fn save(_: *anyopaque, _: ?*anyopaque, _: u32) anyerror!api.prompt_cache.Saved { + return error.Unused; + } + fn restore(_: *anyopaque, _: ?*anyopaque, _: api.prompt_cache.Saved) anyerror!void {} + fn drop(_: *anyopaque, _: api.prompt_cache.Saved) void {} + }; + var unused: u8 = 0; + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const prompt = try arena.allocator().alloc(u32, 12780); + for (prompt, 0..) |*t, i| t.* = @intCast(i); + const vtable: api.prompt_cache.Snapshots.VTable = .{ .bytes = Stub.bytes, .save = Stub.save, .restore = Stub.restore, .drop = Stub.drop }; + const copies: api.prompt_cache.Snapshots = .{ .ptr = &unused, .vtable = &vtable }; + // any row: the history mark cuts the cold pass's last chunk to 7 rows, a different kernel path than the fresh pass + var any = api.prompt_cache.Store.init(std.testing.allocator, copies, passRules(Exl3, 4096, 0), 1 << 30); + defer any.deinit(); + const loose = try any.marks(arena.allocator(), prompt, 0, 12773, &.{}, &.{}, &.{}); + try std.testing.expect(std.mem.indexOfScalar(u32, loose, 12773) != null); + var grid = api.prompt_cache.Store.init(std.testing.allocator, copies, passRules(Exl3, 4096, 4096), 1 << 30); + defer grid.deinit(); + const exact = try grid.marks(arena.allocator(), prompt, 0, 12773, &.{}, &.{}, &.{}); + try std.testing.expect(exact.len > 0); + for (exact) |m| try std.testing.expectEqual(@as(u32, 0), m % 4096); +} diff --git a/zig/src/native/cuda_startup_test.zig b/zig/src/native/cuda_startup_test.zig index 1a7d265a81..d9aea9735b 100644 --- a/zig/src/native/cuda_startup_test.zig +++ b/zig/src/native/cuda_startup_test.zig @@ -99,7 +99,17 @@ test "CUDA startup: every other registered family retains its declared pass size test "CUDA startup: Flash Next cache cuts use the loaded cold pass size" { var loaded: flash.native.Loaded = undefined; loaded.prompt_rows = 4096; - const rules = @import("cuda_snaps.zig").passRules(flash.native, startup.rows(flash.native, loaded)); + loaded.snap_grid = 0; + const rules = @import("cuda_snaps.zig").passRules(flash.native, startup.rows(flash.native, loaded), @import("cuda_snaps.zig").loadedGrid(loaded)); try std.testing.expectEqual(@as(u32, 4096), rules.chunk); - try std.testing.expect(!rules.chunk_starts and rules.lookahead == 1); + try std.testing.expect(!rules.chunk_starts and rules.lookahead == 1 and !rules.planned); +} + +test "CUDA startup: an EXL3 Flash Next keeps and resumes states on its prefill step only" { + var loaded: flash.native.Loaded = undefined; + loaded.prompt_rows = 4096; + loaded.snap_grid = 4096; + const snaps = @import("cuda_snaps.zig"); + const rules = snaps.passRules(flash.native, startup.rows(flash.native, loaded), snaps.loadedGrid(loaded)); + try std.testing.expect(rules.planned and rules.grid == 4096 and !rules.chunk_starts); } diff --git a/zig/tests/cuda/contract_limits.zig b/zig/tests/cuda/contract_limits.zig index 2a16afa5b6..48bb601d76 100644 --- a/zig/tests/cuda/contract_limits.zig +++ b/zig/tests/cuda/contract_limits.zig @@ -26,6 +26,8 @@ pub fn image(name: []const u8) error{UnknownContractImage}![]const u8 { inline for (.{ "qwen4_projection", "qwen4_hc", "qwen4_gdn", "qwen4_ple" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); inline for (.{ "tree_round", "act_rows", "gdn_prep", "qk_rows", "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "qmm_prefill", "lane_gemv", "prompt16", "qmmf", "ck_experts", "nvfp4_experts", "nvfp4_experts_tile", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "glm_glue", "glm_ops", "glm_index", "nemotron_mamba", "sample", "qwen35", "qwen35_bf16_rows", "qwen35_keep", "qwen35_kv8", "w4a8_prefill", "qwen35_draft", "ck_lane", "gdn_train", "train_q4", "train_attn", "flashnext_train", "train", "attn_train", "qwen27_train", "tree_attention", "block_attention", "window_attention", "window_pages", "moe_glue", "dense_rows", "qmm", "flashnext_fact", "flashnext_hcf", "flashnext_fexp", "experts_rows" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{"flashnext_exl3"}) |n| + if (std.mem.eql(u8, name, n)) return @field(kk, n); if (std.mem.eql(u8, name, "contract_controls")) return if (options.controls) &controls.bytes else &.{}; std.debug.print("FAIL kernel-contract: unknown image name: {s}\n", .{name}); return error.UnknownContractImage; diff --git a/zig/tests/cuda/flashnext_exl3/kv/exl3_kv.zig b/zig/tests/cuda/flashnext_exl3/kv/exl3_kv.zig new file mode 120000 index 0000000000..f5414b3e7b --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/kv/exl3_kv.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/kv/exl3_kv.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_attn_cpu.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_attn_cpu.zig new file mode 120000 index 0000000000..cb95d2d7a0 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_attn_cpu.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_attn_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_cache.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_cache.zig new file mode 120000 index 0000000000..a3cf4c6f46 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_cache.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_cache.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_cpu.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_cpu.zig new file mode 120000 index 0000000000..012c3b18c0 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_cpu.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_fwd_cpu.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_fwd_cpu.zig new file mode 120000 index 0000000000..8e732b9fe6 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_fwd_cpu.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_fwd_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_gdn_cpu.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_gdn_cpu.zig new file mode 120000 index 0000000000..3c9a8a52a3 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_gdn_cpu.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_gdn_cpu.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_hc.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_hc.zig new file mode 120000 index 0000000000..47b748a982 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_hc.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_hc.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_moe.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_moe.zig new file mode 120000 index 0000000000..cbbbe19f51 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_moe.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_moe.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_ple2_tests.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_ple2_tests.zig new file mode 100644 index 0000000000..a3e338829d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_ple2_tests.zig @@ -0,0 +1,598 @@ +//! exl3-ple2: an EXL3 pack's n-gram path on v1.1.0's Flash Next family (cuda_ple.Ngram with `format = .exl3`). 1. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const st = core.safetensors; +const ple = flash.ple; +const x3 = flash.exl3_ple; +const Gpu = check.Gpu; + +const py_ids_crc: u32 = 0xd04d5584; +const py_raw_crc: u32 = 0x2d57eafe; +const py_decoded_crc: u32 = 0xdcf520dc; + +const streams = 4; +const hidden = 2560; +const wide = streams * hidden; +const taps = 4; +const dilation = 3; +const tail_rows = (taps - 1) * dilation; +const eps: f32 = 1e-6; +const prefix = "model.language_model.layers.1.ple."; + +// ---- RSS ---------------------------------------------------------------------------------------------------------- + +const Rss = struct { total: u64 = 0, file: u64 = 0, anon: u64 = 0 }; // kB + +fn rss() Rss { + var buf: [1 << 14]u8 = undefined; + const fd = std.c.open("/proc/self/status", .{ .ACCMODE = .RDONLY }); + if (fd < 0) return .{}; + defer _ = std.c.close(fd); + var len: usize = 0; + while (len < buf.len) { + const n = std.c.read(fd, buf[len..].ptr, buf.len - len); + if (n <= 0) break; + len += @intCast(n); + } + var out: Rss = .{}; + var it = std.mem.splitScalar(u8, buf[0..len], '\n'); + while (it.next()) |line| { + const fields = [_]struct { key: []const u8, dst: *u64 }{ .{ .key = "VmRSS:", .dst = &out.total }, .{ .key = "RssFile:", .dst = &out.file }, .{ .key = "RssAnon:", .dst = &out.anon } }; + for (fields) |f| { + if (!std.mem.startsWith(u8, line, f.key)) continue; + var parts = std.mem.tokenizeAny(u8, line[f.key.len..], " \tkB"); + if (parts.next()) |n| f.dst.* = std.fmt.parseInt(u64, n, 10) catch 0; + } + } + return out; +} + +fn showRss(what: []const u8) Rss { + const r = rss(); + std.debug.print("RSS {s}: total {d} MiB, file-backed {d} MiB, anon {d} MiB\n", .{ what, r.total >> 10, r.file >> 10, r.anon >> 10 }); + return r; +} + +// ---- the branch's ids, ported from exl3-flashnext-window-wip zig/src/families/flashnext/exl3_ple_ids.zig ---------- + +/// Row ids [tokens.len][16] after `history` (incremental form; a token equal to eos starts a segment after itself). +fn oursIds(h: ple.Hash, history: [2]i64, tokens: []const u32, out: [][16]u64) void { + var seq: [3]i64 = undefined; + @memcpy(seq[0..2], &history); + var seg: usize = 0; + for (history) |t| seg = if (t == h.eos) 0 else @min(seg + 1, 2); + for (tokens, out[0..tokens.len]) |tok, *row| { + seq[2] = tok; + for (row, 0..) |*id, head| { + var mixed: u64 = 0; + for (0..2 + head / h.per) |q| { + const t = if (seg >= q) seq[2 - q] else h.eos; + const term = @as(u64, @bitCast(t)) *% @as(u64, @bitCast(h.mult[q])); + mixed = if (q == 0) term else mixed ^ term; + } + id.* = @as(u64, @intCast(h.offsets[head])) + mixed % @as(u64, @intCast(h.sizes[head])); + } + seg = if (tok == h.eos) 0 else @min(seg + 1, 2); + seq[0] = seq[1]; + seq[1] = seq[2]; + } +} + +// ---- weights -------------------------------------------------------------------------------------------------------- + +const Index = struct { + parsed: std.json.Parsed(std.json.Value), + files: std.StringHashMapUnmanaged(st.File) = .empty, + dir: []const u8, + + fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Index { + const path = try std.fs.path.join(gpa, &.{ dir, "model.safetensors.index.json" }); + defer gpa.free(path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); + defer gpa.free(text); + return .{ .parsed = try std.json.parseFromSlice(std.json.Value, gpa, text, .{}), .dir = dir }; + } + + fn deinit(x: *Index, gpa: std.mem.Allocator, io: std.Io) void { + var it = x.files.valueIterator(); + while (it.next()) |f| f.close(io); + x.files.deinit(gpa); + x.parsed.deinit(); + } + + fn get(x: *Index, gpa: std.mem.Allocator, io: std.Io, name: []const u8) !st.Tensor { + const rel = (x.parsed.value.object.get("weight_map").?.object.get(name) orelse return error.MissingTensor).string; + if (x.files.get(rel)) |f| return f.get(name) orelse error.MissingTensor; + const path = try std.fs.path.join(gpa, &.{ x.dir, rel }); + defer gpa.free(path); + const f = try st.File.open(gpa, io, path); + try x.files.put(gpa, rel, f); + return f.get(name) orelse error.MissingTensor; + } +}; + +fn widen(gpa: std.mem.Allocator, t: st.Tensor, offset: f32) ![]f32 { + const out = try gpa.alloc(f32, t.numel()); + for (out, 0..) |*o, i| { + const w = std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little); + const v: f32 = switch (t.dtype) { + .f16 => @floatCast(@as(f16, @bitCast(w))), + .bf16 => @bitCast(@as(u32, w) << 16), + else => return error.UnexpectedTensor, + }; + o.* = v + offset; + } + return out; +} + +fn bf(x: f32) f32 { + return @bitCast(@as(u32, ple.toBf16(x)) << 16); +} + +fn bfBits(x: f32) u16 { + return ple.toBf16(x); +} + +fn bfVal(b: u16) f32 { + return @bitCast(@as(u32, b) << 16); +} + +// ---- references of the injection (from exl3-flashnext-window-wip exl3_ple_cpu.zig) --------------------------------- + +const Weights = struct { + key: []const f32, // [wide][hidden] + value: []const f32, // [hidden][hidden] + norm_key: []const f32, // [wide], 1 + w + norm_query: []const f32, + norm_conv: []const f32, + conv: []const f32, // [wide][taps], as stored (F16) +}; + +fn dot32(w: []const f32, x: []const f32) f32 { + var acc: f32 = 0; + for (w, x) |a, b| acc += a * b; + return acc; +} + +/// The keys and values the reference's F16 dense projections store (fp32 sums, bf16 outputs). +fn project(w: Weights, emb: []const f32, keys: []f32, vals: []f32) void { + for (0..wide) |o| keys[o] = bf(dot32(w.key[o * hidden ..][0..hidden], emb)); + for (0..hidden) |o| vals[o] = bf(dot32(w.value[o * hidden ..][0..hidden], emb)); +} + +/// The reference's roundings (glue.py `_ple_gate` / `_ple_conv`) on the host, one token at a time. +const Branch = struct { + w: Weights, + conv_bf: []f32, + gated: []f32, + nrow: []f32, + tail: []f32, + + fn init(gpa: std.mem.Allocator, w: Weights) !Branch { + const conv_bf = try gpa.alloc(f32, w.conv.len); + for (conv_bf, w.conv) |*d, s| d.* = bf(s); + const tail = try gpa.alloc(f32, tail_rows * wide); + @memset(tail, 0); + return .{ .w = w, .conv_bf = conv_bf, .gated = try gpa.alloc(f32, wide), .nrow = try gpa.alloc(f32, wide), .tail = tail }; + } + + fn step(b: *Branch, keys: []const f32, vals: []const f32, h: []const f32, out: []f32) void { + const w = b.w; + const d: f32 = hidden; + for (0..streams) |s| { + const k = keys[s * hidden ..][0..hidden]; + const q = h[s * hidden ..][0..hidden]; + var kss: f32 = 0; + var qss: f32 = 0; + for (k, q) |a, c| { + kss += a * a; + qss += c * c; + } + const krinv = 1.0 / @sqrt(kss / d + eps); + const qrinv = 1.0 / @sqrt(qss / d + eps); + var dot: f32 = 0; + for (0..hidden) |i| dot += bf(bf(k[i] * krinv * w.norm_key[s * hidden + i]) * bf(q[i] * qrinv * w.norm_query[s * hidden + i])); + var gate = bf(dot / @sqrt(d)); + const mag = @sqrt(@max(@abs(gate), 1e-6)); + gate = bf(if (gate > 0) mag else if (gate < 0) -mag else 0); + const sg = bf(1.0 / (1.0 + @exp(-gate))); + var gss: f32 = 0; + for (0..hidden) |i| { + const g = bf(sg * vals[i]); + b.gated[s * hidden + i] = g; + gss += g * g; + } + const rinv = 1.0 / @sqrt(gss / d + eps); + for (0..hidden) |i| b.nrow[s * hidden + i] = bf(b.gated[s * hidden + i] * rinv * w.norm_conv[s * hidden + i]); + } + for (0..wide) |e| { + var acc: f32 = 0; + for (0..taps) |t| { + const at = t * dilation; + acc += b.conv_bf[e * taps + t] * (if (at < tail_rows) b.tail[at * wide + e] else b.nrow[e]); + } + const cv = bf(acc); + out[e] = bf(b.gated[e] + bf(cv / (1.0 + @exp(-cv)))); + } + std.mem.copyForwards(f32, b.tail[0 .. (tail_rows - 1) * wide], b.tail[wide..]); + @memcpy(b.tail[(tail_rows - 1) * wide ..], b.nrow); + } +}; + +/// model.py `PLELayer` in float64 with no intermediate rounding (emb as decoded fp16, weights as stored). +const Oracle = struct { + w: Weights, + tail: []f64, + keys: []f64, + vals: []f64, + gated: []f64, + normed: []f64, + + fn init(gpa: std.mem.Allocator, w: Weights) !Oracle { + const tail = try gpa.alloc(f64, tail_rows * wide); + @memset(tail, 0); + return .{ .w = w, .tail = tail, .keys = try gpa.alloc(f64, wide), .vals = try gpa.alloc(f64, hidden), .gated = try gpa.alloc(f64, wide), .normed = try gpa.alloc(f64, wide) }; + } + + fn centred(x: []const f64, scale: []const f32, out: []f64) void { + var ss: f64 = 0; + for (x) |v| ss += v * v; + const r = 1.0 / @sqrt(ss / @as(f64, @floatFromInt(x.len)) + @as(f64, eps)); + for (out, x, scale) |*o, v, s| o.* = v * r * s; + } + + fn step(o: *Oracle, emb: []const f32, h: []const f32, out: []f64) void { + const w = o.w; + for (0..wide) |r| { + var acc: f64 = 0; + for (w.key[r * hidden ..][0..hidden], emb) |a, b| acc += @as(f64, a) * b; + o.keys[r] = acc; + } + for (0..hidden) |r| { + var acc: f64 = 0; + for (w.value[r * hidden ..][0..hidden], emb) |a, b| acc += @as(f64, a) * b; + o.vals[r] = acc; + } + var kn: [hidden]f64 = undefined; + var qn: [hidden]f64 = undefined; + var hs: [hidden]f64 = undefined; + for (0..streams) |s| { + centred(o.keys[s * hidden ..][0..hidden], w.norm_key[s * hidden ..][0..hidden], &kn); + for (&hs, h[s * hidden ..][0..hidden]) |*d, v| d.* = v; + centred(&hs, w.norm_query[s * hidden ..][0..hidden], &qn); + var dot: f64 = 0; + for (kn, qn) |a, b| dot += a * b; + var gate = dot / @sqrt(@as(f64, hidden)); + const sgn: f64 = if (gate > 0) 1 else if (gate < 0) -1 else 0; + gate = sgn * @sqrt(@max(@abs(gate), 1e-6)); + const sg = 1.0 / (1.0 + @exp(-gate)); + for (0..hidden) |i| o.gated[s * hidden + i] = sg * o.vals[i]; + centred(o.gated[s * hidden ..][0..hidden], w.norm_conv[s * hidden ..][0..hidden], o.normed[s * hidden ..][0..hidden]); + } + for (0..wide) |e| { + var acc: f64 = 0; + for (0..taps) |t| { + const at = t * dilation; + acc += @as(f64, w.conv[e * taps + t]) * (if (at < tail_rows) o.tail[at * wide + e] else o.normed[e]); + } + out[e] = o.gated[e] + acc / (1.0 + @exp(-acc)); + } + std.mem.copyForwards(f64, o.tail[0 .. (tail_rows - 1) * wide], o.tail[wide..]); + @memcpy(o.tail[(tail_rows - 1) * wide ..], o.normed); + } +}; + +const Qbench = struct { rows: []const struct { input_ids: []const u32 } }; + +fn crc(bytes: []const u8) u32 { + return std.hash.Crc32.hash(bytes); +} + +fn bytesOf(comptime T: type, v: []const T) []const u8 { + return std.mem.sliceAsBytes(v); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + const qrow: usize = if (rest.len > 1) try std.fmt.parseInt(usize, rest[1], 10) else 3; + const count: usize = if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 24; + var arena_state = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + const io = gpu.io; + const rss_start = showRss("start"); + + const qpath = try std.fs.path.join(a, &.{ dir, "qbench_prompts.json" }); + const qtext = try std.Io.Dir.cwd().readFileAlloc(io, qpath, a, .limited(1 << 26)); + const parsed = try std.json.parseFromSlice(Qbench, a, qtext, .{ .ignore_unknown_fields = true }); + const qrows = parsed.value.rows; + try check.expect(qrow < qrows.len and count <= qrows[qrow].input_ids.len, "qbench row {d} with {d} tokens", .{ qrow, count }); + const tokens = qrows[qrow].input_ids[0..count]; + + // ---- the table through upstream's Ngram (mapping) ----------------------------------------------------------- + var g = try ple.Ngram.init(gpu.gpa, io, dir); + defer g.deinit(); + try check.expect(g.format == .exl3, "Ngram.init took the EXL3 table", .{}); + const h = g.hash; + const heads = h.heads; + const dh = h.dh; + std.debug.print("hash: n {d}, heads {d} x {d}, eos {d}, {d} shards, {d} rows; table {d}-bit, {d} bytes a row, read {t}, emb {t}\n", .{ h.n, heads, dh, h.eos, h.shards, h.starts[h.shards], g.x3.?.bits, g.x3.?.rowBytes(), g.x3.?.opts.read, g.x3.?.opts.emb }); + check.pass("Ngram.init on the pack: config-derived multipliers, head sizes and offsets equal the pack's shipped tensors; 128 shard row counts equal Hash.starts", .{}); + const rss_open = showRss("after Ngram.init (30.4 GiB table mapped, header only)"); + + // ---- (1) ids --------------------------------------------------------------------------------------------------- + const ids = try a.alloc(i64, count * heads); + { + var hist = [_]i64{ h.eos, h.eos }; + for (tokens, 0..) |tok, t| { + try ple.lastIds(h, &hist, tok, ids[t * heads ..][0..heads]); + hist = .{ hist[1], tok }; + } + const got = crc(bytesOf(i64, ids)); + std.debug.print("ids[0][0..4] = {any}; ids crc32 {x:0>8} (numpy NGram.ids {x:0>8})\n", .{ ids[0..4], got, py_ids_crc }); + if (qrow == 3 and count == 24) try check.expect(got == py_ids_crc, "upstream lastIds crc32 {x:0>8} != numpy {x:0>8}", .{ got, py_ids_crc }); + // every token of every qbench prompt: upstream's lastIds against the branch's incremental ids. + var compared: usize = 0; + var eos_seen: usize = 0; + for (qrows) |r| { + const ours = try gpu.gpa.alloc([16]u64, r.input_ids.len); + defer gpu.gpa.free(ours); + oursIds(h, .{ h.eos, h.eos }, r.input_ids, ours); + var hist2 = [_]i64{ h.eos, h.eos }; + for (r.input_ids, ours) |tok, want| { + if (tok == h.eos) eos_seen += 1; + var up: [16]i64 = undefined; + try ple.lastIds(h, &hist2, tok, up[0..heads]); + hist2 = .{ hist2[1], tok }; + for (up[0..heads], want[0..heads]) |u, o| { + if (u < 0 or @as(u64, @intCast(u)) != o) { + std.debug.print("FAIL ids differ at token {d}: upstream {any} ours {any}\n", .{ tok, up[0..heads], want[0..heads] }); + return error.TestFailed; + } + } + compared += 1; + } + } + check.pass("row ids: upstream Hash/lastIds byte-equal to the numpy NGram.ids crc{s} and to the branch's exl3_ple_ids on all {d} tokens of {d} qbench prompts ({d} eos inside)", .{ if (qrow == 3 and count == 24) "" else " (crc not checked: not row 3 x 24)", compared, qrows.len, eos_seen }); + } + + // ---- (2) raw rows, two readers; decoded rows through Ngram.rows, two readers ------------------------------------ + var pread_t = (try x3.Table.open(gpu.gpa, io, dir, h, .{ .read = .pread })).?; + defer pread_t.destroy(); + const rb = g.x3.?.rowBytes(); + const raw_map = try a.alloc(u8, ids.len * rb); + const raw_ssd = try a.alloc(u8, ids.len * rb); + for (ids, 0..) |id, i| { + const loc = try h.locate(id); + try g.x3.?.raw(loc.shard, loc.local, raw_map[i * rb ..][0..rb]); + try pread_t.raw(loc.shard, loc.local, raw_ssd[i * rb ..][0..rb]); + } + try check.sameBytes("rows: mapping vs pread", raw_ssd, raw_map); + const raw_crc = crc(raw_map); + std.debug.print("{d} rows x {d} bytes: raw crc32 {x:0>8} (numpy {x:0>8})\n", .{ ids.len, rb, raw_crc, py_raw_crc }); + if (qrow == 3 and count == 24) try check.expect(raw_crc == py_raw_crc, "raw rows crc32 {x:0>8} != numpy {x:0>8}", .{ raw_crc, py_raw_crc }); + const rss_map = showRss("after the mapped row reads"); + + const w = g.width(); + const emb_map = try a.alloc(u16, count * w); + const xs_map = try a.alloc(f32, count * w / 32); + { + var hist = [_]i64{ h.eos, h.eos }; + try g.rows(&hist, tokens, emb_map, xs_map); + } + const dec_crc = crc(bytesOf(u16, emb_map)); + var peak_emb: f32 = 0; + for (emb_map) |v| peak_emb = @max(peak_emb, @abs(@as(f32, @floatCast(@as(f16, @bitCast(v)))))); + std.debug.print("Ngram.rows: {d} fp16 values, max |value| {d:.4}, crc32 {x:0>8} (numpy _ple_rows {x:0>8})\n", .{ emb_map.len, peak_emb, dec_crc, py_decoded_crc }); + if (qrow == 3 and count == 24) try check.expect(dec_crc == py_decoded_crc, "decoded crc32 {x:0>8} != numpy {x:0>8}", .{ dec_crc, py_decoded_crc }); + // the same call with the pread reader swapped in behind Ngram + const emb_ssd = try a.alloc(u16, count * w); + const xs_ssd = try a.alloc(f32, count * w / 32); + { + const mapped_t = g.x3.?; + g.x3 = pread_t; + defer g.x3 = mapped_t; + var hist = [_]i64{ h.eos, h.eos }; + try g.rows(&hist, tokens, emb_ssd, xs_ssd); + } + try check.sameBytes("Ngram.rows: mapping vs pread", bytesOf(u16, emb_ssd), bytesOf(u16, emb_map)); + try check.sameBytes("Ngram.rows group sums: mapping vs pread", bytesOf(f32, xs_ssd), bytesOf(f32, xs_map)); + check.pass("rows: {d} ids, stored bytes and Ngram.rows fp16 values byte-identical through the mapping and through pread + hot cache{s}", .{ ids.len, if (qrow == 3 and count == 24) "; raw and decoded crc32 equal the numpy references" else "" }); + + // ---- (3) hot cache over real text; RSS ------------------------------------------------------------------------- + { + var fresh = (try x3.Table.open(gpu.gpa, io, dir, h, .{ .read = .pread })).?; + defer fresh.destroy(); + const r0 = rss(); + var buf: [256]u8 = undefined; + var ntok: usize = 0; + for (qrows) |r| { + var hist = [_]i64{ h.eos, h.eos }; + for (r.input_ids) |tok| { + var row_ids: [16]i64 = undefined; + try ple.lastIds(h, &hist, tok, row_ids[0..heads]); + hist = .{ hist[1], tok }; + for (row_ids[0..heads]) |id| { + const loc = try h.locate(id); + try fresh.raw(loc.shard, loc.local, buf[0..rb]); + } + ntok += 1; + } + } + const s = fresh.hotStats(); + const r1 = rss(); + std.debug.print("pread + 4096-row hot cache over all {d} qbench tokens: {d} lookups, {d} hits ({d:.2}%); RSS file-backed {d} -> {d} MiB, anon {d} -> {d} MiB\n", .{ ntok, s.lookups, s.lookups - s.misses, 100.0 * @as(f64, @floatFromInt(s.lookups - s.misses)) / @as(f64, @floatFromInt(s.lookups)), r0.file >> 10, r1.file >> 10, r0.anon >> 10, r1.anon >> 10 }); + try check.expect(r1.file <= r0.file + 16 * 1024, "pread reads grew file-backed RSS by {d} MiB", .{(r1.file - r0.file) >> 10}); + } + const rss_rows = showRss("after rows (both readers) and the hot-cache pass"); + const table_mib = g.x3.?.file.?.map.memory.len >> 20; + std.debug.print("table {d} MiB; file-backed RSS: start {d}, after open {d}, after mapped reads {d}, now {d} MiB\n", .{ table_mib, rss_start.file >> 10, rss_open.file >> 10, rss_map.file >> 10, rss_rows.file >> 10 }); + try check.expect(rss_map.file <= rss_open.file + 128 * 1024 and rss_rows.file <= 512 * 1024, "the table became resident", .{}); + check.pass("30.4 GiB table stays out of memory: file-backed RSS {d} MiB after the mapped reads (+{d} MiB), anon {d} MiB", .{ rss_rows.file >> 10, (rss_map.file - @min(rss_map.file, rss_open.file)) >> 10, rss_rows.anon >> 10 }); + + // ---- (4) layer 1's injection through upstream's fn_ple_gate / fn_ple_conv --------------------------------------- + try injection(gpu, a, dir, tokens, emb_map, h.eos); +} + +fn injection(gpu: Gpu, a: std.mem.Allocator, dir: []const u8, tokens: []const u32, emb_bits: []const u16, eos: i64) !void { + _ = eos; + const io = gpu.io; + const n = tokens.len; + var index = try Index.open(gpu.gpa, io, dir); + defer index.deinit(gpu.gpa, io); + const key_t = try index.get(gpu.gpa, io, prefix ++ "key_proj.weight"); + const val_t = try index.get(gpu.gpa, io, prefix ++ "value_proj.weight"); + const conv_t = try index.get(gpu.gpa, io, prefix ++ "conv1d.weight"); + try check.expect(key_t.is(.f16, &.{ wide, hidden }) and val_t.is(.f16, &.{ hidden, hidden }) and conv_t.is(.f16, &.{ wide, 1, taps }), "PLE key/value/conv dtypes or shapes (want F16)", .{}); + var norms: [3][]f32 = undefined; + for ([_][]const u8{ "norm_key", "norm_query", "norm_conv" }, 0..) |nm, i| { + var nb: [96]u8 = undefined; + const t = try index.get(gpu.gpa, io, try std.fmt.bufPrint(&nb, prefix ++ "{s}.weight", .{nm})); + try check.expect(t.is(.bf16, &.{wide}), "{s} dtype/shape (want bf16 [10240])", .{nm}); + const raw = try widen(a, t, 0); + var mean: f64 = 0; + for (raw) |v| mean += v; + mean /= wide; + std.debug.print("{s} stored mean {d:.4} (gamma - 1 storage: scale 1 + w)\n", .{ nm, mean }); + try check.expect(@abs(mean) < 0.5, "{s} does not look gamma - 1", .{nm}); + norms[i] = try widen(a, t, 1.0); + } + const weights: Weights = .{ .key = try widen(a, key_t, 0), .value = try widen(a, val_t, 0), .norm_key = norms[0], .norm_query = norms[1], .norm_conv = norms[2], .conv = try widen(a, conv_t, 0) }; + // what upstream's loader (cuda_prompt_load.zig with_ple) would upload as CW: the raw bytes, read as bf16 + var raw_as_bf: usize = 0; + for (0..weights.conv.len) |i| { + const bits = std.mem.readInt(u16, conv_t.bytes[2 * i ..][0..2], .little); + if (bfVal(bits) != bf(weights.conv[i])) raw_as_bf += 1; + } + std.debug.print("conv1d.weight is F16: {d} of {d} values change if its bytes are read as bf16 (upstream's loader does), so it is converted (RNE) as exl3.py does\n", .{ raw_as_bf, weights.conv.len }); + + // inputs: the decoded rows (fp16), the layer input stand-in (token embedding tiled over 4 streams, bf16) + const embed_t = try index.get(gpu.gpa, io, "model.language_model.embed_tokens.weight"); + try check.expect(embed_t.dtype == .bf16 and embed_t.dim(1) == hidden, "embed_tokens", .{}); + const emb = try a.alloc(f32, n * hidden); + for (emb, emb_bits) |*e, b| e.* = @floatCast(@as(f16, @bitCast(b))); + const hq = try a.alloc(f32, n * wide); + const hq_bits = try a.alloc(u16, n * wide); + for (tokens, 0..) |tok, t| for (0..streams) |s| for (0..hidden) |d| { + const bits = std.mem.readInt(u16, embed_t.bytes[(@as(usize, tok) * hidden + d) * 2 ..][0..2], .little); + hq_bits[t * wide + s * hidden + d] = bits; + hq[t * wide + s * hidden + d] = bfVal(bits); + }; + // the projections the reference's F16 dense stores (bf16) + const keys = try a.alloc(f32, n * wide); + const vals = try a.alloc(f32, n * hidden); + for (0..n) |t| project(weights, emb[t * hidden ..][0..hidden], keys[t * wide ..][0..wide], vals[t * hidden ..][0..hidden]); + const keys_bits = try a.alloc(u16, keys.len); + for (keys_bits, keys) |*b, v| b.* = bfBits(v); + const vals_bits = try a.alloc(u16, vals.len); + for (vals_bits, vals) |*b, v| b.* = bfBits(v); + const cw_bits = try a.alloc(u16, weights.conv.len); + for (cw_bits, weights.conv) |*b, v| b.* = bfBits(v); + + // host references + const want = try a.alloc(f64, n * wide); + const emu = try a.alloc(f32, n * wide); + { + var oracle = try Oracle.init(a, weights); + var branch = try Branch.init(a, weights); + for (0..n) |t| { + oracle.step(emb[t * hidden ..][0..hidden], hq[t * wide ..][0..wide], want[t * wide ..][0..wide]); + branch.step(keys[t * wide ..][0..wide], vals[t * hidden ..][0..hidden], hq[t * wide ..][0..wide], emu[t * wide ..][0..wide]); + } + } + + // the device: upstream's glue kernels + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + const D = cuda.DeviceBuffer; + var d_keys = try D.fromHost(gpu.d, bytesOf(u16, keys_bits)); + defer d_keys.free(); + var d_vals = try D.fromHost(gpu.d, bytesOf(u16, vals_bits)); + defer d_vals.free(); + var d_h = try D.fromHost(gpu.d, bytesOf(u16, hq_bits)); + defer d_h.free(); + var d_nk = try D.fromHost(gpu.d, bytesOf(f32, weights.norm_key)); + defer d_nk.free(); + var d_nq = try D.fromHost(gpu.d, bytesOf(f32, weights.norm_query)); + defer d_nq.free(); + var d_nc = try D.fromHost(gpu.d, bytesOf(f32, weights.norm_conv)); + defer d_nc.free(); + var d_cw = try D.fromHost(gpu.d, bytesOf(u16, cw_bits)); + defer d_cw.free(); + const zeros = try a.alloc(u8, n * wide * 2); + @memset(zeros, 0); + var d_zero = try D.fromHost(gpu.d, zeros); // conv's H: zero, so HOUT is the injection itself + defer d_zero.free(); + var d_tail = try D.fromHost(gpu.d, zeros[0 .. tail_rows * wide * 2]); + defer d_tail.free(); + var d_gated = try D.alloc(gpu.d, n * wide * 2); + defer d_gated.free(); + var d_pss = try D.alloc(gpu.d, n * streams * 4); + defer d_pss.free(); + var d_out = try D.alloc(gpu.d, n * wide * 2); + defer d_out.free(); + var d_nrow = try D.alloc(gpu.d, n * wide * 2); + defer d_nrow.free(); + try gpu.ctx.synchronize(); + + // whole prompt: one launch each, the conv reading its own earlier rows + try glue.pleGate(s, d_keys.ptr, d_vals.ptr, d_h.ptr, d_nk.ptr, d_nq.ptr, d_gated.ptr, d_pss.ptr, eps, n, hidden, streams); + try glue.pleConv(s, d_gated.ptr, d_pss.ptr, d_nc.ptr, d_tail.ptr, d_cw.ptr, d_zero.ptr, d_out.ptr, d_nrow.ptr, eps, n, hidden, streams, taps, dilation); + try s.synchronize(); + const prompt_bits = try a.alloc(u16, n * wide); + try d_out.download(0, std.mem.sliceAsBytes(prompt_bits)); + + // token by token: the tail carried on the host from each launch's NROW, as shiftWindow keeps it + const tail = try a.alloc(u16, tail_rows * wide); + @memset(tail, 0); + const decode_bits = try a.alloc(u16, n * wide); + const nrow = try a.alloc(u16, wide); + for (0..n) |t| { + try d_tail.upload(0, bytesOf(u16, tail)); + try gpu.ctx.synchronize(); + try glue.pleGate(s, d_keys.ptr + t * wide * 2, d_vals.ptr + t * hidden * 2, d_h.ptr + t * wide * 2, d_nk.ptr, d_nq.ptr, d_gated.ptr, d_pss.ptr, eps, 1, hidden, streams); + try glue.pleConv(s, d_gated.ptr, d_pss.ptr, d_nc.ptr, d_tail.ptr, d_cw.ptr, d_zero.ptr, d_out.ptr, d_nrow.ptr, eps, 1, hidden, streams, taps, dilation); + try s.synchronize(); + try d_out.download(0, std.mem.sliceAsBytes(decode_bits[t * wide ..][0..wide])); + try d_nrow.download(0, std.mem.sliceAsBytes(nrow)); + std.mem.copyForwards(u16, tail[0 .. (tail_rows - 1) * wide], tail[wide..]); + @memcpy(tail[(tail_rows - 1) * wide ..], nrow); + } + + var worst: f64 = 0; + var peak: f64 = 0; + var worst_rel_rms: f64 = 0; + var emu_worst: f64 = 0; + var emu_off: usize = 0; + var emu_vs_oracle: f64 = 0; + for (0..n) |t| { + var se: f64 = 0; + var sw: f64 = 0; + for (prompt_bits[t * wide ..][0..wide], want[t * wide ..][0..wide], emu[t * wide ..][0..wide]) |gb, o, e| { + const gv: f64 = bfVal(gb); + const err = gv - o; + worst = @max(worst, @abs(err)); + peak = @max(peak, @abs(o)); + se += err * err; + sw += o * o; + emu_worst = @max(emu_worst, @abs(gv - e)); + if (gv != e) emu_off += 1; + emu_vs_oracle = @max(emu_vs_oracle, @abs(@as(f64, e) - o)); + } + worst_rel_rms = @max(worst_rel_rms, @sqrt(se / sw)); + } + const same_modes = std.mem.eql(u16, prompt_bits, decode_bits); + std.debug.print("injection over {d} tokens: GPU max |gpu - f64| {e:.3} at peak {e:.3} ({d:.3}% of peak), worst per-token rms rel {e:.3}\n", .{ n, worst, peak, 100 * worst / peak, worst_rel_rms }); + std.debug.print(" GPU vs the reference's roundings emulated on the host: max |diff| {e:.3}, {d} of {d} values differ; emulation vs f64 max {e:.3}\n", .{ emu_worst, emu_off, n * wide, emu_vs_oracle }); + std.debug.print(" whole-prompt launch vs token-by-token launches: {s}\n", .{if (same_modes) "bit-identical" else "DIFFER"}); + try check.expect(peak > 0 and worst <= 0.03 * peak and worst_rel_rms <= 0.02, "injection beyond bf16 rounding of the f64 oracle", .{}); + try check.expect(same_modes, "prompt and decode launches differ", .{}); + check.pass("layer-1 PLE injection through upstream's fn_ple_gate/fn_ple_conv on real EXL3 weights: max |err| {e:.2} <= 3% of peak {e:.2}, rms rel {e:.2} <= 2e-2; prompt == decode bits", .{ worst, peak, worst_rel_rms }); +} diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_route.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_route.zig new file mode 120000 index 0000000000..587fabf2a4 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_route.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_route.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/exl3_router.zig b/zig/tests/cuda/flashnext_exl3/model/exl3_router.zig new file mode 120000 index 0000000000..cd9eb177b7 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/exl3_router.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/model/exl3_router.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig new file mode 100644 index 0000000000..15236b0bad --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_engine_tests.zig @@ -0,0 +1,536 @@ +//! flashnext-exl3-engine: the Qwen3.8-Flash-Next EXL3 pack through the Flash Next family's own Engine +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const Gpu = check.Gpu; +const st = core.safetensors; +const fmt = core.exl3_format; +const prompt = flash.prompt; +const xpass = flash.exl3.pass; +const vocab = prompt.vocabulary; +const dims: usize = 2560; +const cap_default: usize = 32768; +/// Rows of logits a teacher-forced pass returns at most. +const out_rows: usize = 64; + +fn promote(bits: u16) f32 { + return @bitCast(@as(u32, bits) << 16); +} + +fn ms(t0: i96, t1: i96) f64 { + return @as(f64, @floatFromInt(t1 - t0)) / 1e6; +} + +const Eng = struct { + gpu: Gpu, + stream: cuda.Stream, + glue: flash.Glue, + e: *prompt.Engine, + seq: prompt.Seq, + out: cuda.DeviceBuffer, + cap: usize, + attached: xpass.Attached, + load_s: f64, + + fn init(gpu: Gpu, dir: []const u8, lines: ?usize, cap: usize) !*Eng { + const g = try gpu.gpa.create(Eng); + errdefer gpu.gpa.destroy(g); + g.gpu = gpu; + g.cap = cap; + g.stream = try cuda.Stream.init(gpu.d, true); + g.glue = try flash.Glue.load(gpu.d); + const t0 = check.now(gpu.io); + g.e = try prompt.Engine.initRank(gpu.gpa, gpu.io, @constCast(gpu.d), &g.stream, dir, 1, 0); + const t1 = check.now(gpu.io); + // as cuda_lanes_load.open: the pass's scratch for a full chunk and the block scores, then the stream's caches + _ = try g.e.reserve(@min(cap, g.e.chunk)); + _ = try g.e.reserveScores(cap, @min(cap, g.e.chunk)); + g.seq = try g.e.newSeq(cap); + g.out = try cuda.DeviceBuffer.alloc(gpu.d, out_rows * vocab * 2); + var b = try xpass.budgetFromEnv(0); + if (lines) |n| b.lines = n; + g.attached = try xpass.attach(gpu.gpa, gpu.io, gpu.d, g.stream, g.e.layers[0..g.e.n_layers], g.e.chunk, b, null); + const t2 = check.now(gpu.io); + g.load_s = ms(t0, t2) / 1e3; + var free: usize = 0; + var total: usize = 0; + try gpu.d.check(gpu.d.api.cuMemGetInfo_v2(&free, &total), "cuMemGetInfo"); + std.debug.print("engine: loaded in {d:.1} s (Engine.initRank {d:.1} s incl. the pinned bank's copies {d:.1} s; pools {d:.1} s); {d} lines a layer x {d} layers ({d:.2} GiB), {d:.2} GiB free after, chunk {d}, context {d}\n", .{ g.load_s, ms(t0, t1) / 1e3, @as(f64, @floatFromInt(g.attached.bank_ns)) / 1e9, ms(t1, t2) / 1e3, g.attached.lines, g.attached.layers, @as(f64, @floatFromInt(g.attached.device_bytes)) / (1 << 30), @as(f64, @floatFromInt(free)) / (1 << 30), g.e.chunk, cap }); + return g; + } + + fn deinit(g: *Eng) void { + g.out.free(); + g.seq.deinit(g.gpu.gpa); + g.e.deinit(); + g.glue.deinit(); + g.stream.deinit(); + g.gpu.gpa.destroy(g); + } + + fn tri(g: *Eng) flash.Tri { + return .{ .own = &g.glue, .s = g.stream }; + } + + /// The sequence as Engine.newSeq makes it, without a new allocation (the pools took the free memory). + fn reset(g: *Eng) !void { + const s = &g.seq; + for ([_][]cuda.DeviceBuffer{ s.conv, s.state, s.kc, s.vc, s.ikc, s.pooled }) |bs| for (bs) |b| try b.fill8(0, g.stream.handle); + try s.tail.fill8(0, g.stream.handle); + try g.stream.synchronize(); + s.fed = 0; + s.history = .{ g.e.eos, g.e.eos }; + } + + fn stats(g: *Eng) flash.exl3_pool.Stats { + return xpass.poolStats(g.e.layers[0..g.e.n_layers]); + } +}; + +const Golden = struct { + parsed: std.json.Parsed(std.json.Value), + + fn open(gpu: Gpu, path: []const u8) !Golden { + const text = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpu.gpa, .limited(1 << 26)); + defer gpu.gpa.free(text); + return .{ .parsed = try std.json.parseFromSlice(std.json.Value, gpu.gpa, text, .{}) }; + } + + fn prompts(g: *const Golden) []std.json.Value { + return g.parsed.value.object.get("prompts").?.array.items; + } +}; + +fn ints(a: std.mem.Allocator, v: std.json.Value) ![]u32 { + const items = v.array.items; + const out = try a.alloc(u32, items.len); + for (items, out) |x, *o| o.* = @intCast(x.integer); + return out; +} + +fn num(v: std.json.Value) f64 { + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => std.math.nan(f64), + }; +} + +fn top5(row: []const u16) [5]u32 { + var ids: [5]u32 = @splat(0); + var vals: [5]f32 = @splat(-std.math.inf(f32)); + for (row, 0..) |b, i| { + const v = promote(b); + if (!(v > vals[4])) continue; + var at: usize = 4; + while (at > 0 and v > vals[at - 1]) : (at -= 1) { + vals[at] = vals[at - 1]; + ids[at] = ids[at - 1]; + } + vals[at] = v; + ids[at] = @intCast(i); + } + return ids; +} + +// -- the EXL3 head in float64 ------------------------------------------------------------------------------------ + +const Raw = struct { + files: [7]st.File = undefined, + n: usize = 0, + + fn open(gpu: Gpu, dir: []const u8) !Raw { + var r: Raw = .{}; + var buf: [512]u8 = undefined; + for (0..7) |i| { + r.files[i] = try st.File.open(gpu.gpa, gpu.io, try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-00007.safetensors", .{ dir, i + 1 })); + r.n += 1; + } + return r; + } + + fn close(r: *Raw, io: std.Io) void { + for (r.files[0..r.n]) |*f| f.close(io); + } + + fn get(r: *const Raw, name: []const u8) !st.Tensor { + for (r.files[0..r.n]) |*f| if (f.get(name)) |t| return t; + return error.MissingTensor; + } +}; + +/// lm_head as stored: trellis int16 [K/16, N/16, 8 * halves] (bytes), suh fp16 [K], svh fp16 [N]. +const Head = struct { + trellis: []const u8, + suh: []const u8, + svh: []const u8, + k: usize, + n: usize, + halves: u16, + + fn words(h: Head) usize { + return @as(usize, h.halves) * 8; + } + + fn of(trellis: st.Tensor, suh: st.Tensor, svh: st.Tensor, mul1: ?st.Tensor) !Head { + if (trellis.dtype != .i16 or trellis.rank != 3 or trellis.dim(2) % 8 != 0) return error.UnexpectedTensor; + const h: Head = .{ .trellis = trellis.bytes, .suh = suh.bytes, .svh = svh.bytes, .k = trellis.dim(0) * 16, .n = trellis.dim(1) * 16, .halves = @intCast(trellis.dim(2) / 8) }; + if (suh.dtype != .f16 or suh.numel() != h.k or svh.dtype != .f16 or svh.numel() != h.n) return error.UnexpectedTensor; + if (h.k % fmt.hadamard_dim != 0 or h.n % fmt.hadamard_dim != 0 or h.halves < 2 or h.halves > 16) return error.UnexpectedTensor; + if (mul1) |m| { + if (m.bytes.len != 4 or std.mem.readInt(u32, m.bytes[0..4], .little) != fmt.Mul1Mul) return error.UnexpectedTensor; + } + return h; + } +}; + +const HeadWork = struct { + head: Head, + x: []const f64, // [rows, K] + rows: usize, + out: []f64, // [rows, N] + next: std.atomic.Value(usize) = .init(0), + failed: ?anyerror = null, +}; + +fn headWorker(w: *HeadWork, gpa: std.mem.Allocator) void { + headWork(w, gpa) catch |err| { + w.failed = err; + }; +} + +/// logits[r][v] = sum_k x[r][k] * W[k][v], W decoded in float64 a 128-column block at a time (exl3_format.dequantize). +fn headWork(w: *HeadWork, gpa: std.mem.Allocator) !void { + const h = w.head; + const bn: usize = fmt.hadamard_dim; + const tiles_k = h.k / 16; + const tiles_n = h.n / 16; + const per = bn / 16; + const words = h.words(); + const sub = try gpa.alloc(i16, tiles_k * per * words); + defer gpa.free(sub); + const vals = try gpa.alloc(f64, h.k * bn); + defer gpa.free(vals); + const scratch = try gpa.alloc(f64, 2 * h.k * bn); + defer gpa.free(scratch); + const suh = try gpa.alloc(f16, h.k); + defer gpa.free(suh); + for (suh, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u16, h.suh[2 * i ..][0..2], .little)); + var svh: [fmt.hadamard_dim]f16 = undefined; + while (true) { + const nb = w.next.fetchAdd(1, .monotonic); + if (nb >= h.n / bn) return; + for (0..tiles_k) |tk| for (0..per) |j| { + const src = ((tk * tiles_n) + nb * per + j) * words * 2; + const dst = sub[(tk * per + j) * words ..][0..words]; + for (dst, 0..) |*d, i| d.* = std.mem.readInt(i16, h.trellis[src + 2 * i ..][0..2], .little); + }; + for (&svh, 0..) |*o, i| o.* = @bitCast(std.mem.readInt(u16, h.svh[2 * (nb * bn + i) ..][0..2], .little)); + fmt.dequantize(.{ .halves = h.halves }, .mul1, sub, tiles_k, per, suh, &svh, vals, scratch); + for (0..w.rows) |r| { + const xr = w.x[r * h.k ..][0..h.k]; + var acc: [fmt.hadamard_dim]f64 = @splat(0); + for (0..h.k) |k| { + const xv = xr[k]; + const row = vals[k * bn ..][0..bn]; + for (&acc, row) |*a, v| a.* += xv * v; + } + @memcpy(w.out[r * h.n + nb * bn ..][0..bn], &acc); + } + } +} + +fn headF64(gpa: std.mem.Allocator, raw: *const Raw, x: []const f64, rows: usize) ![]f64 { + const head = try Head.of(try raw.get("lm_head.trellis"), try raw.get("lm_head.suh"), try raw.get("lm_head.svh"), raw.get("lm_head.mul1") catch null); + const out = try gpa.alloc(f64, rows * head.n); + errdefer gpa.free(out); + var w: HeadWork = .{ .head = head, .x = x, .rows = rows, .out = out }; + var pool: [64]std.Thread = undefined; + const n = @min(pool.len, @max(1, std.Thread.getCpuCount() catch 1)); + var spawned: usize = 0; + for (0..n) |_| { + pool[spawned] = try std.Thread.spawn(.{}, headWorker, .{ &w, gpa }); + spawned += 1; + } + for (pool[0..spawned]) |t| t.join(); + if (w.failed) |err| return err; + return out; +} + +// -- modes --------------------------------------------------------------------------------------------------------- + +fn tfMode(gpu: Gpu, dir: []const u8, golden_path: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const p = golden.prompts()[0]; + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + const tf = p.object.get("teacher_forced").?.array.items; + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + try g.reset(); + const n = ids.len; + const ops = g.e.layers[0].exl3.?.ops; + const before = ops.fingerprint(); + const t0 = check.now(gpu.io); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, .{ .logits = .all, .logits_out = g.out.ptr }); + try g.stream.synchronize(); + const t1 = check.now(gpu.io); + const logits = try a.alloc(u16, n * vocab); + defer a.free(logits); + try g.out.download(0, std.mem.sliceAsBytes(logits)); + // the head rows the pass fed lm_head (the final mixer's read-out, bf16 [n, dims]) + const hrows = try a.alloc(u16, n * dims); + defer a.free(hrows); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(hrows.ptr, g.e.batch_s.?.mixed, n * dims * 2), "cuMemcpyDtoH"); + std.debug.print("teacher-forced pass: {d} rows in {d:.1} ms (first pass: includes every miss's copy)\n", .{ n, ms(t0, t1) }); + // one decode round too, then every EXL3 scratch buffer must be where Ops.reserve put it + { + var tok = prompt.argmax(g.seq.logits); + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, (&tok)[0..1], g.out.ptr); + tok = prompt.argmax(g.seq.logits); + } + const after = ops.fingerprint(); + try check.expect(std.mem.eql(u8, std.mem.sliceAsBytes(&before), std.mem.sliceAsBytes(&after)), "an EXL3 scratch buffer grew inside a pass", .{}); + check.pass("EXL3 scratch reserved up front: the {d} scratch buffers kept their addresses and sizes through a {d}-row prompt pass and a decode round", .{ blk: { + var k: usize = 0; + for (after) |e| k += @intFromBool(e[1] != 0); + break :blk k; + }, n }); + + std.debug.print(" pos input ours top1 (logit) golden top1 (logit) match top5 overlap max|dlogit| on golden top5\n", .{}); + var top1_ok: usize = 0; + var top5_same: usize = 0; + var worst: f64 = 0; + var sum_abs: f64 = 0; + var count: usize = 0; + for (tf, 0..) |row, pos| { + const r = logits[pos * vocab ..][0..vocab]; + const ours = top5(r); + const g5 = try ints(a, row.object.get("top5_ids").?); + defer a.free(g5); + const gl = row.object.get("top5_logits").?.array.items; + var overlap: usize = 0; + for (ours) |o| for (g5) |x| { + overlap += @intFromBool(o == x); + }; + var dmax: f64 = 0; + for (g5, gl) |id, lv| { + const d = @abs(@as(f64, promote(r[id])) - num(lv)); + dmax = @max(dmax, d); + sum_abs += d; + count += 1; + } + worst = @max(worst, dmax); + const ok = ours[0] == g5[0]; + top1_ok += @intFromBool(ok); + top5_same += @intFromBool(std.mem.eql(u32, &ours, g5)); + std.debug.print("{d:4} {d:7} {d:9} ({d:7.3}) {d:9} ({d:7.3}) {s:5} {d}/5 {s} {d:.4}\n", .{ pos, ids[pos], ours[0], promote(r[ours[0]]), g5[0], num(gl[0]), if (ok) "yes" else "NO", overlap, if (std.mem.eql(u32, &ours, g5)) "(same order)" else " ", dmax }); + } + std.debug.print("TF top1 {d}/{d}, top5 identical in order {d}/{d}, |our logit - golden bf16 logit| on the golden's top-5 ids: max {d:.4}, mean {d:.4}\n", .{ top1_ok, tf.len, top5_same, tf.len, worst, sum_abs / @as(f64, @floatFromInt(count)) }); + + // the EXL3 lm_head itself, on the same head rows, against an exact float64 product over all 248,320 columns + var raw = try Raw.open(gpu, dir); + defer raw.close(gpu.io); + const xf = try a.alloc(f64, n * dims); + defer a.free(xf); + for (xf, hrows) |*o, b| o.* = promote(b); + const t2 = check.now(gpu.io); + const exact = try headF64(a, &raw, xf, n); + defer a.free(exact); + const t3 = check.now(gpu.io); + var head_abs: f64 = 0; + var head_rel: f64 = 0; + var head_top1: usize = 0; + for (0..n) |r| { + var peak: f64 = 0; + var best: usize = 0; + var dmax: f64 = 0; + for (0..vocab) |v| { + const want = exact[r * vocab + v]; + peak = @max(peak, @abs(want)); + if (want > exact[r * vocab + best]) best = v; + dmax = @max(dmax, @abs(@as(f64, promote(logits[r * vocab + v])) - want)); + } + head_abs = @max(head_abs, dmax); + head_rel = @max(head_rel, dmax / peak); + const ours = top5(logits[r * vocab ..][0..vocab])[0]; + // a bf16 tie at the top is no disagreement + head_top1 += @intFromBool(ours == best or promote(logits[r * vocab + ours]) == promote(logits[r * vocab + best])); + } + std.debug.print("EXL3 lm_head (Dense, bf16 out) vs exact float64 on the pass's {d} head rows, {d} columns ({d:.1} s on the host): max |d| {d:.4} (rel to the row peak {e:.2}), top-1 equal on {d}/{d}\n", .{ n, vocab, ms(t2, t3) / 1e3, head_abs, head_rel, head_top1, n }); + // bf16 output rounding is 2^-9 of a value (0.06 at a logit of 30); the fp16 rotated input adds ~2.4e-4 of the peak + try check.expect(head_rel < 4e-3 and head_top1 == n, "lm_head within bf16 rounding of float64", .{}); + check.pass("lm_head: EXL3 hook through exl3.Dense, {d} rows, max |d| {d:.4} (rel {e:.2}) vs float64, top-1 equal on every row", .{ n, head_abs, head_rel }); + try check.expect(top1_ok == tf.len, "teacher-forced top-1 equals the golden on every position ({d}/{d})", .{ top1_ok, tf.len }); + check.pass("teacher-forced top-1 {d}/{d} vs the golden; max |dlogit| {d:.4} on the golden's top-5", .{ top1_ok, tf.len, worst }); + printPool(g, "after tf"); +} + +fn printPool(g: *Eng, what: []const u8) void { + const s = g.stats(); + const routed = s.hit_slots + s.miss_slots; + std.debug.print("pool {s}: {d} rounds, {d} passes, hit rate {d:.1}% ({d}/{d} routed slots), {d} copies, {d} evictions; host cut {d:.1} ms over {d} blocks, plan {d:.1} ms\n", .{ what, s.rounds, s.passes, 100 * @as(f64, @floatFromInt(s.hit_slots)) / @as(f64, @floatFromInt(@max(routed, 1))), s.hit_slots, routed, s.misses, s.evictions, @as(f64, @floatFromInt(xpass.pool_time.cut_ns)) / 1e6, xpass.pool_time.blocks, @as(f64, @floatFromInt(xpass.pool_time.plan_ns)) / 1e6 }); +} + +/// Greedy `steps` tokens after `ids`: the prefill's last row, then a decode round a token. +fn greedyRun(g: *Eng, ids: []const u32, out: []u32, step_ms: ?*f64) !void { + try g.reset(); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids, .{ .logits = .last }); + out[0] = prompt.argmax(g.seq.logits); + const t0 = check.now(g.gpu.io); + for (1..out.len) |i| { + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, out[i - 1 ..][0..1], g.out.ptr); + out[i] = prompt.argmax(g.seq.logits); + } + if (step_ms) |p| p.* = ms(t0, check.now(g.gpu.io)) / @as(f64, @floatFromInt(@max(out.len - 1, 1))); +} + +/// Each golden prompt followed by the golden's own 24 greedy ids, teacher-forced in one pass: our top-1 at every +fn tfcMode(gpu: Gpu, dir: []const u8, golden_path: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + var all_ok = true; + for (golden.prompts()) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + const cont = try ints(a, p.object.get("greedy_ids").?); + defer a.free(cont); + const steps = p.object.get("greedy_steps").?.array.items; + const all = try a.alloc(u32, ids.len + cont.len - 1); + defer a.free(all); + @memcpy(all[0..ids.len], ids); + @memcpy(all[ids.len..], cont[0 .. cont.len - 1]); + if (all.len > out_rows) return error.UnexpectedTensor; + try g.reset(); + try g.e.forward(flash.Tri, g.tri(), &g.seq, all, .{ .logits = .all, .logits_out = g.out.ptr }); + try g.stream.synchronize(); + const logits = try a.alloc(u16, all.len * vocab); + defer a.free(logits); + try g.out.download(0, std.mem.sliceAsBytes(logits[0 .. all.len * vocab])); + std.debug.print("TFC {s}\n step golden id ours top1 our margin golden margin max|dlogit| golden top5\n", .{p.object.get("name").?.string}); + var ok_n: usize = 0; + for (cont, 0..) |want, k| { + const r = logits[(ids.len - 1 + k) * vocab ..][0..vocab]; + const t = top5(r); + const margin = promote(r[t[0]]) - promote(r[t[1]]); + const g5 = try ints(a, steps[k].object.get("top5_ids").?); + defer a.free(g5); + var dmax: f64 = 0; + for (g5, steps[k].object.get("top5_logits").?.array.items) |id, lv| dmax = @max(dmax, @abs(@as(f64, promote(r[id])) - num(lv))); + const ok = t[0] == want; + ok_n += @intFromBool(ok); + std.debug.print(" {d:4} {d:10} {d:9} {s} {d:10.4} {d:14.4} {d:10.4}\n", .{ k, want, t[0], if (ok) " " else "*", margin, num(steps[k].object.get("margin12").?), dmax }); + } + std.debug.print("TFC {s}: top-1 equals the golden's greedy id on {d}/{d} continuation positions\n", .{ p.object.get("name").?.string, ok_n, cont.len }); + if (ok_n != cont.len) all_ok = false; + } + printPool(g, "after tfc"); + try check.expect(all_ok, "teacher-forced continuation top-1 equals the golden everywhere", .{}); + check.pass("teacher-forced continuation: top-1 equals the golden's greedy id at every step of all {d} prompts", .{golden.prompts().len}); +} + +fn genMode(gpu: Gpu, dir: []const u8, golden_path: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + var golden = try Golden.open(gpu, golden_path); + defer golden.parsed.deinit(); + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + var all_ok = true; + for (golden.prompts()) |p| { + const ids = try ints(a, p.object.get("ids").?); + defer a.free(ids); + const want = try ints(a, p.object.get("greedy_ids").?); + defer a.free(want); + const got = try a.alloc(u32, want.len); + defer a.free(got); + var step: f64 = 0; + try greedyRun(g, ids, got, &step); + var first_diff: ?usize = null; + for (got, want, 0..) |x, y, i| if (x != y and first_diff == null) { + first_diff = i; + }; + std.debug.print("GEN {s}: ", .{p.object.get("name").?.string}); + for (got) |x| std.debug.print("{d},", .{x}); + std.debug.print(" {s} ({d:.1} ms/step decode)\n", .{ if (first_diff != null) "DIFFERS" else "equal to the golden", step }); + if (first_diff) |d| { + std.debug.print(" first difference at step {d}: ours {d}, golden {d}; golden margin there {d:.4}\n", .{ d, got[d], want[d], num(p.object.get("greedy_steps").?.array.items[d].object.get("margin12").?) }); + all_ok = false; + } + } + printPool(g, "after gen"); + try check.expect(all_ok, "greedy 24 equals the golden on every prompt", .{}); + check.pass("greedy 24 tokens equal the golden on all {d} prompts", .{golden.prompts().len}); +} + +fn benchMode(gpu: Gpu, dir: []const u8, lines: ?usize) !void { + const a = gpu.gpa; + const g = try Eng.init(gpu, dir, lines, cap_default); + defer g.deinit(); + const max_n: usize = 2048; + const ids = try a.alloc(u32, max_n); + defer a.free(ids); + // text-like ids: a fixed pseudo-random walk over the common part of the vocabulary + var prng = std.Random.DefaultPrng.init(0x5eed); + for (ids) |*x| x.* = 200 + prng.random().uintLessThan(u32, 60000); + // warm-up: one pass to load kernels, settle the pool + try g.reset(); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids[0..64], .{ .logits = .last }); + for ([_]usize{ 16, 256, 2048 }) |n| { + var best: f64 = std.math.inf(f64); + for (0..2) |_| { + try g.reset(); + const before = g.stats(); + const t0 = check.now(gpu.io); + try g.e.forward(flash.Tri, g.tri(), &g.seq, ids[0..n], .{ .logits = .last }); + const t = ms(t0, check.now(gpu.io)); + const after = g.stats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print(" prefill {d:5} tokens: {d:8.1} ms = {d:7.1} rows/s (hit {d:.1}%, {d} copies)\n", .{ n, t, @as(f64, @floatFromInt(n)) / (t / 1e3), 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))), after.misses - before.misses }); + best = @min(best, t); + } + std.debug.print("BENCH prefill {d} tokens: best {d:.1} ms, {d:.1} rows/s\n", .{ n, best, @as(f64, @floatFromInt(n)) / (best / 1e3) }); + // decode after this prompt: 32 tokens, greedy + var tok = prompt.argmax(g.seq.logits); + const before = g.stats(); + const t0 = check.now(gpu.io); + const steps: usize = 32; + for (0..steps) |_| { + const q = [_]*prompt.Seq{&g.seq}; + try g.e.round(flash.Tri, g.tri(), &q, (&tok)[0..1], g.out.ptr); + tok = prompt.argmax(g.seq.logits); + } + const t = ms(t0, check.now(gpu.io)) / steps; + const after = g.stats(); + const hits = after.hit_slots - before.hit_slots; + const all = hits + after.miss_slots - before.miss_slots; + std.debug.print("BENCH decode after {d} tokens: {d:.2} ms/token over {d} steps (hit {d:.1}%, {d:.1} copies a token)\n", .{ n, t, steps, 100 * @as(f64, @floatFromInt(hits)) / @as(f64, @floatFromInt(@max(all, 1))), @as(f64, @floatFromInt(after.misses - before.misses)) / steps }); + } + printPool(g, "after bench"); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const mode = rest[1]; + const lineArg = struct { + fn f(r: []const [:0]const u8, i: usize) !?usize { + return if (r.len > i) try std.fmt.parseInt(usize, r[i], 10) else null; + } + }.f; + if (std.mem.eql(u8, mode, "tf")) return tfMode(gpu, dir, if (rest.len > 2) rest[2] else return error.MissingArgument, try lineArg(rest, 3)); + if (std.mem.eql(u8, mode, "tfc")) return tfcMode(gpu, dir, if (rest.len > 2) rest[2] else return error.MissingArgument, try lineArg(rest, 3)); + if (std.mem.eql(u8, mode, "gen")) return genMode(gpu, dir, if (rest.len > 2) rest[2] else return error.MissingArgument, try lineArg(rest, 3)); + if (std.mem.eql(u8, mode, "bench")) return benchMode(gpu, dir, try lineArg(rest, 2)); + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_layer_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_layer_tests.zig new file mode 100644 index 0000000000..cbc2cb518d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_layer_tests.zig @@ -0,0 +1,386 @@ +//! flashnext-exl3: one decoder layer through the family pass against the float64 layer (gate 5). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const fwd = @import("exl3_fwd_cpu.zig"); +const hc = @import("exl3_hc.zig"); +const moe = @import("exl3_moe.zig"); +const router_cpu = @import("exl3_router.zig"); +const route_cpu = @import("exl3_route.zig"); +const scan = @import("../pack/exl3_scan.zig"); +const place = @import("../pack/exl3_place.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); + +const X = flash.exl3; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const batch = X.batch; +const Gpu = check.Gpu; +const st = core.safetensors; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const topk = slots - 1; +const experts_n = wt.experts_n; +const shared = @import("flashnext_exl3_rig.zig"); +const typed = shared.typed; +const valueOf = shared.valueOf; +const f64Of = shared.f64Of; +const f32Of = shared.f32Of; +const f16Of = shared.f16Of; +const qmat = shared.qmat; +const bf16Bits = shared.bf16Bits; +const bfValue = shared.bfValue; +const Err = shared.Err; +const errOf = shared.errOf; +const download = shared.download; +const uploadTo = shared.uploadTo; +const RawPack = shared.RawPack; +const Pass = shared.Pass; +const Rig = shared.Rig; +const bf16Rows = shared.bf16Rows; +const default_ids = shared.default_ids; +const projections = @import("flashnext_exl3_proj_tests.zig").projections; + +// -- the float64 layer -------------------------------------------------------------------------------------------- + +const HcHost = struct { + norm: []f32, + down: []f32, + up: []f32, + inject: ?[]f32, + + fn load(a: std.mem.Allocator, raw: *const RawPack, comptime base: []const u8, args: anytype, with_inject: bool) !HcHost { + return .{ + .norm = try f32Of(a, try raw.get(base ++ ".hc_norm.weight", args)), + .down = try f32Of(a, try raw.get(base ++ ".input_mix_weight_down.weight", args)), + .up = try f32Of(a, try raw.get(base ++ ".input_mix_weight_up.weight", args)), + .inject = if (with_inject) try f32Of(a, try raw.get(base ++ ".block_inject_weight.weight", args)) else null, + }; + } + + fn weights(h: HcHost) hc.Weights { + return .{ .norm = h.norm, .down = h.down, .up = h.up, .inject = h.inject }; + } +}; + +const ExpertCtx = struct { + gpa: std.mem.Allocator, + raw: *const RawPack, + layer: usize, + + fn load(ctx: *anyopaque, e: usize, out: *[3]fwd.QMat) anyerror!void { + const self: *ExpertCtx = @ptrCast(@alignCast(ctx)); + const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + var made: usize = 0; + errdefer for (0..made) |i| out[i].deinit(self.gpa); + while (made < 3) : (made += 1) { + var nb: [160]u8 = undefined; + out[made] = try qmat(self.gpa, self.raw, try std.fmt.bufPrint(&nb, "model.language_model.layers.{d}.mlp.experts.{d}.{s}", .{ self.layer, e, projs[made] }), 1); + } + } +}; + +/// The oracle's layer (exl3_fwd_cpu.LayerW) from the pack's stored tensors, decoded by the Zig f64 decoder. +const Oracle = struct { + attn_hc: HcHost, + mlp_hc: HcHost, + gdn: fwd.GdnW = undefined, + attn: fwd.AttnW = undefined, + moe: fwd.MoeW, + ctx: ExpertCtx, + state: fwd.LayerState, + /// The mixer's state for the staged comparison (fed the device's own inputs), apart from the chained one. + staged: fwd.LayerState, + linear: bool, + + fn create(a: std.mem.Allocator, raw: *const RawPack, l: usize, offset: f32, cap: usize) !*Oracle { + const o = try a.create(Oracle); + o.linear = wt.linearLayer(l); + o.attn_hc = try HcHost.load(a, raw, "model.language_model.layers.{d}.attn_hyper_connection", .{l}, true); + o.mlp_hc = try HcHost.load(a, raw, "model.language_model.layers.{d}.mlp_hyper_connection", .{l}, true); + o.ctx = .{ .gpa = a, .raw = raw, .layer = l }; + var nb: [4][160]u8 = undefined; + o.moe = .{ + .gate16 = try f16Of(a, try raw.get("model.language_model.layers.{d}.mlp.gate.weight", .{l})), + .sgate16 = try f16Of(a, try raw.get("model.language_model.layers.{d}.mlp.shared_expert_gate.weight", .{l})), + .shared = .{ + try qmat(a, raw, try std.fmt.bufPrint(&nb[0], "model.language_model.layers.{d}.mlp.shared_expert.gate_proj", .{l}), 48), + try qmat(a, raw, try std.fmt.bufPrint(&nb[1], "model.language_model.layers.{d}.mlp.shared_expert.up_proj", .{l}), 48), + try qmat(a, raw, try std.fmt.bufPrint(&nb[2], "model.language_model.layers.{d}.mlp.shared_expert.down_proj", .{l}), 48), + }, + .source = .{ .ctx = &o.ctx, .load = ExpertCtx.load }, + }; + if (o.linear) { + const stem = "model.language_model.layers.{d}.linear_attn"; + o.gdn = .{ + .qkv = try qmat(a, raw, try std.fmt.bufPrint(&nb[0], stem ++ ".in_proj_qkv", .{l}), 48), + .z = try qmat(a, raw, try std.fmt.bufPrint(&nb[1], stem ++ ".in_proj_z", .{l}), 48), + .out = try qmat(a, raw, try std.fmt.bufPrint(&nb[2], stem ++ ".out_proj", .{l}), 48), + .a_w = try f64Of(a, try raw.get(stem ++ ".in_proj_a.weight", .{l})), + .b_w = try f64Of(a, try raw.get(stem ++ ".in_proj_b.weight", .{l})), + .params = .{ + .conv_w = try f64Of(a, try raw.get(stem ++ ".conv1d.weight", .{l})), + .a_log = try f64Of(a, try raw.get(stem ++ ".A_log", .{l})), + .dt_bias = try f64Of(a, try raw.get(stem ++ ".dt_bias", .{l})), + .norm_w = try f64Of(a, try raw.get(stem ++ ".norm.weight", .{l})), + }, + }; + o.state = .{ .gdn = try gdn_cpu.State.init(a, gdn_cpu.qwen) }; + o.staged = .{ .gdn = try gdn_cpu.State.init(a, gdn_cpu.qwen) }; + } else { + const stem = "model.language_model.layers.{d}.self_attn"; + var sc: [4][]f32 = undefined; + for ([_][]const u8{ "q_norm", "k_norm", "indexer.q_layernorm", "indexer.k_layernorm" }, 0..) |nm, i| { + var buf: [160]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.self_attn.{s}.weight", .{ l, nm }); + sc[i] = try f32Of(a, try raw.get("{s}", .{name})); + for (sc[i]) |*v| v.* += offset; + } + o.attn = .{ + .q = try qmat(a, raw, try std.fmt.bufPrint(&nb[0], stem ++ ".q_proj", .{l}), 48), + .k = try qmat(a, raw, try std.fmt.bufPrint(&nb[1], stem ++ ".k_proj", .{l}), 48), + .v = try qmat(a, raw, try std.fmt.bufPrint(&nb[2], stem ++ ".v_proj", .{l}), 48), + .o = try qmat(a, raw, try std.fmt.bufPrint(&nb[3], stem ++ ".o_proj", .{l}), 48), + .idx = try qmat(a, raw, try std.fmt.bufPrint(&nb[0], stem ++ ".indexer.index_qk_proj", .{l}), 48), + .norms = .{ .q = sc[0], .k = sc[1], .iq = sc[2], .ik = sc[3], .inv = attn_cpu.invFreq() }, + }; + o.state = .{ .attn = try attn_cpu.Cache.init(a, cap) }; + o.staged = .{ .attn = try attn_cpu.Cache.init(a, cap) }; + } + return o; + } + + fn weights(o: *Oracle) fwd.LayerW { + return .{ .attn_hc = o.attn_hc.weights(), .mlp_hc = o.mlp_hc.weights(), .mixer = if (o.linear) .{ .gdn = &o.gdn } else .{ .attn = &o.attn }, .moe = &o.moe }; + } +}; + +// -- one decoder layer through the family's pass (gate 5) ---------------------------------------------------------- + +const Stage = struct { mixer: f64 = 0, streams: f64 = 0, moe: f64 = 0, out: f64 = 0, picks_same: usize = 0, rows: usize = 0 }; + +/// What the device fed each block: the read-outs' mixed rows and inject gates (bf16), attention's then the MoE's. +pub const Taps = struct { attn_mixed: []u16, attn_inj: []u16, mlp_mixed: []u16, mlp_inj: []u16 }; + +/// The pass's per-layer sequence (cuda_batch.zig Pass.chunk's loop body), stepped with read-backs between blocks. +pub fn deviceLayer(gpu: Gpu, p: *Pass, lw: *const pl.Layer, l: usize, ids: []const u32, pos0: usize, h_in: []const u16, tr: fwd.Trace, h_out: []u16, taps: Taps) !void { + const s = p.s; + const n = ids.len; + p.rows = n; + p.pos0 = pos0; + p.pending = null; + try p.stage(ids); + try uploadTo(gpu, p.stream, s.h, std.mem.sliceAsBytes(h_in)); + try p.writeback(0, n); + try p.readout(lw.attn_hc, 0, n); + try tap(gpu, p, s.mixed, taps.attn_mixed[0 .. n * dims]); + try tap(gpu, p, s.inj, taps.attn_inj[0 .. n * streams]); + if (lw.linear) try p.gdnBlock(lw, l, wt.linearIndex(l), .{}) else try p.attnBlock(lw, wt.attnIndex(l)); + const branch = try download(f32, gpu, p.stream, s.part, n * dims); + defer gpu.gpa.free(branch); + for (tr.attn_branch[0 .. n * dims], branch) |*d, v| d.* = v; + try p.tri.hcWriteback(s.h, s.h, s.pss, s.part, "*fp32", s.inj, s.dummy, "*bf16", s.dummy, "*bf16", n * dims, n, dims, streams, 4, 1, 1, wt.out_sk); + const mid = try download(u16, gpu, p.stream, s.h, n * wide); + defer gpu.gpa.free(mid); + for (tr.h_mid[0 .. n * wide], mid) |*d, v| d.* = bfValue(v); + try p.readout(lw.mlp_hc, 0, n); + try tap(gpu, p, s.mixed, taps.mlp_mixed[0 .. n * dims]); + try tap(gpu, p, s.inj, taps.mlp_inj[0 .. n * streams]); + try p.moeBlock(lw); + const picks = try download(i32, gpu, p.stream, s.picks, n * slots); + defer gpu.gpa.free(picks); + const wts = try download(f32, gpu, p.stream, s.wts, n * slots); + defer gpu.gpa.free(wts); + const y = try download(f32, gpu, p.stream, s.y, n * slots * dims); + defer gpu.gpa.free(y); + for (0..n) |r| { + for (0..topk) |k| tr.picks[r * topk + k] = @intCast(picks[r * slots + k]); + for (0..dims) |d| { + var acc: f64 = 0; + for (0..slots) |k| acc += @as(f64, wts[r * slots + k]) * y[(r * slots + k) * dims + d]; + tr.mlp_branch[r * dims + d] = acc; + } + } + p.pending = l; + try p.writeback(0, n); + const out = try download(u16, gpu, p.stream, s.h, n * wide); + defer gpu.gpa.free(out); + @memcpy(h_out, out); +} + +fn tap(gpu: Gpu, p: *Pass, ptr: u64, out: []u16) !void { + const got = try download(u16, gpu, p.stream, ptr, out.len); + defer gpu.gpa.free(got); + @memcpy(out, got); +} + +fn toF64(a: std.mem.Allocator, bits: []const u16) ![]f64 { + const out = try a.alloc(f64, bits.len); + for (out, bits) |*o, b| o.* = bfValue(b); + return out; +} + +/// Each block against the float64 oracle fed the DEVICE's own input to that block: its error alone. +const Staged = struct { attn_mixed: f64 = 0, attn_inj: f64 = 0, mixer: f64 = 0, mlp_mixed: f64 = 0, mlp_inj: f64 = 0, moe: f64 = 0, writeback: f64 = 0 }; + +pub fn layerCheck(gpu: Gpu, r: *Rig, raw: *const RawPack, lw: *const pl.Layer, l: usize, tokens: []const u32) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const prefill = 8; + var t0 = check.now(gpu.io); + const o = try Oracle.create(a, raw, l, r.store.norm_offset, r.cap); + std.debug.print(" float64 layer {d} decoded in {d:.0} ms (offset {d})\n", .{ l, @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6, r.store.norm_offset }); + try r.reset(); + var p = r.pass((lw)[0..1]); + const R = prefill; + const mk = struct { + fn trace(al: std.mem.Allocator) !fwd.Trace { + return .{ .attn_branch = try al.alloc(f64, R * dims), .mlp_branch = try al.alloc(f64, R * dims), .picks = try al.alloc(u16, R * topk), .h_mid = try al.alloc(f64, R * wide) }; + } + }; + const dt = try mk.trace(a); + const ot = try mk.trace(a); + const table = try raw.get("model.language_model.embed_tokens.weight", .{}); + const calls = [_][2]usize{ .{ 0, prefill }, .{ 8, 1 }, .{ 9, 1 }, .{ 10, 1 }, .{ 11, 1 } }; + var worst: Stage = .{}; + var alone: Staged = .{}; + var device_ms: f64 = 0; + const taps: Taps = .{ .attn_mixed = try a.alloc(u16, R * dims), .attn_inj = try a.alloc(u16, R * streams), .mlp_mixed = try a.alloc(u16, R * dims), .mlp_inj = try a.alloc(u16, R * streams) }; + for (calls) |c| { + const toks = tokens[c[0] .. c[0] + c[1]]; + const n = c[1]; + // the layer's input: each token's embedding row in all four streams (bf16, as the pass's expand writes it) + const hb = try a.alloc(u16, n * wide); + const ho = try a.alloc(f64, n * wide); + for (toks, 0..) |tk, t| for (0..dims) |i| { + const bits = std.mem.readInt(u16, table.bytes[(tk * dims + i) * 2 ..][0..2], .little); + for (0..streams) |sidx| { + hb[t * wide + sidx * dims + i] = bits; + ho[t * wide + sidx * dims + i] = bfValue(bits); + } + }; + const hg = try a.alloc(u16, n * wide); + t0 = check.now(gpu.io); + try deviceLayer(gpu, &p, lw, l, toks, c[0], hb, dt, hg, taps); + device_ms += @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6; + // staged: every block from the device's own input (before the chained oracle updates `ho` in place) + { + const w = o.weights(); + const off = r.store.norm_offset; + const mixed = try a.alloc(f64, n * dims); + const inj = try a.alloc(f64, n * streams); + try fwd.hcPreRows(a, w.attn_hc, off, ho, n, mixed, inj); + const e_am = errOf(@as([]const u16, taps.attn_mixed[0 .. n * dims]), mixed); + const e_ai = errOf(@as([]const u16, taps.attn_inj[0 .. n * streams]), inj); + const dev_mixed = try toF64(a, taps.attn_mixed[0 .. n * dims]); + const branch = try a.alloc(f64, n * dims); + switch (w.mixer) { + .gdn => |g| try fwd.gdnBranch(a, g, o.staged.gdn, dev_mixed, n, branch), + .attn => |at| try fwd.attnBranch(a, at, &o.staged.attn, c[0], dev_mixed, n, branch), + } + const e_mx = errOf(dt.attn_branch[0 .. n * dims], branch); + const dev_mid = dt.h_mid[0 .. n * wide]; + try fwd.hcPreRows(a, w.mlp_hc, off, dev_mid, n, mixed, inj); + const e_mm = errOf(@as([]const u16, taps.mlp_mixed[0 .. n * dims]), mixed); + const e_mi = errOf(@as([]const u16, taps.mlp_inj[0 .. n * streams]), inj); + const dev_mlp = try toF64(a, taps.mlp_mixed[0 .. n * dims]); + const moe_out = try a.alloc(f64, n * dims); + const picks = try a.alloc(u16, n * topk); + try fwd.moeBranch(a, w.moe, dev_mlp, n, moe_out, picks); + var same_picks = true; + for (0..n) |t| { + var x = dt.picks[t * topk ..][0..topk].*; + var y = picks[t * topk ..][0..topk].*; + std.mem.sort(u16, &x, {}, std.sort.asc(u16)); + std.mem.sort(u16, &y, {}, std.sort.asc(u16)); + same_picks = same_picks and std.mem.eql(u16, &x, &y); + } + const e_moe = errOf(dt.mlp_branch[0 .. n * dims], moe_out); + // the write-back alone: the device's streams after the mixer plus its MoE branch times its inject gates + const wb = try a.alloc(f64, n * wide); + for (0..n) |t| for (0..streams) |sidx| for (0..dims) |d| { + const at = t * wide + sidx * dims + d; + wb[at] = dev_mid[at] + dt.mlp_branch[t * dims + d] * bfValue(taps.mlp_inj[t * streams + sidx]); + }; + const e_wb = errOf(@as([]const u16, hg), wb); + std.debug.print(" staged (device input to each block): attn readout mixed rel {e:.3}, inject rel {e:.3}; mixer rel {e:.3}; mlp readout mixed rel {e:.3}, inject rel {e:.3}; MoE rel {e:.3} (picks {s}); write-back rel {e:.3}\n", .{ e_am.rel, e_ai.rel, e_mx.rel, e_mm.rel, e_mi.rel, e_moe.rel, if (same_picks) "equal" else "DIFFER", e_wb.rel }); + alone.attn_mixed = @max(alone.attn_mixed, e_am.rel); + alone.attn_inj = @max(alone.attn_inj, e_ai.rel); + alone.mixer = @max(alone.mixer, e_mx.rel); + alone.mlp_mixed = @max(alone.mlp_mixed, e_mm.rel); + alone.mlp_inj = @max(alone.mlp_inj, e_mi.rel); + if (same_picks) alone.moe = @max(alone.moe, e_moe.rel); + alone.writeback = @max(alone.writeback, e_wb.rel); + } + try fwd.layerRows(a, o.weights(), &o.state, c[0], r.store.norm_offset, ho, n, ot); + const e_mix = errOf(dt.attn_branch[0 .. n * dims], ot.attn_branch[0 .. n * dims]); + const e_mid = errOf(dt.h_mid[0 .. n * wide], ot.h_mid[0 .. n * wide]); + var same: usize = 0; + for (0..n) |t| { + var x = dt.picks[t * topk ..][0..topk].*; + var y = ot.picks[t * topk ..][0..topk].*; + std.mem.sort(u16, &x, {}, std.sort.asc(u16)); + std.mem.sort(u16, &y, {}, std.sort.asc(u16)); + same += @intFromBool(std.mem.eql(u16, &x, &y)); + } + const e_moe = errOf(dt.mlp_branch[0 .. n * dims], ot.mlp_branch[0 .. n * dims]); + const e_out = errOf(@as([]const u16, hg), ho); + std.debug.print(" rows {d}..{d}: mixer branch rel {e:.3} (peak {d:.3}); streams after mixer rel {e:.3}; MoE picks same set {d}/{d}, MoE branch rel {e:.3} (peak {d:.3}); layer output streams rel {e:.3} (peak {d:.2}, max_abs {e:.2})\n", .{ c[0], c[0] + n - 1, e_mix.rel, e_mix.peak, e_mid.rel, same, n, e_moe.rel, e_moe.peak, e_out.rel, e_out.peak, e_out.max_abs }); + worst.mixer = @max(worst.mixer, e_mix.rel); + worst.streams = @max(worst.streams, e_mid.rel); + worst.picks_same += same; + worst.rows += n; + if (same == n) { + worst.moe = @max(worst.moe, e_moe.rel); + worst.out = @max(worst.out, e_out.rel); + } + } + std.debug.print(" device time for the 5 calls (read-backs between blocks included): {d:.0} ms\n", .{device_ms}); + std.debug.print(" worst staged: attn readout mixed {e:.3} inject {e:.3}; mixer {e:.3}; mlp readout mixed {e:.3} inject {e:.3}; MoE {e:.3}; write-back {e:.3}\n", .{ alone.attn_mixed, alone.attn_inj, alone.mixer, alone.mlp_mixed, alone.mlp_inj, alone.moe, alone.writeback }); + std.debug.print(" worst chained: mixer {e:.3}, streams {e:.3}, MoE {e:.3}, output {e:.3}\n", .{ worst.mixer, worst.streams, worst.moe, worst.out }); + // Each block alone, fed the device's own input: its bf16 output and the roundings inside it. + const staged = [_]f64{ alone.attn_mixed, alone.attn_inj, alone.mixer, alone.mlp_mixed, alone.mlp_inj, alone.moe, alone.writeback }; + for (staged) |v| try check.expect(v <= 0.01, "layer {d}: a block alone differs from float64 by {d:.4} of its peak (> 0.01)", .{ l, v }); + // Chained from the embedding: the old engine's gates (mixer 1% GDN / 3% attention, streams 1%, MoE 2%) hold; its + const mixer_tol: f64 = if (wt.linearLayer(l)) 0.01 else 0.03; + try check.expect(worst.mixer <= mixer_tol, "layer {d}: mixer branch differs from the oracle ({d:.4} > {d})", .{ l, worst.mixer, mixer_tol }); + try check.expect(worst.streams <= 0.01, "layer {d}: streams after the mixer differ ({d:.4})", .{ l, worst.streams }); + try check.expect(worst.moe <= 0.02, "layer {d}: MoE branch differs from the oracle ({d:.4})", .{ l, worst.moe }); + try check.expect(worst.out <= 0.02, "layer {d}: layer output streams differ ({d:.4})", .{ l, worst.out }); + if (worst.out > 0.01) std.debug.print(" NOTE layer {d}: chained output {e:.3} is above the old engine's 1% gate (f32 streams there)\n", .{ l, worst.out }); + check.pass("layer {d} ({s}) through the family's pass (writeback, readout, {s}, write-back, readout, EXL3 moeBlock, write-back) vs float64: mixer rel {e:.3}, streams rel {e:.3}, MoE rel {e:.3}, output rel {e:.3}; picks equal on {d}/{d} rows", .{ l, if (wt.linearLayer(l)) "GDN" else "attention", if (wt.linearLayer(l)) "gdnBlock" else "attnBlock", worst.mixer, worst.streams, worst.moe, worst.out, worst.picks_same, worst.rows }); +} + +pub fn layerMode(gpu: Gpu, dir: []const u8, which: []const usize) !void { + var rig: Rig = undefined; + try rig.init(gpu, dir, 8, 64); + defer rig.deinit(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var prng = std.Random.DefaultPrng.init(0xe713_c03e); + var failed: usize = 0; + for (which) |l| { + const t0 = check.now(gpu.io); + var lw = try pl.loadLayerRank(gpu.gpa, @constCast(gpu.d), &rig.stream, &rig.store, rig.pack_fn, l, 1, 0, null); + defer lw.free(gpu.gpa); + std.debug.print("== layer {d} ({s}): the family's loader (cuda_prompt_load.loadLayerRank) in {d:.0} ms, every expert resident\n", .{ l, if (wt.linearLayer(l)) "GDN" else "attention", @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6 }); + try check.expect(lw.exl3 != null, "layer {d}: experts did not load as EXL3", .{l}); + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + // every layer runs, a failed check reported and counted + projections(gpu, &rig, &raw, &lw, l, arena.allocator(), prng.random()) catch |err| if (err == error.TestFailed) { + failed += 1; + } else return err; + layerCheck(gpu, &rig, &raw, &lw, l, &default_ids) catch |err| if (err == error.TestFailed) { + failed += 1; + } else return err; + } + try check.expect(failed == 0, "flashnext-exl3 layer: {d} checks failed", .{failed}); +} diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_names_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_names_tests.zig new file mode 100644 index 0000000000..49a8d78eb7 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_names_tests.zig @@ -0,0 +1,374 @@ +//! flashnext-exl3: tensor names, norm statistics and the router against the float64 oracle. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const fwd = @import("exl3_fwd_cpu.zig"); +const hc = @import("exl3_hc.zig"); +const moe = @import("exl3_moe.zig"); +const router_cpu = @import("exl3_router.zig"); +const route_cpu = @import("exl3_route.zig"); +const scan = @import("../pack/exl3_scan.zig"); +const place = @import("../pack/exl3_place.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); + +const X = flash.exl3; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const batch = X.batch; +const Gpu = check.Gpu; +const st = core.safetensors; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const topk = slots - 1; +const experts_n = wt.experts_n; +const shared = @import("flashnext_exl3_rig.zig"); +const typed = shared.typed; +const valueOf = shared.valueOf; +const f64Of = shared.f64Of; +const f32Of = shared.f32Of; +const f16Of = shared.f16Of; +const qmat = shared.qmat; +const bf16Bits = shared.bf16Bits; +const bfValue = shared.bfValue; +const Err = shared.Err; +const errOf = shared.errOf; +const download = shared.download; +const uploadTo = shared.uploadTo; +const RawPack = shared.RawPack; +const Pass = shared.Pass; +const Rig = shared.Rig; +const bf16Rows = shared.bf16Rows; +const default_ids = shared.default_ids; + +// -- names ---------------------------------------------------------------------------------------------------------- + +const Names = struct { + gpa: std.mem.Allocator, + store: *wt.Store, + claimed: std.StringHashMapUnmanaged(void) = .empty, + missing: usize = 0, + matrices: usize = 0, + dense: usize = 0, + widened: usize = 0, + by_halves: [17]usize = @splat(0), + + fn claim(n: *Names, requested: []const u8) !void { + var buf: [512]u8 = undefined; + const key = n.store.resolveKey(requested, &buf) catch { + n.missing += 1; + if (n.missing <= 20) std.debug.print(" missing: {s}\n", .{requested}); + return; + }; + if (n.claimed.contains(key)) return; + try n.claimed.put(n.gpa, try n.gpa.dupe(u8, key), {}); + } + + /// An EXL3 matrix: its four parts claimed and the matrix admitted by the family's own reader (marker, widths). + fn matrix(n: *Names, comptime fmt: []const u8, args: anytype) !void { + var buf: [256]u8 = undefined; + const prefix = try std.fmt.bufPrint(&buf, fmt, args); + var pb: [280]u8 = undefined; + for ([_][]const u8{ "trellis", "suh", "svh", "mul1" }) |part| try n.claim(try std.fmt.bufPrint(&pb, "{s}.{s}", .{ prefix, part })); + if (!X.load.has(n.store, prefix)) { + n.missing += 1; + std.debug.print(" not EXL3: {s}\n", .{prefix}); + return; + } + const m = try X.load.mat(n.store, prefix); + n.by_halves[m.halves] += 1; + n.matrices += 1; + } + + /// A dense tensor as the loaders read it: through Store.tensor, so an F16 `.weight` is widened (exact or refused). + fn tensor(n: *Names, comptime fmt: []const u8, args: anytype) !void { + var buf: [256]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + try n.claim(name); + if (!n.store.has(name)) return; + const before = n.store.widened.items.len; + _ = try n.store.tensor(name); + if (n.store.widened.items.len > before) n.widened += 1; + n.dense += 1; + } + + fn hcw(n: *Names, base: []const u8, inject: bool) !void { + try n.tensor("{s}.hc_norm.weight", .{base}); + try n.tensor("{s}.input_mix_weight_down.weight", .{base}); + try n.tensor("{s}.input_mix_weight_up.weight", .{base}); + if (inject) try n.tensor("{s}.block_inject_weight.weight", .{base}); + } + + /// What cuda_prompt_load.loadLayerAtRank (and the EXL3 hooks it reaches) read for one decoder layer at `base`. + fn layer(n: *Names, base: []const u8, linear: bool, with_ple: bool) !void { + var hb: [128]u8 = undefined; + try n.hcw(try std.fmt.bufPrint(&hb, "{s}.attn_hyper_connection", .{base}), true); + try n.hcw(try std.fmt.bufPrint(&hb, "{s}.mlp_hyper_connection", .{base}), true); + try n.tensor("{s}.mlp.gate.weight", .{base}); + try n.tensor("{s}.mlp.shared_expert_gate.weight", .{base}); + if (linear) { + for ([_][]const u8{ "in_proj_qkv", "in_proj_z", "out_proj" }) |p| try n.matrix("{s}.linear_attn.{s}", .{ base, p }); + for ([_][]const u8{ "in_proj_b.weight", "in_proj_a.weight", "conv1d.weight", "A_log", "dt_bias", "norm.weight" }) |p| try n.tensor("{s}.linear_attn.{s}", .{ base, p }); + } else { + for ([_][]const u8{ "q_proj", "k_proj", "v_proj", "indexer.index_qk_proj", "o_proj" }) |p| try n.matrix("{s}.self_attn.{s}", .{ base, p }); + for ([_][]const u8{ "q_norm", "k_norm", "indexer.q_layernorm", "indexer.k_layernorm" }) |p| try n.tensor("{s}.self_attn.{s}.weight", .{ base, p }); + } + if (with_ple) { + for ([_][]const u8{ "key_proj", "value_proj", "norm_key", "norm_query", "norm_conv", "conv1d" }) |p| try n.tensor("{s}.ple.{s}.weight", .{ base, p }); + } + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |p| try n.matrix("{s}.mlp.shared_expert.{s}", .{ base, p }); + for (0..experts_n) |e| { + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }) |p| try n.matrix("{s}.mlp.experts.{d}.{s}", .{ base, e, p }); + if (e % 64 == 63) n.store.release(); + } + n.store.release(); + } +}; + +pub fn names(gpu: Gpu, dir: []const u8) !void { + const gpa = gpu.gpa; + var store = try wt.Store.open(gpa, gpu.io, dir); + defer store.deinit(); + var n: Names = .{ .gpa = gpa, .store = &store }; + defer { + var it = n.claimed.keyIterator(); + while (it.next()) |k| gpa.free(k.*); + n.claimed.deinit(gpa); + } + const t0 = check.now(gpu.io); + // the engine (cuda_prompt.zig Engine.initRank): embedding, the final mixer, the head (EXL3 here: plan step 9) + try n.tensor("language_model.model.embed_tokens.weight", .{}); + try n.hcw("language_model.model.hyper_connection_mixer", false); + try n.matrix("language_model.lm_head", .{}); + for (0..wt.layers_n) |l| { + var bb: [64]u8 = undefined; + try n.layer(try std.fmt.bufPrint(&bb, "language_model.model.layers.{d}", .{l}), wt.linearLayer(l), l == wt.ple_layer); + } + // the MTP head (cuda_mtp.zig Head.init): norms, two EXL3 fc projections, its mixer (the patch file), one layer + try n.tensor("language_model.mtp.pre_fc_norm_embedding.weight", .{}); + try n.tensor("language_model.mtp.pre_fc_norm_hidden.weight", .{}); + try n.matrix("language_model.mtp.fc_embedding", .{}); + try n.matrix("language_model.mtp.fc_hidden", .{}); + try n.hcw("language_model.mtp.hyper_connection_mixer", false); + try n.layer("language_model.mtp.layers.0", false, false); + // everything the index (and the patch) names: claimed, or the vision tower; and the header rules' homes agree + const wm = store.parsed.value.object.get("weight_map").?.object; + var vision: usize = 0; + var unclaimed: usize = 0; + var homes: [8]usize = @splat(0); + var disagree: usize = 0; + var it = wm.iterator(); + while (it.next()) |kv| { + const name = kv.key_ptr.*; + const home = place.home(.{ .name = name, .rank = if (std.mem.endsWith(u8, name, ".trellis")) 3 else 1, .file = kv.value_ptr.string }); + homes[@backingInt(home)] += 1; + if (n.claimed.contains(name)) { + // a routed expert's part must be the routed bank's, a shared expert's part the resident one's + if (std.mem.indexOf(u8, name, ".mlp.experts.") != null and home != .routed_host) disagree += 1; + if (std.mem.indexOf(u8, name, ".shared_expert.") != null and home != .shared_resident) disagree += 1; + continue; + } + if (std.mem.startsWith(u8, name, "model.visual.")) { + vision += 1; + continue; + } + unclaimed += 1; + if (unclaimed <= 20) std.debug.print(" unclaimed: {s}\n", .{name}); + } + const ms = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6; + std.debug.print("names: {d} indexed (+ the mixer patch), {d} claimed, {d} vision, {d} unclaimed, {d} missing; {d} EXL3 matrices admitted (mul1 marker, widths from the trellis), {d} dense tensors read, {d} F16 weights widened to BF16 exactly; {d:.0} ms\n", .{ wm.count(), n.claimed.count(), vision, unclaimed, n.missing, n.matrices, n.dense, n.widened, ms }); + std.debug.print(" EXL3 widths (halves: count):", .{}); + for (n.by_halves, 0..) |c, h| if (c > 0) std.debug.print(" {d}:{d}", .{ h, c }); + std.debug.print("\n header homes (exl3_place):", .{}); + for (homes, 0..) |c, i| if (c > 0) std.debug.print(" {s}={d}", .{ @tagName(@as(place.Home, @fromBackingInt(@intCast(i)))), c }); + std.debug.print("\n", .{}); + try check.expect(n.missing == 0, "names: {d} tensors the family reads do not resolve", .{n.missing}); + try check.expect(unclaimed == 0, "names: {d} indexed tensors no family loader reads (vision excepted)", .{unclaimed}); + try check.expect(disagree == 0, "names: {d} expert tensors whose exl3_place home disagrees", .{disagree}); + // the family's device sizing (cuda_exl3_memory.zig) on the same pack + const gib = 1024.0 * 1024.0 * 1024.0; + const on_host = try X.memory.deviceBytes(gpa, gpu.io, dir, .{}); + const resident = try X.memory.deviceBytes(gpa, gpu.io, dir, .{ .routed_on_host = false }); + std.debug.print(" device bytes: {d:.2} GiB with routed experts on the host, {d:.2} GiB all resident\n", .{ @as(f64, @floatFromInt(on_host)) / gib, @as(f64, @floatFromInt(resident)) / gib }); + check.pass("names: every tensor the family reads resolves ({d}), nothing unclaimed but {d} vision tensors; {d} EXL3 matrices admitted; {d} F16 weights widened exactly", .{ n.claimed.count(), vision, n.matrices, n.widened }); +} + +// -- norms ---------------------------------------------------------------------------------------------------------- + +fn meanOf(t: st.Tensor) f64 { + var s: f64 = 0; + const count = t.bytes.len / 2; + for (0..count) |i| s += valueOf(t, i); + return s / @as(f64, @floatFromInt(count)); +} + +pub fn norms(gpu: Gpu, dir: []const u8) !void { + const gpa = gpu.gpa; + var store = try wt.Store.open(gpa, gpu.io, dir); + defer store.deinit(); + try store.detectNorms(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var means: [wt.layers_n]f64 = undefined; + for (&means, 0..) |*m, l| m.* = meanOf(try raw.get("model.language_model.layers.{d}.attn_hyper_connection.hc_norm.weight", .{l})); + var sorted = means; + std.mem.sort(f64, &sorted, {}, std.sort.asc(f64)); + std.debug.print("norms: 48 attn hc_norm means min {d:.4} median {d:.4} max {d:.4} -> offset {d} (cuda_norms.offset via Store.detectNorms)\n", .{ sorted[0], (sorted[23] + sorted[24]) / 2, sorted[47], store.norm_offset }); + try check.expect(store.norm_offset == 1, "norms: the pack stores gamma - 1; the convention came out {d}", .{store.norm_offset}); + // the independent cross-check: the indexer norms' own 24 means locate the same storage + var idx: [24]f32 = undefined; + var c: usize = 0; + for (0..wt.layers_n) |l| { + if (wt.linearLayer(l)) continue; + for ([_][]const u8{ "q_layernorm", "k_layernorm" }) |nm| { + idx[c] = @floatCast(meanOf(try raw.get("model.language_model.layers.{d}.self_attn.indexer.{s}.weight", .{ l, nm }))); + c += 1; + } + } + const indexer = try scan.centredOffset(&idx); + std.debug.print("norms: indexer q/k_layernorm means over the 12 attention layers -> exl3_scan.centredOffset {d}\n", .{indexer}); + try check.expect(indexer == store.norm_offset, "norms: the indexer norms locate offset {d}, the hc_norm means {d}", .{ indexer, store.norm_offset }); + // every centred norm the family scales, MTP's included, loads through Store.normScale with that one offset + const mtp_names = [_][]const u8{ "language_model.mtp.pre_fc_norm_embedding.weight", "language_model.mtp.pre_fc_norm_hidden.weight", "language_model.mtp.hyper_connection_mixer.hc_norm.weight", "language_model.mtp.layers.0.attn_hyper_connection.hc_norm.weight", "language_model.mtp.layers.0.mlp_hyper_connection.hc_norm.weight", "language_model.mtp.layers.0.self_attn.q_norm.weight", "language_model.mtp.layers.0.self_attn.k_norm.weight", "language_model.mtp.layers.0.self_attn.indexer.q_layernorm.weight", "language_model.mtp.layers.0.self_attn.indexer.k_layernorm.weight", "language_model.model.layers.1.ple.norm_key.weight", "language_model.model.layers.1.ple.norm_query.weight", "language_model.model.layers.1.ple.norm_conv.weight", "language_model.model.layers.3.self_attn.q_norm.weight", "language_model.model.layers.3.self_attn.k_norm.weight", "language_model.model.hyper_connection_mixer.hc_norm.weight" }; + for (mtp_names) |name| { + const scales = try store.normScale(gpa, try store.tensor(name)); + defer gpa.free(scales); + var sum: f64 = 0; + var lo: f32 = std.math.inf(f32); + for (scales) |v| { + sum += v; + lo = @min(lo, v); + } + std.debug.print(" {s:<72} stored mean {d:>7.4} -> scale mean {d:.4} (min {d:.4})\n", .{ name, sum / @as(f64, @floatFromInt(scales.len)) - store.norm_offset, sum / @as(f64, @floatFromInt(scales.len)), lo }); + } + store.release(); + check.pass("norms: one convention (offset {d}) from the 48 hc_norm means, confirmed by the indexer norms; every centred norm incl. MTP's scales finite", .{store.norm_offset}); +} + +// -- router --------------------------------------------------------------------------------------------------------- + +pub fn router(gpu: Gpu, dir: []const u8, layers: []const usize, rows: usize) !void { + const gpa = gpu.gpa; + var store = try wt.Store.open(gpa, gpu.io, dir); + defer store.deinit(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var glue = try flash.Glue.load(gpu.d); + defer glue.deinit(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0x5eed_2026); + const rng = prng.random(); + var worst_rows_diff: usize = 0; + for (layers) |l| { + // the family's router buffer, as cuda_prompt_load builds it: mlp.gate's 512 rows then the shared gate. + var nb: [128]u8 = undefined; + const gate = try store.tensor(try std.fmt.bufPrint(&nb, "language_model.model.layers.{d}.mlp.gate.weight", .{l})); + try check.expect(gate.is(.bf16, &.{ experts_n, dims }), "router: layer {d} gate is not the loaders' bf16 [512, 2560]", .{l}); + const w = try gpa.alloc(u8, (experts_n + 1) * dims * 2); + defer gpa.free(w); + @memcpy(w[0 .. experts_n * dims * 2], gate.bytes); + const sg = try store.tensor(try std.fmt.bufPrint(&nb, "language_model.model.layers.{d}.mlp.shared_expert_gate.weight", .{l})); + @memcpy(w[experts_n * dims * 2 ..], sg.bytes); + store.release(); + var wd = try cuda.DeviceBuffer.fromHost(gpu.d, w); + defer wd.free(); + // the oracle's weights: the pack's F16 bits, never the widened copy + const g16 = try f16Of(gpa, try raw.get("model.language_model.layers.{d}.mlp.gate.weight", .{l})); + defer gpa.free(g16); + const s16 = try f16Of(gpa, try raw.get("model.language_model.layers.{d}.mlp.shared_expert_gate.weight", .{l})); + defer gpa.free(s16); + // rows: N(0, sigma) per row, sigma in {0.25, 1, 4}, rounded to bf16 as the pass's `mixed` is + const xb = try gpa.alloc(u16, rows * dims); + defer gpa.free(xb); + for (0..rows) |r| { + const sigma: f32 = switch (r % 3) { + 0 => 0.25, + 1 => 1, + else => 4, + }; + for (0..dims) |i| xb[r * dims + i] = bf16Bits(rng.floatNorm(f32) * sigma); + } + var xd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xb)); + defer xd.free(); + var ld = try cuda.DeviceBuffer.alloc(gpu.d, rows * (experts_n + 1) * 4); + defer ld.free(); + var pd = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer pd.free(); + var wdd = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); + defer wdd.free(); + try glue.routerMany(s, xd.ptr, wd.ptr, ld.ptr, dims, rows); + try glue.topkMany(s, ld.ptr, pd.ptr, wdd.ptr, rows); + const gl = try download(f32, gpu, s, ld.ptr, rows * (experts_n + 1)); + defer gpa.free(gl); + const gp = try download(i32, gpu, s, pd.ptr, rows * slots); + defer gpa.free(gp); + const gw = try download(f32, gpu, s, wdd.ptr, rows * slots); + defer gpa.free(gw); + var diff_rows: usize = 0; + var logit_err: f64 = 0; + var min_gap: f64 = std.math.inf(f64); + var worst_w: f64 = 0; + var worst_sw: f64 = 0; + var f32_agree: usize = 0; + const x = try gpa.alloc(f64, dims); + defer gpa.free(x); + const xf = try gpa.alloc(f32, dims); + defer gpa.free(xf); + var logits: [experts_n + 1]f64 = undefined; + for (0..rows) |r| { + for (x, xf, 0..) |*v, *vf, i| { + v.* = bfValue(xb[r * dims + i]); + vf.* = @floatCast(v.*); + } + try moe.logitsF16(g16, experts_n, s16, x, &logits); + var picks: [topk]u16 = undefined; + var wts: [topk]f64 = undefined; + var sw: f64 = 0; + try moe.routeFromLogits(&logits, topk, &picks, &wts, &sw); + for (0..experts_n + 1) |i| logit_err = @max(logit_err, @abs(@as(f64, gl[r * (experts_n + 1) + i]) - logits[i])); + var sorted = logits[0..experts_n].*; + std.mem.sort(f64, &sorted, {}, std.sort.desc(f64)); + const gap = sorted[topk - 1] - sorted[topk]; + min_gap = @min(min_gap, gap); + var a: [topk]u16 = undefined; + for (&a, 0..) |*v, k| v.* = @intCast(gp[r * slots + k]); + var b = picks; + std.mem.sort(u16, &a, {}, std.sort.asc(u16)); + std.mem.sort(u16, &b, {}, std.sort.asc(u16)); + if (!std.mem.eql(u16, &a, &b)) { + diff_rows += 1; + if (diff_rows <= 10) std.debug.print(" layer {d} row {d}: picks differ; f64 logit gap between the 10th and 11th {e:.3}\n", .{ l, r, gap }); + } else { + for (0..topk) |k| { + // the same expert's weight: the device's bf16 against the float64 softmax + const id = gp[r * slots + k]; + for (picks, wts) |pk, wv| if (pk == id) { + worst_w = @max(worst_w, @abs(@as(f64, gw[r * slots + k]) - wv)); + }; + } + try check.expect(gp[r * slots + topk] == experts_n, "router: the shared slot holds {d}", .{gp[r * slots + topk]}); + worst_sw = @max(worst_sw, @abs(@as(f64, gw[r * slots + topk]) - sw)); + } + // the f32 host reference (exl3_router.launchDtype(.f16)) on the same row: its picks as a set + var l32: [experts_n + 1]f32 = undefined; + var p32: [topk + 1]u16 = undefined; + var w32: [topk + 1]f32 = undefined; + var scratch: [experts_n]f32 = undefined; + try router_cpu.launchDtype(.f16, 0, std.mem.sliceAsBytes(g16), xf, std.mem.sliceAsBytes(s16), xf, topk, &l32, &p32, &w32, &scratch); + var c32 = p32[0..topk].*; + std.mem.sort(u16, &c32, {}, std.sort.asc(u16)); + f32_agree += @intFromBool(std.mem.eql(u16, &c32, &b)); + } + std.debug.print("router layer {d}: {d} rows, picks equal to float64 on {d}/{d} (f32 host reference equal on {d}/{d}); max |device logit - f64| {e:.3}; smallest f64 10th/11th gap {e:.3}; weights max |bf16 - f64| {e:.3}, shared {e:.3}\n", .{ l, rows, rows - diff_rows, rows, f32_agree, rows, logit_err, min_gap, worst_w, worst_sw }); + worst_rows_diff = @max(worst_rows_diff, diff_rows); + try check.expect(diff_rows == 0 or min_gap <= 4 * logit_err, "router layer {d}: {d} rows pick differently with every gap above the logit error", .{ l, diff_rows }); + } + try check.expect(worst_rows_diff == 0, "router: picks differ from float64 on {d} rows (see the gaps above)", .{worst_rows_diff}); + check.pass("router: the exactly widened F16 router's top-10 equals the float64 oracle on every row of layers {any}", .{layers}); +} diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_proj_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_proj_tests.zig new file mode 100644 index 0000000000..0ed5e8c8d8 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_proj_tests.zig @@ -0,0 +1,269 @@ +//! flashnext-exl3: the projections (gate 4) and the PLE key/value projections. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const fwd = @import("exl3_fwd_cpu.zig"); +const hc = @import("exl3_hc.zig"); +const moe = @import("exl3_moe.zig"); +const router_cpu = @import("exl3_router.zig"); +const route_cpu = @import("exl3_route.zig"); +const scan = @import("../pack/exl3_scan.zig"); +const place = @import("../pack/exl3_place.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); + +const X = flash.exl3; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const batch = X.batch; +const Gpu = check.Gpu; +const st = core.safetensors; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const topk = slots - 1; +const experts_n = wt.experts_n; +const shared = @import("flashnext_exl3_rig.zig"); +const typed = shared.typed; +const valueOf = shared.valueOf; +const f64Of = shared.f64Of; +const f32Of = shared.f32Of; +const f16Of = shared.f16Of; +const qmat = shared.qmat; +const bf16Bits = shared.bf16Bits; +const bfValue = shared.bfValue; +const Err = shared.Err; +const errOf = shared.errOf; +const download = shared.download; +const uploadTo = shared.uploadTo; +const RawPack = shared.RawPack; +const Pass = shared.Pass; +const Rig = shared.Rig; +const bf16Rows = shared.bf16Rows; +const default_ids = shared.default_ids; + +// -- projections (gate 4) -------------------------------------------------------------------------------------------- + +/// One Projection through Pass.lane(): `sk` 1 is bf16 [m, n]; `sk` > 1 fp32 [m, n] in the first slice, others zeroed. +fn laneOut(gpu: Gpu, p: *Pass, w: wt.Q4, x: u64, m: usize, sk: usize) !struct { bf: ?[]u16 = null, f: ?[]f32 = null, rest_zero: bool = true } { + const s = p.s; + if (sk == 1) { + _ = try p.lane(x, s.dummy, w, s.proj, s.dummy, m, 1); + return .{ .bf = try download(u16, gpu, p.stream, s.proj, m * w.n) }; + } + // the slices in a buffer of their own: a 12,288-wide part's exceeds the pass's out_sk x dims `part` + var slices = try cuda.DeviceBuffer.alloc(gpu.d, sk * m * w.n * 4); + defer slices.free(); + try slices.fill8(0xff, p.stream.handle); + _ = try p.lane(x, s.dummy, w, s.dummy, slices.ptr, m, sk); + const all = try download(f32, gpu, p.stream, slices.ptr, sk * m * w.n); + defer gpu.gpa.free(all); + var zero = true; + for (all[m * w.n ..]) |v| zero = zero and v == 0; + return .{ .f = try gpu.gpa.dupe(f32, all[0 .. m * w.n]), .rest_zero = zero }; +} + +/// The f64 product of `rows` bf16 rows through the EXL3 matrix `q`. +fn exactRows(a: std.mem.Allocator, q: *const fwd.QMat, xb: []const u16, rows: usize) ![]f64 { + const y = try a.alloc(f64, rows * q.n); + const x = try a.alloc(f64, q.k); + defer a.free(x); + const tmp = try a.alloc(f64, q.k); + defer a.free(tmp); + for (0..rows) |r| { + for (x, 0..) |*v, i| v.* = bfValue(xb[r * q.k + i]); + q.apply(x, y[r * q.n ..][0..q.n], tmp); + } + return y; +} + +/// The f64 product of `rows` bf16 rows through a dense row-major [n, k] weight. +fn denseRows(a: std.mem.Allocator, w: []const f64, n: usize, k: usize, xb: []const u16, rows: usize) ![]f64 { + const y = try a.alloc(f64, rows * n); + for (0..rows) |r| for (0..n) |j| { + var acc: f64 = 0; + for (0..k) |i| acc += bfValue(xb[r * k + i]) * w[j * k + i]; + y[r * n + j] = acc; + }; + return y; +} + +/// The fp16 rotated-input contract: x * suh after the Hadamard is rounded to fp16 (2^-11 of each term), so a column's +const proj_f32_tol = 4e-4; +const proj_bf16_tol = 4e-3; // bf16 output rounding: half an ulp is 2^-9 of a value, at most 1.95e-3 of the peak + +pub fn projections(gpu: Gpu, r: *Rig, raw: *const RawPack, lw: *const pl.Layer, l: usize, a: std.mem.Allocator, rng: std.Random) !void { + var p = r.pass((lw)[0..1]); + const linear = wt.linearLayer(l); + const kind = if (linear) "linear_attn" else "self_attn"; + const e = lw.proj.exl3 orelse return error.NotExl3; + try check.expect(lw.out.exl3 != null and lw.proj.native == null and lw.out.native == null, "layer {d}: projections did not load as EXL3", .{l}); + var worst32: f64 = 0; + var worst16: f64 = 0; + for ([_]usize{ 1, 5, 8 }) |m| { + const xb = try bf16Rows(a, rng, m * dims, 1); + var xd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xb)); + defer xd.free(); + // the stitched input projection (bf16, as the pass reads it) and each of its parts alone (fp32) + const whole = try laneOut(gpu, &p, lw.proj, xd.ptr, m, 1); + for (e.parts) |pt| { + const want = switch (pt.kind) { + .exl3 => blk: { + var names_buf: [128]u8 = undefined; + const name = partName(&names_buf, l, kind, pt.col, linear); + var q = try qmat(a, raw, name, 48); + defer q.deinit(a); + const y = try exactRows(a, &q, xb, m); + // the part alone in fp32: Projection.exl3 over one part + var one = e.*; + var single = [1]X.load.Part{pt}; + single[0].col = 0; + one.parts = &single; + one.n = pt.n; + var view = wt.Q4.empty(gpu.d); + view.n = pt.n; + view.k = dims; + view.exl3 = &one; + const got = try laneOut(gpu, &p, view, xd.ptr, m, 2); + const ef = errOf(got.f.?, y); + try check.expect(got.rest_zero, "layer {d}: {s} left the other K slice unzeroed", .{ l, name }); + std.debug.print(" L{d} {s:<50} rows={d} fp32 out: max|err| {e:.3} peak {d:.3} rel {e:.3}\n", .{ l, name, m, ef.max_abs, ef.peak, ef.rel }); + worst32 = @max(worst32, ef.rel); + break :blk y; + }, + .native => blk: { + // in_proj_b then in_proj_a: the F16 rows the Store widened, the oracle reads the stored F16 + const b = try f64Of(a, try raw.get("model.language_model.layers.{d}.linear_attn.in_proj_b.weight", .{l})); + const av = try f64Of(a, try raw.get("model.language_model.layers.{d}.linear_attn.in_proj_a.weight", .{l})); + const w = try std.mem.concat(a, f64, &.{ b, av }); + break :blk try denseRows(a, w, pt.n, dims, xb, m); + }, + }; + // the part's columns of the stitched bf16 row + const got = try a.alloc(u16, m * pt.n); + for (0..m) |row| @memcpy(got[row * pt.n ..][0..pt.n], whole.bf.?[row * e.n + pt.col ..][0..pt.n]); + const eb = errOf(got, want); + std.debug.print(" L{d} stitched in-proj columns {d}..{d} ({s}) rows={d} bf16 out: rel {e:.3}\n", .{ l, pt.col, pt.col + pt.n, @tagName(pt.kind), m, eb.rel }); + worst16 = @max(worst16, eb.rel); + } + // the output projection as the pass runs it: fp32 into the first of out_sk K slices + const ok = lw.out.k; + const ob = try bf16Rows(a, rng, m * ok, 1); + var od = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(ob)); + defer od.free(); + var nb: [128]u8 = undefined; + const oname = try std.fmt.bufPrint(&nb, "model.language_model.layers.{d}.{s}.{s}", .{ l, kind, if (linear) "out_proj" else "o_proj" }); + var q = try qmat(a, raw, oname, 48); + defer q.deinit(a); + const want = try exactRows(a, &q, ob, m); + const got = try laneOut(gpu, &p, lw.out, od.ptr, m, wt.out_sk); + const eo = errOf(got.f.?, want); + try check.expect(got.rest_zero, "layer {d}: the output projection left its other K slices unzeroed", .{l}); + std.debug.print(" L{d} {s:<50} rows={d} fp32 out (slice 0 of {d}): max|err| {e:.3} peak {d:.3} rel {e:.3}\n", .{ l, oname, m, wt.out_sk, eo.max_abs, eo.peak, eo.rel }); + worst32 = @max(worst32, eo.rel); + } + try check.expect(worst32 <= proj_f32_tol, "layer {d}: an EXL3 projection's fp32 output is {e:.3} of peak from float64 (> {e})", .{ l, worst32, proj_f32_tol }); + try check.expect(worst16 <= proj_bf16_tol, "layer {d}: the stitched bf16 input projection is {e:.3} of peak from float64 (> {e})", .{ l, worst16, proj_bf16_tol }); + check.pass("layer {d} projections through Pass.lane(): fp32 out worst {e:.3} of peak, stitched bf16 in-proj worst {e:.3}, rows 1/5/8", .{ l, worst32, worst16 }); +} + +/// The EXL3 matrix behind a stitched part, from its first column. +fn partName(buf: []u8, l: usize, kind: []const u8, col: usize, linear: bool) []const u8 { + const name: []const u8 = if (linear) switch (col) { + 0 => "in_proj_qkv", + else => "in_proj_z", + } else switch (col) { + 0 => "q_proj", + 12288 => "k_proj", + 12800 => "v_proj", + else => "indexer.index_qk_proj", + }; + return std.fmt.bufPrint(buf, "model.language_model.layers.{d}.{s}.{s}", .{ l, kind, name }) catch unreachable; +} + +/// The PLE key/value projections of an EXL3 pack take the n-gram rows as fp16 (cuda_exl3_ple.zig's default): through +pub fn pleMode(gpu: Gpu, dir: []const u8) !void { + var rig: Rig = undefined; + try rig.init(gpu, dir, 8, 64); + defer rig.deinit(); + var raw = try RawPack.open(gpu, dir); + defer raw.close(gpu); + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const d: *cuda.Driver = @constCast(gpu.d); + // an EXL3 pack's store holds the shared Ops once its first EXL3 matrix loads; this check loads none + (try X.load.opsOf(&rig.store, gpu.d)).release(); + var key = try X.projection_load.load(a, d, &rig.store, "language_model.model.layers.1.ple.key_proj", wide, dims); + defer key.free(); + var value = try X.projection_load.load(a, d, &rig.store, "language_model.model.layers.1.ple.value_proj", dims, dims); + defer value.free(); + try X.load.pleRows(a, d, &rig.store, &key, &value); + try check.expect(key.exl3 != null and key.exl3.?.x == .f16 and value.exl3.?.x == .f16, "ple: key/value did not take fp16 rows", .{}); + var p = rig.pass(&.{}); + const kw = try f64Of(a, try raw.get("model.language_model.layers.1.ple.key_proj.weight", .{})); + const vw = try f64Of(a, try raw.get("model.language_model.layers.1.ple.value_proj.weight", .{})); + var prng = std.Random.DefaultPrng.init(0x91e_f16); + const rng = prng.random(); + var worst_key: f64 = 0; + var worst_value: f64 = 0; + var worst_bf16: f64 = 0; + for ([_]usize{ 1, 5, 8 }) |m| { + const xh = try a.alloc(f16, m * dims); + const xv = try a.alloc(f64, m * dims); + const xr = try a.alloc(u16, m * dims); // the same rows rounded to bf16 + for (xh, xv, xr) |*h, *v, *b| { + h.* = @floatCast(rng.floatNorm(f32) * 0.5); + v.* = @floatCast(h.*); + b.* = bf16Bits(@floatCast(h.*)); + } + var xd = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer xd.free(); + const want_k = try a.alloc(f64, m * wide); + const want_v = try a.alloc(f64, m * dims); + for (0..m) |r| { + for (0..wide) |j| { + var acc: f64 = 0; + for (0..dims) |i| acc += xv[r * dims + i] * kw[j * dims + i]; + want_k[r * wide + j] = acc; + } + for (0..dims) |j| { + var acc: f64 = 0; + for (0..dims) |i| acc += xv[r * dims + i] * vw[j * dims + i]; + want_v[r * dims + j] = acc; + } + } + const got_k = try laneOut(gpu, &p, key, xd.ptr, m, 1); + const ek = errOf(got_k.bf.?, want_k); + // the value projection as pleBlock reads it: the sum of its 4 fp32 K slices + var slices = try cuda.DeviceBuffer.alloc(gpu.d, 4 * m * dims * 4); + defer slices.free(); + try slices.fill8(0xff, p.stream.handle); + _ = try p.lane(xd.ptr, p.s.dummy, value, p.s.dummy, slices.ptr, m, 4); + const all = try download(f32, gpu, p.stream, slices.ptr, 4 * m * dims); + defer gpu.gpa.free(all); + const sum = try a.alloc(f64, m * dims); + for (sum, 0..) |*o, i| o.* = @as(f64, all[i]) + all[m * dims + i] + all[2 * m * dims + i] + all[3 * m * dims + i]; + const ev = errOf(sum, want_v); + // the bf16-row alternative, on the host in float64: rows rounded to bf16, weights exact + const alt = try a.alloc(f64, m * dims); + for (0..m) |r| for (0..dims) |j| { + var acc: f64 = 0; + for (0..dims) |i| acc += bfValue(xr[r * dims + i]) * vw[j * dims + i]; + alt[r * dims + j] = acc; + }; + const eb = errOf(alt, want_v); + std.debug.print(" ple rows={d}: key_proj bf16 out rel {e:.3}; value_proj fp32 slices rel {e:.3}; (bf16-rounded rows would give {e:.3})\n", .{ m, ek.rel, ev.rel, eb.rel }); + worst_key = @max(worst_key, ek.rel); + worst_value = @max(worst_value, ev.rel); + worst_bf16 = @max(worst_bf16, eb.rel); + } + // key: bf16 output rounding (2^-9 of a value); value: fp32 sums of exact products + try check.expect(worst_key <= 4e-3, "ple: key_proj from fp16 rows is {e:.3} of peak from float64", .{worst_key}); + try check.expect(worst_value <= 1e-5, "ple: value_proj from fp16 rows is {e:.3} of peak from float64", .{worst_value}); + check.pass("ple: fp16 n-gram rows through key/value (exact bf16 split): key {e:.3}, value {e:.3} of peak (bf16-rounded rows: {e:.3})", .{ worst_key, worst_value, worst_bf16 }); +} diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_rig.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_rig.zig new file mode 100644 index 0000000000..55f85d788f --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_rig.zig @@ -0,0 +1,251 @@ +//! flashnext-exl3: the raw-pack reader, f64 conversions, error metrics and the one-layer pass rig the gates share. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); +const fwd = @import("exl3_fwd_cpu.zig"); +const hc = @import("exl3_hc.zig"); +const moe = @import("exl3_moe.zig"); +const router_cpu = @import("exl3_router.zig"); +const route_cpu = @import("exl3_route.zig"); +const scan = @import("../pack/exl3_scan.zig"); +const place = @import("../pack/exl3_place.zig"); +const gdn_cpu = @import("exl3_gdn_cpu.zig"); +const attn_cpu = @import("exl3_attn_cpu.zig"); + +const X = flash.exl3; +const wt = X.prompt_wt; +const pl = X.prompt_load; +const batch = X.batch; +const Gpu = check.Gpu; +const st = core.safetensors; + +const dims = wt.dims; +const wide = wt.wide; +const streams = wt.streams; +const slots = wt.slots; +const topk = slots - 1; +const experts_n = wt.experts_n; + +/// "The capital of France is Paris. It is known for the Eiffel Tower." (exl3_fwd_tests.zig's default ids) +pub const default_ids = [_]u32{ 760, 6511, 314, 9338, 369, 11751, 13, 1049, 369, 3750, 364, 279, 242476, 300, 21262, 13 }; + +// -- the raw pack (the oracle's view: tensors exactly as stored, no widening) ------------------------------------- + +pub const RawPack = struct { + files: std.ArrayList(st.File) = .empty, + + pub fn open(gpu: Gpu, dir: []const u8) !RawPack { + var p: RawPack = .{}; + var buf: [512]u8 = undefined; + for (0..7) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-00007.safetensors", .{ dir, i + 1 }); + try p.files.append(gpu.gpa, try st.File.open(gpu.gpa, gpu.io, path)); + } + return p; + } + + pub fn close(p: *RawPack, gpu: Gpu) void { + for (p.files.items) |*f| f.close(gpu.io); + p.files.deinit(gpu.gpa); + } + + pub fn get(p: *const RawPack, comptime fmt: []const u8, args: anytype) !st.Tensor { + var buf: [256]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + for (p.files.items) |*f| if (f.get(name)) |t| return t; + std.debug.print("raw pack: no {s}\n", .{name}); + return error.MissingTensor; + } +}; + +pub fn typed(comptime T: type, a: std.mem.Allocator, bytes: []const u8) ![]T { + const out = try a.alloc(T, bytes.len / @sizeOf(T)); + @memcpy(std.mem.sliceAsBytes(out), bytes[0 .. out.len * @sizeOf(T)]); + return out; +} + +pub fn valueOf(t: st.Tensor, i: usize) f64 { + const h = std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little); + return switch (t.dtype) { + .f16 => @floatCast(@as(f16, @bitCast(h))), + .bf16 => @as(f32, @bitCast(@as(u32, h) << 16)), + else => unreachable, + }; +} + +pub fn f64Of(a: std.mem.Allocator, t: st.Tensor) ![]f64 { + if (t.dtype != .f16 and t.dtype != .bf16) return error.UnexpectedTensor; + const out = try a.alloc(f64, t.bytes.len / 2); + for (out, 0..) |*o, i| o.* = valueOf(t, i); + return out; +} + +pub fn f32Of(a: std.mem.Allocator, t: st.Tensor) ![]f32 { + if (t.dtype != .f16 and t.dtype != .bf16) return error.UnexpectedTensor; + const out = try a.alloc(f32, t.bytes.len / 2); + for (out, 0..) |*o, i| o.* = @floatCast(valueOf(t, i)); + return out; +} + +pub fn f16Of(a: std.mem.Allocator, t: st.Tensor) ![]f16 { + if (t.dtype != .f16) return error.UnexpectedTensor; + return typed(f16, a, t.bytes); +} + +/// The EXL3 matrix `prefix` as the float64 oracle's QMat (W = diag(suh) H Wq H diag(svh), exl3_format's decoder). +pub fn qmat(a: std.mem.Allocator, raw: *const RawPack, prefix: []const u8, threads: usize) !fwd.QMat { + const t = try raw.get("{s}.trellis", .{prefix}); + const trellis = try typed(i16, a, t.bytes); + defer a.free(trellis); + const suh = try typed(f16, a, (try raw.get("{s}.suh", .{prefix})).bytes); + defer a.free(suh); + const svh = try typed(f16, a, (try raw.get("{s}.svh", .{prefix})).bytes); + defer a.free(svh); + return fwd.QMat.init(a, @intCast(t.dim(2) / 8), trellis, t.dim(0), t.dim(1), suh, svh, threads); +} + +pub fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +pub fn bfValue(bits: u16) f64 { + return @as(f32, @bitCast(@as(u32, bits) << 16)); +} + +pub const Err = struct { max_abs: f64 = 0, peak: f64 = 0, rel: f64 = 0 }; + +pub fn errOf(got: anytype, want: []const f64) Err { + var e: Err = .{}; + for (want, 0..) |w, i| { + const g: f64 = switch (@TypeOf(got)) { + []const u16, []u16 => bfValue(got[i]), + else => @floatCast(got[i]), + }; + e.max_abs = @max(e.max_abs, @abs(g - w)); + e.peak = @max(e.peak, @abs(w)); + } + e.rel = if (e.peak > 0) e.max_abs / e.peak else 0; + return e; +} + +pub fn download(comptime T: type, gpu: Gpu, s: cuda.Stream, ptr: u64, n: usize) ![]T { + try s.synchronize(); + const out = try gpu.gpa.alloc(T, n); + const b: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = ptr, .len = n * @sizeOf(T) }; + try b.download(0, std.mem.sliceAsBytes(out)); + return out; +} + +/// A pageable copy into the pass's scratch: the legacy-stream copy is ordered against nothing, so drain both sides. +pub fn uploadTo(gpu: Gpu, s: cuda.Stream, ptr: u64, bytes: []const u8) !void { + try s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = ptr, .len = bytes.len }; + try b.upload(0, bytes); + try gpu.ctx.synchronize(); +} + +// -- the pass rig ----------------------------------------------------------------------------------------------------- + +pub const Pass = batch.Pass(flash.Tri); + +pub const Rig = struct { + gpu: Gpu, + stream: cuda.Stream, + glue: flash.Glue, + k: batch.Kernels, + s: batch.Scratch, + store: wt.Store, + pack_mod: cuda.Module, + pack_fn: cuda.Function, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + conv: [wt.linear_n]cuda.DeviceBuffer = undefined, + state: [wt.linear_n]cuda.DeviceBuffer = undefined, + kc: [wt.attn_n]cuda.DeviceBuffer = undefined, + vc: [wt.attn_n]cuda.DeviceBuffer = undefined, + ikc: [wt.attn_n]cuda.DeviceBuffer = undefined, + pooled: [wt.attn_n]cuda.DeviceBuffer = undefined, + tail: cuda.DeviceBuffer = undefined, + history: [2]i64 = .{ 0, 0 }, + scores: cuda.DeviceBuffer = undefined, + cap: usize, + + pub fn init(r: *Rig, gpu: Gpu, dir: []const u8, rows: usize, cap: usize) !void { + r.* = .{ .gpu = gpu, .stream = try cuda.Stream.init(gpu.d, true), .glue = undefined, .k = undefined, .s = undefined, .store = undefined, .pack_mod = undefined, .pack_fn = undefined, .cap = cap }; + r.glue = try flash.Glue.load(gpu.d); + r.k = try batch.Kernels.load(gpu.d); + r.s = try batch.Scratch.alloc(gpu.gpa, gpu.d, r.stream, rows); + r.store = try wt.Store.open(gpu.gpa, gpu.io, dir); + try r.store.detectNorms(); + r.pack_mod = try cuda.Module.load(gpu.d, cuda.kernels.experts_pack); + r.pack_fn = try r.pack_mod.function(wt.pack_symbol); + const gdn = flash.gdn; + const z: cuda.DeviceBuffer = .{ .d = gpu.d, .ptr = 0, .len = 0 }; + r.conv = @splat(z); + r.state = @splat(z); + r.kc = @splat(z); + r.vc = @splat(z); + r.ikc = @splat(z); + r.pooled = @splat(z); + // every layer's state, as Engine.newSeq makes them (a check may step any layer) + for (&r.conv, &r.state) |*c, *x| { + c.* = try r.zeroed((gdn.taps - 1) * gdn.conv_dim * 2); + x.* = try r.zeroed(wt.state_n * 4); + } + for (&r.kc, &r.vc, &r.ikc, &r.pooled) |*kc, *vc, *ikc, *pooled| { + kc.* = try r.zeroed(cap * wt.kv_heads * wt.head_dim * 2); + vc.* = try r.zeroed(cap * wt.kv_heads * wt.head_dim * 2); + ikc.* = try r.zeroed(cap * wt.index_dim * 2); + pooled.* = try r.zeroed(flash.prompt.pooledBytes(cap)); + } + r.tail = try r.zeroed(wt.ple_tail_n * wide * 2); + r.scores = try r.zeroed(batch.scoreBytes(cap, rows)); + } + + fn zeroed(r: *Rig, bytes: usize) !cuda.DeviceBuffer { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, bytes); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0, r.stream.handle); + try r.stream.synchronize(); + return b; + } + + pub fn reset(r: *Rig) !void { + for (r.bufs.items) |b| try b.fill8(0, r.stream.handle); + try r.stream.synchronize(); + } + + pub fn deinit(r: *Rig) void { + r.stream.synchronize() catch {}; + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(r.gpu.gpa); + r.pack_mod.unload(); + r.store.deinit(); + r.s.free(r.gpu.gpa); + r.k.deinit(); + r.glue.deinit(); + r.stream.deinit(); + } + + pub fn pass(r: *Rig, layers: []const pl.Layer) Pass { + return .{ + .gpa = r.gpu.gpa, + .io = r.gpu.io, + .stream = r.stream, + .tri = .{ .own = &r.glue, .s = r.stream }, + .k = &r.k, + .s = &r.s, + .m = .{ .table_w = 0, .table_s = 0, .table_b = 0, .layers = layers, .mixer = undefined, .head = undefined, .ngram = undefined, .collective = null }, + .seq = .{ .conv = &r.conv, .state = &r.state, .kc = &r.kc, .vc = &r.vc, .ikc = &r.ikc, .pooled = &r.pooled, .tail = r.tail, .history = &r.history, .cap = r.cap, .scores = r.scores.ptr }, + .hooks = .{}, + }; + } +}; + +pub fn bf16Rows(a: std.mem.Allocator, rng: std.Random, n: usize, sigma: f32) ![]u16 { + const v = try a.alloc(u16, n); + for (v) |*x| x.* = bf16Bits(rng.floatNorm(f32) * sigma); + return v; +} diff --git a/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_tests.zig b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_tests.zig new file mode 100644 index 0000000000..1713d06354 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/model/flashnext_exl3_tests.zig @@ -0,0 +1,50 @@ +//! flashnext-exl3: the Flash Next CUDA family on a real EXL3 pack against our float64 oracle; this hub parses the mode. +const std = @import("std"); +const check = @import("../../check.zig"); +const shared = @import("flashnext_exl3_rig.zig"); +const names_tests = @import("flashnext_exl3_names_tests.zig"); +const proj_tests = @import("flashnext_exl3_proj_tests.zig"); +const layer_tests = @import("flashnext_exl3_layer_tests.zig"); + +const Gpu = check.Gpu; + +pub const default_ids = shared.default_ids; +pub const RawPack = shared.RawPack; +pub const bfValue = shared.bfValue; +pub const Err = shared.Err; +pub const errOf = shared.errOf; +pub const download = shared.download; +pub const uploadTo = shared.uploadTo; +pub const Pass = shared.Pass; +pub const Rig = shared.Rig; +pub const Taps = layer_tests.Taps; +pub const deviceLayer = layer_tests.deviceLayer; +pub const layerCheck = layer_tests.layerCheck; + +fn parseList(gpa: std.mem.Allocator, text: []const u8) ![]usize { + var list: std.ArrayList(usize) = .empty; + var it = std.mem.splitScalar(u8, text, ','); + while (it.next()) |tok| if (tok.len > 0) try list.append(gpa, try std.fmt.parseInt(usize, tok, 10)); + return list.toOwnedSlice(gpa); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const mode = rest[1]; + if (std.mem.eql(u8, mode, "names")) return names_tests.names(gpu, dir); + if (std.mem.eql(u8, mode, "norms")) return names_tests.norms(gpu, dir); + if (std.mem.eql(u8, mode, "ple")) return proj_tests.pleMode(gpu, dir); + if (std.mem.eql(u8, mode, "router")) { + const which = try parseList(gpu.gpa, if (rest.len > 2) rest[2] else "0,3,17"); + defer gpu.gpa.free(which); + const rows = if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4096; + return names_tests.router(gpu, dir, which, rows); + } + if (std.mem.eql(u8, mode, "layer")) { + const which = try parseList(gpu.gpa, if (rest.len > 2) rest[2] else "0,3"); + defer gpu.gpa.free(which); + return layer_tests.layerMode(gpu, dir, which); + } + return error.UnknownMode; +} diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_bench_tests.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_bench_tests.zig new file mode 100644 index 0000000000..a23f3fcb55 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_bench_tests.zig @@ -0,0 +1,255 @@ +//! exl3-pool: gate 4, miss cost, copy/compute overlap and the decode pipeline. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const st = core.safetensors; +const exl3e = cuda.exl3_experts; +const fpool = flash.exl3_pool; +const policy = fpool.offload; +const Gpu = check.Gpu; + +const data = @import("exl3_pool_data.zig"); +const dims = data.dims; +const width = data.width; +const routed_n = data.routed_n; +const table_n = data.table_n; +const top_k = data.top_k; +const slots = data.slots; +const max_rows = data.max_rows; +const shard_count = data.shard_count; +const projs = data.projs; +const Pack = data.Pack; +const LayerData = data.LayerData; +const Zipf = data.Zipf; +const bf16Bits = data.bf16Bits; +const Io = data.Io; +const call = data.call; +const poolRound = data.poolRound; +const genRound = data.genRound; + +// ---------------------------------------------------------------------------------------------------- gate 4 + +fn ms(t0: i96, t1: i96) f64 { + return @as(f64, @floatFromInt(t1 - t0)) / 1e6; +} + +fn median(xs: []f64) f64 { + std.mem.sort(f64, xs, {}, std.sort.asc(f64)); + return xs[xs.len / 2]; +} + +/// A row whose first `k` routed picks are not held by the pool and the rest are. +fn rowWithMisses(pool: *const fpool.Pool, rng: std.Random, k: usize, out: []i32) void { + var n: usize = 0; + // held experts first (up to 10 - k) + for (pool.slots_tab, pool.line_tab) |id, state| { + if (n == top_k - k) break; + if (state == .resident) { + out[k + n] = id; + n += 1; + } + } + var m: usize = 0; + while (m < k) { + const id: u16 = rng.intRangeLessThan(u16, 0, routed_n); + if (policy.lineOf(pool.v, pool.opt.layer, id) != null) continue; + if (std.mem.indexOfScalar(i32, out[0..m], id) != null) continue; + out[m] = id; + m += 1; + } + out[top_k] = 512; +} + +fn missCost(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, src: fpool.Source, direct: bool, label: []const u8) !void { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, src, ld.scales(), .{ .layer = @intCast(ld.layer), .lines = 64, .direct = direct, .double_buffer = false }, null); + defer pool.deinit(); + var vp = try Io.init(gpu, pool.tables()); + defer vp.free(); + var prng = std.Random.DefaultPrng.init(0xbe7c4); + const rng = prng.random(); + var x: [dims]u16 = undefined; + for (&x) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + const wts: [slots]f32 = @splat(0.1); + var picks: [slots]i32 = undefined; + // warm: fill the pool with 64 experts + for (0..7) |_| { + rowWithMisses(&pool, rng, top_k, &picks); + try vp.set(&x, &picks, &wts); + _ = try poolRound(&pool, ex, s, vp, &picks, 1, false); + try s.synchronize(); + } + var c0 = try cuda.Event.init(gpu.d, true); + defer c0.deinit(); + var c1 = try cuda.Event.init(gpu.d, true); + defer c1.deinit(); + var k0 = try cuda.Event.init(gpu.d, true); + defer k0.deinit(); + var k1 = try cuda.Event.init(gpu.d, true); + defer k1.deinit(); + std.debug.print(" [{s}] one decode row (10 routed + shared), 64 lines, rounds of k misses (median of 9; idle card):\n", .{label}); + var base: f64 = 0; + for ([_]usize{ 0, 1, 2, 4, 10 }) |k| { + var wall: [9]f64 = undefined; + var copy: [9]f64 = undefined; + var comp: [9]f64 = undefined; + var stage_host: [9]f64 = undefined; + for (0..9) |it| { + rowWithMisses(&pool, rng, k, &picks); + try vp.set(&x, &picks, &wts); + try s.synchronize(); + const t0 = check.now(gpu.io); + const r = try pool.plan(&picks, 1); + const t1 = check.now(gpu.io); + try c0.record(pool.lines.copy); + try pool.stage(r, 0, s); + try c1.record(pool.lines.copy); + const t2 = check.now(gpu.io); + try k0.record(s); + try call(ex, s, pool.tables(), vp, pool.picksFor(r, 0, vp.picks.ptr), 1, false); + try k1.record(s); + try pool.release(r, 0, s); + try s.synchronize(); + try pool.lines.copy.synchronize(); + const t3 = check.now(gpu.io); + wall[it] = ms(t0, t3); + copy[it] = try c0.elapsedMs(c1); + comp[it] = try k0.elapsedMs(k1); + stage_host[it] = ms(t1, t2); + } + const w = median(&wall); + if (k == 0) base = w; + std.debug.print(" k={d:>2}: round wall {d:.3} ms (+{d:.3} vs all-hit), copy stream {d:.3} ms ({d:.3} ms/expert), experts compute {d:.3} ms (measured after the stream's wait on the copies), host stage {d:.3} ms\n", .{ + k, w, w - base, median(©), if (k > 0) median(©) / @as(f64, @floatFromInt(k)) else 0, median(&comp), median(&stage_host), + }); + } + const ls = pool.lines.stats; + if (!direct) std.debug.print(" host staging memcpy (mmap -> pinned): {d:.3} ms/expert average over {d} copies ({d:.2} GB/s)\n", .{ @as(f64, @floatFromInt(ls.stage_ns)) / 1e6 / @as(f64, @floatFromInt(ls.copies)), ls.copies, @as(f64, @floatFromInt(ls.staged_bytes)) / @as(f64, @floatFromInt(ls.stage_ns)) }); +} + +/// A 16-row chunk through a 24-line pool (several passes): copies of pass j+1 overlap compute of pass j or not. +fn overlapChunk(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, src: fpool.Source, serial: bool) !struct { wall: f64, copy: f64, comp: f64, passes: usize, copies: usize } { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, src, ld.scales(), .{ .layer = @intCast(ld.layer), .lines = 24, .direct = true }, null); + defer pool.deinit(); + var vp = try Io.init(gpu, pool.tables()); + defer vp.free(); + const z = Zipf.init(0x0e1a9 + ld.layer, 0.8); + var prng = std.Random.DefaultPrng.init(0x0e1a9); + const rng = prng.random(); + var x: [max_rows * dims]u16 = undefined; + var picks: [max_rows * slots]i32 = undefined; + var wts: [max_rows * slots]f32 = undefined; + genRound(&z, rng, max_rows, &x, &picks, &wts); + try vp.set(&x, &picks, &wts); + var cs: [64]cuda.Event = undefined; + for (&cs) |*e| e.* = try cuda.Event.init(gpu.d, true); + defer for (&cs) |*e| e.deinit(); + var ks: [64]cuda.Event = undefined; + for (&ks) |*e| e.* = try cuda.Event.init(gpu.d, true); + defer for (&ks) |*e| e.deinit(); + try s.synchronize(); + const t0 = check.now(gpu.io); + const r = try pool.plan(&picks, max_rows); + if (r.n_passes > 32) return error.TooManyPasses; + const l = pool.tables(); + for (0..r.n_passes) |j| { + try cs[2 * j].record(pool.lines.copy); + try pool.stage(r, j, s); + try cs[2 * j + 1].record(pool.lines.copy); + if (serial) try pool.lines.copy.synchronize(); + try ks[2 * j].record(s); + try call(ex, s, l, vp, pool.picksFor(r, j, vp.picks.ptr), max_rows, false); + try ks[2 * j + 1].record(s); + try pool.release(r, j, s); + if (serial) try s.synchronize(); + } + try s.synchronize(); + try pool.lines.copy.synchronize(); + const wall = ms(t0, check.now(gpu.io)); + var copy: f64 = 0; + var comp: f64 = 0; + for (0..r.n_passes) |j| { + copy += try cs[2 * j].elapsedMs(cs[2 * j + 1]); + comp += try ks[2 * j].elapsedMs(ks[2 * j + 1]); + } + return .{ .wall = wall, .copy = copy, .comp = comp, .passes = r.n_passes, .copies = r.n_copies }; +} + +/// 60 decode rounds (one row each, skewed picks known ahead), sequential (sync every round) vs pipelined (the host +fn decodePipe(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, src: fpool.Source, lines: usize, pipelined: bool, double: bool) !struct { wall: f64, hit: f64, release_waits: u64, protected: u64 } { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, src, ld.scales(), .{ .layer = @intCast(ld.layer), .lines = lines, .direct = true, .double_buffer = double }, null); + defer pool.deinit(); + const rounds = 60; + var ios: [rounds]Io = undefined; + var made: usize = 0; + defer for (ios[0..made]) |*v| v.free(); + const z = Zipf.init(0xdec0 + ld.layer, 1.0); + var prng = std.Random.DefaultPrng.init(0xdec0); + const rng = prng.random(); + var picks: [rounds][slots]i32 = undefined; + var x: [dims]u16 = undefined; + var wts: [slots]f32 = undefined; + for (0..rounds) |i| { + ios[i] = try Io.init(gpu, pool.tables()); + made += 1; + genRound(&z, rng, 1, &x, &picks[i], &wts); + try ios[i].set(&x, &picks[i], &wts); + } + try s.synchronize(); + const t0 = check.now(gpu.io); + for (0..rounds) |i| { + _ = try poolRound(&pool, ex, s, ios[i], &picks[i], 1, false); + if (!pipelined) try s.synchronize(); + } + try s.synchronize(); + const wall = ms(t0, check.now(gpu.io)); + const ps = pool.stats; + return .{ .wall = wall, .hit = @as(f64, @floatFromInt(ps.hit_slots)) / @as(f64, @floatFromInt(ps.hit_slots + ps.miss_slots)), .release_waits = pool.lines.stats.release_waits, .protected = ps.protected }; +} + +pub fn benchRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const gpa = gpu.gpa; + // (a) miss cost, staged from the mapped checkpoint + try missCost(gpu, ex, s, ld, ld.source(), false, "staged: mmap -> pinned staging -> H2D"); + // the same experts in one pinned bank (what keeping the routed trellis in pinned host RAM gives): direct H2D + var total: usize = 0; + for (ld.routed) |e| for (e) |b| { + total += b.len; + }; + var bank = try cuda.HostBuffer.alloc(gpu.d, total); + defer bank.free(); + const pinned = try gpa.alloc([3][]const u8, routed_n); + defer gpa.free(pinned); + var at: usize = 0; + for (ld.routed, pinned) |e, *o| for (e, 0..) |b, p| { + @memcpy(bank.bytes[at..][0..b.len], b); + o[p] = bank.bytes[at..][0..b.len]; + at += b.len; + }; + std.debug.print(" pinned bank for layer {d}: {d:.0} MiB of routed trellis copied into page-locked host memory\n", .{ ld.layer, @as(f64, @floatFromInt(total)) / 1048576.0 }); + const psrc: fpool.Source = .{ .routed = pinned, .shared = ld.shared }; + try missCost(gpu, ex, s, ld, psrc, true, "pinned bank: direct H2D"); + + // (b) a 16-row chunk split into passes: serial vs overlapped copies + const ser = try overlapChunk(gpu, ex, s, ld, psrc, true); + const ovl = try overlapChunk(gpu, ex, s, ld, psrc, false); + const hidden_ms = (ser.wall - ovl.wall); + std.debug.print(" [16-row chunk, 24 lines, pinned bank] {d} passes, {d} copies: serial {d:.3} ms (copy {d:.3} + compute {d:.3}), overlapped {d:.3} ms; overlap hides {d:.3} ms = {d:.0}% of the smaller of copy/compute\n", .{ + ser.passes, ser.copies, ser.wall, ser.copy, ser.comp, ovl.wall, hidden_ms, 100.0 * hidden_ms / @min(ser.copy, ser.comp), + }); + // (c) decode rounds: pipelined staging, double buffering on/off + for ([_]usize{ 24, 48 }) |lines| { + const seq = try decodePipe(gpu, ex, s, ld, psrc, lines, false, true); + const pipe1 = try decodePipe(gpu, ex, s, ld, psrc, lines, true, false); + const pipe2 = try decodePipe(gpu, ex, s, ld, psrc, lines, true, true); + std.debug.print(" [60 decode rounds, {d} lines, pinned bank, hit {d:.3}] sequential {d:.2} ms; pipelined {d:.2} ms single-buffered ({d} copies waited for the previous round's compute), {d:.2} ms double-buffered ({d} waited, {d} rounds spared the previous round's lines)\n", .{ + lines, seq.hit, seq.wall, pipe1.wall, pipe1.release_waits, pipe2.wall, pipe2.release_waits, pipe2.protected, + }); + } + check.pass("layer {d}: miss cost, chunk overlap and decode pipelining measured (see lines above)", .{ld.layer}); +} diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_data.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_data.zig new file mode 100644 index 0000000000..783bbc8a76 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_data.zig @@ -0,0 +1,284 @@ +//! exl3-pool: the pack reader, one layer's experts and the synthetic routing the pool gates share. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const st = core.safetensors; +const exl3e = cuda.exl3_experts; +const fpool = flash.exl3_pool; +const policy = fpool.offload; +const Gpu = check.Gpu; + +pub const dims = 2560; +pub const width = 640; +pub const routed_n = 512; +pub const table_n = 513; +pub const top_k = 10; +pub const slots = 11; +pub const max_rows = 16; +pub const shard_count = 7; +pub const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + +pub const Pack = struct { + files: [shard_count]st.File = undefined, + open_n: usize = 0, + + pub fn open(p: *Pack, gpu: Gpu, dir: []const u8) !void { + var buf: [512]u8 = undefined; + for (0..shard_count) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-{d:0>5}.safetensors", .{ dir, i + 1, shard_count }); + p.files[i] = try st.File.open(gpu.gpa, gpu.io, path); + p.open_n += 1; + } + } + + pub fn close(p: *Pack, io: std.Io) void { + for (p.files[0..p.open_n]) |*f| f.close(io); + } + + pub fn get(p: *const Pack, comptime fmt: []const u8, args: anytype) !st.Tensor { + var buf: [200]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + for (p.files[0..p.open_n]) |*f| if (f.get(name)) |t| return t; + std.debug.print("missing tensor {s}\n", .{name}); + return error.MissingTensor; + } +}; + +// ---------------------------------------------------------------------------------------------------- one layer + +/// One layer's experts: host trellis slices (mmap), the stacked scales and the fully resident reference on the device. +pub const LayerData = struct { + layer: usize, + routed: [routed_n][3][]const u8, + shared: [3][]const u8, + scale_bufs: [6]cuda.DeviceBuffer, + res_mem: cuda.DeviceBuffer, + res_tab: cuda.DeviceBuffer, + res: exl3e.Layer, + + pub fn load(gpu: Gpu, pack: *const Pack, layer: usize) !*LayerData { + const gpa = gpu.gpa; + const ld = try gpa.create(LayerData); + errdefer gpa.destroy(ld); + ld.layer = layer; + // stacked scales [E, ยท]: gate suh [D] svh [I], up suh [D] svh [I], down suh [I] svh [D] + const sizes = [6]usize{ dims, dims, width, width, width, dims }; // suh_g suh_u svh_g svh_u suh_d svh_d + var host: [6][]u8 = undefined; + for (&host, sizes) |*h, n| h.* = try gpa.alloc(u8, table_n * n * 2); + defer for (host) |h| gpa.free(h); + var total: usize = 0; + for (0..table_n) |e| { + for (projs, 0..) |pn, p| { + const tr = if (e < routed_n) + try pack.get("model.language_model.layers.{d}.mlp.experts.{d}.{s}.trellis", .{ layer, e, pn }) + else + try pack.get("model.language_model.layers.{d}.mlp.shared_expert.{s}.trellis", .{ layer, pn }); + const su = if (e < routed_n) + try pack.get("model.language_model.layers.{d}.mlp.experts.{d}.{s}.suh", .{ layer, e, pn }) + else + try pack.get("model.language_model.layers.{d}.mlp.shared_expert.{s}.suh", .{ layer, pn }); + const sv = if (e < routed_n) + try pack.get("model.language_model.layers.{d}.mlp.experts.{d}.{s}.svh", .{ layer, e, pn }) + else + try pack.get("model.language_model.layers.{d}.mlp.shared_expert.{s}.svh", .{ layer, pn }); + const k: usize = if (p == 2) width else dims; + const n: usize = if (p == 2) dims else width; + if (tr.rank != 3 or tr.dim(0) * 16 != k or tr.dim(1) * 16 != n) return error.BadShape; + if (su.bytes.len != k * 2 or sv.bytes.len != n * 2) return error.BadShape; + if (e < routed_n) ld.routed[e][p] = tr.bytes else ld.shared[p] = tr.bytes; + total += (tr.bytes.len + 255) / 256 * 256; + const su_i: usize = switch (p) { + 0 => 0, + 1 => 1, + else => 4, + }; + const sv_i: usize = switch (p) { + 0 => 2, + 1 => 3, + else => 5, + }; + @memcpy(host[su_i][e * k * 2 ..][0 .. k * 2], su.bytes); + @memcpy(host[sv_i][e * n * 2 ..][0 .. n * 2], sv.bytes); + } + } + var made: usize = 0; + errdefer for (ld.scale_bufs[0..made]) |*b| b.free(); + for (&ld.scale_bufs, host) |*b, h| { + b.* = try cuda.DeviceBuffer.fromHost(gpu.d, h); + made += 1; + } + // the fully resident reference: every expert's trellis back to back, its own tables + ld.res_mem = try cuda.DeviceBuffer.alloc(gpu.d, total); + errdefer ld.res_mem.free(); + var ptrs: [3][table_n]u64 = undefined; + var k2s: [3][table_n]i32 = undefined; + var at: usize = 0; + const tiles = (dims / 16) * (width / 16); + for (0..table_n) |e| for (0..3) |p| { + const bytes = if (e < routed_n) ld.routed[e][p] else ld.shared[p]; + try ld.res_mem.upload(at, bytes); + ptrs[p][e] = ld.res_mem.ptr + at; + k2s[p][e] = @intCast(bytes.len / (tiles * 16)); + at += (bytes.len + 255) / 256 * 256; + }; + const tl = cuda.exl3_pool.Tables.init(table_n); + const tab = try gpa.alloc(u8, tl.bytes); + defer gpa.free(tab); + @memset(tab, 0); + for (0..3) |p| { + @memcpy(tab[tl.ptr[p]..][0 .. table_n * 8], std.mem.sliceAsBytes(&ptrs[p])); + @memcpy(tab[tl.k2[p]..][0 .. table_n * 4], std.mem.sliceAsBytes(&k2s[p])); + } + ld.res_tab = try cuda.DeviceBuffer.fromHost(gpu.d, tab); + ld.res = .{ + .gate = ld.res_tab.ptr + tl.ptr[0], + .up = ld.res_tab.ptr + tl.ptr[1], + .down = ld.res_tab.ptr + tl.ptr[2], + .gate_k2 = ld.res_tab.ptr + tl.k2[0], + .up_k2 = ld.res_tab.ptr + tl.k2[1], + .down_k2 = ld.res_tab.ptr + tl.k2[2], + .suh_g = ld.scale_bufs[0].ptr, + .suh_u = ld.scale_bufs[1].ptr, + .svh_g = ld.scale_bufs[2].ptr, + .svh_u = ld.scale_bufs[3].ptr, + .suh_d = ld.scale_bufs[4].ptr, + .svh_d = ld.scale_bufs[5].ptr, + .experts = table_n, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = 1, + }; + std.debug.print(" layer {d}: 513 experts mapped, fully resident reference {d:.0} MiB; routed k2 gate/up/down {d}/{d}/{d}, shared {d}/{d}/{d}\n", .{ + layer, @as(f64, @floatFromInt(total)) / 1048576.0, k2s[0][0], k2s[1][0], k2s[2][0], k2s[0][512], k2s[1][512], k2s[2][512], + }); + return ld; + } + + pub fn free(ld: *LayerData, gpa: std.mem.Allocator) void { + ld.res_tab.free(); + ld.res_mem.free(); + for (&ld.scale_bufs) |*b| b.free(); + gpa.destroy(ld); + } + + pub fn scales(ld: *const LayerData) fpool.Scales { + return .{ .suh_g = ld.scale_bufs[0].ptr, .suh_u = ld.scale_bufs[1].ptr, .svh_g = ld.scale_bufs[2].ptr, .svh_u = ld.scale_bufs[3].ptr, .suh_d = ld.scale_bufs[4].ptr, .svh_d = ld.scale_bufs[5].ptr }; + } + + pub fn source(ld: *const LayerData) fpool.Source { + return .{ .routed = &ld.routed, .shared = ld.shared }; + } +}; + +// ---------------------------------------------------------------------------------------------------- inputs + +/// Skewed expert choice: Zipf(s) over a fixed random permutation of the 512 ids. +pub const Zipf = struct { + perm: [routed_n]u16, + cdf: [routed_n]f64, + + pub fn init(seed: u64, s: f64) Zipf { + var z: Zipf = undefined; + var prng = std.Random.DefaultPrng.init(seed); + for (&z.perm, 0..) |*p, i| p.* = @intCast(i); + prng.random().shuffle(u16, &z.perm); + var sum: f64 = 0; + for (&z.cdf, 0..) |*c, i| { + sum += 1.0 / std.math.pow(f64, @floatFromInt(i + 1), s); + c.* = sum; + } + for (&z.cdf) |*c| c.* /= sum; + return z; + } + + /// One row: 10 distinct routed ids then the shared expert. + pub fn row(z: *const Zipf, r: std.Random, out: []i32) void { + var n: usize = 0; + while (n < top_k) { + const u = r.float(f64); + var i: usize = 0; + while (i + 1 < routed_n and z.cdf[i] < u) i += 1; + const id: i32 = z.perm[i]; + if (std.mem.indexOfScalar(i32, out[0..n], id) != null) continue; + out[n] = id; + n += 1; + } + out[top_k] = 512; + } +}; + +pub fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +/// The device buffers of one engine side (pool or resident): x, picks, weights, scratch, out. +pub const Io = struct { + x: cuda.DeviceBuffer, + picks: cuda.DeviceBuffer, + wts: cuda.DeviceBuffer, + out: cuda.DeviceBuffer, + sbuf: cuda.DeviceBuffer, + sc: exl3e.Scratch, + + pub fn init(gpu: Gpu, l: exl3e.Layer) !Io { + var v: Io = undefined; + v.x = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * dims * 2); + v.picks = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * slots * 4); + v.wts = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * slots * 4); + v.out = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * dims * 4); + v.sbuf = try cuda.DeviceBuffer.alloc(gpu.d, try exl3e.scratchBytes(l, max_rows, slots)); + v.sc = try exl3e.scratch(v.sbuf.ptr, l, max_rows, slots); + return v; + } + + pub fn free(v: *Io) void { + v.x.free(); + v.picks.free(); + v.wts.free(); + v.out.free(); + v.sbuf.free(); + } + + pub fn set(v: *Io, x: []const u16, picks: []const i32, wts: []const f32) !void { + try v.x.upload(0, std.mem.sliceAsBytes(x)); + try v.picks.upload(0, std.mem.sliceAsBytes(picks)); + try v.wts.upload(0, std.mem.sliceAsBytes(wts)); + } + + pub fn y(v: Io, gpa: std.mem.Allocator, rows: usize) ![]u8 { + const out = try gpa.alloc(u8, rows * slots * dims * 4); + try v.sbuf.d.check(v.sbuf.d.api.cuMemcpyDtoH_v2(out.ptr, v.sc.y, out.len), "cuMemcpyDtoH"); + return out; + } +}; + +pub fn call(ex: *const exl3e.Experts, s: cuda.Stream, l: exl3e.Layer, v: Io, picks: u64, rows: usize, combine: bool) !void { + if (combine) { + try ex.run(s, l, v.x.ptr, dims, .bf16, picks, v.wts.ptr, rows, v.out.ptr, v.sc); + } else { + try ex.runSlots(s, l, v.x.ptr, dims, .bf16, picks, rows, v.sc); + } +} + +/// One round through the pool: plan, then each pass staged, computed and released. +pub fn poolRound(pool: *fpool.Pool, ex: *const exl3e.Experts, s: cuda.Stream, v: Io, picks: []const i32, rows: usize, combine: bool) !*const fpool.Round { + const r = try pool.plan(picks, rows); + const l = pool.tables(); + for (0..r.n_passes) |j| { + try pool.stage(r, j, s); + try call(ex, s, l, v, pool.picksFor(r, j, v.picks.ptr), rows, combine); + try pool.release(r, j, s); + } + return r; +} + +pub fn genRound(z: *const Zipf, rng: std.Random, rows: usize, x: []u16, picks: []i32, wts: []f32) void { + for (x[0 .. rows * dims]) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + for (0..rows) |i| z.row(rng, picks[i * slots ..][0..slots]); + for (wts[0 .. rows * slots]) |*w| w.* = rng.float(f32) * 0.2; +} diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_equal_tests.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_equal_tests.zig new file mode 100644 index 0000000000..2dcd93790d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_equal_tests.zig @@ -0,0 +1,151 @@ +//! exl3-pool: gate 2, the pool against the resident experts. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const st = core.safetensors; +const exl3e = cuda.exl3_experts; +const fpool = flash.exl3_pool; +const policy = fpool.offload; +const Gpu = check.Gpu; + +const data = @import("exl3_pool_data.zig"); +const dims = data.dims; +const width = data.width; +const routed_n = data.routed_n; +const table_n = data.table_n; +const top_k = data.top_k; +const slots = data.slots; +const max_rows = data.max_rows; +const shard_count = data.shard_count; +const projs = data.projs; +const Pack = data.Pack; +const LayerData = data.LayerData; +const Zipf = data.Zipf; +const bf16Bits = data.bf16Bits; +const Io = data.Io; +const call = data.call; +const poolRound = data.poolRound; +const genRound = data.genRound; + +// ---------------------------------------------------------------------------------------------------- gate 2 + +fn verifyPool(gpu: Gpu, pool: *fpool.Pool, ld: *const LayerData) !usize { + const gpa = gpu.gpa; + try pool.lines.drain(); + const lb = pool.lines.layout.line_bytes; + const got = try gpa.alloc(u8, lb); + defer gpa.free(got); + // tables + const tl = pool.lines.tlay; + const tab = try gpa.alloc(u8, tl.bytes); + defer gpa.free(tab); + try pool.lines.tables.download(0, tab); + var checked: usize = 0; + var resident: [routed_n]?usize = @splat(null); + for (pool.slots_tab, pool.line_tab, 0..) |id, state, i| { + if (state == .empty) continue; + resident[id] = i; + try pool.lines.mem.download(i * lb, got); + for (0..3) |p| { + const want = ld.routed[id][p]; + try check.sameBytes("pool line vs mapped checkpoint", got[pool.lines.layout.off[p]..][0..want.len], want); + } + checked += 1; + } + for (0..table_n) |e| for (0..3) |p| { + const ptr = std.mem.readInt(u64, tab[tl.ptr[p] + e * 8 ..][0..8], .little); + const k2 = std.mem.readInt(i32, tab[tl.k2[p] + e * 4 ..][0..4], .little); + if (e == 512) { + try check.expect(ptr != 0 and k2 > 0, "shared expert entry missing", .{}); + } else if (resident[e]) |line| { + try check.expect(ptr == pool.lines.at(line, p) and k2 == pool.k2[e][p], "expert {d} proj {d}: table points at {x}, line {d} is at {x}", .{ e, p, ptr, line, pool.lines.at(line, p) }); + } else { + try check.expect(ptr == 0 and k2 == 0, "absent expert {d} has a stale table entry", .{e}); + } + }; + return checked; +} + +const Script = struct { rows: usize, combine: bool }; + +pub fn equalRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData, lines: usize, prefetch_lines: usize, seed: u64) !void { + const gpa = gpu.gpa; + var pool = try fpool.Pool.init(gpa, gpu.io, gpu.d, ld.source(), ld.scales(), .{ .layer = @intCast(ld.layer), .lines = lines, .prefetch_bytes = 0 }, null); + defer pool.deinit(); + pool.opt.prefetch_bytes = prefetch_lines * pool.lineBytes(); + var vp = try Io.init(gpu, pool.tables()); + defer vp.free(); + var vr = try Io.init(gpu, ld.res); + defer vr.free(); + const z = Zipf.init(0x5eed + ld.layer, 1.1); + var prng = std.Random.DefaultPrng.init(seed); + const rng = prng.random(); + // decode rows, two 16-row prompt chunks, a 4-row chunk; runSlots and run (with combine) alternate + var script: [40]Script = undefined; + for (&script, 0..) |*sc, i| sc.* = .{ .rows = 1, .combine = i % 2 == 1 }; + script[6] = .{ .rows = max_rows, .combine = false }; + script[17] = .{ .rows = max_rows, .combine = true }; + script[25] = .{ .rows = 4, .combine = false }; + var x: [max_rows * dims]u16 = undefined; + var picks: [max_rows * slots]i32 = undefined; + var wts: [max_rows * slots]f32 = undefined; + var ledger_staged: u64 = 0; + var ledger_h2d: u64 = 0; + var ledger_copies: u64 = 0; + var split_rounds: usize = 0; + var max_passes: usize = 0; + var slot_bytes: usize = 0; + for (script, 0..) |sc, ri| { + genRound(&z, rng, sc.rows, &x, &picks, &wts); + const n = sc.rows * slots; + try vp.set(x[0 .. sc.rows * dims], picks[0..n], wts[0..n]); + try vr.set(x[0 .. sc.rows * dims], picks[0..n], wts[0..n]); + try vp.sbuf.fill8(0xff, s.handle); // a slot no pass wrote shows up + const r = try poolRound(&pool, ex, s, vp, picks[0..n], sc.rows, sc.combine); + for (r.copies[0..r.n_copies]) |c| { + const b = ld.routed[c.id]; + ledger_staged += b[0].len + b[1].len + b[2].len; + ledger_h2d += pool.lines.layout.off[2] + b[2].len; + ledger_copies += 1; + } + if (r.masked()) split_rounds += 1; + max_passes = @max(max_passes, r.n_passes); + try call(ex, s, ld.res, vr, vr.picks.ptr, sc.rows, sc.combine); + try s.synchronize(); + const yp = try vp.y(gpa, sc.rows); + defer gpa.free(yp); + const yr = try vr.y(gpa, sc.rows); + defer gpa.free(yr); + check.sameBytes("pool y vs resident y", yp, yr) catch |e| { + std.debug.print(" round {d} ({d} rows, {d} passes) differs\n", .{ ri, sc.rows, r.n_passes }); + return e; + }; + slot_bytes += yp.len; + if (sc.combine) { + const op = try gpa.alloc(u8, sc.rows * dims * 4); + defer gpa.free(op); + const or_ = try gpa.alloc(u8, sc.rows * dims * 4); + defer gpa.free(or_); + try vp.out.download(0, op); + try vr.out.download(0, or_); + try check.sameBytes("pool combined out vs resident", op, or_); + } + } + const lines_checked = try verifyPool(gpu, &pool, ld); + const ps = pool.stats; + const ls = pool.lines.stats; + const hit = @as(f64, @floatFromInt(ps.hit_slots)) / @as(f64, @floatFromInt(ps.hit_slots + ps.miss_slots)); + std.debug.print(" [{d} lines, prefetch {d}/swap round] {d} rounds ({d} split into passes, up to {d}): routed slots hit {d} miss {d} ({d:.3} hit); {d} copies ({d} misses + {d} prefetch), {d} evictions, {d} protected rounds; H2D {d:.1} MiB, staged {d:.1} MiB (host memcpy {d:.2} ms/expert), staging waits {d}, release waits {d}, gate waits {d}; y {d} bytes compared\n", .{ + lines, prefetch_lines, ps.rounds, split_rounds, max_passes, ps.hit_slots, ps.miss_slots, hit, ls.copies, ps.misses, ps.prefetches, ps.evictions, ps.protected, + @as(f64, @floatFromInt(ls.bytes)) / 1048576.0, @as(f64, @floatFromInt(ls.staged_bytes)) / 1048576.0, @as(f64, @floatFromInt(ls.stage_ns)) / 1e6 / @as(f64, @floatFromInt(@max(ls.copies, 1))), ls.staging_waits, ls.release_waits, ls.gate_waits, slot_bytes, + }); + try check.expect(ls.copies == ledger_copies and ls.copies == ps.misses + ps.prefetches, "copy count: device {d}, plans {d}, misses+prefetch {d}", .{ ls.copies, ledger_copies, ps.misses + ps.prefetches }); + try check.expect(ls.staged_bytes == ledger_staged, "staged bytes {d} vs plan ledger {d}", .{ ls.staged_bytes, ledger_staged }); + try check.expect(ls.bytes == ledger_h2d, "H2D bytes {d} vs plan ledger {d}", .{ ls.bytes, ledger_h2d }); + try check.expect(ps.evictions > 0 and ps.miss_slots > 0 and ps.hit_slots > 0 and split_rounds > 0, "a path was not exercised (evictions {d}, misses {d}, hits {d}, split {d})", .{ ps.evictions, ps.miss_slots, ps.hit_slots, split_rounds }); + if (prefetch_lines > 0) try check.expect(ps.prefetches > 0, "no prefetch happened", .{}); + check.pass("layer {d}, {d} lines: every round's runSlots/run output is byte-identical to the fully resident layer; {d} lines and 513 table entries match the checkpoint; staged bytes match the plans", .{ ld.layer, lines, lines_checked }); +} diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_hits_tests.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_hits_tests.zig new file mode 100644 index 0000000000..fc36a6d404 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_hits_tests.zig @@ -0,0 +1,166 @@ +//! exl3-pool: hit rates of the cache policy replayed over recorded routing. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const flash = @import("flashnext"); +const check = @import("../../check.zig"); + +const st = core.safetensors; +const exl3e = cuda.exl3_experts; +const fpool = flash.exl3_pool; +const policy = fpool.offload; +const Gpu = check.Gpu; + +const data = @import("exl3_pool_data.zig"); +const dims = data.dims; +const width = data.width; +const routed_n = data.routed_n; +const table_n = data.table_n; +const top_k = data.top_k; +const slots = data.slots; +const max_rows = data.max_rows; +const shard_count = data.shard_count; +const projs = data.projs; +const Pack = data.Pack; +const LayerData = data.LayerData; +const Zipf = data.Zipf; +const bf16Bits = data.bf16Bits; +const Io = data.Io; +const call = data.call; +const poolRound = data.poolRound; +const genRound = data.genRound; + +// ---------------------------------------------------------------------------------------------------- hit rates + +const line_counts = [_]usize{ 10, 24, 48, 96, 160 }; + +const Sim = struct { + hit: u64 = 0, + miss: u64 = 0, + dec_hit: u64 = 0, + dec_miss: u64 = 0, + copies: u64 = 0, +}; + +/// Host-only replay of a pick sequence through one layer's policy. +fn simulate(gpa: std.mem.Allocator, layer: u8, picks: []const u16, tokens: usize, prompt: usize, lines: usize, prefetch_lines: usize, counts: *[policy.shape.layers][policy.shape.experts]u32) !Sim { + const slots_tab = try gpa.alloc(u16, lines); + defer gpa.free(slots_tab); + const line_tab = try gpa.alloc(policy.Line, lines); + defer gpa.free(line_tab); + const stamp = try gpa.alloc(u32, lines); + defer gpa.free(stamp); + const lay = try gpa.alloc(u8, lines); + defer gpa.free(lay); + var v = policy.Vram.init(slots_tab, line_tab, stamp, lay); + var r = try fpool.Round.init(gpa, lines, max_rows * slots); + defer r.deinit(gpa); + var buf: [max_rows * slots]i32 = undefined; + var sim: Sim = .{}; + var prev: [routed_n]u16 = undefined; + var prev_n: usize = 0; + var t: usize = 0; + var round: u32 = 0; + while (t < tokens) { + const rows: usize = if (t < prompt) @min(max_rows, prompt - t) else 1; + for (0..rows) |i| { + for (0..top_k) |k| buf[i * slots + k] = picks[(t + i) * top_k + k]; + buf[i * slots + top_k] = 512; + } + round += 1; + try fpool.round.plan(&v, layer, buf[0 .. rows * slots], counts, round, round, .{ .prefetch_bytes = prefetch_lines, .line_bytes = 1, .protect = prev[0..prev_n] }, &r); + @memcpy(prev[0..r.n_ids], r.ids[0..r.n_ids]); + prev_n = r.n_ids; + // prefetch lines land before the next round (the device has a whole layer stack of time) + for (r.prefetch()) |c| policy.reapAt(&v, layer, c.id, round) catch {}; + sim.hit += r.hit_slots; + sim.miss += r.miss_slots; + sim.copies += r.n_copies; + if (t >= prompt) { + sim.dec_hit += r.hit_slots; + sim.dec_miss += r.miss_slots; + } + t += rows; + } + return sim; +} + +fn rate(h: u64, m: u64) f64 { + return if (h + m == 0) 0 else @as(f64, @floatFromInt(h)) / @as(f64, @floatFromInt(h + m)); +} + +pub fn hitRun(gpu: Gpu, files: []const [:0]const u8) !void { + const gpa = gpu.gpa; + const counts = try gpa.create([policy.shape.layers][policy.shape.experts]u32); + defer gpa.destroy(counts); + // synthetic: 200 tokens of Zipf(1.1) picks, decode rows only + { + const z = Zipf.init(0x200, 1.1); + var prng = std.Random.DefaultPrng.init(0x200); + var row: [slots]i32 = undefined; + var picks: [200 * top_k]u16 = undefined; + for (0..200) |t| { + z.row(prng.random(), &row); + for (0..top_k) |k| picks[t * top_k + k] = @intCast(row[k]); + } + std.debug.print(" synthetic Zipf(1.1) over 512 experts, 200 decode tokens, one layer:\n", .{}); + for (line_counts) |lines| { + counts.* = @splat(@splat(0)); + const a = try simulate(gpa, 0, &picks, 200, 0, lines, 0, counts); + counts.* = @splat(@splat(0)); + const b = try simulate(gpa, 0, &picks, 200, 0, lines, 4, counts); + std.debug.print(" {d:>3} lines ({d:>5.1} MiB/layer at 1.76 MiB a line): hit {d:.3} (with 4-line prefetch every 4 rounds {d:.3})\n", .{ lines, @as(f64, @floatFromInt(lines)) * 1.7578, rate(a.hit, a.miss), rate(b.hit, b.miss) }); + } + } + for (files) |path| { + const blob = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, gpa, .limited(1 << 30)); + defer gpa.free(blob); + const hdr = std.mem.bytesAsSlice(u32, blob[0..12]); + const tokens: usize = hdr[0]; + const prompt: usize = hdr[1]; + const layers: usize = hdr[2]; + const raw = try gpa.alloc(u16, tokens * layers * top_k); + defer gpa.free(raw); + @memcpy(std.mem.sliceAsBytes(raw), blob[12..][0 .. raw.len * 2]); + std.debug.print(" real router picks {s}: {d} tokens ({d} prompt as 16-row chunks, {d} decode rows) x {d} layers\n", .{ path, tokens, prompt, tokens - prompt, layers }); + const per = try gpa.alloc(u16, tokens * top_k); + defer gpa.free(per); + // distinct experts a layer over the decode part (an upper bound on what any pool must hold) + var distinct_sum: usize = 0; + for (0..layers) |l| { + var seen: [routed_n]bool = @splat(false); + for (prompt..tokens) |t| for (0..top_k) |k| { + seen[raw[(t * layers + l) * top_k + k]] = true; + }; + for (seen) |b| distinct_sum += @intFromBool(b); + } + std.debug.print(" distinct experts a layer over the {d} decode tokens: {d:.1} on average\n", .{ tokens - prompt, @as(f64, @floatFromInt(distinct_sum)) / @as(f64, @floatFromInt(layers)) }); + for (line_counts) |lines| { + for ([_]usize{ 0, 4 }) |pf| { + counts.* = @splat(@splat(0)); + var tot: Sim = .{}; + var lo: f64 = 1; + var hi: f64 = 0; + var copies: u64 = 0; + for (0..layers) |l| { + for (0..tokens) |t| @memcpy(per[t * top_k ..][0..top_k], raw[(t * layers + l) * top_k ..][0..top_k]); + const s = try simulate(gpa, @intCast(l), per, tokens, prompt, lines, pf, counts); + tot.hit += s.hit; + tot.miss += s.miss; + tot.dec_hit += s.dec_hit; + tot.dec_miss += s.dec_miss; + copies += s.copies; + const dr = rate(s.dec_hit, s.dec_miss); + lo = @min(lo, dr); + hi = @max(hi, dr); + } + const dec_tokens = tokens - prompt; + const miss_per_tok = @as(f64, @floatFromInt(tot.dec_miss)) / @as(f64, @floatFromInt(dec_tokens)); + std.debug.print(" {d:>3} lines{s}: decode hit {d:.3} (layers {d:.3}..{d:.3}), {d:.1} routed misses a decode token over 48 layers; all rows hit {d:.3}; pool {d:.2} GiB for 48 layers\n", .{ + lines, if (pf > 0) " +prefetch" else " ", rate(tot.dec_hit, tot.dec_miss), lo, hi, miss_per_tok, rate(tot.hit, tot.miss), @as(f64, @floatFromInt(lines * 48)) * 1.7578 / 1024.0, + }); + } + } + } + check.pass("hit rates printed (host simulation of the pool's own planner)", .{}); +} diff --git a/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_tests.zig b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_tests.zig new file mode 100644 index 0000000000..53a2a3e594 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/offload/exl3_pool_tests.zig @@ -0,0 +1,57 @@ +//! exl3-pool: the Flash Next EXL3 expert offload pool (cuda_exl3_pool.zig) on a real pack; this hub parses the mode. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("../../check.zig"); +const data = @import("exl3_pool_data.zig"); +const equal = @import("exl3_pool_equal_tests.zig"); +const bench = @import("exl3_pool_bench_tests.zig"); +const hits = @import("exl3_pool_hits_tests.zig"); + +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; +const Pack = data.Pack; +const LayerData = data.LayerData; + +fn parseLayers(spec: []const u8, out: []usize) !usize { + var n: usize = 0; + var it = std.mem.splitScalar(u8, spec, ','); + while (it.next()) |tok| { + if (n == out.len) return error.TooManyLayers; + out[n] = try std.fmt.parseInt(usize, tok, 10); + n += 1; + } + return n; +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + var mode: []const u8 = "all"; + var layer_spec: []const u8 = "0,17"; + if (rest.len > 1) { + if (std.ascii.isDigit(rest[1][0])) layer_spec = rest[1] else mode = rest[1]; + } + if (std.mem.eql(u8, mode, "hit")) return hits.hitRun(gpu, rest[2..]); + if (rest.len > 2) layer_spec = rest[2]; + var layers: [8]usize = undefined; + const n_layers = try parseLayers(layer_spec, &layers); + var pack: Pack = .{}; + defer pack.close(gpu.io); + try pack.open(gpu, dir); + var ex = try exl3e.Experts.load(gpu.d); + defer ex.unload(); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + for (layers[0..n_layers]) |layer| { + std.debug.print("== layer {d}\n", .{layer}); + const ld = try LayerData.load(gpu, &pack, layer); + defer ld.free(gpu.gpa); + const all = std.mem.eql(u8, mode, "all"); + if (all or std.mem.eql(u8, mode, "equal")) { + try equal.equalRun(gpu, &ex, s, ld, 10, 0, 0x10 + layer); + try equal.equalRun(gpu, &ex, s, ld, 24, 4, 0x24 + layer); + try equal.equalRun(gpu, &ex, s, ld, 64, 0, 0x64 + layer); + } + if (std.mem.eql(u8, mode, "bench")) try bench.benchRun(gpu, &ex, s, ld); + } +} diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_format.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_format.zig new file mode 120000 index 0000000000..0618faed80 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_format.zig @@ -0,0 +1 @@ +../../../../src/core/exl3_format.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_pack_tests.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_pack_tests.zig new file mode 100644 index 0000000000..a7c195b25d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_pack_tests.zig @@ -0,0 +1,522 @@ +//! exl3-pack: v1.1.0's EXL3 operators (exl3.Dense, exl3_experts.Experts) on a real Flash Next EXL3 pack's shapes. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("../../check.zig"); +const exl3fmt = core.exl3_format; +const st = core.safetensors; +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; + +const shard_count = 7; + +const Pack = struct { + files: [shard_count]st.File = undefined, + open_n: usize = 0, + + fn open(p: *Pack, gpu: Gpu, dir: []const u8) !void { + var buf: [512]u8 = undefined; + for (0..shard_count) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-{d:0>5}.safetensors", .{ dir, i + 1, shard_count }); + p.files[i] = try st.File.open(gpu.gpa, gpu.io, path); + p.open_n += 1; + } + } + + fn close(p: *Pack, io: std.Io) void { + for (p.files[0..p.open_n]) |*f| f.close(io); + } + + fn get(p: *const Pack, name: []const u8) ?st.Tensor { + for (p.files[0..p.open_n]) |*f| if (f.get(name)) |t| return t; + return null; + } +}; + +/// One EXL3 matrix copied to aligned host memory: trellis int16 [K/16, N/16, 8 * halves], fp16 suh [K], svh [N]. +const Mat = struct { + trellis: []i16, + suh: []f16, + svh: []f16, + tiles_k: usize, + tiles_n: usize, + halves: u16, + + fn k(m: Mat) usize { + return m.tiles_k * 16; + } + fn n(m: Mat) usize { + return m.tiles_n * 16; + } +}; + +fn copyAs(comptime T: type, a: std.mem.Allocator, bytes: []const u8) ![]T { + const out = try a.alloc(T, bytes.len / @sizeOf(T)); + @memcpy(std.mem.sliceAsBytes(out), bytes[0 .. out.len * @sizeOf(T)]); + return out; +} + +fn mat(p: *const Pack, a: std.mem.Allocator, base: []const u8) !Mat { + var buf: [256]u8 = undefined; + const t = p.get(try std.fmt.bufPrint(&buf, "{s}.trellis", .{base})) orelse return error.MissingTensor; + const su = p.get(try std.fmt.bufPrint(&buf, "{s}.suh", .{base})) orelse return error.MissingTensor; + const sv = p.get(try std.fmt.bufPrint(&buf, "{s}.svh", .{base})) orelse return error.MissingTensor; + if (p.get(try std.fmt.bufPrint(&buf, "{s}.mul1", .{base})) == null) return error.NotMul1; + if (t.rank != 3 or t.dtype != .i16) return error.UnexpectedTensor; + return .{ + .trellis = try copyAs(i16, a, t.bytes), + .suh = try copyAs(f16, a, su.bytes), + .svh = try copyAs(f16, a, sv.bytes), + .tiles_k = t.dim(0), + .tiles_n = t.dim(1), + .halves = @intCast(t.dim(2) / 8), + }; +} + +fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +fn bfValue(bits: u16) f64 { + return @floatCast(@as(f32, @bitCast(@as(u32, bits) << 16))); +} + +/// y [rows, N] = x [rows, K] ยท diag(suh) H Wq H diag(svh), all in float64 (x taken as given, no fp16 step). +fn exact(a: std.mem.Allocator, m: Mat, x: []const f64, rows: usize, y: []f64) !void { + const k = m.k(); + const n = m.n(); + const wq = try a.alloc(f16, k * n); + defer a.free(wq); + exl3fmt.unpack(.{ .halves = m.halves }, .mul1, m.trellis, m.tiles_k, m.tiles_n, wq); + const xs = try a.alloc(f64, rows * k); + defer a.free(xs); + for (0..rows) |r| for (0..k) |i| { + xs[r * k + i] = x[r * k + i] * @as(f64, @floatCast(m.suh[i])); + }; + exl3fmt.rotate(xs, rows, k, 1); + @memset(y, 0); + for (0..k) |ki| { + const row_w = wq[ki * n ..][0..n]; + for (0..rows) |r| { + const xv = xs[r * k + ki]; + const yr = y[r * n ..][0..n]; + for (yr, row_w) |*o, w| o.* += xv * @as(f64, @floatCast(w)); + } + } + exl3fmt.rotate(y, rows, n, 1); + for (0..rows) |r| for (0..n) |c| { + y[r * n + c] *= @as(f64, @floatCast(m.svh[c])); + }; +} + +const Err = struct { max_abs: f64 = 0, peak: f64 = 0, worst_rel: f64 = 0 }; + +fn compareBf16(got: []const u16, want: []const f64) Err { + var e: Err = .{}; + for (got, want) |g, w| { + const d = @abs(bfValue(g) - w); + e.max_abs = @max(e.max_abs, d); + e.peak = @max(e.peak, @abs(w)); + } + e.worst_rel = e.max_abs / @max(e.peak, 1e-30); + return e; +} + +fn compareF32(got: []const f32, want: []const f64) Err { + var e: Err = .{}; + for (got, want) |g, w| { + const d = @abs(@as(f64, g) - w); + e.max_abs = @max(e.max_abs, d); + e.peak = @max(e.peak, @abs(w)); + } + e.worst_rel = e.max_abs / @max(e.peak, 1e-30); + return e; +} + +const Bench = struct { + gpu: Gpu, + s: cuda.Stream, + ev0: cuda.Event, + ev1: cuda.Event, + reps: usize, + total_ms: f64 = 0, + + fn time(b: *Bench, ctx: anytype, comptime f: fn (@TypeOf(ctx)) anyerror!void) !f64 { + for (0..3) |_| try f(ctx); + try b.s.synchronize(); + var best: f64 = std.math.inf(f64); + for (0..5) |_| { + try b.ev0.record(b.s); + for (0..b.reps) |_| try f(ctx); + try b.ev1.record(b.s); + try b.ev1.synchronize(); + const ms = try b.ev0.elapsedMs(b.ev1); + best = @min(best, @as(f64, ms) / @as(f64, @floatFromInt(b.reps))); + } + return best; + } +}; + +const Call = struct { + de: *exl3.Dense, + s: cuda.Stream, + w: exl3.Weight, + x: u64, + y: u64, + rows: usize, + fn go(c: Call) anyerror!void { + try c.de.decode(c.s, c.w, c.x, .bf16, c.rows, c.y, .bf16); + } +}; + +fn randX(a: std.mem.Allocator, rng: std.Random, n: usize) !struct { bits: []u16, vals: []f64 } { + const bits = try a.alloc(u16, n); + const vals = try a.alloc(f64, n); + for (bits, vals) |*b, *v| { + b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + v.* = bfValue(b.*); + } + return .{ .bits = bits, .vals = vals }; +} + +fn linear(b: *Bench, de: *exl3.Dense, pack: *const Pack, a: std.mem.Allocator, rng: std.Random, name: []const u8, impl: @TypeOf(de.implementation)) !void { + const m = mat(pack, a, name) catch |e| { + std.debug.print("{s}: skipped ({t})\n", .{ name, e }); + return; + }; + const k = m.k(); + const n = m.n(); + const strides = exl3.strides(k, n, m.halves, false); + var tr = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.trellis)); + defer tr.free(); + var su = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.suh)); + defer su.free(); + var sv = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.svh)); + defer sv.free(); + const w: exl3.Weight = .{ .trellis = tr.ptr, .suh = su.ptr, .svh = sv.ptr, .k = @intCast(k), .n = @intCast(n), .halves = m.halves, .codebook = .mul1, .stride_k = strides.k, .stride_nb = strides.nb }; + if (impl == .tiles) { + _ = exl3.Dense.candidateWeight(w) catch |e| { + std.debug.print("{s:<58} tiles: refused ({t})\n", .{ name, e }); + return; + }; + } + de.implementation = impl; + defer de.implementation = .reference; + for ([_]usize{ 1, 5 }) |rows| { + const x = try randX(a, rng, rows * k); + var xd = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(x.bits)); + defer xd.free(); + var yd = try cuda.DeviceBuffer.alloc(b.gpu.d, rows * n * 2); + defer yd.free(); + try yd.fill8(0xff, b.s.handle); + const call: Call = .{ .de = de, .s = b.s, .w = w, .x = xd.ptr, .y = yd.ptr, .rows = rows }; + try call.go(); + try b.s.synchronize(); + const got = try a.alloc(u16, rows * n); + try yd.download(0, std.mem.sliceAsBytes(got)); + // the reference row count: lm_head's 636M products a row, so only its first row + const ref_rows: usize = if (n > 100_000) 1 else rows; + const want = try a.alloc(f64, ref_rows * n); + try exact(a, m, x.vals[0 .. ref_rows * k], ref_rows, want); + const e = compareBf16(got[0 .. ref_rows * n], want); + // a second call's bytes equal the first's (deterministic) + try call.go(); + try b.s.synchronize(); + const again = try a.alloc(u16, rows * n); + try yd.download(0, std.mem.sliceAsBytes(again)); + const same = std.mem.eql(u16, got, again); + const ms = try b.time(call, Call.go); + if (rows == 1) b.total_ms += ms; + std.debug.print("{s:<58} {s:<9} K={d:<5} N={d:<6} halves={d:<2} rows={d} ms={d:>8.4} max|err|={e:.3} peak={d:.3} rel={e:.2} repeat={s}\n", .{ name, @tagName(impl), k, n, m.halves, rows, ms, e.max_abs, e.peak, e.worst_rel, if (same) "same" else "DIFF" }); + if (rows == 1) { + // the same row with fp32 outputs: the accumulation's own error, without bf16's output rounding + var yf = try cuda.DeviceBuffer.alloc(b.gpu.d, n * 4); + defer yf.free(); + try de.decode(b.s, w, xd.ptr, .bf16, 1, yf.ptr, .f32); + try b.s.synchronize(); + const gf = try a.alloc(f32, n); + try yf.download(0, std.mem.sliceAsBytes(gf)); + const ef = compareF32(gf, want[0..n]); + std.debug.print("{s:<58} {s:<9} fp32 out rows=1: max|err|={e:.3} peak={d:.3} rel={e:.2}\n", .{ name, @tagName(impl), ef.max_abs, ef.peak, ef.worst_rel }); + } + } +} + +const ExpCall = struct { + ex: *const exl3e.Experts, + s: cuda.Stream, + l: exl3e.Layer, + x: u64, + pick: u64, + wts: u64, + out: u64, + sc: exl3e.Scratch, + fn go(c: ExpCall) anyerror!void { + try c.ex.run(c.s, c.l, c.x, c.l.dims, .bf16, c.pick, c.wts, 1, c.out, c.sc); + } +}; + +/// Ten experts of `layer` as E = 10, one row picking all ten (weights 0.1): fp32 out vs float64. +fn experts(b: *Bench, ex: *const exl3e.Experts, pack: *const Pack, a: std.mem.Allocator, rng: std.Random, layer: usize, act_mode: i32, mapped: usize) !void { + const count = 10; + var buf: [160]u8 = undefined; + var g: [count]Mat = undefined; + var u: [count]Mat = undefined; + var dn: [count]Mat = undefined; + for (0..count) |e| { + g[e] = try mat(pack, a, try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.experts.{d}.gate_proj", .{ layer, e })); + u[e] = try mat(pack, a, try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.experts.{d}.up_proj", .{ layer, e })); + dn[e] = try mat(pack, a, try std.fmt.bufPrint(&buf, "model.language_model.layers.{d}.mlp.experts.{d}.down_proj", .{ layer, e })); + } + const dims = g[0].k(); + const width = g[0].n(); + var bufs: std.ArrayList(cuda.DeviceBuffer) = .empty; + defer { + for (bufs.items) |*x| x.free(); + bufs.deinit(b.gpu.gpa); + } + const up = struct { + fn f(gpu: Gpu, list: *std.ArrayList(cuda.DeviceBuffer), bytes: []const u8) !u64 { + const d = try cuda.DeviceBuffer.fromHost(gpu.d, bytes); + try list.append(gpu.gpa, d); + return d.ptr; + } + }.f; + var hosts: std.ArrayList(cuda.HostBuffer) = .empty; + defer { + for (hosts.items) |*h| h.free(); + hosts.deinit(b.gpu.gpa); + } + const hostUp = struct { + fn f(gpu: Gpu, list: *std.ArrayList(cuda.HostBuffer), bytes: []const u8) !u64 { + var h = try cuda.HostBuffer.allocMapped(gpu.d, bytes.len); + errdefer h.free(); + @memcpy(h.bytes[0..bytes.len], bytes); + try list.append(gpu.gpa, h); + return try h.device(); + } + }.f; + var gp: [count]u64 = undefined; + var upp: [count]u64 = undefined; + var dp: [count]u64 = undefined; + var gk2: [count]i32 = undefined; + var uk2: [count]i32 = undefined; + var dk2: [count]i32 = undefined; + var suh_g: std.ArrayList(f16) = .empty; + var suh_u: std.ArrayList(f16) = .empty; + var svh_g: std.ArrayList(f16) = .empty; + var svh_u: std.ArrayList(f16) = .empty; + var suh_d: std.ArrayList(f16) = .empty; + var svh_d: std.ArrayList(f16) = .empty; + for (0..count) |e| { + if (e < mapped) { + gp[e] = try hostUp(b.gpu, &hosts, std.mem.sliceAsBytes(g[e].trellis)); + upp[e] = try hostUp(b.gpu, &hosts, std.mem.sliceAsBytes(u[e].trellis)); + dp[e] = try hostUp(b.gpu, &hosts, std.mem.sliceAsBytes(dn[e].trellis)); + } else { + gp[e] = try up(b.gpu, &bufs, std.mem.sliceAsBytes(g[e].trellis)); + upp[e] = try up(b.gpu, &bufs, std.mem.sliceAsBytes(u[e].trellis)); + dp[e] = try up(b.gpu, &bufs, std.mem.sliceAsBytes(dn[e].trellis)); + } + gk2[e] = g[e].halves; + uk2[e] = u[e].halves; + dk2[e] = dn[e].halves; + try suh_g.appendSlice(a, g[e].suh); + try suh_u.appendSlice(a, u[e].suh); + try svh_g.appendSlice(a, g[e].svh); + try svh_u.appendSlice(a, u[e].svh); + try suh_d.appendSlice(a, dn[e].suh); + try svh_d.appendSlice(a, dn[e].svh); + } + const l: exl3e.Layer = .{ + .gate = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&gp)), + .up = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&upp)), + .down = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&dp)), + .gate_k2 = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&gk2)), + .up_k2 = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&uk2)), + .down_k2 = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&dk2)), + .suh_g = try up(b.gpu, &bufs, std.mem.sliceAsBytes(suh_g.items)), + .suh_u = try up(b.gpu, &bufs, std.mem.sliceAsBytes(suh_u.items)), + .svh_g = try up(b.gpu, &bufs, std.mem.sliceAsBytes(svh_g.items)), + .svh_u = try up(b.gpu, &bufs, std.mem.sliceAsBytes(svh_u.items)), + .suh_d = try up(b.gpu, &bufs, std.mem.sliceAsBytes(suh_d.items)), + .svh_d = try up(b.gpu, &bufs, std.mem.sliceAsBytes(svh_d.items)), + .experts = count, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = act_mode, + }; + const x = try randX(a, rng, dims); + var pick: [count]i32 = undefined; + for (&pick, 0..) |*p, i| p.* = @intCast((i * 7) % count); + var wts: [count]f32 = @splat(0.1); + const sbytes = try exl3e.scratchBytes(l, 1, count); + var sbuf = try cuda.DeviceBuffer.alloc(b.gpu.d, sbytes); + defer sbuf.free(); + var outd = try cuda.DeviceBuffer.alloc(b.gpu.d, dims * 4); + defer outd.free(); + try outd.fill8(0xff, b.s.handle); + const call: ExpCall = .{ + .ex = ex, + .s = b.s, + .l = l, + .x = try up(b.gpu, &bufs, std.mem.sliceAsBytes(x.bits)), + .pick = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&pick)), + .wts = try up(b.gpu, &bufs, std.mem.sliceAsBytes(&wts)), + .out = outd.ptr, + .sc = try exl3e.scratch(sbuf.ptr, l, 1, count), + }; + try call.go(); + try b.s.synchronize(); + const got = try a.alloc(f32, dims); + try outd.download(0, std.mem.sliceAsBytes(got)); + // float64: out = sum_e 0.1 * down_e(silu(gate_e(x)) * up_e(x)) + const want = try a.alloc(f64, dims); + @memset(want, 0); + const gy = try a.alloc(f64, width); + const uy = try a.alloc(f64, width); + const dy = try a.alloc(f64, dims); + for (pick) |pe| { + const e: usize = @intCast(pe); + try exact(a, g[e], x.vals, 1, gy); + try exact(a, u[e], x.vals, 1, uy); + for (gy, uy) |*gv, uv| gv.* = gv.* / (1.0 + @exp(-gv.*)) * uv; + try exact(a, dn[e], gy, 1, dy); + for (want, dy) |*o, v| o.* += 0.1 * v; + } + const err = compareF32(got, want); + const ms = try b.time(call, ExpCall.go); + // with a 64 MiB memset between calls (evicts the 32 MiB L2, so mapped host bytes come over PCIe each call). + var flush = try cuda.DeviceBuffer.alloc(b.gpu.d, 64 << 20); + defer flush.free(); + const Flushed = struct { + c: ExpCall, + f: u64, + with_call: bool, + fn go(fl: @This()) anyerror!void { + try fl.c.s.d.check(fl.c.s.d.api.cuMemsetD8Async(fl.f, 0, 64 << 20, fl.c.s.handle), "cuMemsetD8Async"); + if (fl.with_call) try fl.c.go(); + } + }; + const with = try b.time(Flushed{ .c = call, .f = flush.ptr, .with_call = true }, Flushed.go); + const alone = try b.time(Flushed{ .c = call, .f = flush.ptr, .with_call = false }, Flushed.go); + std.debug.print(" cold L2 (64 MiB memset between calls, its time subtracted): {d:.4} ms\n", .{with - alone}); + std.debug.print("routed experts L{d} (1 row x 10 picks, act_mode {d}, {d} of 10 from mapped host memory) K={d} I={d} halves={d}/{d}/{d} ms={d:>8.4} max|err|={e:.3} peak={d:.3} rel={e:.2}\n", .{ layer, act_mode, mapped, dims, width, g[0].halves, u[0].halves, dn[0].halves, ms, err.max_abs, err.peak, err.worst_rel }); + if (mapped > 0) { + // the copy-engine alternative: the same experts' trellis bytes, pinned host to device. + var total: usize = 0; + for (0..mapped) |e| total += (g[e].trellis.len + u[e].trellis.len + dn[e].trellis.len) * 2; + var src = try cuda.HostBuffer.alloc(b.gpu.d, total); + defer src.free(); + var dst = try cuda.DeviceBuffer.alloc(b.gpu.d, total); + defer dst.free(); + const Copy = struct { + d: *const cuda.Driver, + s: cuda.Stream, + src: []u8, + dst: u64, + fn go(c: @This()) anyerror!void { + try c.d.check(c.d.api.cuMemcpyHtoDAsync_v2(c.dst, c.src.ptr, c.src.len, c.s.handle), "cuMemcpyHtoDAsync"); + } + }; + const c: Copy = .{ .d = b.gpu.d, .s = b.s, .src = src.bytes[0..total], .dst = dst.ptr }; + const cms = try b.time(c, Copy.go); + std.debug.print(" pinned H2D copy of those {d} experts' trellis ({d} bytes): {d:.4} ms ({d:.1} GB/s)\n", .{ mapped, total, cms, @as(f64, @floatFromInt(total)) / cms / 1e6 }); + } +} + +/// Rows=1 timing only, warm and with a 64 MiB memset between calls (the memset's own time subtracted). +fn coldLinear(b: *Bench, de: *exl3.Dense, pack: *const Pack, a: std.mem.Allocator, rng: std.Random, name: []const u8) !void { + const m = try mat(pack, a, name); + const k = m.k(); + const n = m.n(); + const strides = exl3.strides(k, n, m.halves, false); + var tr = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.trellis)); + defer tr.free(); + var su = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.suh)); + defer su.free(); + var sv = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(m.svh)); + defer sv.free(); + const w: exl3.Weight = .{ .trellis = tr.ptr, .suh = su.ptr, .svh = sv.ptr, .k = @intCast(k), .n = @intCast(n), .halves = m.halves, .codebook = .mul1, .stride_k = strides.k, .stride_nb = strides.nb }; + const x = try randX(a, rng, k); + var xd = try cuda.DeviceBuffer.fromHost(b.gpu.d, std.mem.sliceAsBytes(x.bits)); + defer xd.free(); + var yd = try cuda.DeviceBuffer.alloc(b.gpu.d, n * 2); + defer yd.free(); + var flush = try cuda.DeviceBuffer.alloc(b.gpu.d, 64 << 20); + defer flush.free(); + const call: Call = .{ .de = de, .s = b.s, .w = w, .x = xd.ptr, .y = yd.ptr, .rows = 1 }; + const warm = try b.time(call, Call.go); + const F = struct { + c: Call, + f: u64, + with_call: bool, + fn go(fl: @This()) anyerror!void { + try fl.c.s.d.check(fl.c.s.d.api.cuMemsetD8Async(fl.f, 0, 64 << 20, fl.c.s.handle), "cuMemsetD8Async"); + if (fl.with_call) try fl.c.go(); + } + }; + const with = try b.time(F{ .c = call, .f = flush.ptr, .with_call = true }, F.go); + const alone = try b.time(F{ .c = call, .f = flush.ptr, .with_call = false }, F.go); + const bytes = m.trellis.len * 2; + std.debug.print("{s:<58} rows=1 bytes={d:>10} warm {d:.4} ms cold {d:.4} ms ({d:.0} GB/s cold)\n", .{ name, bytes, warm, with - alone, @as(f64, @floatFromInt(bytes)) / (with - alone) / 1e6 }); +} + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + const reps: usize = if (rest.len > 1) try std.fmt.parseInt(usize, rest[1], 10) else 20; + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var pack: Pack = .{}; + defer pack.close(gpu.io); + try pack.open(gpu, dir); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + var b: Bench = .{ .gpu = gpu, .s = s, .ev0 = try cuda.Event.init(gpu.d, true), .ev1 = try cuda.Event.init(gpu.d, true), .reps = reps }; + defer b.ev0.deinit(); + defer b.ev1.deinit(); + var de = try exl3.Dense.load(gpu.d); + defer de.unload(); + var ex = try exl3e.Experts.load(gpu.d); + defer ex.unload(); + var prng = std.Random.DefaultPrng.init(0x5eed); + const rng = prng.random(); + + const names = [_][]const u8{ + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "model.language_model.layers.0.linear_attn.in_proj_z", + "model.language_model.layers.0.linear_attn.out_proj", + "model.language_model.layers.3.self_attn.q_proj", + "model.language_model.layers.3.self_attn.k_proj", + "model.language_model.layers.3.self_attn.v_proj", + "model.language_model.layers.3.self_attn.o_proj", + "model.language_model.layers.3.mlp.shared_expert.gate_proj", + "model.language_model.layers.3.mlp.shared_expert.down_proj", + "model.language_model.layers.3.self_attn.indexer.index_qk_proj", + "model.language_model.layers.0.mlp.experts.0.gate_proj", + "model.language_model.layers.0.mlp.experts.0.down_proj", + "lm_head", + }; + if (rest.len > 2 and std.mem.eql(u8, rest[2], "cold")) { + for (names) |nm| try coldLinear(&b, &de, &pack, a, rng, nm); + try experts(&b, &ex, &pack, a, rng, 17, 1, 0); + return; + } + if (rest.len > 2 and std.mem.eql(u8, rest[2], "experts")) { + for ([_]usize{ 0, 1, 4, 10 }) |mapped| try experts(&b, &ex, &pack, a, rng, 0, 1, mapped); + return; + } + for (names) |nm| try linear(&b, &de, &pack, a, rng, nm, .reference); + std.debug.print("sum of rows=1 ms over the projections above (reference Dense): {d:.3}\n", .{b.total_ms}); + b.total_ms = 0; + for (names) |nm| try linear(&b, &de, &pack, a, rng, nm, .tiles); + try experts(&b, &ex, &pack, a, rng, 0, 1, 0); + try experts(&b, &ex, &pack, a, rng, 0, 0, 0); + try experts(&b, &ex, &pack, a, rng, 17, 1, 0); + check.pass("exl3-pack: v1.1.0 EXL3 operators ran on the pack's shapes (see the error and timing lines)", .{}); +} diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_place.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_place.zig new file mode 120000 index 0000000000..0a791de308 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_place.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/pack/exl3_place.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_qmat.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_qmat.zig new file mode 120000 index 0000000000..e690658604 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_qmat.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/pack/exl3_qmat.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_scan.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_scan.zig new file mode 120000 index 0000000000..640d68021d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_scan.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/pack/exl3_scan.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/pack/exl3_shape.zig b/zig/tests/cuda/flashnext_exl3/pack/exl3_shape.zig new file mode 120000 index 0000000000..e4dccac406 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/pack/exl3_shape.zig @@ -0,0 +1 @@ +../../../../src/families/flashnext/exl3/pack/exl3_shape.zig \ No newline at end of file diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_copy.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_copy.zig new file mode 100644 index 0000000000..0aeb73061d --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_copy.zig @@ -0,0 +1,157 @@ +//! exl3-prefill copy: the expert line copies of the offload pool, timed and checked on one layer. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("../../check.zig"); +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; +const packs = @import("exl3_prefill_pack.zig"); +const experts = @import("exl3_prefill_experts.zig"); + +const dims = 2560; +const width = 640; +const routed_n = 512; +const table_n = 513; +const slots = 11; + +const LayerData = packs.LayerData; +const bf16Bits = packs.bf16Bits; +const genPicks = packs.genPicks; +const SlotsCall = experts.SlotsCall; + +// ---------------------------------------------------------------------------------------------------- copies + +pub fn copyRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const a = gpu.gpa; + // the layer's routed trellis in one pinned bank (what S9 keeps), lines laid out as Lines does + var total: usize = 0; + for (ld.routed) |e| for (e) |b| { + total += b.len; + }; + var bank = try cuda.HostBuffer.alloc(gpu.d, total); + defer bank.free(); + const pinned = try a.alloc([3][]const u8, routed_n); + defer a.free(pinned); + var at: usize = 0; + for (ld.routed, pinned) |e, *o| for (e, 0..) |b, p| { + @memcpy(bank.bytes[at..][0..b.len], b); + o[p] = bank.bytes[at..][0..b.len]; + at += b.len; + }; + const n_lines: usize = 64; + var bytes: [3]usize = .{ 0, 0, 0 }; + var k2: [routed_n][3]i32 = undefined; + const tiles = (dims / 16) * (width / 16); + for (pinned, &k2) |e, *w| for (0..3) |p| { + bytes[p] = @max(bytes[p], e[p].len); + w[p] = @intCast(e[p].len / (tiles * 16)); + }; + var lines = try cuda.exl3_pool.Lines.init(a, gpu.d, .{ .experts = table_n, .bytes = bytes, .lines = n_lines, .staging = 0 }); + defer lines.deinit(); + var c0 = try cuda.Event.init(gpu.d, true); + defer c0.deinit(); + var c1 = try cuda.Event.init(gpu.d, true); + defer c1.deinit(); + var raw = try cuda.DeviceBuffer.alloc(gpu.d, n_lines * lines.layout.line_bytes); + defer raw.free(); + const moved = blk: { + var m: usize = 0; + for (pinned[0..n_lines]) |e| m += e[0].len + e[1].len + e[2].len; + break :blk m; + }; + std.debug.print(" {d} expert lines ({d:.2} MiB, {d:.3} MiB a line) from the pinned bank, best of 5:\n", .{ n_lines, @as(f64, @floatFromInt(moved)) / 1048576.0, @as(f64, @floatFromInt(moved)) / 1048576.0 / n_lines }); + var best_load: f64 = std.math.inf(f64); + var best_raw: f64 = std.math.inf(f64); + var best_raw1: f64 = std.math.inf(f64); + for (0..5) |it| { + try lines.drain(); + try c0.record(lines.copy); + for (0..n_lines) |i| { + const id = (it * n_lines + i) % routed_n; + const prev = (id + routed_n - n_lines) % routed_n; + try lines.load(i, id, pinned[id], k2[id], if (it == 0) null else prev, true, gpu.io); + } + try c1.record(lines.copy); + try c1.synchronize(); + best_load = @min(best_load, try c0.elapsedMs(c1)); + // the same bytes as three H2D a line, no table writes + try c0.record(lines.copy); + for (0..n_lines) |i| for (0..3) |p| { + const e = pinned[i]; + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr + i * lines.layout.line_bytes + lines.layout.off[p], e[p].ptr, e[p].len, lines.copy.handle), "cuMemcpyHtoDAsync"); + }; + try c1.record(lines.copy); + try c1.synchronize(); + best_raw = @min(best_raw, try c0.elapsedMs(c1)); + // one H2D of the bank's first n_lines experts (contiguous in the bank) + try c0.record(lines.copy); + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr, bank.bytes.ptr, moved, lines.copy.handle), "cuMemcpyHtoDAsync"); + try c1.record(lines.copy); + try c1.synchronize(); + best_raw1 = @min(best_raw1, try c0.elapsedMs(c1)); + } + const gbs = struct { + fn f(b: usize, ms: f64) f64 { + return @as(f64, @floatFromInt(b)) / ms / 1e6; + } + }.f; + std.debug.print(" Lines.load (3 H2D + 6-9 table memsets + event a line): {d:.3} ms = {d:.3} ms a line, {d:.2} GB/s\n", .{ best_load, best_load / n_lines, gbs(moved, best_load) }); + std.debug.print(" 3 bare H2D a line: {d:.3} ms = {d:.3} ms a line, {d:.2} GB/s\n", .{ best_raw, best_raw / n_lines, gbs(moved, best_raw) }); + std.debug.print(" one H2D of the same bytes: {d:.3} ms, {d:.2} GB/s\n", .{ best_raw1, gbs(moved, best_raw1) }); + + // copies beside compute: a 1024-row experts call on the resident layer while the copy stream moves 64 lines + const rows: usize = 1024; + var prng = std.Random.DefaultPrng.init(0xc0b1); + const rng = prng.random(); + const xh = try a.alloc(u16, rows * dims); + defer a.free(xh); + for (xh) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer x.free(); + const picks = try a.alloc(i32, rows * slots); + defer a.free(picks); + genPicks(rng, 0, rows, picks); + var pk = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(picks)); + defer pk.free(); + var sbuf = try cuda.DeviceBuffer.alloc(gpu.d, try exl3e.scratchBytes(ld.res, rows, slots)); + defer sbuf.free(); + const sc = try exl3e.scratch(sbuf.ptr, ld.res, rows, slots); + const call: SlotsCall = .{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }; + var k0 = try cuda.Event.init(gpu.d, true); + defer k0.deinit(); + var k1 = try cuda.Event.init(gpu.d, true); + defer k1.deinit(); + var alone_c: f64 = std.math.inf(f64); + var alone_k: f64 = std.math.inf(f64); + var both_c: f64 = std.math.inf(f64); + var both_k: f64 = std.math.inf(f64); + var both_w: f64 = std.math.inf(f64); + for (0..5) |_| { + try k0.record(s); + try call.go(); + try k1.record(s); + try k1.synchronize(); + alone_k = @min(alone_k, try k0.elapsedMs(k1)); + try c0.record(lines.copy); + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr, bank.bytes.ptr, moved, lines.copy.handle), "cuMemcpyHtoDAsync"); + try c1.record(lines.copy); + try c1.synchronize(); + alone_c = @min(alone_c, try c0.elapsedMs(c1)); + // both at once + try s.synchronize(); + const t0 = check.now(gpu.io); + try c0.record(lines.copy); + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(raw.ptr, bank.bytes.ptr, moved, lines.copy.handle), "cuMemcpyHtoDAsync"); + try c1.record(lines.copy); + try k0.record(s); + try call.go(); + try k1.record(s); + try k1.synchronize(); + try c1.synchronize(); + both_w = @min(both_w, @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e6); + both_c = @min(both_c, try c0.elapsedMs(c1)); + both_k = @min(both_k, try k0.elapsedMs(k1)); + } + std.debug.print(" overlap: {d}-row experts call alone {d:.3} ms, 64-line H2D alone {d:.3} ms; together: experts {d:.3} ms, H2D {d:.3} ms, wall {d:.3} ms (serial would be {d:.3})\n", .{ rows, alone_k, alone_c, both_k, both_c, both_w, alone_k + alone_c }); + check.pass("layer {d}: copy cost and copy/compute overlap measured", .{ld.layer}); +} diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_dense.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_dense.zig new file mode 100644 index 0000000000..833f380083 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_dense.zig @@ -0,0 +1,171 @@ +//! exl3-prefill dense: the prompt tiles of a dense EXL3 projection against a float64 reference. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("../../check.zig"); +const exl3fmt = core.exl3_format; +const exl3 = cuda.exl3; +const Gpu = check.Gpu; +const packs = @import("exl3_prefill_pack.zig"); + +const Pack = packs.Pack; +const Timer = packs.Timer; +const bf16Bits = packs.bf16Bits; +const bfValue = packs.bfValue; + +// ---------------------------------------------------------------------------------------------------- dense + +const DenseCall = struct { + de: *exl3.Dense, + s: cuda.Stream, + w: exl3.Weight, + x: u64, + y: u64, + rows: usize, + pub fn go(c: DenseCall) anyerror!void { + try c.de.prompt(c.s, c.w, c.x, .bf16, c.rows, c.y, .bf16); + } +}; + +/// y [rows, N] = x [rows, K] . diag(suh) H Wq H diag(svh) in float64. +fn exact(a: std.mem.Allocator, trellis: []const i16, suh: []const f16, svh: []const f16, tk: usize, tn: usize, halves: u16, x: []const f64, rows: usize, y: []f64) !void { + const k = tk * 16; + const n = tn * 16; + const wq = try a.alloc(f16, k * n); + defer a.free(wq); + exl3fmt.unpack(.{ .halves = halves }, .mul1, trellis, tk, tn, wq); + const xs = try a.alloc(f64, rows * k); + defer a.free(xs); + for (0..rows) |r| for (0..k) |i| { + xs[r * k + i] = x[r * k + i] * @as(f64, @floatCast(suh[i])); + }; + exl3fmt.rotate(xs, rows, k, 1); + @memset(y, 0); + for (0..k) |ki| { + const row_w = wq[ki * n ..][0..n]; + for (0..rows) |r| { + const xv = xs[r * k + ki]; + const yr = y[r * n ..][0..n]; + for (yr, row_w) |*o, w| o.* += xv * @as(f64, @floatCast(w)); + } + } + exl3fmt.rotate(y, rows, n, 1); + for (0..rows) |r| for (0..n) |c| { + y[r * n + c] *= @as(f64, @floatCast(svh[c])); + }; +} + +fn copyAs(comptime T: type, a: std.mem.Allocator, bytes: []const u8) ![]T { + const out = try a.alloc(T, bytes.len / @sizeOf(T)); + @memcpy(std.mem.sliceAsBytes(out), bytes[0 .. out.len * @sizeOf(T)]); + return out; +} + +pub fn denseRun(gpu: Gpu, pack: *const Pack, s: cuda.Stream) !void { + const a = gpu.gpa; + var de = try exl3.Dense.load(gpu.d); + defer de.unload(); + var t = try Timer.init(gpu.d); + defer t.deinit(); + const names = [_][]const u8{ + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "model.language_model.layers.0.linear_attn.in_proj_z", + "model.language_model.layers.0.linear_attn.out_proj", + "model.language_model.layers.3.self_attn.q_proj", + "model.language_model.layers.3.self_attn.k_proj", + "model.language_model.layers.3.self_attn.o_proj", + "model.language_model.layers.3.mlp.shared_expert.gate_proj", + "model.language_model.layers.3.mlp.shared_expert.down_proj", + }; + const row_counts = [_]usize{ 16, 128, 256, 1024, 2048 }; + const max_rows: usize = 2048; + var prng = std.Random.DefaultPrng.init(0xde45e); + const rng = prng.random(); + var sum_ref: [row_counts.len]f64 = @splat(0); + var sum_tiles: [row_counts.len]f64 = @splat(0); + for (names) |name| { + const tr = try pack.get("{s}.trellis", .{name}); + const su = try pack.get("{s}.suh", .{name}); + const sv = try pack.get("{s}.svh", .{name}); + const tk = tr.dim(0); + const tn = tr.dim(1); + const halves: u16 = @intCast(tr.dim(2) / 8); + const k = tk * 16; + const n = tn * 16; + const trellis = try copyAs(i16, a, tr.bytes); + defer a.free(trellis); + const suh = try copyAs(f16, a, su.bytes); + defer a.free(suh); + const svh = try copyAs(f16, a, sv.bytes); + defer a.free(svh); + var trd = try cuda.DeviceBuffer.fromHost(gpu.d, tr.bytes); + defer trd.free(); + var sud = try cuda.DeviceBuffer.fromHost(gpu.d, su.bytes); + defer sud.free(); + var svd = try cuda.DeviceBuffer.fromHost(gpu.d, sv.bytes); + defer svd.free(); + const strides = exl3.strides(k, n, halves, false); + const w: exl3.Weight = .{ .trellis = trd.ptr, .suh = sud.ptr, .svh = svd.ptr, .k = @intCast(k), .n = @intCast(n), .halves = halves, .codebook = .mul1, .stride_k = strides.k, .stride_nb = strides.nb }; + const tiles_ok = if (exl3.Dense.candidateWeight(w)) |_| true else |_| false; + const xh = try a.alloc(u16, max_rows * k); + defer a.free(xh); + for (xh) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer x.free(); + var yr = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * n * 2); + defer yr.free(); + var yt = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * n * 2); + defer yt.free(); + // f64 of rows 0 and 1 + const xf = try a.alloc(f64, 2 * k); + defer a.free(xf); + for (xf, xh[0 .. 2 * k]) |*o, b| o.* = bfValue(b); + const want = try a.alloc(f64, 2 * n); + defer a.free(want); + try exact(a, trellis, suh, svh, tk, tn, halves, xf, 2, want); + std.debug.print(" {s} K={d} N={d} halves={d}{s}\n", .{ name, k, n, halves, if (tiles_ok) "" else " (tiles refuse it)" }); + for (row_counts, 0..) |rows, ri| { + de.implementation = .reference; + const rc: DenseCall = .{ .de = &de, .s = s, .w = w, .x = x.ptr, .y = yr.ptr, .rows = rows }; + const ms_ref = try t.best(s, 3, rc, DenseCall.go); + sum_ref[ri] += ms_ref; + var ms_tiles: f64 = std.math.nan(f64); + var diff: f64 = std.math.nan(f64); + var err_ref: f64 = 0; + var err_tiles: f64 = std.math.nan(f64); + var peak: f64 = 0; + const got_r = try a.alloc(u16, rows * n); + defer a.free(got_r); + try yr.download(0, std.mem.sliceAsBytes(got_r)); + for (got_r[0 .. 2 * n], want) |g, wv| { + err_ref = @max(err_ref, @abs(bfValue(g) - wv)); + peak = @max(peak, @abs(wv)); + } + if (tiles_ok) { + de.implementation = .tiles; + const tc: DenseCall = .{ .de = &de, .s = s, .w = w, .x = x.ptr, .y = yt.ptr, .rows = rows }; + ms_tiles = try t.best(s, 3, tc, DenseCall.go); + sum_tiles[ri] += ms_tiles; + const got_t = try a.alloc(u16, rows * n); + defer a.free(got_t); + try yt.download(0, std.mem.sliceAsBytes(got_t)); + diff = 0; + var rpeak: f64 = 0; + for (got_r, got_t) |g, h| { + diff = @max(diff, @abs(bfValue(g) - bfValue(h))); + rpeak = @max(rpeak, @abs(bfValue(g))); + } + diff /= rpeak; + err_tiles = 0; + for (got_t[0 .. 2 * n], want) |g, wv| err_tiles = @max(err_tiles, @abs(bfValue(g) - wv)); + de.implementation = .reference; + } else sum_tiles[ri] += ms_ref; + std.debug.print(" rows {d:5}: reference {d:8.4} ms ({d:6.2} us/row) tiles {d:8.4} ms ({d:6.2} us/row) max|ref-tiles|/peak {e:.2} rows 0-1 vs f64: ref {e:.2} tiles {e:.2} (rel to peak {d:.3})\n", .{ rows, ms_ref, 1e3 * ms_ref / @as(f64, @floatFromInt(rows)), ms_tiles, 1e3 * ms_tiles / @as(f64, @floatFromInt(rows)), diff, err_ref / peak, err_tiles / peak, peak }); + } + } + std.debug.print(" sum over the {d} shapes (ms): ", .{names.len}); + for (row_counts, sum_ref, sum_tiles) |r, x, y| std.debug.print("rows {d}: reference {d:.3} best {d:.3}; ", .{ r, x, y }); + std.debug.print("\n", .{}); + check.pass("dense prompt kernels timed", .{}); +} diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_experts.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_experts.zig new file mode 100644 index 0000000000..2dcca86822 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_experts.zig @@ -0,0 +1,118 @@ +//! exl3-prefill experts: the grouped, slot and expert-major masked passes timed on one layer. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("../../check.zig"); +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; +const packs = @import("exl3_prefill_pack.zig"); + +const dims = 2560; +const width = 640; +const routed_n = 512; +const table_n = 513; +const slots = 11; +const max_call_rows: usize = 1408; + +const LayerData = packs.LayerData; +const Timer = packs.Timer; +const bf16Bits = packs.bf16Bits; +const genPicks = packs.genPicks; + +pub const SlotsCall = struct { + ex: *const exl3e.Experts, + s: cuda.Stream, + l: exl3e.Layer, + x: u64, + picks: u64, + rows: usize, + sc: exl3e.Scratch, + pub fn go(c: SlotsCall) anyerror!void { + try c.ex.runSlots(c.s, c.l, c.x, dims, .bf16, c.picks, c.rows, c.sc); + } +}; + +/// Expert-major passes over one block: pass j runs the block's picks of experts [j*per, (j+1)*per) (others -1). +const PassesCall = struct { + ex: *const exl3e.Experts, + s: cuda.Stream, + l: exl3e.Layer, + x: u64, + masked: u64, // int32 [passes, rows * slots] + rows: usize, + passes: usize, + sc: exl3e.Scratch, + pub fn go(c: PassesCall) anyerror!void { + for (0..c.passes) |j| try c.ex.runSlots(c.s, c.l, c.x, dims, .bf16, c.masked + j * c.rows * slots * 4, c.rows, c.sc); + } +}; + +pub fn expertsRun(gpu: Gpu, ex: *const exl3e.Experts, s: cuda.Stream, ld: *const LayerData) !void { + const a = gpu.gpa; + var t = try Timer.init(gpu.d); + defer t.deinit(); + const max_rows = max_call_rows; + var prng = std.Random.DefaultPrng.init(0x5eed); + const rng = prng.random(); + const xh = try a.alloc(u16, max_rows * dims); + defer a.free(xh); + for (xh) |*b| b.* = bf16Bits((rng.float(f32) - 0.5) * 2.0); + var x = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(xh)); + defer x.free(); + const picks = try a.alloc(i32, max_rows * slots); + defer a.free(picks); + var pk = try cuda.DeviceBuffer.alloc(gpu.d, max_rows * slots * 4); + defer pk.free(); + var sbuf = try cuda.DeviceBuffer.alloc(gpu.d, try exl3e.scratchBytes(ld.res, max_rows, slots)); + defer sbuf.free(); + const sc = try exl3e.scratch(sbuf.ptr, ld.res, max_rows, slots); + std.debug.print(" runSlots, one resident layer (10 routed + shared a row), ms a call (best of 5 x 5):\n", .{}); + for ([_]f64{ 0.8, 0 }) |zs| { + genPicks(rng, zs, max_rows, picks); + try pk.upload(0, std.mem.sliceAsBytes(picks)); + for ([_]usize{ 16, 64, 256, 512, 1024, max_rows }) |rows| { + var seen: [table_n]bool = @splat(false); + var uniq: usize = 0; + for (picks[0 .. rows * slots]) |p| if (!seen[@intCast(p)]) { + seen[@intCast(p)] = true; + uniq += 1; + }; + const ms = try t.best(s, 5, SlotsCall{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }, SlotsCall.go); + std.debug.print(" {s:<8} rows {d:5}: {d:8.3} ms {d:7.2} us/row {d:6.0} rows/s/layer ({d} experts incl. shared, {d:.1} member rows each)\n", .{ if (zs == 0) "uniform" else "zipf0.8", rows, ms, 1e3 * ms / @as(f64, @floatFromInt(rows)), @as(f64, @floatFromInt(rows)) / ms * 1e3, uniq, @as(f64, @floatFromInt(rows * slots)) / @as(f64, @floatFromInt(uniq)) }); + } + } + // expert-major passes over one 1024-row block (uniform picks): bytes equal to one call, and the time + const rows: usize = 1024; + genPicks(rng, 0, rows, picks); + try pk.upload(0, std.mem.sliceAsBytes(picks[0 .. rows * slots])); + const one_bytes = try a.alloc(u8, rows * slots * dims * 4); + defer a.free(one_bytes); + try (SlotsCall{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }).go(); + try s.synchronize(); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(one_bytes.ptr, sc.y, one_bytes.len), "cuMemcpyDtoH"); + const base_ms = try t.best(s, 5, SlotsCall{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .picks = pk.ptr, .rows = rows, .sc = sc }, SlotsCall.go); + std.debug.print(" expert-major passes over one {d}-row block (uniform picks; pass j: experts [j*512/P, (j+1)*512/P) plus the shared one in pass 0):\n", .{rows}); + for ([_]usize{ 1, 2, 4, 8, 16, 32 }) |passes| { + const per = routed_n / passes; + const host = try a.alloc(i32, passes * rows * slots); + defer a.free(host); + for (0..passes) |j| for (picks[0 .. rows * slots], host[j * rows * slots ..][0 .. rows * slots]) |p, *o| { + const in = if (p == 512) j == 0 else @as(usize, @intCast(p)) / per == j; + o.* = if (in) p else -1; + }; + var masked = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(host)); + defer masked.free(); + try gpu.d.check(gpu.d.api.cuMemsetD8_v2(sc.y, 0xff, rows * slots * dims * 4), "cuMemsetD8"); + const call: PassesCall = .{ .ex = ex, .s = s, .l = ld.res, .x = x.ptr, .masked = masked.ptr, .rows = rows, .passes = passes, .sc = sc }; + try call.go(); + try s.synchronize(); + const got = try a.alloc(u8, one_bytes.len); + defer a.free(got); + try gpu.d.check(gpu.d.api.cuMemcpyDtoH_v2(got.ptr, sc.y, got.len), "cuMemcpyDtoH"); + const same = std.mem.eql(u8, got, one_bytes); + const ms = try t.best(s, 3, call, PassesCall.go); + std.debug.print(" {d:2} passes of {d:3} experts: {d:8.3} ms (+{d:.3} ms vs one call, {d:.3} ms a pass) y {s}\n", .{ passes, per, ms, ms - base_ms, ms / @as(f64, @floatFromInt(passes)), if (same) "bit-identical to one call" else "DIFFERS" }); + try check.expect(same, "masked passes reproduce one call's bytes", .{}); + } + check.pass("layer {d}: runSlots timed at 16..{d} rows; expert-major masked passes are bit-identical to one call", .{ ld.layer, max_rows }); +} diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_pack.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_pack.zig new file mode 100644 index 0000000000..366f189eca --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_pack.zig @@ -0,0 +1,218 @@ +//! exl3-prefill fixtures: the checkpoint pack, a timer, one layer's experts on the device and random picks. + +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const check = @import("../../check.zig"); +const st = core.safetensors; +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; + +const dims = 2560; +const width = 640; +const routed_n = 512; +const table_n = 513; +const top_k = 10; +const slots = 11; +const shard_count = 7; +const projs = [3][]const u8{ "gate_proj", "up_proj", "down_proj" }; + +pub const Pack = struct { + files: [shard_count]st.File = undefined, + open_n: usize = 0, + + pub fn open(p: *Pack, gpu: Gpu, dir: []const u8) !void { + var buf: [512]u8 = undefined; + for (0..shard_count) |i| { + const path = try std.fmt.bufPrint(&buf, "{s}/model-{d:0>5}-of-{d:0>5}.safetensors", .{ dir, i + 1, shard_count }); + p.files[i] = try st.File.open(gpu.gpa, gpu.io, path); + p.open_n += 1; + } + } + + pub fn close(p: *Pack, io: std.Io) void { + for (p.files[0..p.open_n]) |*f| f.close(io); + } + + pub fn get(p: *const Pack, comptime fmt: []const u8, args: anytype) !st.Tensor { + var buf: [200]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, fmt, args); + for (p.files[0..p.open_n]) |*f| if (f.get(name)) |t| return t; + std.debug.print("missing tensor {s}\n", .{name}); + return error.MissingTensor; + } +}; + +pub fn bf16Bits(v: f32) u16 { + const b: u32 = @bitCast(v); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +pub fn bfValue(bits: u16) f64 { + return @floatCast(@as(f32, @bitCast(@as(u32, bits) << 16))); +} + +pub const Timer = struct { + a: cuda.Event, + b: cuda.Event, + + pub fn init(d: *const cuda.Driver) !Timer { + return .{ .a = try cuda.Event.init(d, true), .b = try cuda.Event.init(d, true) }; + } + pub fn deinit(t: *Timer) void { + t.a.deinit(); + t.b.deinit(); + } + /// Best of 5 batches of `reps` calls, ms a call (after 2 warm-up calls). + pub fn best(t: *Timer, s: cuda.Stream, reps: usize, ctx: anytype, comptime f: fn (@TypeOf(ctx)) anyerror!void) !f64 { + for (0..2) |_| try f(ctx); + try s.synchronize(); + var b: f64 = std.math.inf(f64); + for (0..5) |_| { + try t.a.record(s); + for (0..reps) |_| try f(ctx); + try t.b.record(s); + try t.b.synchronize(); + b = @min(b, @as(f64, try t.a.elapsedMs(t.b)) / @as(f64, @floatFromInt(reps))); + } + return b; + } +}; + +/// One layer's 513 experts: the mmapped trellis, the stacked scales and every expert resident on the device. +pub const LayerData = struct { + layer: usize, + routed: [routed_n][3][]const u8, + shared: [3][]const u8, + scale_bufs: [6]cuda.DeviceBuffer, + res_mem: cuda.DeviceBuffer, + res_tab: cuda.DeviceBuffer, + res: exl3e.Layer, + line_bytes: usize, + + pub fn load(gpu: Gpu, pack: *const Pack, layer: usize) !*LayerData { + const gpa = gpu.gpa; + const ld = try gpa.create(LayerData); + errdefer gpa.destroy(ld); + ld.layer = layer; + const sizes = [6]usize{ dims, dims, width, width, width, dims }; // suh_g suh_u svh_g svh_u suh_d svh_d + var host: [6][]u8 = undefined; + for (&host, sizes) |*h, n| h.* = try gpa.alloc(u8, table_n * n * 2); + defer for (host) |h| gpa.free(h); + var total: usize = 0; + ld.line_bytes = 0; + for (0..table_n) |e| { + var line: usize = 0; + for (projs, 0..) |pn, p| { + const base = "model.language_model.layers.{d}.mlp."; + const tr = if (e < routed_n) try pack.get(base ++ "experts.{d}.{s}.trellis", .{ layer, e, pn }) else try pack.get(base ++ "shared_expert.{s}.trellis", .{ layer, pn }); + const su = if (e < routed_n) try pack.get(base ++ "experts.{d}.{s}.suh", .{ layer, e, pn }) else try pack.get(base ++ "shared_expert.{s}.suh", .{ layer, pn }); + const sv = if (e < routed_n) try pack.get(base ++ "experts.{d}.{s}.svh", .{ layer, e, pn }) else try pack.get(base ++ "shared_expert.{s}.svh", .{ layer, pn }); + const k: usize = if (p == 2) width else dims; + const n: usize = if (p == 2) dims else width; + if (tr.rank != 3 or tr.dim(0) * 16 != k or tr.dim(1) * 16 != n) return error.BadShape; + if (e < routed_n) ld.routed[e][p] = tr.bytes else ld.shared[p] = tr.bytes; + total += (tr.bytes.len + 255) / 256 * 256; + line += (tr.bytes.len + 255) / 256 * 256; + const su_i: usize = switch (p) { + 0 => 0, + 1 => 1, + else => 4, + }; + const sv_i: usize = switch (p) { + 0 => 2, + 1 => 3, + else => 5, + }; + @memcpy(host[su_i][e * k * 2 ..][0 .. k * 2], su.bytes); + @memcpy(host[sv_i][e * n * 2 ..][0 .. n * 2], sv.bytes); + } + if (e < routed_n) ld.line_bytes = @max(ld.line_bytes, line); + } + var made: usize = 0; + errdefer for (ld.scale_bufs[0..made]) |*b| b.free(); + for (&ld.scale_bufs, host) |*b, h| { + b.* = try cuda.DeviceBuffer.fromHost(gpu.d, h); + made += 1; + } + ld.res_mem = try cuda.DeviceBuffer.alloc(gpu.d, total); + errdefer ld.res_mem.free(); + var ptrs: [3][table_n]u64 = undefined; + var k2s: [3][table_n]i32 = undefined; + var at: usize = 0; + const tiles = (dims / 16) * (width / 16); + for (0..table_n) |e| for (0..3) |p| { + const bytes = if (e < routed_n) ld.routed[e][p] else ld.shared[p]; + try ld.res_mem.upload(at, bytes); + ptrs[p][e] = ld.res_mem.ptr + at; + k2s[p][e] = @intCast(bytes.len / (tiles * 16)); + at += (bytes.len + 255) / 256 * 256; + }; + const tl = cuda.exl3_pool.Tables.init(table_n); + const tab = try gpa.alloc(u8, tl.bytes); + defer gpa.free(tab); + @memset(tab, 0); + for (0..3) |p| { + @memcpy(tab[tl.ptr[p]..][0 .. table_n * 8], std.mem.sliceAsBytes(&ptrs[p])); + @memcpy(tab[tl.k2[p]..][0 .. table_n * 4], std.mem.sliceAsBytes(&k2s[p])); + } + ld.res_tab = try cuda.DeviceBuffer.fromHost(gpu.d, tab); + ld.res = .{ + .gate = ld.res_tab.ptr + tl.ptr[0], + .up = ld.res_tab.ptr + tl.ptr[1], + .down = ld.res_tab.ptr + tl.ptr[2], + .gate_k2 = ld.res_tab.ptr + tl.k2[0], + .up_k2 = ld.res_tab.ptr + tl.k2[1], + .down_k2 = ld.res_tab.ptr + tl.k2[2], + .suh_g = ld.scale_bufs[0].ptr, + .suh_u = ld.scale_bufs[1].ptr, + .svh_g = ld.scale_bufs[2].ptr, + .svh_u = ld.scale_bufs[3].ptr, + .suh_d = ld.scale_bufs[4].ptr, + .svh_d = ld.scale_bufs[5].ptr, + .experts = table_n, + .dims = dims, + .width = width, + .codebook = .mul1, + .act_mode = 1, + }; + std.debug.print(" layer {d}: 513 experts resident ({d:.0} MiB), routed k2 gate/up/down {d}/{d}/{d}, a routed line {d:.3} MiB\n", .{ layer, @as(f64, @floatFromInt(total)) / 1048576.0, k2s[0][0], k2s[1][0], k2s[2][0], @as(f64, @floatFromInt(ld.line_bytes)) / 1048576.0 }); + return ld; + } + + pub fn free(ld: *LayerData, gpa: std.mem.Allocator) void { + ld.res_tab.free(); + ld.res_mem.free(); + for (&ld.scale_bufs) |*b| b.free(); + gpa.destroy(ld); + } +}; + +/// Picks for `rows` rows: ten distinct routed ids then the shared expert; +pub fn genPicks(rng: std.Random, s: f64, rows: usize, out: []i32) void { + var perm: [routed_n]u16 = undefined; + for (&perm, 0..) |*p, i| p.* = @intCast(i); + var prng = std.Random.DefaultPrng.init(0x9e3779b9); + prng.random().shuffle(u16, &perm); + var cdf: [routed_n]f64 = undefined; + var sum: f64 = 0; + for (&cdf, 0..) |*c, i| { + sum += if (s == 0) 1.0 else 1.0 / std.math.pow(f64, @floatFromInt(i + 1), s); + c.* = sum; + } + for (&cdf) |*c| c.* /= sum; + for (0..rows) |r| { + const row = out[r * slots ..][0..slots]; + var n: usize = 0; + while (n < top_k) { + const u = rng.float(f64); + var i: usize = 0; + while (i + 1 < routed_n and cdf[i] < u) i += 1; + const id: i32 = perm[i]; + if (std.mem.indexOfScalar(i32, row[0..n], id) != null) continue; + row[n] = id; + n += 1; + } + row[top_k] = 512; + } +} diff --git a/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_tests.zig b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_tests.zig new file mode 100644 index 0000000000..5695cf4cf4 --- /dev/null +++ b/zig/tests/cuda/flashnext_exl3/prefill/exl3_prefill_tests.zig @@ -0,0 +1,37 @@ +//! exl3-prefill: the pieces of a Flash Next EXL3 prefill, timed on the pack's real tensors (one idle card). + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("../../check.zig"); +const exl3e = cuda.exl3_experts; +const Gpu = check.Gpu; +const packs = @import("exl3_prefill_pack.zig"); +const experts = @import("exl3_prefill_experts.zig"); +const copies = @import("exl3_prefill_copy.zig"); +const dense = @import("exl3_prefill_dense.zig"); + +const LayerData = packs.LayerData; +const Pack = packs.Pack; +const expertsRun = experts.expertsRun; +const copyRun = copies.copyRun; +const denseRun = dense.denseRun; + +pub fn run(gpu: Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const mode = rest[1]; + const layer: usize = if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 17; + var pack: Pack = .{}; + defer pack.close(gpu.io); + try pack.open(gpu, dir); + var s = try cuda.Stream.init(gpu.d, true); + defer s.deinit(); + if (std.mem.eql(u8, mode, "dense")) return denseRun(gpu, &pack, s); + var ex = try exl3e.Experts.load(gpu.d); + defer ex.unload(); + const ld = try LayerData.load(gpu, &pack, layer); + defer ld.free(gpu.gpa); + if (std.mem.eql(u8, mode, "experts")) return expertsRun(gpu, &ex, s, ld); + if (std.mem.eql(u8, mode, "copy")) return copyRun(gpu, &ex, s, ld); + return error.UnknownMode; +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 2d37fc734a..14adcb4c19 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -95,6 +95,12 @@ const lane_half_tests = @import("lane_half_tests.zig"); const flashnext_dense_hc_tests = @import("flashnext_dense_hc_tests.zig"); const flashnext_hc_tests = @import("flashnext_hc_tests.zig"); const flashnext_attn_tests = @import("flashnext_attn_tests.zig"); +const exl3_pack_tests = @import("flashnext_exl3/pack/exl3_pack_tests.zig"); +const exl3_ple2_tests = @import("flashnext_exl3/model/exl3_ple2_tests.zig"); +const exl3_pool_tests = @import("flashnext_exl3/offload/exl3_pool_tests.zig"); +const exl3_prefill_tests = @import("flashnext_exl3/prefill/exl3_prefill_tests.zig"); +const flashnext_exl3_tests = @import("flashnext_exl3/model/flashnext_exl3_tests.zig"); +const flashnext_exl3_engine_tests = @import("flashnext_exl3/model/flashnext_exl3_engine_tests.zig"); const ggml_tests = @import("ggml_tests.zig"); const paged_tests = @import("paged_tests.zig"); const glm_loader_gate = @import("glm_loader_gate.zig"); @@ -135,6 +141,8 @@ const usage = \\ lane-half lane half tiles (MXFP8, BF16): bytes equal to the 64-column tile, FP64, cold-weight time \\ flashnext-dense-hc dense connection weights and FP32 normalization against FP64 \\ flashnext-hc Flash Next's own hyper-connection and n-gram glue against float64 (synthetic rows) + \\ flashnext-exl3 DIR names|norms|router [layers] [rows]|layer [layers] Flash Next on an EXL3 pack vs float64 + \\ flashnext-exl3-engine DIR tf|gen GOLDEN [LINES] | bench [LINES] the family's Engine with the offload pool vs the reference golden \\ flashnext-attn Flash Next's own attention glue, pooling and block scores against float64 (synthetic) \\ flashnext-pages the attention glue's caches through a page table: paged == contiguous, float64, controls \\ info driver, device and library versions @@ -223,6 +231,10 @@ const usage = \\ dense-rows bf16 rows times weight rows (routers, heads) against float64 references (synthetic) \\ exl3 [dir] EXL3 decode-row linear: float64 checks, row invariance, the Python engine's fixture (dir none skips it) \\ exl3-tp2 [dir] EXL3 tensor parallel: rank shards vs the Python engine's, all-gathered (dir none skips them) + \\ exl3-pack DIR [REPS] EXL3 Dense and grouped experts on a real EXL3 pack's shapes: float64 error and timing + \\ exl3-ple2 DIR [ROW] [N] EXL3 n-gram table through cuda_ple.Ngram (ids, rows, RSS) and the PLE kernels vs float64 + \\ exl3-prefill DIR experts|copy [LAYER]|dense EXL3 prefill pieces on the pack: grouped experts by rows, expert-major passes, line copies, prompt Dense + \\ exl3-pool DIR [equal|bench|hit FILES] [LAYERS] Flash Next EXL3 expert offload pool on Experts' pointer tables vs fully resident \\ exl3-tiles [DIR] EXL3's tile candidate: codebooks and tiles bit-equal, FP64, rows through qlinear and Dense \\ train-pack MODEL EXAMPLES [COUNT] training passes holding several sequences against one a pass: differences, time \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) @@ -333,6 +345,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "flashnext-dense-hc")) return flashnext_dense_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-hc")) return flashnext_hc_tests.run(gpu); if (std.mem.eql(u8, cmd, "flashnext-attn")) return flashnext_attn_tests.run(gpu); + if (std.mem.eql(u8, cmd, "flashnext-exl3")) return flashnext_exl3_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "flashnext-exl3-engine")) return flashnext_exl3_engine_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "flashnext-pages")) return @import("contract_fn_pages.zig").run(gpu); if (std.mem.eql(u8, cmd, "info")) return info(gpu); if (std.mem.eql(u8, cmd, "smoke")) return runtime_tests.smoke(gpu); @@ -386,6 +400,10 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exlr8-bench")) return exlr8_tests.bench(gpu, try arg(rest, 0), if (rest.len > 1) rest[1] else "home", if (rest.len > 2) try std.fmt.parseInt(u16, rest[2], 10) else 40, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4); if (std.mem.eql(u8, cmd, "exl3")) return exl3_tests.run(gpu, if (rest.len > 0) rest[0] else "none"); + if (std.mem.eql(u8, cmd, "exl3-pack")) return exl3_pack_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-ple2")) return exl3_ple2_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-pool")) return exl3_pool_tests.run(gpu, rest); + if (std.mem.eql(u8, cmd, "exl3-prefill")) return exl3_prefill_tests.run(gpu, rest); if (std.mem.eql(u8, cmd, "exl3-tiles")) return exl3_candidate_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exl3-tp2")) return exl3_tp2_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else "none"); if (std.mem.eql(u8, cmd, "sample")) return sample_tests.draws(gpu);