Speculative Decoding: Rejection Sampling

Hard · Solved in PyTorch · Machine learning coding practice

Problem

Speculative decoding (Leviathan et al., 2022) addresses the memory-bandwidth bottleneck in autoregressive LLM inference by using a "guess-and-verify" strategy with two models: a <tooltip term="Draft model (The Guesser)" definition="A much smaller, faster, and cheaper model is used to quickly generate a sequence of \'draft\' tokens \(for example, the next 4 or 5 tokens\). Because it is small, it requires very little memory bandwidth to run.">draft model</tooltip> and a <tooltip term="Target model (The Verifier)" definition="The massive, slow, and highly accurate main model takes the draft model\'s guessed sequence and evaluates all of them simultaneously in a single parallel forward pass.">target model</tooltip>.

This utilizes idle GPU compute while maintaining the target model's exact output distribution.

Topics: inference, optimization, reject-sampling

The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Speculative Decoding: Rejection Sampling in the browser and run it against the tests.

Related problems