Satyam Agnihotri

Speedrunning Inference Engineering Day 2: First-order inference performance math The goal of today is to create a small Python program that takes in model config, batch size, context length, hardware, etc. and outputs the estimated vLLM performance metrics in prefill and decode. This is more about deriving metrics from first principles to gain intuition than actually profiling a real system. We want to predict: Model-weight memory, KV cache memory, decode FLOPs, decode memory traffic, compute and bandwidth lower bounds, decode iteration time and output throughput (TOPS), and whether a workload is weight, KV, or compute bound. Intrinsically, a load is always going to be either compute OR memory bound. Our goal in inference engineering is to bring this ratio as close to unity as possible such that compute and memory are waiting on one another as little as possible To make these calculations, we need to get model params and GPU params. I am using the common mans Llama 3.1 8B Instruct and a single H100 (lets not overcomplicate things on day2) Llama 3.1 8B model params are hidden_size 4096 (d_model, thus each d_head = 4096 / 32 = 128) intermediate_size 14336 num_hidden_layers 32 num_attention_heads 32 num_key_value_heads 8 (KV heads ≠ Q heads bc GQA model) vocab_size 128256 torch_dtype bfloat16 (precision, no quantization) max_position_embeddings 131072 (context length) H100 params are 80GBs HBM3 700W peak power 132 SMs 3.35 TB/s HBM bandwidth 989 TFLOPs of BF16 compute Obviuosly these are all peak performance numbers and represent an ideal. Most production systems rarely performance at these numbers for any significant duration of time, but our goal is to maximize that time. Now for Llamas exact size to get an idea of where its parameter allocation lies Token embeddings 525,336,576 params 1,050,673,152 bytes Attention weights 1,342,177,280 params 2,684,354,560 bytes MLP weights 5,637,144,576 params 11,274,289,152 bytes Normalization 266,240 params 532,480 bytes Output lm_head 525,336,576 params 1,050,673,152 bytes Total 8,030,261,248 params 16,060,522,496 bytes So we can see tat the vast majority of its weights still lay in the MLP, rather than attention or embedding/un-embedding (in the lm_head) Now to calculate the total KV memory per user. If you remember from the last post, the equation is 2 * num_layers * num_kv_heads * d_head * precision * seq_length, so with the Llama numbers above at BF16 (2 bytes), a single token of KV adds 131KBs to the cache (~4KBs per layer) So across seq_lenghts, we have 128 tokens → 16.0 MiB 1,024 tokens → 128.0 MiB 2,048 tokens → 256.0 MiB 4,096 tokens → 512.0 MiB 8,192 tokens → 1.0 GiB 16,384 tokens → 2.0 GiB 32,768 tokens → 4.0 GiB 65,536 tokens → 8.0 GiB 131,072 tokens → 16.0 GiB ← model's max context Thus at max length, we have at most 32GiBs sitting in HBM. This cleanly fits in an H100, but the better question is how does the growing KV affect bandwidth and SM utilization? Now to derive transformer linear layers FLOPs (the linear layers are just the matmuls that require linear projections, so q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, and down_proj) For the attention projections, to get q_proj, we multiply a (1,4096) tensor by the (4096, 4096) W_Q matrix. Same math for our eventual o_proj. So these each give us d^2 params for each of these matrices per token and thus 2d^2 FLOPs (forward pass FLOPs is just one addition and multiply per param). For the K and V projections, we only multiply a (1, 4096) tensor by the K and V (4096, 1024) matrices. These matrices have smaller ouptuts (1/4th the size of Q) because for every 4 Q heads, we only have 1 K and V head (again, the entire idea behind GQA). Therefore, each K and V projection gives us d * d_head * H_KV params. So total FLOPs for the attention projections is (4 * d^2) + (4 * d * (d_head * H_KV)) Now for the MLP layers, each SwiGLU MLP has three matrices (gate_proj, up_proj, and down_proj). The gate and up projection are in parallel, each taking the embedding (1, 4096) and computing (4096, 14336) taking the embedding now to be a (1, 14336) tensor. When it is time for the activation, the up_proj computes a candidate value for each of the 14336 neurons, while the gate_proj acts as a learned valve on how strongly each neuron fires. So after each of these tensor operations, we run the non-linearity as h = silu(g) * u (we actually run the non-linearity on the gate projection, not on the up projection). Then, we take this product h and down_project it back into d_model space: (1, 14336) * (14336, 4096) -> (1, 4096). Thus we can see for all of these operation weoperate on 3 * d_model * d_ff (intermediate dimension) params and thus 6 * d_model * d_ff FLOPs So in total we have (4 * d^2) + (4 * d * (d_head * H_KV)) + (6 * d_model * d_ff) FLOPs. Then we multiply this all by L layers. The last thing to remember is the final output projection (lm_head calculation) which takes the (1, 4096) embedding tensor and maps it to a vocab_size-d dimension vector (each dimension here representing a token probability). Thus (1, 4096) * (4096, 128256) -> (1, 128256), which is the final vector yielding a probability for every possible next token. We take softmax on this, sample, and then return that next token. This shows we operate on d_model * vocab_size params and thus 2 * d_model * vobab_size FLOPs Aggregating these terms across 32 Layers yields ≈ 15 GFLOPs per forward pass which is what we expect as the rough rule of thumb is 2 * num_params (this is a slighly overestimation as it ignores embedding params which is a row gather, not a matmul) Now to calculate the decode attention FLOPs. For each generated token at context length N, each qeury head computes QK^T across all N (assuming dense attention) and a weighted sum of this dot product over all N value vectors. So o = qk^T and o * v This operates on 2 * num_layers * seq_lenght * Q_h * d_head params, thus 4 * num_layers * seq_length * Q_h * d_head FLOPs. Since Q_h * d_head = d (bc Q_h = d / d_head), we can simplify to 4 * num_layers * seq_length * d. Thus, total decode FLOPs per sequence are F_fixed + (4 * num_layers * seq_length * d). We just derived F_fixed ≈ 15 GFLOPs, so total FLOPS ≈ 15 GLOPs + 4LNd Across context lengths, this gives Context: 128 Fixed FLOPs: 15.01 GF Attention FLOPs: 0.07 GF Total: 15.08 GF Context: 2K Fixed FLOPs: 15.01 GF Attention FLOPs: 1.07 GF Total: 16.08 GF Context: 8K Fixed FLOPs: 15.01 GF Attention FLOPs: 4.29 GF Total: 19.30 GF Context: 16K Fixed FLOPs: 15.01 GF Attention FLOPs: 8.59 GF Total: 23.60 GF Context: 32K Fixed FLOPs: 15.01 GF Attention FLOPs: 17.18 GF Total: 32.19 GF Context: 64K Fixed FLOPs: 15.01 GF Attention FLOPs: 34.36 GF Total: 49.37 GF Context: 128K Fixed FLOPs: 15.01 GF Attention FLOPs: 68.72 GF Total: 83.73 GF Although, everyone knows that bc of the KV cache, decode is memory bound. We aren't really so concerned about FLOPs when most of the SMs are sitting idle waiting for memory. So lets model bandwidth. Per step, from HBM, the GPU needs to pull active weights and all KV cache and push one new KV token per sequence. So with context length = N and batches = B, to a first order approximation, Total Bandwidth = model_weights + (B * N * KV_user) + (B * KV_new_token). Model weights are shared across the batch, in one decode step all sequences pass through the same layer together. KV traffic however is individual to each sequence (this is where continious batching will come in later). So roughly we can model memory byte movement growing as B=1 128: 16.08 2K: 16.33 8K: 17.13 16K: 18.21 32K: 20.36 64K: 24.65 128K: 33.24 B=8 128: 16.20 2K: 18.21 8K: 24.65 16K: 33.24 32K: 50.42 64K: 84.78 128K: OOM B=16 128: 16.33 2K: 20.36 8K: 33.24 16K: 50.42 32K: 84.78 64K: OOM 128K: OOM B=32 128: 16.60 2K: 24.65 8K: 50.42 16K: 84.78 32K: OOM 64K: OOM 128K: OOM B=64 128: 17.14 2K: 33.25 8K: 84.79 16K: OOM 32K: OOM 64K: OOM 128K: OOM B=128 128: 18.22 2K: 50.44 8K: OOM 16K: OOM 32K: OOM 64K: OOM 128K: OOM OOM = crossing 64.7 GiB which saturates H100 total HBM when including weights Ok now with all of this info lets try calculating memory and compute bounds. If we were memory bound (i.e. have ample idle compute), the time per decode step would be memory_per_step / HBM_BW. We just calculated memory_per_step is dependent on context and batching. So to see this in action I had Claude make a chart (first one attached) We calculated that F_step = B * F_decode_per_token = B * (F_fixed+ 4LNd). So if we were compute bound, time per decode step would be F_step / H100_FLOPs (H100 FLOPs = 989TFLOPs). This is represented in the second chart attached. However because there is so much more compute on an H100 than relative memory, we are essentially never compute bound (except at high batch prefill). This shows up clearly comparing the two charts. I think going forward I am going to make these way shorter and more approachable. Or do people like them long and dense like this? Lmk!