File size: 3,747 Bytes
7f02c5a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | #!/usr/bin/env python3
"""Verify the A16W8 trunk (prefill) context binary on real hardware.
The trunk is stateless: one forward per prompt, no KV. So this pushes the per-prompt inputs,
runs qnn-net-run once, pulls the hidden state, and compares against the float reference
(trunk_ref.npz, produced on the AWS box by trunk_ref.py) on:
- cosine of the last-real-position hidden vs float
- the argmax next token through the tied softcapped lm_head
Usage:
python verify_trunk.py --adb-serial <serial> --ref /path/to/trunk_ref.npz
"""
import argparse, pathlib, subprocess, sys, numpy as np
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
import hostlib
BASE = "/data/local/tmp/gemma"
STEP = f"{BASE}/tstep"
OUTR = f"{BASE}/tout/Result_0"
SEQ = 128
NEG = hostlib.NEG
L_ = pathlib.Path("/tmp/gemma_trunk"); L_.mkdir(exist_ok=True)
def adb(a, serial, **kw):
return subprocess.run(["adb"] + (["-s", serial] if serial else []) + a,
capture_output=True, text=True, **kw)
def cmask(L):
j = np.arange(SEQ)[None, :]
i = np.arange(SEQ)[:, None]
return np.where((j <= i) & (j < L), 0.0, NEG).astype(np.float32).reshape(1, 1, SEQ, SEQ)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--adb-serial", required=True)
ap.add_argument("--ref", default="/home/azek/.claude/jobs/db9a6094/tmp/trunk_ref.npz")
args = ap.parse_args()
ref = np.load(args.ref)
n = int(ref["n"][0])
m = hostlib.HostModel()
adb(["shell", f"mkdir -p {STEP} {BASE}/tout"], args.adb_serial)
ok = 0
coss = []
for k in range(n):
ids = [int(x) for x in ref[f"ids_{k}"]]
L = len(ids)
hf = ref[f"h_{k}"]
tf = int(ref[f"tok_{k}"][0])
# host embeddings for the whole padded window
padded = ids + [0] * (SEQ - L)
ie = np.concatenate([m.embeds(t)[0].reshape(1, 1, hostlib.H) for t in padded], axis=1)
ple = np.concatenate([m.embeds(t)[1].reshape(1, 1, hostlib.NL, hostlib.PLD) for t in padded], axis=1)
mk = cmask(L)
files = {
"inputs_embeds": ie.astype(np.float32),
"per_layer_inputs": ple.astype(np.float32),
"position_ids": np.arange(SEQ, dtype=np.int32).reshape(1, SEQ),
"full_mask": mk,
"sliding_mask": mk.copy(),
}
for name, arr in files.items():
p = L_ / f"{name}.raw"
arr.tofile(p)
r = adb(["push", str(p), f"{STEP}/{name}.raw"], args.adb_serial, timeout=300)
if r.returncode != 0:
raise RuntimeError(f"push {name}: {r.stderr}")
r = adb(["shell", f"sh {BASE}/gate_ondevice_trunk.sh"], args.adb_serial, timeout=600)
if "TRUNK_OK" not in r.stdout:
raise RuntimeError(f"trunk net-run failed:\n{r.stdout}\n{r.stderr}")
adb(["pull", f"{OUTR}/hidden.raw", str(L_ / "hidden.raw")], args.adb_serial, timeout=300)
hd = np.fromfile(L_ / "hidden.raw", np.float32).reshape(SEQ, hostlib.H)[L - 1]
cos = float(hd @ hf / (np.linalg.norm(hd) * np.linalg.norm(hf) + 1e-9))
coss.append(cos)
td = int(m.logits(hd).argmax())
hit = (td == tf)
ok += hit
print(f" [{k:2d}] cos={cos:.5f} |hw|={np.linalg.norm(hd):7.2f} |float|={np.linalg.norm(hf):7.2f} "
f"hw={td:6d} {m.decode([td])!r:16s} float={tf:6d} {m.decode([tf])!r:16s} "
f"{'OK' if hit else 'MISMATCH'}", flush=True)
print(f"\n hidden cos: mean={np.mean(coss):.5f} min={np.min(coss):.5f}", flush=True)
print(f" NEXT-TOKEN TOP-1 vs float on HARDWARE: {ok}/{n} ({100.0*ok/n:.1f}%)", flush=True)
print("TRUNK_VERIFY_DONE", flush=True)
if __name__ == "__main__":
main()
|