LLM inference has become increasingly challenging, with the release of 1+ Trillion Parameter models, occupying 100s of GBs in weights as parameters, and requiring 10s-100s of GPUs/ASICs to serve them efficiently.
While traditional inference optimization techniques tend to focus on system-level algorithmic optimizations for decreasing latency (e.g Flash Attention) and increasing throughput (e.g. KV Cache, Paged Attention, Data, Tensor and Pipeline Parallelism). There are other approaches which improve model performance by simplifying the token generation step as part of the model itself.
In this post, I will cover the basic process process of LLM inference and how the process can be accelerated by sampling tokens from a model a bit more optimistically via speculative decoding. I will then cover one of the speculative decoding techniques in limelight, Dspark, by Deepseek.
LLM Decoding Basics
Here I cover the core architecture of a Transformer, and how the architecture of a transformer leads to a logical split of the token generation process into 2 phases. I do not dive deeply into the core architecture of a transformer but cover enough details from an inference engineering perspective.
Autoregressive Transformer
An autogressive transformer is a decoder-only transformer, which at the time of sampling/processing token $ a_t $ at time an arbitrary position $t$, only has access to hidden states of the tokens $\{h_0, h_1, ..., h_{t-1}\}$.
This means that the token $ a_t $ can only view the preceeding tokens and their corresponding states. Any token that comes after $ a_t $ is not visible to the token.
Once token $ a_t $ is generated, the token is appended to the list of generated tokens and then the generation of next token (${a_{t+1}}$) begins.
2 Stages in Autoregressive Generation
When interacting with a model like GPT/Claude, A user typically writes a query in plaintext and sends to the model to generate the next sequence of tokens (say $ N $). All the input tokens are first tokenized, and passed to the model for generating the hidden states corresponding to each token $ {h_0, h_1, ..., h_N} $, which are used to generate the subsequent tokens.
KV Caching
Now the most important point. Since the model is autoregressive, Each token $ a_t $ Only interacts with hidden states of $ a_{\lt t} $. So whenever a new token is generated for time step $ \gt t$, the hidden state of token $a_t$ need not be generated again, set of tokens each generation step looks at does not change for token $ a_t $ for any time step ${\gt t}$. This means that a token never looks at the tokens that come after it in sequence.
This means, all the hidden states that are generated for each token can be “cached” and “reused” once a token is generated.
Because the hidden states once generated for a query can be reused throughought the token generation process, this naturally leads to the separation of the token generation process into 2 steps.
- Prefill Phase - Hidden state of all input tokens is generated simultaneously.
- Decode Phase - Hidden state of the token to be generated is generated, using the hidden state of the previous tokens.
Note: