Spaces:
Sleeping
Sleeping
| """DSSD demo: exact self-speculative decoding with pruned exit heads and context copying. | |
| Local UI work without a GPU: DSSD_DEMO_BACKEND=tiny python app.py | |
| """ | |
| import gradio as gr | |
| from demo.backend import MAX_NEW_TOKENS, MODELS, generate_updates, load_models | |
| from demo.examples import EXAMPLES | |
| from demo.render import legend_html | |
| def gpu_seconds(model, domain_label, prompt, max_new, compare): | |
| # ZeroGPU refuses a call whose requested duration exceeds the visitor's remaining daily quota | |
| # (5 min on free accounts), so ask for what the request needs: engine setup + both decodes. | |
| return int(30 + int(max_new) * (0.25 if compare else 0.12)) | |
| try: # ZeroGPU on Hugging Face Spaces | |
| import spaces | |
| GPU = spaces.GPU(duration=gpu_seconds) | |
| except ImportError: | |
| def GPU(f): | |
| return f | |
| ENGINES = load_models() | |
| DOMAINS = {"Web text": "web", "Math (few-shot)": "gsm8k", "Python code": "code"} | |
| HEADER = """ | |
| <div class="hero"> | |
| <h1>DSSD <span class="hero-sub">exact self-speculative decoding</span></h1> | |
| <p>A tiny exit head on layer 4 of the frozen model, plus tokens copied from the context, propose a | |
| small tree of draft tokens. The full model checks the whole tree in one forward pass and keeps the | |
| longest correct branch. The text is <b>exactly</b> what greedy decoding would write, only faster.</p> | |
| <div class="badges"><span>✓ identical to greedy decoding</span><span>no draft model</span> | |
| <span>≈ 280 MB of heads</span></div> | |
| </div> | |
| """ | |
| HOW = """ | |
| <div class="how"> | |
| <div class="how-step"><b>1 · Draft</b><br>Copy n-grams found earlier in the context<br>+ chain of the | |
| layer-4 head<br>+ its top-4 alternatives at the root</div> | |
| <div class="how-arrow">→</div> | |
| <div class="how-step"><b>2 · Verify</b><br>One forward pass of the full model over the whole tree | |
| (16 rows, row-invariant kernels)</div> | |
| <div class="how-arrow">→</div> | |
| <div class="how-step"><b>3 · Commit</b><br>Longest branch that matches the model's own greedy choice, | |
| plus one token from the model</div> | |
| </div> | |
| <p class="how-note">The exit heads predict over a small, layer-specific vocabulary (≈ 3k tokens at layer 4) | |
| plus a <i>reject</i> class, about 50× cheaper than a full projection. Speedups depend on the GPU and on | |
| whether CUDA graphs are used. The paper's measurements (AMD MI210, greedy, batch 1) are 2.50× / 1.72× / 1.46× on | |
| web text, GSM8K and HumanEval for Qwen3-8B-Base, and 2.51× / 1.80× / 1.57× for Llama-3.1-8B.</p> | |
| """ | |
| CSS = """ | |
| .hero h1 { margin: 0 0 .3rem 0; font-size: 1.9rem; } | |
| .hero-sub { font-weight: 400; opacity: .65; font-size: 1.1rem; margin-left: .4rem; } | |
| .hero p { margin: .2rem 0 .6rem 0; max-width: 60rem; } | |
| .badges span { display: inline-block; padding: .15rem .6rem; margin: 0 .4rem .3rem 0; border-radius: 999px; | |
| font-size: .82rem; background: var(--block-background-fill); border: 1px solid var(--border-color-primary); } | |
| :root { --tok-copy: #cfe4fb; --tok-head: #ffe1b0; --tok-star: #f3d1ec; --tok-draft: #d9f2d0; } | |
| .dark { --tok-copy: #1f4b78; --tok-head: #7a5413; --tok-star: #6b2f60; --tok-draft: #2f5a2a; } | |
| .tok { border-radius: 3px; padding: 0 1px; } | |
| .tok-copy { background: var(--tok-copy); } | |
| .tok-head { background: var(--tok-head); } | |
| .tok-star { background: var(--tok-star); box-shadow: inset 0 -2px 0 #b3479b; } | |
| .tok-draft { background: var(--tok-draft); } | |
| .tok-model { background: transparent; border: 1px dashed var(--border-color-primary); } | |
| .legend { display: flex; flex-wrap: wrap; gap: .9rem; font-size: .88rem; margin: .2rem 0; } | |
| .legend-item { display: inline-flex; align-items: center; gap: .35rem; } | |
| .swatch { font-size: .8rem; padding: 0 .3rem; } | |
| .dssd-text { white-space: pre-wrap; overflow-wrap: anywhere; font-family: var(--font-mono); font-size: .9rem; line-height: 1.55; | |
| min-height: 14rem; max-height: 30rem; overflow-y: auto; padding: .8rem; border-radius: 8px; | |
| background: var(--block-background-fill); border: 1px solid var(--border-color-primary); } | |
| .cursor { display: inline-block; width: .5em; height: 1em; background: var(--body-text-color); | |
| vertical-align: text-bottom; animation: blink 1s steps(1) infinite; opacity: .6; } | |
| @keyframes blink { 50% { opacity: 0; } } | |
| .panel-title { font-weight: 600; margin-bottom: .2rem; } | |
| .stats { display: flex; flex-wrap: wrap; gap: .6rem; margin: .4rem 0 .6rem 0; } | |
| .stat { flex: 1 1 8rem; padding: .6rem .8rem; border-radius: 8px; background: var(--block-background-fill); | |
| border: 1px solid var(--border-color-primary); } | |
| .stat-value { font-size: 1.35rem; font-weight: 700; } | |
| .stat-label { font-size: .78rem; opacity: .7; } | |
| .stat-main .stat-value { font-size: 1.9rem; color: var(--color-accent); } | |
| .stat-ok .stat-value { color: #1a7f37; } .stat-bad .stat-value { color: #cf222e; } | |
| .share-bar { display: flex; height: .7rem; border-radius: 999px; overflow: hidden; | |
| border: 1px solid var(--border-color-primary); } | |
| .share-bar .tok-model { background: var(--border-color-primary); border: 0; } | |
| .share-labels { font-size: .8rem; opacity: .75; margin-top: .25rem; } | |
| .progress { font-size: .9rem; opacity: .75; padding: .4rem 0; } | |
| .replay-clock { font-size: .78rem; opacity: .7; margin-left: .4rem; } | |
| .how { display: flex; flex-wrap: wrap; align-items: center; gap: .5rem; } | |
| .how-step { flex: 1 1 12rem; padding: .7rem; border-radius: 8px; background: var(--block-background-fill); | |
| border: 1px solid var(--border-color-primary); font-size: .88rem; } | |
| .how-arrow { font-size: 1.4rem; opacity: .5; } | |
| .how-note { font-size: .85rem; opacity: .8; margin-top: .6rem; } | |
| @media (max-width: 700px) { .how-arrow { display: none; } } | |
| """ | |
| # Client-side replay of both generations side by side, at the speed measured on the GPU. | |
| HEAD = """ | |
| <script> | |
| window.dssdReplay = function () { | |
| const holder = document.querySelector('#replay-holder .replay-data'); | |
| if (!holder) return; | |
| const data = JSON.parse(holder.dataset.replay); | |
| const el = {dssd: document.querySelector('#dssd-out .dssd-text'), ar: document.querySelector('#ar-out .dssd-text')}; | |
| const esc = s => s.replace(/&/g, '&').replace(/</g, '<').replace(/>/g, '>'); | |
| for (const k in el) if (el[k]) el[k].innerHTML = ''; | |
| const idx = {dssd: 0, ar: 0}, start = performance.now(); | |
| function frame(now) { | |
| const t = (now - start) / 1000; let more = false; | |
| for (const k of ['dssd', 'ar']) { | |
| const arr = data[k] || [], box = el[k]; if (!box) continue; | |
| let html = ''; | |
| while (idx[k] < arr.length && arr[idx[k]][0] <= t) { | |
| const [, txt, src] = arr[idx[k]++]; | |
| html += (k === 'dssd' && src !== 'model') ? `<span class="tok tok-${src}">${esc(txt)}</span>` : esc(txt); | |
| } | |
| if (html) box.insertAdjacentHTML('beforeend', html); | |
| if (idx[k] < arr.length) more = true; | |
| else if (arr.length && !box.dataset.done) { | |
| box.dataset.done = 1; | |
| box.insertAdjacentHTML('beforeend', `<span class="replay-clock">⏱ ${arr[arr.length - 1][0].toFixed(2)} s</span>`); | |
| } | |
| } | |
| if (more) requestAnimationFrame(frame); | |
| else for (const k in el) if (el[k]) delete el[k].dataset.done; | |
| } | |
| requestAnimationFrame(frame); | |
| }; | |
| </script> | |
| """ | |
| def run(model, domain_label, prompt, max_new, compare): | |
| try: | |
| for u in generate_updates(ENGINES[model], prompt, DOMAINS[domain_label], max_new, compare): | |
| yield u.dssd, u.ar, u.stats, u.replay, gr.update(visible=u.done and compare) | |
| except ValueError as e: | |
| raise gr.Error(str(e)) | |
| def example(domain_label, i): | |
| ex = EXAMPLES[DOMAINS[domain_label]] | |
| i = (i + 1) % len(ex) | |
| return ex[i], i | |
| with gr.Blocks(title="DSSD: exact self-speculative decoding") as demo: | |
| gr.HTML(HEADER) | |
| with gr.Row(equal_height=False): | |
| with gr.Column(scale=2): | |
| prompt = gr.Textbox(label="Prompt (plain-text completion: these are base models)", lines=6, max_lines=14, | |
| value=EXAMPLES["web"][0]) | |
| with gr.Row(): | |
| ex_btn = gr.Button("Another example", size="sm", variant="secondary") | |
| ex_idx = gr.State(0) | |
| with gr.Column(scale=1): | |
| model = gr.Radio(list(ENGINES), value=next(iter(ENGINES)), label="Model") | |
| domain = gr.Radio(list(DOMAINS), value="Web text", label="Domain (calibrated settings and stop rule)") | |
| max_new = gr.Slider(16, MAX_NEW_TOKENS, value=160, step=16, label="Max new tokens") | |
| compare = gr.Checkbox(value=True, label="Also run plain autoregressive decoding and compare") | |
| with gr.Row(): | |
| go = gr.Button("Generate", variant="primary") | |
| stop = gr.Button("Stop", variant="stop") | |
| gr.HTML(legend_html()) | |
| with gr.Row(equal_height=True): | |
| with gr.Column(): | |
| gr.HTML('<div class="panel-title">DSSD</div>') | |
| out_dssd = gr.HTML(elem_id="dssd-out") | |
| with gr.Column(): | |
| gr.HTML('<div class="panel-title">Autoregressive (greedy)</div>') | |
| out_ar = gr.HTML(elem_id="ar-out") | |
| stats = gr.HTML() | |
| replay_btn = gr.Button("▶ Replay both side by side at the measured speed", visible=False, size="sm") | |
| replay = gr.HTML(elem_id="replay-holder", visible=True) | |
| with gr.Accordion("How it works", open=False): | |
| gr.HTML(HOW) | |
| gr.Markdown(f"Heads: {', '.join(f'[{r}](https://huggingface.co/{r})' for r in MODELS.values())}. " | |
| "Llama-3.1-8B heads: built with Llama.", elem_classes="footer") | |
| gen = go.click(run, [model, domain, prompt, max_new, compare], [out_dssd, out_ar, stats, replay, replay_btn]) | |
| stop.click(None, cancels=[gen]) | |
| replay_btn.click(None, js="() => { window.dssdReplay(); }") | |
| domain.change(lambda d: (EXAMPLES[DOMAINS[d]][0], 0), domain, [prompt, ex_idx]) | |
| ex_btn.click(example, [domain, ex_idx], [prompt, ex_idx]) | |
| demo.queue(default_concurrency_limit=1, max_size=16) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"), css=CSS, head=HEAD) | |