EXL3 Flash Next: deep prefill (3/8)

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 13 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_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))

tools/exl3_serve_gate.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+"""EXL3 serve gate: greedy 24 tokens for each golden prompt through tensorfold-native serve, drafts off
+("draft": false) and on (the server's MTP), compared by the reply's token_sha256 (sha256 of the comma-joined
+decimal ids, server/tokens.zig tokenSha256) against the golden's greedy ids and against each other.
+usage: s9_client.py PORT GOLDEN.json [extra known continuations as name=ids,...]"""
+import hashlib, json, sys, time, urllib.request
+
+port, golden = int(sys.argv[1]), json.load(open(sys.argv[2]))
+extra = {}
+for a in sys.argv[3:]:
+    k, v = a.split("=", 1)
+    extra[k] = [int(x) for x in v.split(",") if x]
+
+def sha(ids):
+    return hashlib.sha256(",".join(str(i) for i in ids).encode()).hexdigest()
+
+def ask(ids, draft, n=24):
+    body = {"prompt": ids, "max_tokens": n, "temperature": 0, "ignore_eos": True}
+    if not draft:
+        body["draft"] = False
+    req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/completions", json.dumps(body).encode(), {"Content-Type": "application/json"})
+    t0 = time.time()
+    r = json.load(urllib.request.urlopen(req, timeout=600))
+    r["_wall"] = time.time() - t0
+    return r
+
+ok = True
+report = []
+for p in golden["prompts"]:
+    want = p["greedy_ids"]
+    row = {"name": p["name"]}
+    for draft in (False, True):
+        r = ask(p["ids"], draft)
+        tf, sp = r["tensorfold"], r.get("speculative", {})
+        got = tf["token_sha256"]
+        known = [k for k, v in extra.items() if sha(v) == got and k.startswith(p["name"][:4])]
+        key = "drafts" if draft else "no_drafts"
+        row[key] = {"sha": got[:12], "equals_golden": got == sha(want), "equals": known, "text": r["choices"][0]["text"],
+                    "decode_tok_s": round(tf["tokens_per_second"], 2), "prefill_s": round(tf["prefill_seconds"], 3),
+                    "drafted": sp.get("drafted"), "accepted": sp.get("accepted"), "acceptance_rate": sp.get("acceptance_rate"),
+                    "tokens_per_round": sp.get("tokens_per_round"), "drafts_flag": tf["drafts"]}
+    row["drafted_equals_undrafted"] = row["drafts"]["sha"] == row["no_drafts"]["sha"]
+    row["golden_sha"] = sha(want)[:12]
+    ok &= row["drafted_equals_undrafted"]
+    report.append(row)
+    print(json.dumps(row, indent=1))
+print("DRAFTED==UNDRAFTED on every prompt:", ok)
+# EXL3_GATE_BENCH=N: one longer greedy reply (the 16-id prompt, N tokens), drafts off then on, twice each
+n = int(__import__("os").environ.get("EXL3_GATE_BENCH", "0"))
+if n:
+    ids = golden["prompts"][0]["ids"]
+    for draft in (False, True, False, True):
+        r = ask(ids, draft, n)
+        tf, sp = r["tensorfold"], r.get("speculative", {})
+        print(f"BENCH {n} tokens drafts={draft}: decode {tf['tokens_per_second']:.2f} tok/s ({1000 / tf['tokens_per_second']:.1f} ms/token), "
+              f"prefill {tf['prefill_seconds'] * 1000:.0f} ms, accepted {sp.get('accepted')}/{sp.get('drafted')}, sha {tf['token_sha256'][:12]}")

tools/exl3_serve_tp2.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+# Launch one TP=2 rank of the EXL3 Qwen3.8-Flash-Next server on two RTX 5060 Ti (16 GB), drafts on at every stream count.
+# Usage: tools/exl3_serve_tp2.sh MODEL_DIR RANK PORT MASTER_PORT [extra serve args]; start rank 1 first, then rank 0.
+# The MTP head is loaded by default (rank 0 holds it); `--no-drafts` in the extra arguments serves without it.
+# Environment: LINES (expert lines a layer, default 38), RESERVE_MB (planner reserve), MEMORY_RESERVE_GIB, CONTEXT,
+# PARALLEL (streams, default 2), TENSORFOLD_BIN.
+set -euo pipefail
+model=$1 rank=$2 port=$3 master=$4
+shift 4
+bin=${TENSORFOLD_BIN:-./zig-out/native/bin/tensorfold-native}
+export TENSORFOLD_CHECKPOINT_SHA256=${TENSORFOLD_CHECKPOINT_SHA256:-$(cat "$model/model.safetensors.index.json" "$model/config.json" | sha256sum | cut -c1-64)}
+export TF_FLASHNEXT_EXPERT_LINES=${LINES:-38}
+[ -n "${RESERVE_MB:-}" ] && export TF_FLASHNEXT_EXPERT_RESERVE_MB=$RESERVE_MB
+[ -n "${MEMORY_RESERVE_GIB:-}" ] && export TENSORFOLD_MEMORY_RESERVE_GIB=$MEMORY_RESERVE_GIB
+CUDA_VISIBLE_DEVICES=$rank exec "$bin" serve "$model" --backend cuda --context "${CONTEXT:-32768}" --tp 2 --rank "$rank" \
+  --master 127.0.0.1 --master-port "$master" --port "$port" --no-update-check --parallel "${PARALLEL:-2}" --prompt-cache-gib 0 "$@"