prefill

This is prefill:

Prefill is the part of LLM inferencing where the prompt is run through the model. The products of prefill are:

1. K and V vectors for every layer of attention and every token in the prompt
2. The final hidden state of the model prior to output tokens being generatedok

Because the whole input prompt is passed in, it is a dense computation and is therefore compute-limited. This contrasts with the next phase, decode.

Specifically, the input of prefill is $[N\times d_{\text{model}}]$, where $N$ is the number of input tokens and $d_{\text{model}}$ is the hidden dimension. To compute the attention part of prefill,

1. $K$, $Q$, and $V$ projections are computed by multiplying the input matrix by three different weight matrices ($W_K$, $W_Q$, $W_V$). The resulting $K$ and $V$ projections are then saved in KV cache.
2. Attention itself is computed using these computed $K$, $Q$, and $V$ according to $\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d}}\right)V$. This is not cacheable.

The projection part turns into GEMMs (matrix-matrix multiplication). As long as $N$ is large enough to fill a tensor core tile, each GEMM takes longer than loading the inputs from HBM, so it is compute-bound.

Prefill’s compute scales as

• $O(S\cdot d^2)$ when computing projections
• $O(S^2\cdot d)$ when computing attention