spec-decode for BitNet: subNorm+f32-KV batched verify, ?spec + ?bench=spec

index.html

@@ -50,6 +50,9 @@ let m = pick ? (MODELS.find((x) => new RegExp(pick, "i").test(x.name)) || MODELS
 }
 // ?verify=gpu → re-derive each BLAKE3 weight-block κ ENTIRELY on the GPU (2.74 GB/s) instead of pure-JS BLAKE3.
 if (params.get("verify") === "gpu") globalThis.__gpuVerify = true;
+// ?spec → speculative decode (n-gram draft + batched-K verify). Byte-identical to greedy; big wins on echo-heavy
+// text (code/quote/retrieval), no gain on free-form chat. ?bench=spec runs the A/B measurement harness after load.
+if (params.get("spec") || params.get("bench") === "spec") globalThis.__spec = true;
 
 // GROUND the model as on-device Q (a base/instruct model has NO self-knowledge — without this it confabulates
 // a generic "I run on OpenAI/AWS cloud servers" identity, which is false). Injected as the SYSTEM turn.
@@ -129,6 +132,52 @@ send.onclick = onSend;
 input.onkeydown = (e) => { if (e.key === "Enter" && !e.shiftKey) { e.preventDefault(); onSend(); } };
 input.oninput = () => { input.style.height = "auto"; input.style.height = Math.min(140, input.scrollHeight) + "px"; };
 
