LLM Inference Optimization: KV Cache, Quantization & GPUs

The GPU That's Idle While You Pay For It

You provisioned a single A100 for your 8B-parameter chat model. The back-of-envelope math said 16 GB of weights against 80 GB of VRAM — comfortable headroom. Then the profiling report comes back and the GPU sits at 12% utilization while your users wait forty seconds per reply.

Nothing is broken. The GPU is doing exactly what the workload asks of it. The first profiling report I ever read for a chat model showed tensor cores nearly idle while the HBM interface ran saturated — that report changed how I think about every serving decision since. The mismatch isn't hardware, it's that LLM generation is two different workloads wearing the same coat.

Prefill vs. Decode: The Two-Phase Tax

Autoregressive inference splits into two regimes with opposite bottlenecks:

• Prefill is compute-bound. The prompt's tokens are scored simultaneously in one big parallel matmul, which is exactly what tensor cores love. Given a big prompt, the GPU hums.
• Decode is memory-bandwidth-bound. Output tokens are produced one at a time. Each step re-reads the entire weight matrix from HBM into on-chip SRAM to emit a single token. The FLOPs are trivial; the memory traffic is enormous, and arithmetic intensity collapses.

Everything that follows is a way of paying this tax less often.

The KV Cache Is the Real Memory Problem

During decode, attention needs the key and value projections of every previous token. Recomputing them each step would turn $O(N)$ decoding into $O(N^2)$; instead, serving engines cache them. That cache is not free. Its size grows with sequence length, batch size, and hidden dimensionality, and the per-token footprint is easy to derive:

$\text{KV cache size per token}=2\cdot N_{\text{layers}}\cdot N_{\text{heads}}\cdot d_{\text{head}}\cdot \text{bytes per element}$

Plug in an 8B model with 32 layers, 32 heads, a head dimension of 128, and a 32K-token context in FP16: $2\times 32\times 32\times 128\times 2$ bytes per token — 512 KiB of cache per token. At 32K tokens that's 17.2 GB, larger than the 16 GB of weights the model occupies. The thing every deployment guide forgets: past a few thousand tokens of context, the KV cache, not the model, is what fills the GPU. That's why context windows cost so much, and it's why cache management turned into a research field of its own.

PagedAttention: Virtual Memory for the Cache

Traditional serving frameworks pre-allocate contiguous memory for the maximum possible context length (say 2048 or 4096 tokens) per request. Real sequences end at wildly different points, so most of that reservation sits unused — the PagedAttention paper measured 60–80% of VRAM wasted to such fragmentation on typical workloads.

PagedAttention imports the operating system's answer: virtual memory paging. The cache is carved into fixed-size physical blocks, allocated on demand, and logical token sequences point to physical blocks through a block table. Blocks don't need to be contiguous, so fragmentation disappears; blocks can be shared, so multiple generation streams from the same prompt — parallel sampling, beam search, speculative verification trees — reuse one physical copy of the shared prefix instead of each holding their own. That sharing alone cuts memory overheads by up to 55% on parallel workloads, and the whole package buys up to 2.2× throughput on memory-constrained workloads. The cost is a small indirection penalty per block lookup — easily worth it.

Static Caching: Predictability as an Optimization

There's a second, simpler axis: instead of growing cache arrays on demand, pre-allocate contiguous memory for the full sequence length up front. The payoff is kernel friendliness — fixed shapes let execution frameworks fuse kernels and keep the GPU busy without dynamic bookkeeping. PyTorch reports up to a 4× forward-pass speedup for static-cache workflows, a number that depends heavily on model and batch size, so treat it as an upper bound, not a promise.

Continuous Batching: Stop Waiting for the Slowest Request

With request-level batching, the whole batch waits for its slowest sequence to finish generating. In practice a batch of eight answers arrives at eight wildly different times, and the GPU idles through the stragglers.

Orca replaced this with iteration-level scheduling: instead of executing a static batch to completion, the engine invokes the model one autoregressive step at a time, across whatever requests are live. The moment one sequence emits <eos>, its KV blocks are released and a queued request takes its place in the active batch. Orca also introduced selective batching: the matrix multiplications (which benefit from being batched across all active sequences) are executed together, while attention is evaluated per sequence, since each sequence attends over its own context. The combination reports up to 36.9× throughput over request-level engines like FasterTransformer at comparable latency. If you run any production API and haven't enabled continuous batching, this is the cheapest, highest-leverage fix on the list — most engines have it as a default now, but it's worth confirming your config actually uses it.

The Prefill Sneak Attack: Sarathi-Serve

Continuous batching has a hidden failure mode. A long prompt prefill is a compute-heavy chunk that hogs the GPU; inject one into an active batch of memory-bound decodes and every in-flight sequence suffers a latency spike — time-between-tokens jumps from 40 ms to 2 seconds for everyone. Sarathi-Serve fixes the asymmetry with chunked prefill: split long prompts into uniform token chunks and piggyback them onto ongoing decode iterations. Every iteration becomes roughly equal-length, so batch members see steady, predictable timing. It also keeps pipeline-parallel workers uniformly busy, and reports 2.6×–5.6× more serving capacity under tail-latency bounds than the standard scheme. The lesson worth internalizing: in serving, who you batch with is as important as how many you batch.

