LLM Inference Internals: KV Cache, Sampling, and Serving at Scale

Every time you call an LLM API, inference runs in two distinct phases with radically different performance characteristics.

Prefill processes your entire prompt at once. Because all input tokens are known upfront, the GPU can compute attention across all of them in parallel — this is essentially a giant matrix multiply. A 1,000-token prompt typically takes ~50–200 ms on a modern A100.

Decode generates one token at a time. Each step attends over the entire context (prompt + tokens generated so far), produces a probability distribution, samples one token, and feeds it back as input for the next step. This is inherently serial — you cannot generate token n+1 until you have token n. Generating 200 tokens might take 2–5 seconds on the same hardware.

The consequence: latency-sensitive applications care more about decode throughput, while cost-sensitive batch jobs care about prefill efficiency. Almost every optimization in LLM serving targets one of these two phases.

Prefill fills the KV cache in one shot; every decode step reads from it and appends one new row.

Transformer attention computes Q (query), K (key), and V (value) matrices for every token at every layer. During decode, the new token's Q attends over the K and V of all previous tokens. Without caching, you'd recompute every past token's K and V on every step — O(n²) wasted work.

The KV cache stores K and V tensors for every past token, layer by layer. Each decode step only computes K/V for the one new token, then appends them to the cache.

Memory cost is real. For a 70B parameter model with 80 layers, 8 KV heads, and head dimension 128, each token costs:

$\text{bytes} = 2 \times \text{layers} \times \text{heads} \times d_{head} \times \text{bytes\_per\_element}$

With fp16: $2 \times 80 \times 8 \times 128 \times 2 = 327{,}680$ bytes ≈ 320 KB per token. A 4,096-token context consumes ~1.3 GB — just for the KV cache, not model weights. This is why serving long contexts is expensive and why context length is a hard engineering constraint, not an arbitrary cap.

After the forward pass, the model outputs a logit vector over its vocabulary (~32k–128k tokens). Sampling turns that into a single discrete token.

Temperature scales logits before softmax: $\text{logit}_i' = \text{logit}_i / T$. Low T (0.1–0.4) sharpens the distribution — the model almost always picks the highest-probability token. High T (1.2–2.0) flattens it, producing more creative but less coherent output. temperature=0 is greedy decoding: always argmax.

Top-p (nucleus sampling) sorts tokens by probability, accumulates until the cumulative sum reaches p, then samples only from that nucleus. With top_p=0.9, you sample from whatever subset of tokens accounts for 90% of the mass — typically 10–500 tokens depending on confidence. This adapts dynamically: when the model is certain, the nucleus is tiny; when unsure, it's larger.

A common production default: temperature=0.7, top_p=0.9. Setting both to 1.0 gives unfiltered multinomial sampling.

Top-k is the classic third filter: keep only the k highest-probability tokens. Its problem is that k is context-blind. When the model is 99% confident about a token, keeping the top 50 alternatives still lets junk through. When the model is genuinely uncertain, top-50 might cut off good options.

min-p (introduced by Nguyen et al., 2023) sets a relative floor: keep any token whose probability is at least min_p × p_max, where p_max is the probability of the top token. With min_p=0.05, if the top token has probability 0.6, any token with prob ≥ 0.03 is kept; if the top token has prob 0.1, any token ≥ 0.005 is kept.

This scales naturally with model confidence. It tends to produce tighter, more coherent output than top-k while avoiding the truncation artifacts of a fixed k. llama.cpp, Kobold, and several inference servers now support it:

# llama-cpp-python example
llm(prompt, min_p=0.05, temperature=0.8)

min-p and top-p can be combined — apply min-p first (relative floor), then top-p on the surviving tokens.

The fundamental bottleneck of decode is that you need one forward pass per token. Speculative decoding (Leviathan et al., 2022; Chen et al., 2023) breaks this by using a cheap draft model to speculate several tokens ahead, then verifying them in a single large-model forward pass.

1. The draft model (e.g., a 7B) autoregressively generates k candidate tokens (k=4–8 is typical).
2. The target model (e.g., a 70B) runs one forward pass over the prompt + all k draft tokens simultaneously — this is parallelizable like prefill.
3. Each draft token is accepted or rejected using a rejection sampling scheme that guarantees the output distribution matches the target model exactly.
4. On rejection, the target model's own token is used and drafting restarts.

