EXL3 Flash Next: draft lane guard for offloaded experts (8/8)

docs/flash-next-exl3-graphs.md

@@ -0,0 +1,32 @@
+# Flash Next EXL3: CUDA graphs and the device planner
+
+With the expert offload pool, a decode round of Flash Next EXL3 plans its experts on the device
+(`flashnext_exl3_dplan.cu`) instead of waiting for the host, and can replay the whole round as a CUDA graph. Acceptance
+stays in the core lane engine; the device code proposes and plans.
+
+## Switches
+
+| Variable | Default | Meaning |
+| --- | --- | --- |
+| `TF_FLASHNEXT_EXL3_GRAPH` | `0` | `1` plans decode rounds on the device and replays graphs for decode rounds, verify windows, the MTP draft chain and prefill layers. Output is identical to `0`. |
+| `TF_FLASHNEXT_EXL3_DPLAN` | `0` | `1` plans on the device without graphs (each planned round launches kernel by kernel). Implied by `GRAPH=1`. This is the eager arm of an eager-vs-graph comparison. |
+| `TF_FLASHNEXT_EXL3_GRAPH_ASYNC` | `1` | `0` instantiates a captured graph on the step's thread; `1` on a worker thread while that step runs eagerly. |
+
+Everything else is fixed in the code: a shape is captured at its second sighting, a device-planned round holds at most
+16 rows (the MTP window), and the graph memory model (reservation, headroom, per-graph charges) follows the stream
+count. `tools/exl3_serve_tp2.sh` sets `GRAPH=1` and `GRAPH_ASYNC=0`.
+
+## A graph that cannot be built
+
+A capture or an instantiation that fails runs its round eagerly and the process keeps its other graphs. The failed
+shape backs off (64, 128, 256 rounds) and is tried at most three more times before it stays eager. The prefill layer
+shape and the draft chain do the same.
+
+On two ranks the capture outcome is agreed over the rank wire (`graphAgree`) before either rank launches: a graph is
+kept only when both ranks hold one, otherwise both run the round eagerly, so the communicator is never left
+half-launched. A communicator error (a lost peer, a timeout) still ends the step, as before.
+
+## State the engine owns
+
+The device planner, the plan active during a round and the MTP head's pool belong to the engine
+(`Engine.exl3_dplan`), so two engines in one process share none of them.

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 launches (`flashnext_exl3_pending.zig`) carry only an nvcc 13.4 receipt: they stay pending under
+13.0 to 13.3 and take their 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.

docs/multi-rank-receipts.md

@@ -13,7 +13,8 @@ Each path places its ranks its own way:
 | Native ROCm build | `--tp`, `--rank`, `--master` | [Native ROCm placement](../zig/src/native/hip.zig) |
 | The `fabric/tp` transport's harness | `TF_TP_WORLD`, `TF_TP_RANK`, `TF_TP_MASTER` | [The transport's README](../zig/src/fabric/tp/README.md) |
 
-The native CUDA build serves neither `--tp` nor `--rank`: a family's own variables place its ranks. `tf-tp-test` is
+The native CUDA build serves neither `--tp` nor `--rank` for a checkpoint other than a Flash Next EXL3 pack, which they
+place across two GPUs (`--tp 2 --rank N --master HOST`): a family's own variables place its ranks. `tf-tp-test` is
 the transport's harness rather than a server. `tensorfold benchmark` accepts no group placement flags and records 1
 for a managed run. Its child inherits the parent's rank environment, so start managed runs from a single-rank
 configuration. Measure a group as an attached server: start the group, then point the client at the rank that serves.

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)

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__)

tools/exl3_nsys_step_breakdown.py

@@ -0,0 +1,41 @@
+"""Per-decode-step kernel time by category from an nsys SQLite export of a rank (counts fn_exl3_dplan launches / 48 as steps). usage: exl3_nsys_step_breakdown.py EXPORT_PREFIX."""
+import sqlite3, collections, re, subprocess, os, sys
+f = sys.argv[1]  # the export's path without its .sqlite / .nsys-rep suffix
+if not os.path.exists(f + ".sqlite"):
+    subprocess.run(["nsys", "stats", "--report", "cuda_gpu_kern_sum", "-o", "/dev/null", f + ".nsys-rep"], capture_output=True)
+c = sqlite3.connect(f + ".sqlite")
+names = dict(c.execute("select id,value from StringIds"))
+rows = c.execute("select start,end,demangledName from CUPTI_ACTIVITY_KIND_KERNEL order by start").fetchall()
+dp = [r for r in rows if "fn_exl3_dplan" in names[r[2]]]
+steps = len(dp) // 48
+t0, t1 = dp[0][0], dp[-1][1]
+print("dplan launches", len(dp), "decode steps", steps, "window s", round((t1 - t0) / 1e9, 2), "ms/step wall", round((t1 - t0) / 1e6 / steps, 2))
+def cat(n):
+    if "AllReduce" in n or "AllGather" in n: return "nccl all-reduce/gather"
+    if "fn_exl3_dcopy" in n: return "expert miss copy (side stream)"
+    if "tf_exl3_linear" in n or "tf_exl3_grouped" in n or "tf_exl3_group" in n or "tf_exl3_gateup" in n or "tf_exl3_down" in n or "tf_exl3_rot" in n: return "expert/dense exl3 matmuls"
+    if "qmmf" in n: return "qmmf matmuls"
+    if "fn_gdn" in n: return "DeltaNet (gdn)"
+    if "fn_hc_" in n: return "hyper-connection"
+    if "attn" in n or "fn_qsa" in n: return "attention"
+    if "tf_route" in n or "router" in n: return "router"
+    if "dplan" in n: return "device planner"
+    return "other"
+tot = collections.Counter(); cnt = collections.Counter(); other = collections.Counter()
+for s, e, n in rows:
+    if s < t0 or e > t1: continue
+    nm = names[n]; k = cat(nm); tot[k] += e - s; cnt[k] += 1
+    if k == "other": other[nm[:70]] += e - s
+for k, v in tot.most_common():
+    print("%-34s %7.2f ms/step  %5d launches/step" % (k, v / 1e6 / steps, cnt[k] // steps))
+print("sum of kernel time (all streams) ms/step", round(sum(tot.values()) / 1e6 / steps, 2))
+print("top 'other':")
+for k, v in other.most_common(8): print("  %6.2f ms/step %s" % (v / 1e6 / steps, k))
+# main stream busy fraction: union over all non-dcopy kernels
+iv = sorted((s, e) for s, e, n in rows if t0 <= s and e <= t1 and "dcopy" not in names[n])
+busy = 0; cs, ce = iv[0]
+for s, e in iv[1:]:
+    if s > ce: busy += ce - cs; cs, ce = s, e
+    else: ce = max(ce, e)
+busy += ce - cs
+print("busy fraction (non-copy kernels)", round(busy / (t1 - t0), 3))

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("<Q", h.read(8))[0]; hdr = json.loads(h.read(n))
+    for name, m in hdr.items():
+        if name == "__metadata__": continue
+        a, b = m["data_offsets"]; g[cat(name)] += b - a
+tot = sum(g.values())
+for k, v in g.most_common(): print("%8.2f GB %5.1f%%  %s" % (v / 1e9, 100 * v / tot, k))