Inference server: cache the keys and values
A batched attention decoder generates tokens for many sequences at once. The naive step recomputes every past token's key and value projection, for every sequence, on every step, so the work grows with the sequence length. Cache the projections and reuse them to avoid repeated K/V projections; attention still grows with context length.
This project is available to study. Hosted execution is not currently available.
Languages: CUDA C++
Lesson
The trace shows that each decoding step becomes more expensive as the sequence grows. This decoder already batches its sequences, so batching may amortize launch overhead, while the actual bottleneck must be measured. In attention, step t attends the current token over all previous tokens, but the naive path recomputes every past token's key and value projection on every step. Past tokens do not change after they are generated. Identify which intermediate results can be computed once and reused by later decoding steps.
Compute each token K/V projection once. Attention still reads cached history; attention work and cache storage grow with context. Separate projection work, attention work and total step time when interpreting the timeline.
What you will learn
- Recognize redundant recomputation: the same K/V projections recomputed every decode step as sequences grow, even in a batched decoder.
- Introduce a key/value cache so each token is projected once and reused on later steps.
- Measure how caching turns O(seq) per-step projection work into O(1) once the batch has made the step compute-bound.