Speedup depends entirely on the draft acceptance rate. On-domain text (code, formal writing) where the draft model is confident often achieves 2–3× wall-clock speedup. Off-domain or creative text may yield only 1.2×. Google's PaLM 2 paper reported ~2× on coding benchmarks using a 2B draft.

The obvious way to serve multiple users is to batch their requests: collect N prompts, run them through the model together, return results. GPUs love batching — matrix multiplies scale efficiently with batch size.

The problem is padding and variable-length sequences. If request A needs 2,000 tokens and request B needs 200, you must pad B to 2,000 tokens or run two separate batches. Padding wastes computation. Waiting to fill a batch adds latency.

Worse: with static batching, all requests in a batch finish together. If A finishes at step 100 and B at step 500, the batch slot for A sits idle for 400 steps. On a busy server, utilization collapses.

Request A: [tok1][tok2]...[tok100][PAD][PAD]...[PAD]  ← 400 wasted steps
Request B: [tok1][tok2]...[tok500]

This is why naive batching struggles to serve both low-latency and high-throughput workloads simultaneously. The solution is continuous batching.

Continuous batching (Yu et al., 2022) decouples sequence completion from batch scheduling. Instead of running fixed batches to completion, the server maintains a running pool of requests. At each decode step, it:

1. Runs one forward pass over all currently-active sequences.
2. Checks which sequences generated an EOS or hit max_tokens.
3. Immediately inserts new waiting requests into freed slots.

This keeps GPU utilization near 100% because slots are never idle — the moment a sequence finishes, a new one starts.

vLLM (Kwon et al., 2023) adds PagedAttention: it manages KV cache memory like an OS page table, allocating cache in fixed-size blocks (pages) rather than contiguous buffers. This eliminates fragmentation and allows the physical KV memory for one sequence to be scattered across non-contiguous blocks — borrowed from any free page in a shared pool. The result: vLLM can pack far more concurrent sequences into GPU memory than naive implementations.

pip install vllm
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Meta-Llama-3-8B-Instruct \
  --max-num-seqs 256

The key insight: treat the decode loop as a scheduler tick, not a batch boundary.

Rough GPU cost breakdown for a typical production deployment (A100 80GB, ~$2–3/hr on cloud):

A practical benchmark: a single A100 running LLaMA-3 70B with vLLM and INT4 quantization (via AWQ) can sustain ~800–1,200 tokens/sec throughput at batch sizes of 32–64. At $2.50/hr, that's roughly **$0.0006–0.0009 per 1K tokens** — an order of magnitude cheaper than GPT-4 API pricing for equivalent quality.

Cost per token = (GPU $/hr) ÷ (tokens/sec × 3600). Throughput is the lever; latency is a separate SLA.

Model weights in fp16 require 2 bytes per parameter. A 70B model needs ~140 GB — more than fits on a single A100 (80 GB). Quantization compresses weights to INT8 (1 byte) or INT4 (0.5 bytes).

GPTQ (post-training quantization) minimizes reconstruction error layer-by-layer, typically achieving <1% perplexity degradation at INT4. AWQ (Activation-aware Weight Quantization, Lin et al., 2023) identifies the 1% of weights that dominate output variance and protects them, reaching similar quality with less tuning.

KV cache quantization is a separate lever — storing K and V in INT8 instead of fp16 halves cache memory at the cost of minor attention precision loss.

Common gotchas:

• INT4 quantization degrades more on math and code tasks than on prose.
• Quantized models can still be memory bandwidth-bound during decode — smaller weights don't help if the bottleneck is reading them off HBM.
• Always benchmark on your actual workload distribution, not just perplexity on Wikitext.

A production inference stack typically layers these components:

Client → Load balancer → Router (prefix-aware)
       → vLLM workers (continuous batching, PagedAttention)
         ├── Draft model (speculative decoding)
         ├── KV cache (paged, INT8)
         └── Quantized weights (AWQ INT4)

Prefix caching is a further optimization: if many requests share the same system prompt, their KV cache prefix is computed once and reused across requests. vLLM supports this via --enable-prefix-caching. For chat applications with a fixed system prompt, this can eliminate prefill cost entirely for the shared portion.

The decode throughput equation: $\text{tokens/sec} = \frac{\text{batch\_size} \times \text{1}}{\text{time\_per\_step}}$. Time per step depends on memory bandwidth (reading weights + KV cache), not FLOPs — which is why A100s with 2 TB/s HBM3 outperform cards with similar peak FLOP counts but slower memory.