[CHEATSHEET] Optimization that makes server performance realistic.
Problem: Long prompts (L=32k to 128k) exceed GPU memory if processed in one shot. Additionally, a long prompt will keep other shorter prompts waiting as server will be busy.
Solution: Process the prompt in chunks of size L_chunk sequentially, while building the same KV cache.
- Full input: [B, L, d_model]
- Split into: [B, L_chunk, d_model] Ă— N_chunks
- KV cache after processing all chunks: [B, n_heads, L, d_head]
Masking
- At chunk i, self-attention must only attend up to tokens 0 … i·L_chunk.
- Mask ensures no leakage into “future” chunks.
- Equivalent to full prefill’s causal mask, but applied progressively.
Gains:
- Memory footprint reduced: O(L_chunk²) instead of O(L²).
- Enables extremely long prompts on fixed-GPU HBM capacity.
- Throughput slightly lower due to repeated kernel launches and masking overhead.
- Still compute-bound, much higher efficiency than decode.
