Spaces:

abhimittal
/

speculative-decoding-trace

Running

File size: 15,292 Bytes

<!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";

// ?local=1 → serve models from ./models/ (offline mirrors, CI testing)
if (new URLSearchParams(location.search).get("local")) {
  env.allowRemoteModels = false;
  env.allowLocalModels = true;
  env.localModelPath = "./models/";
}
// pin the small legacy quantized files — plain {dtype:"q8"} resolves to the
// much larger model_quantized.onnx (237 MB) on these repos
const MODEL_OPTS = { model_file_name: "decoder_model_merged", dtype: "q8" };

const $ = id => document.getElementById(id);
const esc = s => s.replace(/&/g,"&amp;").replace(/</g,"&lt;").replace(/>/g,"&gt;");
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;

// mulberry32 — seedable RNG so runs are reproducible
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){ // greedy: one-hot argmax
    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; // [1, seq, vocab]
}
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 { /* wasm fallback */ }
  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"){ // some GPUs fail on q8 — retry on wasm
      $("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;
    // draft proposes gamma tokens (no KV cache in-browser; seqs stay short)
    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}…`;
    }
    // target verifies all positions in one forward pass
    $("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 {
        // resample from the residual norm(max(p − q, 0))
        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>