KV-CACHING + PREFILL & DECODING
JUSTIN ZHANG - SWAN LAB
AGENDA
- 01Book-keeping items
- 02General learning points
- 03Papers + key takeaways
- 04Experiments performed
- 05Action items
BOOK-KEEPING ITEMS
Thank you for following up with IFS. Great to hear they've replied.
Proposal: recreate and update the SWAN lab website to add credibility.
LAB WEBSITE - PROPOSED PLAN
- 01Create a Lab GitHub Organization instead of a personal account.
github.com/swan-carleton/swan-carleton.github.io - 02Host on GitHub with a custom domain.
swan.olgabaysal.com - 03Fix current issues: out of date, not mobile responsive, design could be more appealing.
COMPETE WITHOUT INFINITE COMPUTE
Last meeting we discussed how hard it is to compete in research when large, well-funded labs have access to effectively infinite compute.
I've been watching YC Paper Club - a monthly series where ML researchers present their latest studies. Insightful watch.
YC PAPER CLUB
youtube.com/watch?v=wE1ZgJdt4uM&t=3845s
PRE-TRAINING UNDER INFINITE COMPUTE
Compute spent training LLMs grows ~4x per year. Web data grows only ~1.03x per year. We approach a regime of infinite compute, finite data.
A list of scaling recipes that reduce validation loss, and a characterization of how that loss behaves in the long run when compute is infinite.
As data stops scaling but compute keeps growing, validation loss approaches a horizontal asymptote modeled as a power-law function.
FOUNDATIONAL PAPERS WALK-THROUGH
Starting with foundational papers on KV-caching and prefilling, then Disaggregation for distributed LLM inference.
ATTENTION IS ALL YOU NEED
The self-attention mechanism computes a matrix denoting the importance of each token relative to the others.
Self-attention is foundational to all modern LLMs. The breakthrough here was training speedup, not inference - inference is still autoregressive like traditional methods. Builds a good foundation for understanding KV caching.
FLASHATTENTION
Fast and memory-efficient exact attention with IO-awareness. Achieved via block-matrix multiplication (tiling) in low-level GPU kernels.
Not relevant - I will not be writing low-level GPU kernels. Still valuable for understanding the concept of IO-awareness when reasoning about attention cost.
DISTSERVE - DISAGGREGATING PREFILL & DECODE
Prefill / Decoding disaggregation, introduced in 2024 for distributed LLM inferencing.
PREFILL STAGE
When a user submits an initial prompt, the entire token sequence passes through the network in parallel.
- Attention matrices are computed for large chunks of text at once.
- Parallelized, not autoregressive - all tokens in the prompt processed together.
- Compute-heavy by definition.
- This initial processing is known as the PREFILL STAGE.
DECODE STAGE
- After the prompt is processed, the response is generated autoregressively.
- Output tokens are fed back into the decoder to predict the next token.
- No need to recompute the full attention matrix - only an additional row and column are added.
- Not compute-heavy - memory-heavy. This is the DECODE STAGE.
PREFILL & DECODE - VISUAL OVERVIEW
A visual of the prefill and decode stages.
THE PROBLEM & THE SOLUTION
Running both phases on the same hardware wastes resources. During decode, GPUs are under-saturated - few operations, lots of idle compute.
Split the phases onto different hardware tailored to each workload.
- Cluster A - tailored for compute (Prefill)
- Cluster B - tailored for memory (Decode)
Very relevant to my research direction.
MODELING THE PREFILL-TO-DECODE PENALTY
What is the performance cost of running prefill and decode together on a single GPU (vLLM inference framework)? How does that penalty grow as the workload becomes more prefill-heavy?
TIME BETWEEN TOKENS vs. BATCH COMPOSITION
Single GPU, 8 concurrent requests - how Time Between Tokens changes as the batch becomes prefill-heavy vs. decode-heavy.
BATCH COMPOSITION SETUP
LEARNED POLYNOMIAL - A SCHEDULING HEURISTIC
f(r) = 1.63r² + 1.39r + 14.52- The penalty of using one GPU for prefill + decode can be modeled along the Time-between-Token vs. batch composition axis.
- The gap between
r = 0.0andr = 1.0is the benefit you gain from disaggregation (prefill on GPU A, decode on GPU B). - Fit via ordinary least squares regression.
- Pretty reliable - coefficient of determination R² = 0.90.
- The mechanism of this scheduling heuristic is still being explored.
RESERVING COMPUTE WITH THE SCHOOL
Distributing prefill and decoding requires larger clusters. Target: Carleton RCDC.
AVAILABLE GPUS - GENERAL PURPOSE
| GPU | VRAM | QTY |
|---|---|---|
| NVIDIA H200 | 141 GB | 8 |
| NVIDIA A100 | 80 GB | 8 |
| NVIDIA A100 | 40 GB | 16 |
| NVIDIA L40S | 48 GB | 8 |
| NVIDIA V100 | 16 GB | 2 |
| NVIDIA GTX 1080 Ti | 11 GB | 16 |
| AMD MI210 | 64 GB | 2 |
MISC ACTIONS
Gain access to the lab website GitHub repo.
Set up a gathering of the lab for early fall. Send out forms and coordinate a date with all members.
Do we have budget for a lab lunch?
NEXT PAPERS
- 01GPipe - multi-GPU training and serving.
- 02Megatron-LM - how billion-parameter models fit on GPU clusters.
DISCUSSION
JUSTIN ZHANG - SWAN LAB - MCS MEETING AUG 6