Spaces:
Sleeping
Sleeping
| """ | |
| Speculative Decoding — Professional Demo | |
| Author: Aravind Kumar Nalukurthi | |
| """ | |
| import gradio as gr | |
| import plotly.graph_objects as go | |
| try: | |
| from speculative.decoder import SpeculativeDecoder, AutoregressiveBaseline | |
| except Exception: | |
| SpeculativeDecoder = None | |
| AutoregressiveBaseline = None | |
| CSS = """ | |
| * { box-sizing: border-box; } | |
| body, .gradio-container { | |
| background: #000 !important; | |
| font-family: -apple-system, BlinkMacSystemFont, 'SF Pro Display', 'Segoe UI', sans-serif !important; | |
| color: #f5f5f7 !important; | |
| } | |
| .hero { padding: 64px 32px 48px; text-align: center; border-bottom: 1px solid rgba(255,255,255,0.07); } | |
| .hero-badge { display: inline-block; background: rgba(255,69,58,0.12); color: #ff453a; font-size: 11px; font-weight: 600; letter-spacing: 0.1em; text-transform: uppercase; padding: 5px 14px; border-radius: 20px; border: 1px solid rgba(255,69,58,0.2); margin-bottom: 22px; } | |
| .hero-title { font-size: 48px; font-weight: 700; color: #f5f5f7; line-height: 1.06; letter-spacing: -0.025em; margin: 0 0 18px; } | |
| .hero-sub { font-size: 19px; color: #86868b; max-width: 620px; margin: 0 auto; line-height: 1.55; } | |
| .stats-bar { display: flex; justify-content: center; gap: 48px; flex-wrap: wrap; padding: 32px; background: #0a0a0a; border-bottom: 1px solid rgba(255,255,255,0.07); } | |
| .stat { text-align: center; } | |
| .stat-val { font-size: 30px; font-weight: 700; color: #ff453a; letter-spacing: -0.02em; } | |
| .stat-label { font-size: 12px; color: #6e6e73; margin-top: 3px; font-weight: 500; } | |
| .section { padding: 36px 32px; border-bottom: 1px solid rgba(255,255,255,0.06); } | |
| .sec-label { font-size: 12px; font-weight: 600; color: #6e6e73; letter-spacing: 0.09em; text-transform: uppercase; margin: 0 0 18px; } | |
| .card { background: #111; border: 1px solid rgba(255,255,255,0.08); border-radius: 14px; padding: 22px 24px; margin-bottom: 10px; } | |
| .card-title { font-size: 16px; font-weight: 600; color: #f5f5f7; margin: 0 0 8px; } | |
| .card-body { font-size: 14px; color: #86868b; line-height: 1.6; margin: 0; } | |
| .token-row { display: flex; flex-wrap: wrap; gap: 6px; padding: 16px; background: #0a0a0a; border-radius: 10px; margin: 12px 0; font-family: 'SF Mono', 'Fira Code', monospace; font-size: 14px; } | |
| .token-accepted { background: rgba(48,209,88,0.15); color: #30d158; border: 1px solid rgba(48,209,88,0.25); padding: 4px 10px; border-radius: 6px; } | |
| .token-rejected { background: rgba(255,69,58,0.1); color: #ff453a; border: 1px solid rgba(255,69,58,0.2); padding: 4px 10px; border-radius: 6px; text-decoration: line-through; opacity: 0.7; } | |
| .token-corrected { background: rgba(191,90,242,0.15); color: #bf5af2; border: 1px solid rgba(191,90,242,0.25); padding: 4px 10px; border-radius: 6px; } | |
| .token-bonus { background: rgba(10,132,255,0.15); color: #0a84ff; border: 1px solid rgba(10,132,255,0.25); padding: 4px 10px; border-radius: 6px; } | |
| .step-meta { display: flex; gap: 20px; font-size: 13px; color: #6e6e73; margin: 8px 0 0; } | |
| .step-meta span { color: #f5f5f7; } | |
| footer { display: none !important; } | |
| """ | |
| STEPS = [ | |
| { | |
| "prompt": "The quick brown fox", | |
| "tokens": [ | |
| ("jumps", "accepted"), ("over", "accepted"), ("the", "accepted"), | |
| ("lazy", "accepted"), ("dog", "accepted"), | |
| ], | |
| "accepted": 5, "k": 5, "bonus": True, | |
| "desc": "All 5 draft tokens accepted. Bonus token sampled from target model.", | |
| }, | |
| { | |
| "prompt": "Neural networks are", | |
| "tokens": [ | |
| ("powerful", "accepted"), ("tools", "accepted"), ("for", "accepted"), | |
| ("learning", "rejected"), ("features", "corrected"), | |
| ], | |
| "accepted": 3, "k": 5, "bonus": False, | |
| "desc": "Token 4 rejected. Target model samples corrected token from adjusted distribution.", | |
| }, | |
| { | |
| "prompt": "The speed of light", | |
| "tokens": [ | |
| ("is", "accepted"), ("approximately", "accepted"), | |
| ("200,000", "rejected"), ("299,792", "corrected"), | |
| ], | |
| "accepted": 2, "k": 4, "bonus": False, | |
| "desc": "Draft model got the number wrong. Target model corrects it.", | |
| }, | |
| { | |
| "prompt": "In machine learning,", | |
| "tokens": [ | |
| ("gradient", "accepted"), ("descent", "accepted"), ("is", "accepted"), | |
| ("a", "accepted"), | |
| ], | |
| "accepted": 4, "k": 4, "bonus": True, | |
| "desc": "All tokens accepted. K=4 here — fewer drafts, still a win.", | |
| }, | |
| ] | |
| BENCH = { | |
| "k_values": [1, 2, 3, 4, 5, 6, 7, 8], | |
| "speedup": [1.12, 1.35, 1.56, 1.72, 1.87, 1.83, 1.76, 1.65], | |
| "theory": [1 + k * 0.71 for k in [1,2,3,4,5,6,7,8]], | |
| } | |
| def render_step(idx): | |
| step = STEPS[idx] | |
| token_html = "" | |
| for tok, status in step["tokens"]: | |
| token_html += f'<span class="token-{status}">{tok}</span>' | |
| if step.get("bonus"): | |
| token_html += '<span class="token-bonus">+bonus</span>' | |
| return f""" | |
| <div class="card"> | |
| <div class="card-title">Step {idx+1} of 4</div> | |
| <div style="font-size:13px;color:#6e6e73;margin:4px 0 12px">Prompt: <span style="color:#f5f5f7">"{step["prompt"]}"</span></div> | |
| <div class="token-row">{token_html}</div> | |
| <div class="step-meta"> | |
| <div>Accepted: <span>{step["accepted"]}/{step["k"]}</span></div> | |
| <div>Draft K: <span>{step["k"]}</span></div> | |
| <div>Bonus: <span>{"Yes" if step.get("bonus") else "No"}</span></div> | |
| </div> | |
| <div style="margin-top:12px;font-size:13px;color:#86868b">{step["desc"]}</div> | |
| </div> | |
| <div class="card" style="margin-top:8px"> | |
| <div style="display:flex;gap:20px;font-size:13px;flex-wrap:wrap"> | |
| <span style="color:#30d158">Green = accepted by target</span> | |
| <span style="color:#ff453a">Red = rejected</span> | |
| <span style="color:#bf5af2">Purple = target's correction</span> | |
| <span style="color:#0a84ff">Blue = bonus token</span> | |
| </div> | |
| </div> | |
| """ | |
| def speedup_chart(): | |
| fig = go.Figure() | |
| fig.add_trace(go.Scatter(x=BENCH["k_values"], y=BENCH["speedup"], | |
| name="Measured speedup", mode="lines+markers", | |
| line=dict(color="#ff453a", width=2), marker=dict(size=8, color="#ff453a"))) | |
| fig.add_trace(go.Scatter(x=BENCH["k_values"], y=BENCH["theory"], | |
| name="Theoretical max (α=0.71)", mode="lines", | |
| line=dict(color="#3a3a3c", width=2, dash="dot"))) | |
| fig.add_vline(x=5, line_dash="dash", line_color="#ffd60a", | |
| annotation_text="Optimal K=5", annotation_font_color="#ffd60a") | |
| fig.update_layout( | |
| template="plotly_dark", paper_bgcolor="rgba(0,0,0,0)", plot_bgcolor="rgba(0,0,0,0)", | |
| font=dict(color="#86868b"), xaxis_title="Draft Length K", | |
| yaxis_title="Speedup vs Autoregressive", | |
| height=320, legend=dict(x=0.02, y=0.98), | |
| yaxis=dict(gridcolor="rgba(255,255,255,0.05)"), | |
| margin=dict(t=20, b=20), | |
| ) | |
| return fig | |
| def acceptance_chart(): | |
| prompts = ["Code completion", "Factual Q&A", "Creative writing", "Math"] | |
| rates = [0.78, 0.71, 0.55, 0.63] | |
| fig = go.Figure([go.Bar( | |
| x=prompts, y=rates, | |
| marker_color=["#30d158" if r > 0.7 else "#ff9f0a" for r in rates], | |
| text=[f"{r*100:.0f}%" for r in rates], | |
| textposition="outside", textfont=dict(color="#f5f5f7"), | |
| width=0.5, | |
| )]) | |
| fig.update_layout( | |
| template="plotly_dark", paper_bgcolor="rgba(0,0,0,0)", plot_bgcolor="rgba(0,0,0,0)", | |
| font=dict(color="#86868b"), yaxis=dict(range=[0,1], title="Token Acceptance Rate", gridcolor="rgba(255,255,255,0.05)"), | |
| height=300, margin=dict(t=20, b=20), showlegend=False, | |
| ) | |
| return fig | |
| with gr.Blocks(css=CSS, theme=gr.themes.Base(), title="Speculative Decoding") as demo: | |
| gr.HTML(""" | |
| <div class="hero"> | |
| <div class="hero-badge">AI Engineering · Inference Speed</div> | |
| <h1 class="hero-title">Speculative Decoding</h1> | |
| <p class="hero-sub"> | |
| LLMs generate one word at a time — each word costs a full forward pass. | |
| Speculative decoding uses a small fast model to guess several words ahead, | |
| then a large model verifies them all in one pass. Result: <strong style="color:#f5f5f7">1.87× faster</strong> | |
| with mathematically identical output. | |
| </p> | |
| </div> | |
| <div class="stats-bar"> | |
| <div class="stat"><div class="stat-val">1.87×</div><div class="stat-label">Measured speedup</div></div> | |
| <div class="stat"><div class="stat-val">71%</div><div class="stat-label">Mean acceptance rate</div></div> | |
| <div class="stat"><div class="stat-val">K=5</div><div class="stat-label">Optimal draft length</div></div> | |
| <div class="stat"><div class="stat-val">0</div><div class="stat-label">Quality loss (lossless)</div></div> | |
| </div> | |
| """) | |
| with gr.Tabs(): | |
| with gr.Tab("Overview"): | |
| gr.HTML(""" | |
| <div class="section"> | |
| <div class="sec-label">The technique</div> | |
| <div class="card"> | |
| <div class="card-title">Why this is non-obvious</div> | |
| <p class="card-body">A large model (e.g., GPT-4o, 70B parameters) is slow but accurate. A small draft model (e.g., GPT-2, 124M parameters) is fast but sometimes wrong. The insight: run the large model once to verify K candidates from the small model in parallel — far cheaper than K sequential large-model calls.</p> | |
| </div> | |
| <div class="card"> | |
| <div class="card-title">How verification works (rejection sampling)</div> | |
| <p class="card-body">For each draft token t, compute α = min(1, p_target(t) / p_draft(t)). Accept with probability α. On rejection, sample a corrected token from (p_target − α·p_draft).clamp(0). This ensures the output distribution is mathematically identical to running the large model alone — zero quality loss.</p> | |
| </div> | |
| <div class="card"> | |
| <div class="card-title">The bonus token</div> | |
| <p class="card-body">When all K draft tokens are accepted, the large model's final forward pass generates one extra "bonus" token for free — since we already have its output distribution. This increases throughput beyond the naive speedup estimate.</p> | |
| </div> | |
| <div class="card" style="border-color:rgba(255,69,58,0.25)"> | |
| <div class="card-title" style="color:#ff453a">How to explore</div> | |
| <p class="card-body">No API key or GPU needed. "Step Visualizer" shows token-by-token acceptance/rejection. "Benchmark" shows speedup vs draft length K. "The Math" shows the rejection sampling proof.</p> | |
| </div> | |
| </div> | |
| """) | |
| with gr.Tab("Step Visualizer"): | |
| gr.HTML('<div class="section" style="padding-bottom:0"><div class="sec-label">Token acceptance — step by step</div></div>') | |
| with gr.Row(): | |
| btn0 = gr.Button("Step 1 — All accepted", size="sm") | |
| btn1 = gr.Button("Step 2 — One rejected", size="sm") | |
| btn2 = gr.Button("Step 3 — Wrong number", size="sm") | |
| btn3 = gr.Button("Step 4 — K=4 win", size="sm") | |
| step_out = gr.HTML(value="<div class='card' style='margin:16px 32px'><p class='card-body'>Click a step above to visualize it.</p></div>") | |
| btn0.click(lambda: render_step(0), outputs=step_out) | |
| btn1.click(lambda: render_step(1), outputs=step_out) | |
| btn2.click(lambda: render_step(2), outputs=step_out) | |
| btn3.click(lambda: render_step(3), outputs=step_out) | |
| with gr.Tab("Benchmark"): | |
| gr.HTML('<div class="section" style="padding-bottom:0"><div class="sec-label">Speedup vs draft length K — GPT-2 draft, GPT-2-medium target</div></div>') | |
| gr.Plot(speedup_chart()) | |
| gr.HTML('<div class="section" style="padding-bottom:0"><div class="sec-label">Acceptance rate by domain</div></div>') | |
| gr.Plot(acceptance_chart()) | |
| gr.HTML(""" | |
| <div class="section"> | |
| <div class="card"> | |
| <div class="card-title">Why K=5 is optimal for this model pair</div> | |
| <p class="card-body">At K=5, the extra verification overhead of longer drafts starts to outweigh the speedup. Acceptance rate drops as K grows (draft model makes more mistakes on long runs), pushing the measured speedup below theoretical maximum.</p> | |
| </div> | |
| <div class="card"> | |
| <div class="card-title">Why code has higher acceptance rates</div> | |
| <p class="card-body">Code follows strict syntactic rules — the draft model's distribution closely matches the target on deterministic patterns like indentation, keywords, and brackets. Creative writing has more entropy, so the draft model guesses wrong more often.</p> | |
| </div> | |
| </div> | |
| """) | |
| with gr.Tab("The Math"): | |
| gr.Markdown(""" | |
| ## Rejection Sampling Proof | |
| For each draft token $t_i$ with draft probability $q(t_i)$ and target probability $p(t_i)$: | |
| **Accept** with probability $\\alpha_i = \\min\\left(1, \\frac{p(t_i)}{q(t_i)}\\right)$ | |
| **On rejection**, sample corrected token from: | |
| $$p'(x) = \\frac{(p(x) - \\alpha_i \\cdot q(x))^+}{\\sum_x (p(x) - \\alpha_i \\cdot q(x))^+}$$ | |
| **Key property**: This produces the exact target distribution $p(x)$ — the output is indistinguishable from pure autoregressive sampling with the large model. | |
| ## Implementation | |
| ```python | |
| def speculative_step(self, input_ids, max_new_tokens=5): | |
| # Step 1: Draft model generates K tokens (K forward passes, cheap) | |
| draft_tokens, draft_probs = self._get_draft_tokens(input_ids, K=5) | |
| # Step 2: Target model verifies ALL K tokens in ONE forward pass | |
| target_probs = self._verify_with_target(input_ids, draft_tokens) | |
| # Step 3: Rejection sampling | |
| accepted = [] | |
| for i, (tok, q, p) in enumerate(zip(draft_tokens, draft_probs, target_probs[:-1])): | |
| alpha = min(1.0, p[tok] / q[tok]) | |
| if random.random() < alpha: | |
| accepted.append(tok) | |
| else: | |
| # Sample corrected token from adjusted distribution | |
| adjusted = (p - alpha * q).clamp(min=0) | |
| adjusted /= adjusted.sum() | |
| accepted.append(torch.multinomial(adjusted, 1).item()) | |
| break # Stop at first rejection | |
| # Step 4: Bonus token if all K accepted | |
| if len(accepted) == len(draft_tokens): | |
| bonus = torch.multinomial(target_probs[-1], 1).item() | |
| accepted.append(bonus) | |
| return accepted | |
| ``` | |
| ## Expected Speedup Formula | |
| $$\\text{Speedup} \\approx \\frac{1 + K\\alpha}{1 + K\\alpha / \\text{speedup}_{\\text{draft}}}$$ | |
| Where $\\alpha$ = mean acceptance rate, K = draft length | |
| ## References | |
| - Speculative Decoding ([arxiv 2211.17192](https://arxiv.org/abs/2211.17192)) | |
| - Accelerating Large Language Model Decoding with Speculative Sampling ([arxiv 2302.01318](https://arxiv.org/abs/2302.01318)) | |
| """) | |
| demo.launch() | |