Implement and Benchmark Speculative Decoding

Read the full interview experience this question came from →

Quick 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.

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

|Home/Machine Learning/Spacex
Spacex logo
Spacex
Sep 3, 2026
mediumResearch ScientistTechnical ScreenMachine Learning
0
0

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 Guidance

  • 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 Guidance

  • 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 Guidance

  • 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 Guidance

  • 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 Guidance

  • 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 Guidance

  • 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?
Loading comments...