| <!DOCTYPE html> |
| <html lang="en"> |
| <head> |
| <meta charset="utf-8"> |
| <meta name="viewport" content="width=device-width, initial-scale=1"> |
| <title>Real speculative decoding, in your browser</title> |
| <meta name="description" content="distilgpt2 drafts, gpt2 verifies — real rejection sampling with per-token acceptance math, running fully client-side via transformers.js."> |
| <style> |
| :root{ |
| color-scheme:light; |
| --surface:#fcfcfb; --page:#f9f9f7; --ink:#0b0b0b; --ink2:#52514e; |
| --muted:#898781; --grid:#e1e0d9; --axis:#c3c2b7; --border:rgba(11,11,11,.10); |
| --target:#2a78d6; --draft:#008300; --residual:#e87ba4; |
| --accept:#0ca30c; --reject:#d03b3b; --resample:#eb6834; --bonus:#2a78d6; |
| } |
| @media (prefers-color-scheme:dark){ |
| :root:not([data-theme="light"]){ |
| color-scheme:dark; |
| --surface:#1a1a19; --page:#0d0d0d; --ink:#ffffff; --ink2:#c3c2b7; |
| --muted:#898781; --grid:#2c2c2a; --axis:#383835; --border:rgba(255,255,255,.10); |
| --target:#3987e5; --draft:#008300; --residual:#d55181; |
| --accept:#0ca30c; --reject:#d03b3b; --resample:#d95926; --bonus:#3987e5; |
| } |
| } |
| :root[data-theme="dark"]{ |
| color-scheme:dark; |
| --surface:#1a1a19; --page:#0d0d0d; --ink:#ffffff; --ink2:#c3c2b7; |
| --muted:#898781; --grid:#2c2c2a; --axis:#383835; --border:rgba(255,255,255,.10); |
| --target:#3987e5; --draft:#008300; --residual:#d55181; |
| --accept:#0ca30c; --reject:#d03b3b; --resample:#d95926; --bonus:#3987e5; |
| } |
| *{box-sizing:border-box} |
| body{margin:0;background:var(--page);color:var(--ink); |
| font-family:system-ui,-apple-system,"Segoe UI",sans-serif;line-height:1.55} |
| main{max-width:980px;margin:0 auto;padding:24px 16px 80px} |
| h1{font-size:1.6rem;margin:.2em 0} |
| p{color:var(--ink2);max-width:75ch} |
| a{color:var(--target)} |
| code{font-family:ui-monospace,SFMono-Regular,Menlo,monospace;font-size:.92em} |
| .card{background:var(--surface);border:1px solid var(--border);border-radius:12px; |
| padding:16px;margin:12px 0} |
| button{background:var(--target);color:#fff;border:0;border-radius:8px; |
| padding:8px 14px;font:inherit;font-weight:600;cursor:pointer} |
| button:disabled{opacity:.5;cursor:default} |
| input[type=text],select{font:inherit;background:var(--surface);color:var(--ink); |
| border:1px solid var(--axis);border-radius:8px;padding:6px 8px} |
| input[type=range]{accent-color:var(--target)} |
| label{font-size:.82rem;color:var(--ink2)} |
| .stat{display:inline-block;margin-right:26px} |
| .stat b{font-size:1.5rem;font-weight:700} |
| .stat span{display:block;font-size:.76rem;color:var(--muted)} |
| .chip{display:inline-block;padding:2px 6px;margin:2px;border-radius:4px;color:#fff; |
| font-family:ui-monospace,Menlo,monospace;font-size:.9em} |
| .legend{font-size:.82rem;color:var(--ink2);margin:6px 0} |
| .legend i{display:inline-block;width:10px;height:10px;border-radius:2px; |
| margin:0 4px 0 14px;vertical-align:baseline} |
| .poscard{border:1px solid var(--grid);border-radius:8px;padding:10px 12px;margin:8px 0} |
| .poscard .meta{color:var(--muted);font-size:12px} |
| .bars{display:flex;gap:16px;flex-wrap:wrap;margin-top:8px} |
| .bars>div{flex:1;min-width:230px} |
| .brow{display:flex;align-items:center;gap:6px;margin:2px 0} |
| .brow code{width:92px;text-align:right;overflow:hidden;color:var(--ink2)} |
| .btrack{flex:1;background:var(--grid);border-radius:4px;height:13px} |
| .bfill{height:13px;border-radius:4px} |
| .brow span{width:52px;color:var(--ink2);font-size:12px} |
| #status{font-size:.85rem;color:var(--muted)} |
| progress{width:220px;vertical-align:middle} |
| .warn{color:var(--resample);font-size:.85rem} |
| </style> |
| </head> |
| <body> |
| <main> |
| <h1>Real speculative decoding — running in your browser</h1> |
| <p><b style="color:var(--draft)">distilgpt2</b> drafts γ tokens, <b style="color:var(--target)">gpt2</b> |
| verifies them — quantized ONNX models executed client-side with |
| <a href="https://huggingface.co/docs/transformers.js">transformers.js</a> |
| (WebGPU when available, WASM otherwise). No server, no GPU quota, no cold start: |
| the rejection sampling below is the actual algorithm on actual model logits. |
| <a href="index.html">← back to the interactive explainer</a></p> |
|
|
| <div class="card"> |
| <button id="loadBtn">Load models (~215 MB, cached after first visit)</button> |
| <span id="status"></span> |
| <div id="loadWarn" class="warn"></div> |
| </div> |
|
|
| <div class="card" id="controls" style="display:none"> |
| <div style="display:flex;gap:12px;flex-wrap:wrap;align-items:end"> |
| <div style="flex:2;min-width:260px"><label>Prompt</label><br> |
| <input type="text" id="prompt" style="width:100%" |
| value="The key idea behind speculative decoding is"></div> |
| <div><label>γ = <span id="gammaV">4</span></label><br> |
| <input type="range" id="gamma" min="1" max="8" step="1" value="4"></div> |
| <div><label>temperature = <span id="tempV">0.8</span></label><br> |
| <input type="range" id="temp" min="0" max="1.5" step="0.1" value="0.8"></div> |
| <div><label>max new = <span id="maxV">24</span></label><br> |
| <input type="range" id="maxNew" min="8" max="48" step="4" value="24"></div> |
| <div><label>seed</label><br> |
| <input type="text" id="seed" value="0" style="width:64px"></div> |
| <div><button id="runBtn">Run</button></div> |
| </div> |
| </div> |
|
|
| <div class="card" id="results" style="display:none"> |
| <div id="stats"></div> |
| <div class="legend"> |
| <i style="background:var(--accept);margin-left:0"></i>accepted |
| <i style="background:var(--reject)"></i>rejected |
| <i style="background:var(--resample)"></i>resampled |
| <i style="background:var(--bonus)"></i>bonus |
| </div> |
| <div id="ribbon" style="line-height:2"></div> |
| <h3 style="margin:14px 0 4px">Per-position acceptance math</h3> |
| <div id="cards"></div> |
| </div> |
|
|
| <p class="hint" style="font-size:.8rem;color:var(--muted)"> |
| Everything stays on your device — the models are fetched from the Hugging Face |
| CDN and executed locally. Source: |
| <a href="https://github.com/aabhimittal/attention-trace-differ-for-speculative-decoding">GitHub</a>. |
| The repo also ships a Gradio app (with KV caching, attention-trace diffing and |
| a 70B cloud cross-check) you can run locally with <code>python app.py</code>.</p> |
| </main> |
|
|
| <script type="module"> |
| import { AutoTokenizer, AutoModelForCausalLM, Tensor, env } |
| from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.5.2"; |
| |
| |
| if (new URLSearchParams(location.search).get("local")) { |
| env.allowRemoteModels = false; |
| env.allowLocalModels = true; |
| env.localModelPath = "./models/"; |
| } |
| |
| |
| const MODEL_OPTS = { model_file_name: "decoder_model_merged", dtype: "q8" }; |
| |
| const $ = id => document.getElementById(id); |
| const esc = s => s.replace(/&/g,"&").replace(/</g,"<").replace(/>/g,">"); |
| const vis = t => t.trim() === "" ? t.replace(/ /g,"·") || "·" : t; |
| ["gamma","temp","maxNew"].forEach((id,i) => |
| $(id).oninput = () => $(["gammaV","tempV","maxV"][i]).textContent = $(id).value); |
| |
| let tok, draft, target, VOCAB; |
| |
| |
| function rng(seed){ let a = seed >>> 0; |
| return () => { a |= 0; a = a + 0x6D2B79F5 | 0; |
| let t = Math.imul(a ^ a >>> 15, 1 | a); |
| t = t + Math.imul(t ^ t >>> 7, 61 | t) ^ t; |
| return ((t ^ t >>> 14) >>> 0) / 4294967296; }; } |
| |
| function softmax(row, temperature){ |
| const n = row.length, out = new Float32Array(n); |
| if (temperature <= 0){ |
| let mi = 0; for (let i = 1; i < n; i++) if (row[i] > row[mi]) mi = i; |
| out[mi] = 1; return out; } |
| let mx = -Infinity; |
| for (let i = 0; i < n; i++){ const v = row[i]/temperature; out[i] = v; if (v > mx) mx = v; } |
| let s = 0; |
| for (let i = 0; i < n; i++){ out[i] = Math.exp(out[i]-mx); s += out[i]; } |
| for (let i = 0; i < n; i++) out[i] /= s; |
| return out; |
| } |
| function sampleFrom(p, r){ |
| let u = r(), c = 0; |
| for (let i = 0; i < p.length; i++){ c += p[i]; if (u < c) return i; } |
| return p.length - 1; |
| } |
| function top8(p){ |
| const idx = []; |
| for (let i = 0; i < p.length; i++){ |
| if (idx.length < 8){ idx.push(i); idx.sort((a,b)=>p[b]-p[a]); } |
| else if (p[i] > p[idx[7]]){ idx[7] = i; idx.sort((a,b)=>p[b]-p[a]); } } |
| return idx.map(i => ({ id: i, prob: p[i] })); |
| } |
| async function forward(model, ids){ |
| const t = new Tensor("int64", BigInt64Array.from(ids.map(BigInt)), [1, ids.length]); |
| const mask = new Tensor("int64", BigInt64Array.from(ids.map(()=>1n)), [1, ids.length]); |
| const out = await model({ input_ids: t, attention_mask: mask }); |
| return out.logits; |
| } |
| const rowOf = (logits, r) => |
| logits.data.subarray(r * VOCAB, (r + 1) * VOCAB); |
| |
| $("loadBtn").onclick = async () => { |
| $("loadBtn").disabled = true; |
| const status = p => $("status").textContent = p; |
| const progress = d => { if (d.status === "progress" && d.file?.endsWith(".onnx")) |
| status(`${d.file.split("/").pop()} ${Math.round(d.progress||0)}%`); }; |
| let device = "wasm"; |
| try { if (navigator.gpu && await navigator.gpu.requestAdapter()) device = "webgpu"; } |
| catch { } |
| try { |
| status("loading tokenizer…"); |
| tok = await AutoTokenizer.from_pretrained("Xenova/distilgpt2"); |
| status("loading draft (distilgpt2)…"); |
| draft = await AutoModelForCausalLM.from_pretrained("Xenova/distilgpt2", |
| { ...MODEL_OPTS, device, progress_callback: progress }); |
| status("loading target (gpt2)…"); |
| target = await AutoModelForCausalLM.from_pretrained("Xenova/gpt2", |
| { ...MODEL_OPTS, device, progress_callback: progress }); |
| } catch (e) { |
| if (device === "webgpu"){ |
| $("loadWarn").textContent = "WebGPU load failed, retrying on WASM…"; |
| device = "wasm"; |
| draft = await AutoModelForCausalLM.from_pretrained("Xenova/distilgpt2", |
| { ...MODEL_OPTS, device, progress_callback: progress }); |
| target = await AutoModelForCausalLM.from_pretrained("Xenova/gpt2", |
| { ...MODEL_OPTS, device, progress_callback: progress }); |
| } else { status("load failed: " + e.message); $("loadBtn").disabled = false; return; } |
| } |
| status(`ready · running on ${device}`); |
| $("controls").style.display = "block"; |
| }; |
| |
| function bars(top, cls, hlId){ |
| return top.map(({id, prob}) => |
| `<div class="brow"><code style="${id===hlId?'font-weight:700':''}">${esc(vis(tok.decode([id])))}</code>` + |
| `<div class="btrack"><div class="bfill" style="width:${Math.max(prob*100,.5).toFixed(1)}%;` + |
| `background:var(--${cls})"></div></div><span>${prob.toFixed(3)}</span></div>`).join(""); |
| } |
| |
| $("runBtn").onclick = async () => { |
| $("runBtn").disabled = true; |
| $("results").style.display = "block"; |
| $("stats").innerHTML = ""; $("ribbon").innerHTML = ""; $("cards").innerHTML = ""; |
| const gamma = +$("gamma").value, temp = +$("temp").value, |
| maxNew = +$("maxNew").value, r = rng(+$("seed").value || 0); |
| const enc = await tok($("prompt").value); |
| let ids = Array.from(enc.input_ids.data, Number); |
| const promptLen = ids.length; |
| VOCAB = null; |
| let proposed = 0, acceptedN = 0, rounds = 0, t0 = performance.now(); |
| const chip = (txt, color, extra="") => |
| `<span class="chip" style="background:var(--${color});${extra}" >${esc(vis(txt))}</span>`; |
| |
| while (ids.length - promptLen < maxNew) { |
| rounds++; |
| const roundStart = ids.length; |
| |
| let draftIds = ids.slice(); const qDists = []; |
| for (let k = 0; k < gamma; k++){ |
| const logits = await forward(draft, draftIds); |
| VOCAB ??= logits.dims[2]; |
| const q = softmax(rowOf(logits, draftIds.length - 1), temp); |
| qDists.push(q); |
| draftIds.push(sampleFrom(q, r)); |
| $("status").textContent = `round ${rounds}: drafting ${k+1}/${gamma}…`; |
| } |
| |
| $("status").textContent = `round ${rounds}: target verifying…`; |
| const tLogits = await forward(target, draftIds); |
| const pDist = s => softmax(rowOf(tLogits, ids.length - 1 + s), temp); |
| |
| const props = draftIds.slice(ids.length); |
| let rejected = false; |
| for (let i = 0; i < props.length; i++){ |
| const x = props[i], q = qDists[i], p = pDist(i); |
| const qx = q[x], px = p[x]; |
| const ratio = qx > 0 ? px / qx : Infinity; |
| const aProb = Math.min(1, ratio); |
| const u = r(); |
| const ok = u < aProb; |
| proposed++; |
| const tokStr = tok.decode([x]); |
| let cardTail = "", ribbonAdd = ""; |
| if (ok){ |
| acceptedN++; |
| ribbonAdd = chip(tokStr, "accept"); |
| } else { |
| |
| const res = new Float32Array(VOCAB); let s = 0; |
| for (let j = 0; j < VOCAB; j++){ const v = Math.max(p[j]-q[j], 0); res[j] = v; s += v; } |
| let newTok; |
| if (s > 0){ for (let j = 0; j < VOCAB; j++) res[j] /= s; newTok = sampleFrom(res, r); } |
| else newTok = sampleFrom(p, r); |
| const newStr = tok.decode([newTok]); |
| ribbonAdd = chip(tokStr, "reject", "opacity:.55;text-decoration:line-through") |
| + chip(newStr, "resample"); |
| cardTail = `<div class="bars"><div> |
| <div class="meta">residual norm(max(p−q,0))</div>${bars(top8(res), "residual", newTok)}</div></div>`; |
| ids = ids.concat(props.slice(0, i), [newTok]); |
| rejected = true; |
| } |
| $("ribbon").insertAdjacentHTML("beforeend", ribbonAdd); |
| $("cards").insertAdjacentHTML("beforeend", |
| `<div class="poscard"> |
| <div class="meta">position ${roundStart - promptLen + i} — draft proposed <code>${esc(vis(tokStr))}</code></div> |
| <div>q(x)=<b>${qx.toFixed(4)}</b> · p(x)=<b>${px.toFixed(4)}</b> · |
| p/q=<b>${ratio === Infinity ? "∞" : ratio.toFixed(3)}</b> · |
| accept prob=<b>${aProb.toFixed(3)}</b> · u=<b>${u.toFixed(3)}</b> → |
| <b style="color:var(--${ok ? "accept" : "reject"})">${ok ? "✓ ACCEPTED" : "✗ REJECTED"}</b></div> |
| <div class="bars"> |
| <div><div class="meta">draft q — top 8</div>${bars(top8(q), "draft", x)}</div> |
| <div><div class="meta">target p — top 8</div>${bars(top8(p), "target", x)}</div> |
| </div>${cardTail}</div>`); |
| if (!ok) break; |
| } |
| if (!rejected){ |
| const pb = pDist(gamma); |
| const bonus = sampleFrom(pb, r); |
| ids = ids.concat(props, [bonus]); |
| $("ribbon").insertAdjacentHTML("beforeend", chip(tok.decode([bonus]), "bonus")); |
| } |
| const dt = (performance.now() - t0) / 1000; |
| const gen = ids.length - promptLen; |
| $("stats").innerHTML = |
| `<span class="stat"><b>${proposed ? Math.round(100*acceptedN/proposed) : 0}%</b><span>acceptance rate</span></span>` + |
| `<span class="stat"><b>${acceptedN}/${proposed}</b><span>draft tokens accepted</span></span>` + |
| `<span class="stat"><b>${(gen/rounds).toFixed(2)}×</b><span>tokens per target forward pass</span></span>` + |
| `<span class="stat"><b>${rounds}</b><span>target calls (vs ${gen} autoregressive)</span></span>` + |
| `<span class="stat"><b>${(gen/dt).toFixed(1)}</b><span>tokens/s in-browser</span></span>`; |
| } |
| $("status").textContent = "done · " + JSON.stringify(tok.decode(ids.slice(promptLen))); |
| $("runBtn").disabled = false; |
| }; |
| </script> |
| </body> |
| </html> |
|
|