+// ── SPEC-DECODE A/B HARNESS (?bench=spec) ── measures baseline greedy vs speculative on the operator's real
+// GPU across echo-heavy and free-form prompts: byte-identical check (G1), mean accepted tokens/verify (G2),
+// and decode tok/s for both. One load, one table — the honest verdict on whether spec-decode earns its place.
+const SPEC_BENCH = [
+  { tag: "code / edit (echo-heavy)", text: "Here is a function:\n\nfunction add(a, b) {\n  return a + b;\n}\n\nRewrite it exactly the same but rename add to sum." },
+  { tag: "retrieval / quote", text: "Passage: \"The quick brown fox jumps over the lazy dog near the river bank at dawn.\" Repeat that passage back to me word for word." },
+  { tag: "free-form chat", text: "In one short sentence, why is the sky blue?" },
+];
+const eqArr = (a, b) => a.length === b.length && a.every((x, i) => x === b[i]);
+async function runSpecBench() {
+  log.innerHTML = ""; input.disabled = send.disabled = true;
+  if (!engine.specAvailable) { st.textContent = "spec-decode unavailable for this model"; bubble("a", "This model can't use the batched-verify head (specAvailable=false)."); return; }
+  const N = 96, rep = m.rep ?? 1.3, rows = [];
+  let prev = { windows: 0, drafted: 0, accepted: 0 };
+  for (const b of SPEC_BENCH) {
+    st.textContent = `bench: ${b.tag}…`;
+    const ids = engine.tokenize(frameSystem() + engine.frameTurn(b.text, false));
+    globalThis.__spec = false; engine.reset();
+    const t0 = performance.now(); const r0 = await engine.generate(ids.slice(), { maxNew: N, repPenalty: rep }); const w0 = performance.now() - t0;
+    globalThis.__spec = true; engine.reset();
+    const t1 = performance.now(); const r1 = await engine.generate(ids.slice(), { maxNew: N, repPenalty: rep }); const w1 = performance.now() - t1;
+    globalThis.__spec = false;
+    const cur = (r1.stats && r1.stats.spec) || prev;
+    const dd = { windows: cur.windows - prev.windows, drafted: cur.drafted - prev.drafted, accepted: cur.accepted - prev.accepted }; prev = { windows: cur.windows, drafted: cur.drafted, accepted: cur.accepted };
+    rows.push({
+      tag: b.tag, same: eqArr(r0.outIds, r1.outIds), nB: r0.outIds.length, nS: r1.outIds.length,
+      baseTok: r0.outIds.length / (w0 / 1000), specTok: r1.outIds.length / (w1 / 1000),
+      perVerify: dd.windows ? 1 + dd.accepted / dd.windows : 0, accept: dd.drafted ? dd.accepted / dd.drafted : 0, windows: dd.windows,
+    });
+  }
+  const allSame = rows.every((r) => r.same);
+  const fmt = (x) => x.toFixed(x < 10 ? 1 : 0);
+  const tbl = `<div style="font-family:ui-monospace,monospace;font-size:13px;max-width:900px;margin:0 auto;padding:8px">
+    <div style="font-size:18px;font-weight:700;margin-bottom:4px">Speculative decode — measured on your GPU</div>
+    <div style="color:var(--dim);margin-bottom:12px">BitNet-2B · n-gram draft + batched-K verify · greedy, byte-exact by construction</div>
+    <table style="width:100%;border-collapse:collapse">
+      <tr style="color:var(--dim);text-align:left"><th style="padding:6px 8px">workload</th><th style="padding:6px 8px;text-align:right">baseline</th><th style="padding:6px 8px;text-align:right">spec</th><th style="padding:6px 8px;text-align:right">speedup</th><th style="padding:6px 8px;text-align:right">tok/verify</th><th style="padding:6px 8px;text-align:right">accept</th><th style="padding:6px 8px;text-align:right">byte-exact</th></tr>
+      ${rows.map((r) => `<tr style="border-top:1px solid var(--line)"><td style="padding:6px 8px">${r.tag}</td><td style="padding:6px 8px;text-align:right">${fmt(r.baseTok)} tok/s</td><td style="padding:6px 8px;text-align:right">${fmt(r.specTok)} tok/s</td><td style="padding:6px 8px;text-align:right;color:${r.specTok > r.baseTok * 1.05 ? "#48c26c" : r.specTok < r.baseTok * 0.95 ? "#f0616d" : "var(--dim)"}">${(r.specTok / r.baseTok).toFixed(2)}×</td><td style="padding:6px 8px;text-align:right">${r.perVerify.toFixed(2)}</td><td style="padding:6px 8px;text-align:right">${(r.accept * 100).toFixed(0)}%</td><td style="padding:6px 8px;text-align:right;color:${r.same ? "#48c26c" : "#f0616d"}">${r.same ? "✓ identical" : "✗ DIVERGED"}</td></tr>`).join("")}
+    </table>
+    <div style="margin-top:14px;font-weight:600;color:${allSame ? "#48c26c" : "#f0616d"}">${allSame ? "✓ G1 PASS — spec output is byte-identical to greedy on every prompt." : "✗ G1 FAIL — spec diverged from greedy; not shippable until fixed (see console)."}</div>
+  </div>`;
+  log.innerHTML = tbl;
+  st.textContent = "spec-decode bench · done";
+  console.log("[specbench]", rows);
+}
+
 try {
   if (!navigator.gpu) throw new Error("This browser has no WebGPU — open in Chrome, Edge, or a recent mobile browser.");
   st.textContent = `loading ${m.name} (${m.size})…`;
@@ -141,7 +190,8 @@ try {
   armed = true;
   st.textContent = `${m.name} · ${m.size} · in your browser · ready`;
   input.placeholder = "Message Q…";
-  if (pending) { const w = [...log.querySelectorAll(".a")].reverse().find((x) => x.dataset.pending); if (w) w.remove(); const p = pending; pending = null; generate(p, true); }
+  if (params.get("bench") === "spec") { globalThis.__spec = false; await runSpecBench(); }
+  else if (pending) { const w = [...log.querySelectorAll(".a")].reverse().find((x) => x.dataset.pending); if (w) w.remove(); const p = pending; pending = null; generate(p, true); }
   else await proactiveGreeting();
 } catch (e) { st.textContent = "⚠ " + e.message; bubble("a", "Could not start: " + e.message); }
 </script></body></html>