re-computation. These KV-caches are large, and accounts for the majority of the memory
bandwidth utilisation for the attention mechanism. However, since the KV-cache size does not
change as we increase 𝐾, there is little to no delta in this component.
3.All-reduces: Asmodelsgrowinsize, itsparametersneedtobedividedacrossmultipleaccelerators, leading to communication overheads. With Megatron, this manifests itself as an all-reduce
after every feed-forward and attention layer. Since only the activations for a small number of
tokens are transmitted, this operation is typically latency bound instead of throughput bound
for both sampling and scoring (for small 𝐾). Again, this results in a similar amount of time
spent in the two cases.
Other sources of overhead may exist, depending on the exact transformer implementation. Therefore
it is still possible that the choice of positioning encoding, decoding method (e.g. a sort might be
required for nucleus sampling), hardware limitations etc. can introduce some deltas between scoring
and sampling. However, if the conditions are met such that the above components dominate then
scoring should not be signiﬁcantly slower for small 𝐾.
Modiﬁed Rejection Sampling
We present speculative sampling, an algorithm for accelerating transformer
decoding by enabling the generation of multiple tokens from each transformer
call. Our algorithm relies on the observation that the latency of parallel
scoring of short continuations, generated by a faster but less powerful draft
model, is comparable to that of sampling a single token from the larger target
model. This is combined with a novel modified rejection sampling scheme which
preserves the distribution of the target model within hardware numerics. We
benchmark speculative sampling with Chinchilla, a 70 billion parameter language
model, achieving a 2-2.5x decoding speedup in a distributed setup, without
compromising the sample quality or making modifications to the model itself.
Charlie Chen Sebastian Borgeaud Geoffrey Irving Jean-Baptiste Lespiau Laurent Sifre John Jumper