LLM Serving Capacity: KV Cache Math and Decode Throughput

Preview mode. Browse milestones and edit starter code here. Sign in to run and submit.

0/8EDA

Python

advanced~7h8 milestones

The problem

You are the inference engineer on a team that serves open-weight language models. Finance wants to know how many GPUs next quarter needs, and product wants 32k-token context on the largest model. Both answers come from the same place: the KV cache. Every request keeps a key and a value vector for every token, in every layer, for as long as it is running, and every decode step reads all of them back from GPU memory along with every weight. You have 16,000 benchmark runs across 48 model configurations and 5 GPU classes. Your job is to compute the cache from the model shape, work out what fits, prove the cache is exact, and build a throughput model the capacity planners can trust.

Your task · Profile the Benchmark Runs

Count the runs, classify each model's attention layout, and measure how wide the throughput range is.

Each row is one benchmark run: a decoder-only model (n_layers, d_model, n_heads, n_kv_heads, head_dim, params_b), stored at a precision (weight_bytes_per_param and kv_bytes_per_elem are 2 for 16-bit, 1 for 8-bit, 0.5 for 4-bit), served on one GPU (gpu_class, gpu_mem_gb, hbm_bandwidth_tbps, peak_tflops), holding batch_size concurrent requests of context_len tokens each. The label is decode_tokens_per_s: generated tokens per second summed over the whole batch.

Implement explore_runs(train_df) returning a dict with:

row_count
model_count: distinct values of model
gpu_class_count: distinct values of gpu_class
attention_variants: a dict with keys "MHA", "GQA", "MQA" and the number of ROWS of each. A row is MHA when n_kv_heads == n_heads, MQA when n_kv_heads == 1, and GQA otherwise.
median_context_len and median_batch_size, as floats
median_tokens_per_s: median decode_tokens_per_s, rounded to 2 places
throughput_spread: max decode_tokens_per_s divided by min, rounded to 1 place

The spread is the reason this project grades on relative error. A model that is 50 tokens per second off is excellent on one row and useless on another.

Constraints

▪GB means 10^9 bytes throughout
▪Only 92% of GPU memory is usable; the rest goes to activations, the allocator and the runtime
▪Every run in the data fits in memory: weights plus batch x context x KV bytes per token stay under that 92%
▪Throughput spans more than three orders of magnitude, so the grader uses relative error (MAPE, sMAPE), not only R^2
▪The grading sandbox has NumPy, pandas, SciPy and scikit-learn. No PyTorch.
▪Hidden test runs are scored server-side

Graded on

KV bytes per token, capacity and the KV cache proof match the reference to 1e-4
Roofline model: MAPE <= 0.09 and R^2 >= 0.93 with four fitted coefficients
Throughput model: MAPE <= 0.062 and R^2 >= 0.97
Deployed model: sMAPE <= 0.065 within 2 seconds

Evaluated server-side against a hidden test set.

Skills you'll demonstrate

KV CacheAttention MechanismGrouped-Query AttentionLLM InferenceRoofline ModelingFeature EngineeringCapacity Planning