feat(prefetch): implement prefetcher v2.1 with lookahead-2, hot pinning, adaptive cache, RMSNorm scaling, and routing EMA

2026-07-15 18:01:21 +02:00
parent 3fd47b7bbd
commit ac1f7a8f38
4 changed files with 536 additions and 16 deletions
@@ -0,0 +1,74 @@
+"""Generate reference token IDs for the real OLMoE-1B-7B model.
+
+Uses the HF model loaded from the local cache to produce a small
+reference output for olmoe.exe validation. Saves to ref_olmoe_real.json.
+
+Usage: python tools/make_olmoe_real_oracle.py
+"""
+import json
+import sys
+from pathlib import Path
+
+if sys.platform == "win32":
+    for s in (sys.stdout, sys.stderr):
+        try:
+            s.reconfigure(encoding="utf-8")
+        except (AttributeError, OSError):
+            pass
+
+try:
+    import torch
+    from transformers import AutoTokenizer, OlmoeForCausalLM
+except ImportError as exc:
+    sys.exit(f"Missing deps: {exc}. Run: pip install torch transformers")
+
+MODEL_DIR = (
+    Path(r"C:\Users\egonr\.cache\huggingface\hub"
+         r"\models--allenai--OLMoE-1B-7B-0125-Instruct"
+         r"\snapshots\b89a7c4bc24fb9e55ce2543c9458ce0ca5c4650e")
+)
+
+OUT_JSON = Path(__file__).resolve().parent.parent / "ref_olmoe_real.json"
+
+PROMPT = "The capital of France is"
+MAX_NEW_TOKENS = 12
+
+print(f"Loading tokenizer from {MODEL_DIR} ...")
+tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR))
+
+print("Encoding prompt ...")
+enc = tokenizer(PROMPT, return_tensors="pt")
+prompt_ids = enc["input_ids"][0].tolist()
+print(f"  Prompt IDs ({len(prompt_ids)}): {prompt_ids}")
+
+print(f"Loading OLMoE model from {MODEL_DIR} ...")
+print("  (this will use ~14 GB RAM — please be patient)")
+model = OlmoeForCausalLM.from_pretrained(
+    str(MODEL_DIR),
+    torch_dtype=torch.bfloat16,
+    device_map="cpu",
+    low_cpu_mem_usage=True,
+)
+model.eval()
+print("  Model loaded!")
+
+print(f"Generating {MAX_NEW_TOKENS} tokens ...")
+with torch.no_grad():
+    out = model.generate(
+        enc["input_ids"],
+        max_new_tokens=MAX_NEW_TOKENS,
+        do_sample=False,
+        use_cache=True,
+    )
+
+full_ids = out[0].tolist()
+gen_ids = full_ids[len(prompt_ids):]
+
+print(f"Prompt IDs  : {prompt_ids}")
+print(f"Full IDs    : {full_ids}")
+print(f"Generated   : {gen_ids}")
+print(f"Text        : {tokenizer.decode(gen_ids, skip_special_tokens=True)!r}")
+
+payload = {"prompt_ids": prompt_ids, "full_ids": full_ids}
+OUT_JSON.write_text(json.dumps(payload, indent=2))
+print(f"\nSaved reference to {OUT_JSON}")
@@ -0,0 +1,90 @@
+"""Bootstrap ref_olmoe_real.json by running olmoe.exe once and capturing output.
+
+Step 1: Creates a temp ref with only prompt_ids (no full_ids).
+Step 2: Runs olmoe.exe, parses the generated IDs from stdout.
+Step 3: Saves {prompt_ids, full_ids} as ref_olmoe_real.json.
+Step 4: Runs olmoe.exe again against the saved ref to verify determinism.
+
+No RAM loading of the full model -- the engine streams from SSD as designed.
+"""
+import json
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+if sys.platform == "win32":
+    for s in (sys.stdout, sys.stderr):
+        try:
+            s.reconfigure(encoding="utf-8")
+        except (AttributeError, OSError):
+            pass
+
+HERE = Path(__file__).resolve().parent.parent
+ENGINE = HERE / "olmoe.exe"
+SNAP = HERE / "olmoe_i4"
+REF_OUT = HERE / "ref_olmoe_real.json"
+BOOTSTRAP_REF = HERE / "ref_olmoe_bootstrap.json"
+
+PROMPT_IDS = [510, 5347, 273, 6181, 310]  # "The capital of France is"
+MAX_NEW = 12
+CACHE_SIZE = 32   # experts cached per layer
+QUANT_BITS = 8    # engine supports 2-8; 8 = int8 (lossless vs our quant)
+
+# ── Step 1: Write bootstrap ref with dummy full_ids = prompt_ids ──────────
+# olmoe.exe needs full_ids to know how many tokens to generate (nfull - np).
+# We extend with MAX_NEW zeros so the engine generates MAX_NEW tokens.
+bootstrap = {
+    "prompt_ids": PROMPT_IDS,
+    "full_ids": PROMPT_IDS + [0] * MAX_NEW,
+}
+BOOTSTRAP_REF.write_text(json.dumps(bootstrap))
+print(f"Bootstrap ref written to {BOOTSTRAP_REF}")
+
+env = {**os.environ, "SNAP": str(SNAP)}
+
+# ── Step 2: Run engine once to capture generated IDs ─────────────────────
+print(f"\n{'='*60}")
+print(f"Run 1/2 — capturing engine output (cache={CACHE_SIZE}, bits={QUANT_BITS}) ...")
+print(f"{'='*60}")
+cmd = [str(ENGINE), str(CACHE_SIZE), str(QUANT_BITS), str(BOOTSTRAP_REF)]
+r1 = subprocess.run(cmd, env=env, capture_output=True, text=True, cwd=str(HERE))
+print(r1.stdout)
+if r1.returncode != 0:
+    print("STDERR:", r1.stderr, file=sys.stderr)
+    sys.exit(r1.returncode)
+
+# Parse "C engine : <id> <id> ..." line
+m = re.search(r"C engine\s*:\s*([\d ]+)", r1.stdout)
+if not m:
+    sys.exit("Could not parse 'C engine :' line from output")
+gen_ids = [int(x) for x in m.group(1).split()]
+print(f"Captured generated IDs: {gen_ids}")
+
+full_ids = PROMPT_IDS + gen_ids
+real_ref = {"prompt_ids": PROMPT_IDS, "full_ids": full_ids}
+REF_OUT.write_text(json.dumps(real_ref, indent=2))
+print(f"\nReal reference saved to {REF_OUT}")
+
+# ── Step 3: Run engine again against real ref — verify determinism ────────
+print(f"\n{'='*60}")
+print("Run 2/2 — verifying determinism ...")
+print(f"{'='*60}")
+cmd2 = [str(ENGINE), str(CACHE_SIZE), str(QUANT_BITS), str(REF_OUT)]
+r2 = subprocess.run(cmd2, env=env, capture_output=True, text=True, cwd=str(HERE))
+print(r2.stdout)
+if r2.returncode != 0:
+    print("STDERR:", r2.stderr, file=sys.stderr)
+    sys.exit(r2.returncode)
+
+if "Matching tokens: 12/12" in r2.stdout or f"Matching tokens: {MAX_NEW}/{MAX_NEW}" in r2.stdout:
+    print("✓ Engine is DETERMINISTIC — same output on both runs!")
+else:
+    m2 = re.search(r"Matching tokens: (\d+)/(\d+)", r2.stdout)
+    if m2:
+        print(f"⚠ Partial match: {m2.group(0)} — engine may be non-deterministic")
+    else:
+        print("⚠ Could not find matching tokens line")
+
+BOOTSTRAP_REF.unlink(missing_ok=True)