Speculative Decoding: Spending FLOPs to Save Wall Time

Decode is slow because it's sequential, and each step is memory-bound. Speculative decoding inverts the logic: let a small, fast draft model $M_q$ guess the next $K$ tokens, then have the big model $M_p$ verify all $K$ guesses in a single parallel forward pass. If the guesses are good, one memory-bound step produces several tokens. The full algorithm, in order:

1. Sample $K$ candidate tokens from the draft model's distribution $q$.
2. Run the target model once and evaluate its distribution $p$ over all $K$ positions at once.
3. Walk the candidates: accept token $x$ when $p(x)\ge q(x)$, otherwise accept it with probability $P(\text{accept})=\min \left(1,\frac{p(x)}{q(x)}\right)$.
4. If a token is rejected at index $i$, discard every candidate after it and resample a replacement from the corrected distribution:

$p^{\prime }(x)=\frac{\max (0,p(x)-q(x))}{{\sum }_y\max (0,p(y)-q(y))}$

That correction is what makes the scheme exact: the final output distribution is mathematically identical to sampling from the target model alone, and the guarantee costs you nothing in quality. Every pass emits at least one verified token (the worst case where the very first draft token is rejected) and at most $K+1$ (when all $K$ guesses are accepted, you also grab a bonus token from the target distribution). Reported results land at 2×–3× latency reductions.

The honest caveat: it only pays off when the draft model is genuinely good at predicting the target. If your acceptance rate drops below roughly 0.6, the verification overhead eats the gains.

Draft Strategies Beyond the Little Model

The draft model is the weak link, so the research has attacked it from three directions:

• Hierarchical speculative decoding stacks drafts into a tiered ladder — a tiny model proposes, a mid-sized model validates and proposes onward, the target model verifies at the top — with the schedule itself optimized as a dynamic programming problem.
• Medusa gets rid of the separate draft model entirely: it bolts multiple lightweight feed-forward heads onto the target model's top hidden state, each predicting a future token offset, and fine-tunes them while freezing the backbone. No extra model to host, no extra VRAM for weights.
• Prompt Lookup Decoding skips learned models completely. For retrieval-grounded workloads — RAG summarization, document translation — the output constantly reuses n-grams that already exist in the prompt. It extracts those matches and feeds them to the target for verification, delivering 1.5×–2.5× speedups on exactly the workloads where it applies, with zero additional VRAM.

When One GPU Isn't Enough

Once a model won't fit on a single accelerator, you're choosing how to cut it apart:

Pipeline parallelism cuts the model along depth: GPU 0 holds layers 1–8, GPU 1 layers 9–16, and so on. The catch is the pipeline bubble — downstream GPUs idle waiting for activations to arrive, and naive pipelines can leave hardware asleep most of the time. Two mechanisms close the gap. Micro-batching splits each batch into chunks so workers always have something to do. 1F1B scheduling (one forward, one backward, in strict alternation per worker) interleaves work so the backward pass of an older micro-batch fills the idle slots of newer forward passes. The subtlety that makes 1F1B correct: because a micro-batch's backward pass happens long after its forward pass, each worker must stash the exact weight versions it used at forward time, and replay them for the gradient computation — a technique called weight stashing that doubles as its own memory cost. The follow-up 2BW refinement double-buffers weights and coalesces gradient exchanges, restoring single-GPU update equivalence and shaving up to ~20% off peak memory on giant models. Idle time never fully disappears from pipelines — the bubble is physics, not a bug — but scheduling shrinks it from "most of the time" to "a few percent."

Tensor parallelism splits weight matrices inside each layer: column-parallel MLP slices let GPUs compute $Y_1=X{W}_1$ and $Y_2=X{W}_2$ independently before an all-reduce stitches the result together, and attention heads partition across devices with no inter-head communication. Because every layer needs an all-reduce, TP lives or dies on interconnect: it wants NVLink-grade bandwidth, not commodity Ethernet.

My default rule for a single-node cluster: tensor parallel first, pipeline parallel only when the model no longer fits with TP alone, data parallel for throughput on replicated models. If your GPUs sit in separate machines, the calculus flips — TP's per-layer syncs will drown in network latency.

Quantization: The Easiest Win, Read the Fine Print

Dropping from FP16 to INT8 or INT4 halves or quarters both the VRAM footprint and the bytes that have to stream through HBM each decode step. Post-training quantization (PTQ) gets most of the way there without retraining:

• LLM.int8() was the breakthrough that made 70B models usable on a single node: it detects the ~0.1% of activation outlier channels whose extreme values would destroy INT8 ranges, keeps them in FP16, and quantizes the rest vector-wise. Historically important, and in practice superseded by the 4-bit methods below.
• GPTQ compensates for quantization error as it goes: each column's quantization error is corrected by updating the remaining unquantized weights using a second-order approximation of the inverse Hessian $H^{-1}$. It's excellent for 4-bit weight-only compression, but the calibration pass is computationally heavy and its weights can overfit the calibration set.
• AWQ observed that what matters isn't which weights are large but which channels have large activations — protecting just 0.1–1% of salient channels (identified by activation magnitude, not weight magnitude) beats uniform rounding. Instead of mixed precision, it applies a per-channel scale $s$ to the weights before quantizing, chosen by searching a small grid over $\alpha$ on a calibration set:

