[CHEATSHEET] Use multiple GPUs and pipeline.
Idea: Merge Decode and Prefill phase of different requests (in batch dimension) and operate on this batched input. Use Paged KV cache.
This is Continuous Batching with 4 requests (R1, R2, R3, R4) all executing on 1 GPU.
- The smaller uniform blocks are decode phase.
- Larger 3 phases are prefill phase.
At every instance, decode and prefill phases from different requests are merged in batch (B) dimension. So, one operation is executing at a time. This one operation has merged input of multiple operations.
All requests are padded to the largest sequence length among the requests (as tensors of same shape can be concatenated).
Each request in the batched request needs its KV cache. To access efficiently, paged KV Cache is used. The decode phases can update their respective KV cache efficiently.
Note: this is not pipelining. It is a single operation. Just input is a batch of different phases.
Q: (batch_size, heads, 1, head_dim)
K_full: (batch_size, heads, seq_len, head_dim)
V_full: (batch_size, heads, seq_len, head_dim)
Each request has its own KV cache:
- vLLM uses paged Attention
- Triton used gather based attention kernel.
Note: KV cache is not padded across the batch. Matmul happens by gathering the correct K/V blocks for each request. KV-cache layout is ragged.
vLLM approach:
K/V are stored in blocks called pages:
Each page: (H, block_size, Dh)
Typical block_size = 16 or 32
A request with seq_len 50 uses: ceil(50 / block_size) = 4 pages
Example memory layout:
Request A:
Page 101 → tokens 0–31
Page 207 → tokens 32–63
Page 19 → tokens 64–95
(request ends at position 10 inside page 101)
All K/V for all requests lives in a common pool of pages.
Attention uses indirect indexing. K/V are gathered by index.
The attention kernel receives:
- Q tensor
- list of page pointers (per request)
- per-request sequence lengths
The kernel computes accordingly.
