Skip to content

DESIGN.md

Latest commit

 

History

History
182 lines (124 loc) · 5.54 KB

DESIGN.md

File metadata and controls

182 lines (124 loc) · 5.54 KB
Raw

Design Document — speculative-decoding-real

For results and usage see README.md.

Objective

Empirically measure speculative decoding performance using real transformer models on a single consumer GPU, validate the Leviathan et al. (2023) analytical speedup model, and quantify sources of deviation.

Speculative Decoding Algorithm

Standard Greedy Algorithm

1. Prime both draft and target KV caches on the prompt
2. Draft generates gamma tokens autoregressively
3. Target scores all gamma+1 positions in ONE forward pass
4. Accept draft token i if argmax(target[i]) == draft[i]
5. On first rejection: use target correction, stop accepting
6. If all accepted: append target bonus token
7. Re-sync draft KV cache to accepted prefix
8. Repeat until max_new_tokens reached

Key Property

One target forward pass scores gamma+1 positions simultaneously. If all draft tokens are accepted: gamma+1 tokens generated per step. If none accepted: 1 token generated per step (same as autoregressive). Expected tokens per step: (1 - alpha^(gamma+1)) / (1 - alpha)

Analytical Model (Leviathan et al. 2023)

Expected speedup over autoregressive baseline:

E[speedup] = (1 - alpha^(gamma+1)) / ((1 - alpha) x (1 + gamma/c))

Parameters:

  • alpha: per-token acceptance rate (probability draft matches target)
  • gamma: number of draft tokens per speculative step
  • c: cost ratio = T_target / T_draft (single token decode times)

Break-even condition (speedup = 1):

c_min = (gamma x (1 - alpha)) / (1 - alpha^(gamma+1) - (1-alpha))

The model assumes:

  • Zero KV cache synchronization overhead
  • Zero coordination/comparison overhead
  • Constant alpha across all positions

Corrected Model

Adding measured sync overhead s per step:

E[speedup_corrected] = E[tokens_per_step] x T_target /
                       (gamma x T_draft + T_target + s)

Where:

  • E[tokens_per_step] = (1 - alpha^(gamma+1)) / (1 - alpha)
  • s = KV sync time (measured: ~23ms in Python, ~0.1ms in C++)

This model fits measured data better at high gamma values where sync overhead is amortized over more accepted tokens.

Memory Bandwidth Model

At batch_size=1, decode time is dominated by weight loading:

T_theoretical = model_bytes / gpu_bandwidth

For RTX 2070 (448 GB/s): Qwen2-0.5B: 988 MB / 448 GB/s = 2.2 ms Qwen2-1.5B: 3088 MB / 448 GB/s = 6.9 ms Theoretical ratio: 3.12x

Measured times include fixed overhead per layer: KV cache attention: O(seq_len) memory reads Layer normalization: small but constant CUDA kernel launch: 2-5ms per layer regardless of size

This overhead equalizes decode times across model sizes, producing cost_ratio = 1.178x instead of 3.12x.

Three Phases

Phase 1 — speculative_decoding.py

autoregressive_generate(model, input_ids, max_new_tokens)
    Standard KV-cached decode
    Returns: tokens, wall_time

speculative_generate(draft, target, input_ids, max_new_tokens, gamma)
    Full speculative decode loop with greedy acceptance
    Returns: tokens, wall_time, acceptance_rates_per_step

N_REPEATS=3, median reported to suppress outliers
8 prompts x 4 gamma = 32 configurations

Phase 2 — analysis_deep.py

section1: Memory bandwidth roofline explanation
section2: Binary search for min cost_ratio per (alpha, gamma)
section3: Monte Carlo simulation with real alpha values
section4: Regime map table at gamma=4
section5: Overhead factor quantification
section6: Summary of findings

Phase 3 — phase3_profiling.py

section1: Per-step time breakdown (draft/target/sync), N=20 steps
section2: Alpha vs token position, N=64 tokens
section3: Alpha vs temperature (0, 0.3, 0.5, 0.7, 1.0, 1.5)
section4: Isolated single-token timing, N=50 repetitions

Prompt Selection

8 prompts across 4 categories:

Code (2):          High structural regularity, predictable syntax
Technical (2):     Domain vocabulary, moderate predictability
Creative (1):      Diverse vocabulary, lower predictability
Factual (1):       Reference patterns, moderate predictability
Conversational (1): Common patterns, moderate predictability
Mathematical (1):  Formulaic structure, high local regularity

This spread captures the range of real deployment scenarios.

KV Cache Synchronization

After each speculative step, the draft KV cache must reflect the accepted prefix, not the full gamma-token draft.

Python implementation (this project): Re-run draft forward pass on accepted tokens Cost: ~23ms (one full forward pass)

Production C++ (vLLM): Tensor slicing: copy accepted cache entries directly Cost: ~0.1ms

This 230x difference is the primary source of deviation from the analytical model at low gamma values.

Limitations

  1. Single GPU — cost_ratio would differ on multi-GPU
  2. Greedy acceptance only in Phases 1 and 3 (Phase 3 section 3 tests sampling for alpha measurement)
  3. Python overhead adds ~23ms per step vs ~0.1ms in production
  4. One draft/target pair only
  5. Short sequences (64 tokens) — cost_ratio may change with context length
  6. Both models share training distribution — alpha may differ for other pairs

References

  • Leviathan et al., Fast Inference from Transformers via Speculative Decoding, ICML 2023
  • Chen et al., Accelerating Large Language Model Decoding with Speculative Sampling, 2023
  • vLLM speculative decoding implementation, 2023
  • Williams et al., Roofline: An Insightful Visual Performance Model for Multicore Architectures, 2009