Implement and Benchmark Speculative Decoding
Company: Spacex
Role: Research Scientist
Category: Machine Learning
Difficulty: medium
Interview Round: Technical Screen
# Implement and Benchmark Speculative Decoding
Design two generation functions for an autoregressive language model:
- a normal decoder that generates up to `max_tokens` from `model`, `tokenizer`, and `prompt`;
- a speculative decoder that additionally uses a smaller `draft` model to propose several tokens before the target model verifies them.
Explain the implementation for greedy decoding and describe what must change to preserve the target model's distribution under sampling. Both functions must stop consistently on end-of-sequence, maintain attention and position state correctly, and avoid corrupting either model's key-value cache after partial acceptance.
Finally, design a benchmark that tests the source's performance goal: speculative decoding should be faster than normal decoding on an appropriate workload. Explain why this cannot be assumed for every prompt or hardware setup.
### Clarifying Questions to Ask
- Is output required to match greedy target decoding exactly, or preserve a sampling distribution under a fixed random seed?
- How many tokens may the draft propose per iteration?
- Do both models share a tokenizer and vocabulary?
- Which device, precision, batch size, warmup, and synchronization rules define the speed test?
### Part 1 — Normal Decoding
Describe tokenization, prompt prefill, next-token selection, key-value cache reuse, stopping, and output decoding.
#### What This Part Should Cover
- Selection of the first generated token from the prompt prefill's final-position logits.
- For each later token, one incremental target pass that feeds only the newest generated token with the retained cache.
- Correct masks, positions, EOS handling, and maximum-new-token counting.
- Cache reuse without recomputing the full accepted prefix.
### Part 2 — Speculative Decoding
Describe proposal, batched target verification, acceptance, rejection, fallback-token selection, and cache rollback or slicing.
#### What This Part Should Cover
- Greedy acceptance based on agreement with target choices.
- Distribution-preserving acceptance and residual sampling for stochastic generation.
- Synchronization of target and draft cache state to the final accepted prefix.
### Part 3 — Performance Test
Define a repeatable comparison between normal and speculative decoding.
#### What This Part Should Cover
- Identical prompts, output limits, target model, and stopping policy.
- Warmup, device synchronization, repeated trials, and robust latency statistics.
- Acceptance rate, drafted tokens, verified tokens, and draft overhead as explanatory metrics.
### What a Strong Answer Covers
- Semantic equivalence to the chosen target-decoding contract.
- Correct handling of the first rejected proposal and any extra target token.
- Cache and position consistency after acceptance, rejection, and EOS.
- A benchmark that separates algorithmic speedup from compilation, transfer, or warm-cache artifacts.
- Recognition that low acceptance or an expensive draft model can make speculation slower.
### Follow-up Questions
- How would you choose proposal length adaptively from recent acceptance rates?
- What breaks when target and draft tokenizers differ?
- How do batched requests change the latency-versus-throughput benefit?
Overview: Design normal and speculative language-model decoding with correct proposal verification, rejection, sampling, EOS, positions, and key-value cache state. The answer also builds a fair speed benchmark around warmup, device synchronization, acceptance rate, latency statistics, and workloads where draft overhead can outweigh target-pass savings.
Read the full Spacex Research Scientist interview experience this question came from