#!/usr/bin/env python3 """Chunk-ALIGNED warm/cold identity gate (Nex-N2.5-mini: unpatched server, no draft head). Written for the Agnes MTP prompt-cache patch; here it checks the stock server's own checkpoint restore for this hybrid (gated-delta + attention) model. Why a third gate: llama-server splits every prompt so that context checkpoints land (4 + n_ubatch) and 4 tokens before its end (tools/server/server-context.cpp `checkpoint_offsets`, upstream PR #20288). A warm request restores a checkpoint whose position was fixed by the length of the request that CREATED it. When those lengths differ, the warm tail is processed in different chunks than a cold run of the same prompt, so float rounding differs and a greedy token can flip - with or without a draft head. cachegate2 mixed prompt lengths and hit exactly that. Here every prompt in the run is padded to ONE token length L, so warm and cold see identical chunking. What is left under test is whether the restored checkpoint (attention KV + recurrent state) is exact. Per variant: warm-up A (cache on) -> B warm (cache on, must restore L-(4+ub)) -> B cold (cache off). PASS = every warm B reused the cache at the aligned position AND is byte-identical to its cold twin, over all variants.""" import argparse, hashlib, json, os, sys sys.path.insert(0, "/mnt/models/nex-n2.5-mini") from nex_harness import Server, post, prompt_8k, THINK_OFF # noqa: E402 (reads AGNES_BIN at import) UB = 1024 # agnes_harness.Server passes -ub 1024 WARM = "Summarise this file." INSTR = ["Write a function that lists every tensor name in this file.", "Write a function that counts the model classes registered in this file.", "Write a function that finds the longest method in this file.", "Write a function that returns every regular expression used in this file.", "Write a function that maps each class in this file to its base classes.", "Write a function that extracts all string constants from this file.", "Write a function that reports which imports in this file are unused.", "Write a function that lists every method that raises an exception in this file.", "Write a function that counts the lines of code per class in this file.", "Write a function that finds duplicate method names across classes in this file."] KW = THINK_OFF def plen(port, content): p = post(port, "/apply-template", {"messages": [{"role": "user", "content": content}], "chat_template_kwargs": KW})["prompt"] return len(post(port, "/tokenize", {"content": p, "add_special": True, "parse_special": True})["tokens"]) def pad(port, ctx, instr, L): s = instr n = plen(port, ctx + "\n\n" + s) for filler in (" ok", ".", " x"): while n < L: t = s + filler m = plen(port, ctx + "\n\n" + t) if m > L: break s, n = t, m if n == L: return s raise SystemExit(f"could not pad {instr!r} to {L} (stuck at {n})") def req(port, ctx, instr, n, cache): body = {"messages": [{"role": "user", "content": ctx + "\n\n" + instr}], "max_tokens": n, "temperature": 0, "top_k": 1, "ignore_eos": True, "cache_prompt": cache, "chat_template_kwargs": KW} r = post(port, "/v1/chat/completions", body) tm = r.get("timings", {}) return (r["choices"][0]["message"].get("content") or ""), tm.get("prompt_n"), tm.get("cache_n"), \ tm.get("draft_n"), tm.get("draft_n_accepted") def first_diff(x, y): if x == y: return None return next((i for i, (p, q) in enumerate(zip(x, y)) if p != q), min(len(x), len(y))) def main(a): s = Server(a, a.port) rows = [] try: ctx, _ = prompt_8k("code") instr = INSTR[:a.reps] warm = [f"{WARM} (warm-up {i})" for i in range(len(instr))] # L over the FULL prompt set, whatever --reps is: runs with different --reps then share byte-identical # prompts, so their outputs can be compared across configs (e.g. MTP vs no draft head). allp = INSTR + [f"{WARM} (warm-up {i})" for i in range(len(INSTR))] L = max(plen(a.port, ctx + "\n\n" + t) for t in allp) + 1 instr = [pad(a.port, ctx, t, L) for t in instr] warm = [pad(a.port, ctx, t, L) for t in warm] expect_cache = L - (4 + UB) for i, (wa, vb) in enumerate(zip(warm, instr)): req(a.port, ctx, wa, 16, True) tw, pw, cw, dw, aw = req(a.port, ctx, vb, 192, True) tc, pc, cc, dc, ac = req(a.port, ctx, vb, 192, False) rows.append({"variant": i, "L": L, "warm_prompt_n": pw, "warm_cache_n": cw, "cold_prompt_n": pc, "cold_cache_n": cc, "aligned": (cw == expect_cache and pc == L), "warm_draft": [aw, dw], "cold_draft": [ac, dc], "identical": tw == tc, "first_diff_char": first_diff(tw, tc), "warm_sha": hashlib.sha256(tw.encode()).hexdigest()[:12], "cold_sha": hashlib.sha256(tc.encode()).hexdigest()[:12], "instr": vb, "warm_instr": wa, "warm_text": tw, "cold_text": tc}) print(" ", json.dumps({k: v for k, v in rows[-1].items() if not k.endswith("_text")}), flush=True) finally: s.stop() reused = all((r["warm_cache_n"] or 0) > 0 for r in rows) aligned = all(r["aligned"] for r in rows) ident = all(r["identical"] for r in rows) res = {"label": a.label, "gate": "cachegate3-aligned", "n": len(rows), "L": rows[0]["L"] if rows else None, "expect_cache_n": expect_cache if rows else None, "all_reused": reused, "all_aligned": aligned, "identical": sum(r["identical"] for r in rows), "rows": rows, "result": "PASS" if (rows and reused and aligned and ident) else "FAIL"} print(json.dumps({k: v for k, v in res.items() if k != "rows"}), flush=True) with open(a.jsonl, "a") as f: f.write(json.dumps(res) + "\n") sys.exit(0 if res["result"] == "PASS" else 1) if __name__ == "__main__": ap = argparse.ArgumentParser() ap.add_argument("--model", required=True); ap.add_argument("--draft") ap.add_argument("--nmax", type=int, default=4); ap.add_argument("--pmin", type=float, default=0.0) ap.add_argument("--dev", default="ROCm0"); ap.add_argument("--ctx", type=int, default=65536) ap.add_argument("--port", type=int, default=18600); ap.add_argument("--reps", type=int, default=10) ap.add_argument("--label", required=True); ap.add_argument("--jsonl", required=True) ap.add_argument("--serverlog", required=True) ap.add_argument("--mtp-infile", action="store_true"); ap.add_argument("--strict", action="store_true") main(ap.parse_args())