$s={\left(\frac{\max |x|}{\max |w|}\right)}^{\alpha }$

Scaling up a salient channel shrinks its relative quantization step, protecting it without any hardware-unfriendly mixed-precision layout. There's no backprop and no weight reconstruction, which keeps out-of-domain generalization strong. Integrated into execution frameworks like TinyChat, AWQ delivers 3×–4× memory reduction with practical speedups over FP16 baselines on desktop and mobile GPUs.

• AQLM pushes into sub-3-bit territory (2-bit) by reformulating weight compression as a vector quantization problem over learned codebooks — several codebooks sum to approximate each weight vector — rather than per-scalar rounding. It's Pareto-optimal relative to footprint at the extreme end of the compression curve, and the price is a more complex kernel story.
• EXL2 (ExLlamaV2's format) and GGUF (llama.cpp) are the ecosystem formats: GGUF with block quantizations like Q4_K_M targets CPU and edge with dynamic GPU offloading; EXL2 supports variable bitrates per layer, spending more bits on sensitive attention layers and fewer on redundant MLPs.

The honest caveat: 4-bit weights measurably degrade long-tail reasoning and multilingual output, and the degradation is worse the smaller the model. Calibrate with data from your distribution — a calibration set tuned on English chat may surprise you on code or legal text.

The Engine Is a Decision, Not a Default

The big serving engines have converged on the same feature set, so the choice comes down to operations, not checkboxes:

If you don't already know why you need something else, start with vLLM — PagedAttention is a first-class citizen there, the community integrations are the best, and the codebase is the most approachable to debug. TensorRT-LLM earns its keep when you control the whole stack (NVIDIA GPUs with NVLink, a tuned engine build) and need every last token per second, or need exact per-kernel determinism. TGI is the pick when you're already inside the Hugging Face ecosystem and want the smoothest model-loading path.

On the math side, FlashAttention-2 is worth understanding even though it ships in all three engines: it tiles Q, K, V into SRAM blocks and computes exact attention with an online softmax that maintains a running maximum $m^{(i)}$ and exponential sum ${\ell }^{(i)}$, never materializing the quadratic attention matrix in main memory — exact attention in $O(N)$ memory, with kernels that reach 50–70% of theoretical peak FLOPs on modern hardware. Pairing it with PagedAttention and a 4-bit format like AWQ is the combination that dominates most production budgets today.

The Next Bottleneck

The decode tax exists because autoregression is sequential and attention grows with context. Every trick above works around that — speculative decoding guesses ahead, quantization shrinks the bytes per step, parallelism spreads the load. The more interesting question is what happens when the underlying architecture changes: linear attention, state-space models, and KV compression all attack the cache directly, and early serving stacks for them already look very different — the latest systems are disaggregating even multimodal pipelines into separate compute paths with record-and-replay data planes. If decode stops being memory-bound, the entire optimization hierarchy collapses into a new one — and the teams that understood the current bottleneck will be the ones ready to rebuild it. What will you bet that KV caches still dominate GPU budgets five years from now?

References

• Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). Attention Is All You Need. NeurIPS.
• Kwon, W., Li, Z., Zhuang, S., et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICML.
• Yu, G.-I., Jeong, J. S., Kim, G.-W., et al. (2022). Orca: A Distributed Serving System for Transformer-Based Generative Models. OSDI.
• Agrawal, A., Kedia, N., Panwar, A., et al. (2023). Sarathi: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills. arXiv:2303.01469.
• Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. ICML.
• Chen, C., Borgeaud, S., Irving, G., et al. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arXiv:2302.01318.
• Cai, T., Li, Y., Geng, Z., et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. arXiv:2401.10774.
• Huang, Y., Cheng, Y., Bapna, A., et al. (2019). GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism. NeurIPS.
• Shoeybi, M., Patwary, M., Puri, R., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053.
• Narayanan, D., Harlap, A., Phanishayee, A., et al. (2019). PipeDream: Generalized Pipeline Parallelism for DNN Training. SOSP.
• Narayanan, D., Phanishayee, A., Shi, K., et al. (2021). Memory-Efficient Pipeline-Parallel DNN Training. ICML.
• Dettmers, T., Lewis, M., Belkada, Y., & Zettlemoyer, L. (2022). LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale. NeurIPS.
• Frantar, E., Ashkboos, S., Hoefler, T., & Alistarh, D. (2022). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers. ICLR 2023.
• Lin, J., Tang, J., Tang, H., et al. (2024). AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration. MLSys.
• Egiazarian, V., Panferov, A., Kuznedelev, D., et al. (2024). Extreme Compression of Large Language Models via Additive Quantization. ICML.
• A Distributed Serving System for Any-to-Any Multimodal Models. arXiv:2603.12118.
• Prompt Lookup Decoding: community technique, github.com/apoorvumang/prompt-lookup-decoding.