Efficient Inference

Every operation on a GPU has two costs: FLOPs (compute) and bytes moved between HBM and the compute units on the chip (memory transfer). The ratio is arithmetic intensity.

The H100 has 989 TFLOPs (BF16) and 3.35 TB/s HBM bandwidth. Divide: the break-even intensity is 989e12 / 3.35e12 ≈ 295. Any operation with intensity > 295 is compute-bound (good — you're using the silicon). Any operation with intensity < 295 is memory-bound (bad — the GPU is waiting for bytes).

For a dense matrix multiply $X \in \mathbb{R}^{B \times D}$ times $W \in \mathbb{R}^{D \times F}$:

• FLOPs: $2 \cdot B \cdot D \cdot F$
• Bytes: $2 \cdot (B \cdot D + D \cdot F + B \cdot F)$ at BF16

In the limit $D, F \gg B$, intensity simplifies to $B$. The batch dimension is what makes a matmul compute-bound.

Now apply this to LLM inference. During generation, you process one token at a time per sequence ($T=1$), so the MLP layers have intensity $= B$ where $B$ is the number of concurrent requests. For attention layers, the math comes out differently. The K and V matrices grow with sequence length, and they are per-sequence, so batching does not help. Attention during generation has intensity $\approx 1$. Both regimes are far below 295.

This is the fundamental inference problem. Generation is bounded by memory bandwidth on every hardware platform built since 2017. The compute is sitting there. Every technique in this chapter buys you efficiency in one of three ways: reduce bytes-per-token (quantization, MLA, MQA), reduce reads-per-step (KV cache, FlashAttention), or amortize across more tokens (continuous batching, speculative decoding).

The prefill phase (processing the prompt) is different. Prefill has $T = \text{prompt length}$ and is largely parallelizable. The MLP has intensity $B \cdot T$, attention has intensity $S/2$ where $S$ is prompt length. Both are compute-bound for any reasonable prompt. This is why prefill is fast and generation is slow, even though they run the same model.

In another chapter you wrote a transformer that, for each new token, ran a full forward pass over the entire context. For generating token $T$, that means running attention over $T$ keys and $T$ values. The cost is $O(T^2)$ per layer for attention, and you have to do this for every new token, so generating $T$ tokens of output is $O(T^3)$ total. That is intolerable past sequence length ~1000.

The KV cache fixes it. Observation: when you generate token $T$, the keys and values for positions $0..T-1$ depend only on those positions' input embeddings, which don't change as you generate more tokens. So you compute K and V for each position once (when that token is first seen, either during prefill or during generation), and store them in memory. For the next token, you compute the new K and V for position $T$, append to the cache, and run attention against the entire cache.

Per layer, per head, you store one K vector of shape (seq_len, d_k) and one V vector of the same shape. For a typical Llama-2-13B (40 layers, 40 heads, $d_k = 128$), at sequence length 2048, this is:

40 layers × 40 heads × 2048 tokens × 128 dims × 2 (K and V) × 2 bytes (BF16)
= 1.68 GB per sequence

Multiply by batch size: at batanother chapter, that's 107 GB of KV cache. On an 80GB H100, you have no room left for the model weights. This is the practical bottleneck on inference throughput at large context. The KV cache is the dominant memory consumer at non-trivial scale.

The crucial bit in any from-scratch implementation is the difference between prefill (run over all prompt tokens at once, fill the cache) and decode (one new token at a time, append to the cache).

model = AutoModelForCausalLM.from_pretrained("gpt2", torch_dtype=torch.bfloat16)
tok = AutoTokenizer.from_pretrained("gpt2")
inputs = tok("The capital of France is", return_tensors="pt")
# use_cache=True is the default; this IS the KV-cache machinery.
# do_sample=False => greedy, so the output equals decoding with no cache.
out = model.generate(**inputs, max_new_tokens=50, do_sample=False, use_cache=True)

# generate_kv: prefill once over the whole prompt, then decode one token at a time
self.reset_caches()
positions = torch.arange(T, device=idx.device)
logits = self(idx, use_cache=True, positions=positions)        # prefill fills the cache
next_token = logits[:, -1, :].argmax(-1, keepdim=True)
idx = torch.cat([idx, next_token], dim=1)
for _ in range(max_new_tokens - 1):                            # decode: T_new == 1
    pos = torch.tensor([idx.shape[1] - 1], device=idx.device)
    logits = self(idx[:, -1:], use_cache=True, positions=pos)  # only the new token in
    next_token = logits[:, -1, :].argmax(-1, keepdim=True)
    idx = torch.cat([idx, next_token], dim=1)

# inside CausalSelfAttentionWithKV.forward, when use_cache=True:
cl = int(self.cache_len.item())
assert cl + T_new <= self.max_seq_len
self.k_cache[:, :, cl:cl + T_new] = k_new    # in-place slice write, no realloc
self.v_cache[:, :, cl:cl + T_new] = v_new
k = self.k_cache[:, :, :cl + T_new]          # view of the active prefix (all history)
v = self.v_cache[:, :, :cl + T_new]
self.cache_len += T_new
scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k)   # q is the new rows, k is all of history
attn = F.softmax(scores, dim=-1)
out = attn @ v

Two subtleties:

1. The cache is per-layer and per-sequence. A batched implementation has shape (B, n_heads, seq_len, d_k) for both K and V.
2. Concatenating with torch.cat allocates a new tensor every step. Production implementations pre-allocate the maximum-length tensor and write into a slice (cache[:, :, t, :] = k_new). This avoids the realloc overhead and matches how vLLM and TensorRT-LLM do it.

Once you've measured your KV cache and realized it's 50% of your inference memory, the next question is: can you make the cache smaller without breaking the model?

The cheapest answer is to share K and V heads across query heads. Vanilla multi-head attention (MHA) has $n_\text{heads}$ separate K and V projections. Each head's KV is independent.

Multi-Query Attention (MQA, Shazeer 2019): use only one K and V head, shared across all query heads. The KV cache shrinks by a factor of n_heads. Llama-2-70B used MQA experimentally; PaLM-540B used MQA from the start. Accuracy drops slightly but throughput improves a lot.

Grouped-Query Attention (GQA, Ainslie et al. 2023): a compromise. $n_\text{kv\_heads} < n_\text{heads}$, and each KV head is shared by n_heads / n_kv_heads query heads. Llama-2-70B-Chat and Llama-3 use GQA with n_kv_heads = 8 and n_heads = 64, so each KV head is shared by 8 query heads, shrinking the cache by 8×. Accuracy is essentially identical to MHA.

The shape math:

Multi-head Latent Attention (MLA, DeepSeek 2024) goes further. Instead of sharing K/V across heads, it stores a low-rank latent representation of K/V and reconstructs the full K/V on the fly via a learned up-projection. The cache holds vectors of dimension $C \ll n_\text{heads} \cdot d_k$. DeepSeek-V2 uses $C = 512$ for $n_\text{heads} \cdot d_k = 16384$ — a 32× reduction over MHA, and it actually outperforms MHA on the benchmarks they ran.

# Minimal GQA forward (no cache for clarity)
def gqa_forward(x, W_q, W_k, W_v, n_heads, n_kv_heads, d_k):
    B, T, C = x.shape
    q = (x @ W_q.T).view(B, T, n_heads, d_k).transpose(1, 2)        # (B, h, T, d_k)
    k = (x @ W_k.T).view(B, T, n_kv_heads, d_k).transpose(1, 2)     # (B, g, T, d_k)
    v = (x @ W_v.T).view(B, T, n_kv_heads, d_k).transpose(1, 2)     # (B, g, T, d_k)
    # Repeat KV heads to match Q heads
    n_groups = n_heads // n_kv_heads
    k = k.repeat_interleave(n_groups, dim=1)                        # (B, h, T, d_k)
    v = v.repeat_interleave(n_groups, dim=1)
    scores = q @ k.transpose(-2, -1) / (d_k ** 0.5)
    attn = scores.softmax(dim=-1)
    return (attn @ v).transpose(1, 2).contiguous().view(B, T, C)

The repeat_interleave is a logical operation — in practice the kernel reads each KV head and uses it for all n_groups query heads without actually materializing the expanded tensor. FlashAttention implementations handle this natively.

Quantization reduces the precision of numbers in the model. Less precision means fewer bytes per parameter, which means less memory bandwidth, which means faster inference. The question is: how much precision can you lose before accuracy drops?

The hierarchy of numerical formats:

Going from BF16 to INT8 halves your KV cache and your weight memory. From BF16 to INT4 quarters them. On a memory-bound workload (which inference is), that translates directly into speedup, as long as you have a kernel that operates on the quantized type.

The basic linear quantization recipe for a tensor:

where $s$ (scale) and $z$ (zero-point) are chosen per tensor (or per channel, or per group) to map the floating-point range onto the integer range. The choice of "what to quantize together" — per-tensor, per-channel, per-group — is called granularity and matters a lot for accuracy.

Post-training quantization (PTQ) quantizes a trained FP16/BF16 model in one shot, possibly using a small calibration dataset to pick scales. Cheap. Works well for INT8 on most models. Struggles at INT4.

Quantization-aware training (QAT) simulates quantization during training (forward pass uses quantized values, backward uses the floats). More expensive but recovers more accuracy at aggressive bit-widths.

For LLMs, the workflow is almost always PTQ because the models are huge and full QAT is impractical. The question is which PTQ algorithm to use.

Tim Dettmers' LLM.int8 (2022) is the paper that made INT8 quantization actually work on transformers at scale. The problem it solved:

Naive INT8 PTQ on a small BERT works fine. On a 175B-parameter OPT, accuracy collapses. Dettmers identified why: at scale, transformers develop "emergent outlier features" — a small number of hidden dimensions where activation values are 100× larger than the rest. If you quantize these in INT8, the outlier dominates the scale, and the other 99.9% of dimensions get squished to ~zero. Information loss is total in those dimensions.

The fix is mixed-precision decomposition. For each matrix multiply $XW$:

1. Identify the columns of $X$ that contain outliers (any value > 6 in magnitude, typically 0.1% of columns).
2. Compute the matrix multiply for those columns in FP16: $X_\text{outlier} W_\text{outlier}$.
3. Quantize the other 99.9% of columns to INT8, do the multiply in INT8 with vector-wise scaling, dequantize.
4. Sum the two contributions.

The result is essentially full-precision accuracy at near-INT8 memory cost. The catch is the INT8 multiply is slower than the FP16 multiply on most hardware because the kernels are less optimized. LLM.int8 is great for memory savings (lets you fit 175B on 8×A100 instead of 16×A100) but is 15-23% slower than FP16 inference on the same hardware. Different tradeoff than you might expect.

The bitsandbytes library implements this:

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quant_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,   # outlier detection threshold
    llm_int8_has_fp16_weight=False,
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=quant_config,
    device_map="auto",
)
# Now uses ~3.5GB for weights instead of 13GB. Inference works, slightly slower.

The emergent features finding is, in retrospect, a mech-interp result hiding in a quantization paper. Dettmers' observation that 6 specific feature dimensions in a 6.7B model carry the outlier mass, and that these dimensions are stable across layers and inputs — is the kind of finding that the Anthropic mech-interp team picked up later. There's a connection between "outlier features" and "high-importance circuits" that I don't think is fully understood yet.

INT8 is comfortable. INT4 is where the real memory savings (and the real engineering challenges) live. The two algorithms that dominate INT4 quantization in 2024-2025:

GPTQ (Frantar et al. 2022) treats weight quantization as a per-layer reconstruction problem. Given calibration data, find quantized weights $\hat{W}$ that minimize $\|XW - X\hat{W}\|^2$. GPTQ does this column by column, using the Hessian of the squared error to choose which columns to quantize first (the most informative ones). The math is the closed-form update from Optimal Brain Quantization (OBQ); GPTQ is OBQ scaled up to LLM-sized matrices. With ~128 calibration samples, GPTQ quantizes OPT-175B to 3-4 bits in a few GPU-hours, with <1% perplexity loss. (Perplexity is the standard intrinsic language-model quality metric: the exponential of the per-token cross-entropy — lower is better, and a 1% rise means the quantized model is barely worse at predicting held-out text.)

The downside: GPTQ quantizes weights but does nothing to activations. Activations are still FP16 at runtime. The win is purely on weight memory. If your hardware supports INT4 × FP16 matmul (modern NVIDIA GPUs do via Marlin or ExLlamaV2 kernels), you also get a speedup.

AWQ (Activation-aware Weight Quantization; Lin et al. 2023) takes a different angle. Observation: not all weights are equally important. The weights in channels that get hit by large activations matter more. AWQ identifies the ~1% of weight channels that correspond to large activation magnitudes (using calibration data), keeps those in higher precision via a per-channel scaling trick, and quantizes the rest to INT4 normally. The scaling is mathematically equivalent (you scale the weight up and the corresponding activation down), so no extra runtime cost.

AWQ tends to outperform GPTQ on instruction-tuned models, where the "important" channels are more concentrated. GPTQ tends to be slightly more general-purpose. Both produce 4× smaller models than FP16 with <1% perplexity hit on most workloads.

The combined recipe most people use in production:

1. Quantize weights with AWQ or GPTQ.
2. Use an INT4×FP16 kernel at inference (Marlin or ExLlamaV2, both for NVIDIA Ampere/Ada) — without one, the quantized weights save memory but give no speedup.
3. Keep the KV cache in FP16 or INT8 separately.
4. Serve from vLLM, which supports the loading and serving glue.

# AWQ in HuggingFace
from transformers import AutoModelForCausalLM
from awq import AutoAWQForCausalLM

# Load the FP16 model, quantize with AWQ
model = AutoAWQForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
quant_config = {"zero_point": True, "q_group_size": 128, "w_bit": 4}
model.quantize(tokenizer, quant_config=quant_config,
               calib_data="mit-han-lab/pile-val-backup")
model.save_quantized("./llama-2-7b-awq")
# Now ~3.5GB on disk, can serve from vLLM with --quantization awq

The standard attention implementation reads Q, K, V from HBM, computes scores (an $T \times T$ matrix) and writes it back to HBM, reads it again for the softmax, writes the softmax result, reads it again to multiply by V. For long sequences, this materializes a giant $T \times T$ tensor in HBM that we then immediately throw away.

FlashAttention (Dao et al. 2022) rewrites this as an IO-aware algorithm. Tile Q, K, V into blocks that fit in SRAM (the GPU's on-chip cache, ~100KB per SM). Compute the attention for each Q block by iterating over K and V blocks, accumulating a running softmax using the online softmax trick:

For each block of Q:
    Initialize running max m = -inf and running denominator l = 0
    For each block of K, V:
        Compute scores S = Q_block @ K_block.T / sqrt(d_k)
        m_new = max(m, max(S))
        # Rescale previous accumulation: subtract new max
        l = l * exp(m - m_new) + sum(exp(S - m_new))
        O = O * exp(m - m_new) + exp(S - m_new) @ V_block
        m = m_new
    O = O / l
    Write O to HBM

The point: you never materialize the full $T \times T$ score matrix in HBM. Tiles exist in SRAM long enough to update the online softmax and output accumulator. Standard attention moves $\Theta(T^2)$ score-sized data to/from HBM; FlashAttention's paper analyzes HBM accesses as $\Theta(T^2 d^2/M)$ for head dimension $d$ and SRAM capacity $M$ over its supported regime. That is a large, hardware-dependent reduction—not generally $O(T)$. The wall-clock gain depends on shape, dtype, hardware, and kernel implementation, so benchmark the workload rather than translating the complexity result into a fixed “week to day” promise.

FlashAttention-2 (Dao 2023) tweaked the kernel for better warp-level parallelism (a warp is the pack of 32 threads a GPU schedules together as one unit). FlashAttention-3 (Shah et al. 2024) added FP8 support and explicit use of TMA (tensor memory accelerator) on H100. FlashDecoding (Hong et al. 2023) re-derived the algorithm for the decoding case (one query, long KV) and got another 2-8× speedup. All three are now production-standard for transformer inference.

You almost never write FlashAttention yourself. You call it:

import torch
import torch.nn.functional as F

# PyTorch 2.0+ has it built in
q = torch.randn(2, 12, 1024, 64, device="cuda", dtype=torch.bfloat16)
k = torch.randn(2, 12, 1024, 64, device="cuda", dtype=torch.bfloat16)
v = torch.randn(2, 12, 1024, 64, device="cuda", dtype=torch.bfloat16)
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
# Behind the scenes: uses FlashAttention 2 on Ampere, FA3 on Hopper.

The thing to take away is the online softmax pattern. It's an algorithm that lets you compute a softmax over a sequence in one streaming pass without materializing the whole sequence. The same pattern appears in the chunked attention of Ring Attention, in tree-based attention reductions for very long context, and elsewhere. Worth understanding the math.

Most of the time, torch.compile and F.scaled_dot_product_attention are fast enough. But sometimes you need to fuse a custom op (e.g., RMSNorm + a residual add + a Swish, all in one kernel). The two options are CUDA (hard to write, hard to debug) and Triton (Python-like syntax, JIT-compiled to PTX).

Triton is a domain-specific language by OpenAI's Philippe Tillet. It looks like Python with a few annotations and primitives (tl.load, tl.store, tl.arange). It compiles to PTX (GPU assembly). The compiler handles thread scheduling within a block; you only specify what the block does, not how the threads in it cooperate.

Here is a fused softmax kernel — the same one that ships with Triton's tutorials. It reads a 2D matrix, computes max-stabilized softmax along the rows, writes the result, all in one launch:

import triton
import triton.language as tl
import torch


@triton.jit
def softmax_kernel(
    x_ptr, y_ptr,
    x_row_stride, y_row_stride,
    num_cols,
    BLOCK_SIZE: tl.constexpr,
):
    row_idx = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)

    # Load one row (with masking for the tail)
    x_start = x_ptr + row_idx * x_row_stride
    x_row = tl.load(x_start + col_offsets,
                    mask=col_offsets < num_cols,
                    other=float("-inf"))

    # Max-subtract for stability
    x_row = x_row - tl.max(x_row, axis=0)
    numerator = tl.exp(x_row)
    denominator = tl.sum(numerator, axis=0)
    y_row = numerator / denominator

    # Store
    y_start = y_ptr + row_idx * y_row_stride
    tl.store(y_start + col_offsets, y_row, mask=col_offsets < num_cols)


def triton_softmax(x: torch.Tensor) -> torch.Tensor:
    y = torch.empty_like(x)
    M, N = x.shape
    BLOCK_SIZE = triton.next_power_of_2(N)
    softmax_kernel[(M,)](
        x, y, x.stride(0), y.stride(0), N, BLOCK_SIZE=BLOCK_SIZE,
    )
    return y

What's happening: the kernel runs $M$ thread blocks (one per row). Each block loads the row into SRAM, computes softmax in-block, writes back. The whole thing is one HBM read and one HBM write — the theoretical minimum. The PyTorch softmax does 5 reads and 3 writes because it's composed of multiple element-wise operations that each launch a kernel.

The same pattern applies to fused ops in LLM inference: RMSNorm-and-add, RoPE application, KV cache write + attention. The big serving systems (vLLM, TensorRT-LLM, sglang) ship custom Triton kernels for these. You can read them and learn from them; you can write your own when needed.

Triton is best for kernels that are memory-bound and that PyTorch can't fuse automatically. For pure matmul, NVIDIA's CUBLAS is still faster; that's the case where the existing kernels are already at peak performance.

Generation is sequential because token $T+1$ depends on token $T$. The classic insight that breaks this: evaluating $N$ tokens in parallel is cheaper than generating them one at a time, because evaluation is compute-bound (prefill-like) and generation is memory-bound.

Speculative decoding (Leviathan et al. 2023; Chen et al. 2023): use a small "draft model" $p$ to propose $K$ candidate tokens at once. Then use the target model $q$ to evaluate all $K$ in parallel and decide which to accept.

The acceptance rule is rejection sampling. For each draft token $x_i$ with draft probability $p(x_i)$ and target probability $q(x_i)$:

• Accept with probability $\min(1, q(x_i) / p(x_i))$
• If rejected, sample a new token from the residual distribution $\max(0, q - p) / \|\max(0, q - p)\|_1$ and discard the rest of the draft.

This is the modified rejection sampling that gives you an exact sample from $q$. The proof (Leviathan §3) is one page of algebra. The intuition: if the draft model is correct often, you accept many tokens per target-model call. If the draft is wrong, the residual distribution corrects it. In expectation, the output sequence is distributed identically to $q$'s output.

In practice: a 1B draft model with a 70B target gives 2-3× speedup on most workloads. A draft model very close to the target (e.g., a distilled version) gives 3-5×. The math is exact; the speedup is empirical.

Variants:

• Medusa (Cai et al. 2024): instead of a separate draft model, train extra "heads" on the target model that predict the next 2-4 tokens directly. No separate model to maintain.
• EAGLE (Li et al. 2024): predict not the next token but the next hidden state, then decode with the target's head. Higher acceptance rate at lower compute.
• Self-speculative: use a few layers of the target model as the draft (skip later layers for the draft pass).

# Pseudo-code for one speculative decoding step
def speculative_step(target_model, draft_model, prefix, K):
    # 1. Draft model proposes K tokens autoregressively
    draft_tokens = []
    draft_probs = []
    ctx = prefix.clone()
    for _ in range(K):
        logits = draft_model(ctx)[:, -1, :]
        probs = logits.softmax(dim=-1)
        tok = torch.multinomial(probs, 1)
        draft_tokens.append(tok)
        draft_probs.append(probs.gather(-1, tok))
        ctx = torch.cat([ctx, tok], dim=1)

    # 2. Target model evaluates all K tokens in parallel (one forward pass)
    target_logits = target_model(ctx)[:, -K-1:-1, :]   # logits for each draft position
    target_probs = target_logits.softmax(dim=-1)

    # 3. Acceptance loop
    accepted = []
    for i, tok in enumerate(draft_tokens):
        q_i = target_probs[0, i].gather(-1, tok)
        p_i = draft_probs[i]
        if torch.rand(1) < (q_i / p_i).clamp(max=1.0):
            accepted.append(tok)
        else:
            # Sample from residual q - p, return immediately
            residual = (target_probs[0, i] - draft_probs[i]).clamp(min=0)
            residual = residual / residual.sum()
            new_tok = torch.multinomial(residual, 1)
            return prefix.new(accepted + [new_tok])
    # All K accepted: bonus sample from target's logits beyond position K
    bonus = torch.multinomial(target_probs[0, K], 1)
    return prefix.new(accepted + [bonus])

The speedup pattern: if the draft accepts $k$ out of $K$ tokens per round, you generate $k+1$ tokens per target forward pass (the bonus token at the end). If $k$ averages 3 with $K=5$, you go 4× faster. The "acceptance rate" is the headline metric for any spec-decode implementation.

Training works on dense batches: $B$ sequences of the same length, processed together. Inference doesn't. Requests arrive at different times, prompts have different lengths, and generation stops at different lengths (when each request hits its EOS or max-tokens).

The old way (static batching): collect a batch of $B$ requests, pad them to the longest, run the model, return all the responses. Problems: requests wait for the batch to fill (high TTFT — time to first token, how long a user waits before any output appears), short responses sit idle waiting for long responses to finish (low GPU utilization).

The new way (continuous batching, from the Orca paper, Yu et al. 2022): batch at the iteration level, not the request level. Every generation step, check what requests are in flight. If new requests arrived during the last step, add them to the batch (their prefill runs alongside the others' decode). If a request finished, drop it from the batch and free its KV cache. This is the technique that makes vLLM, TensorRT-LLM, and TGI more efficient than naive serving.

The implementation needs:

1. A batch where different requests are at different positions in their generation. Some are decoding, some are still prefilling.
2. Selective batching: for attention layers, process each request's KV separately (because each has its own cache). For non-attention layers, concatenate all sequences and run as one big matmul.
3. A scheduler that picks the next requests to admit based on a queue policy.

Continuous batching alone gets 2-4× higher throughput than static batching on typical workloads. TTFT improves dramatically because new requests do not wait. They are injected on the next step.

You don't implement continuous batching from scratch. You use a serving stack (vLLM is the most popular open-source one). What you do need to understand is the throughput-latency tradeoff: throughput is total tokens/second the server emits across all requests, latency is how long any one request waits (its TTFT plus per-token generation time). High batch sizes give high throughput but slow per-request generation. Small batch sizes are the opposite. Whatever serving system you pick has a "max concurrent requests" knob that controls this tradeoff. The headline efficiency number people quote is MFU (model FLOPs utilization, actual FLOP/s ÷ the hardware's promised FLOP/s); pushing batch size up is mostly about dragging this number off the floor where memory-bound generation leaves it.

vLLM's signature contribution. The problem: the KV cache for each request is a contiguous block of GPU memory of size $S_\text{max} \cdot d \cdot 2$ bytes. If you pre-allocate the maximum length, most of that memory is wasted for requests that finish early. If you allocate dynamically and resize, you fragment the GPU memory (just like a malloc-heavy program fragments RAM).

The vLLM solution: page the KV cache the way an OS pages process memory. Divide the cache into fixed-size blocks (e.g., 16 tokens per block). Allocate blocks as needed. Maintain a per-request "block table" that maps logical token positions to physical block indices.

Request 1's KV cache (logical):
  [block0][block1][block2]...
              |
              v
Physical KV memory:
  [block_42][block_7][block_113]...

A request's KV cache no longer needs to be contiguous in GPU memory. New requests can grab free blocks from anywhere. When a request finishes, its blocks return to the free pool. Fragmentation: gone.

The really cool part: prefix caching. Two requests start with the same prompt (a common system prompt, say). Their first $N$ blocks of KV are identical. Paged attention can share those blocks across both requests (copy-on-write semantics if anyone ever writes, which doesn't happen for KV). One copy of the system prompt's KV serves all requests.

This is huge for production. If you have a 2000-token system prompt and serve 1000 concurrent requests, naive caching uses 1000× the system prompt's KV memory. Paged attention with prefix sharing uses 1×. The math says you can serve 4-10× more concurrent requests on the same hardware just from prefix sharing alone.

The cost: the attention kernel has to do a small indirection on every read — look up which physical block to read from. vLLM ships custom CUDA kernels (and now FlashAttention-3 supports paged KV natively) to keep this overhead minimal. The math says ~5% slower per kernel; the practice says you get 4× more concurrent capacity, so the net is +400% throughput.

# vLLM API. PagedAttention happens behind the scenes.
from vllm import LLM, SamplingParams

llm = LLM(model="meta-llama/Llama-2-7b-hf",
          gpu_memory_utilization=0.9,
          enable_prefix_caching=True,
          max_num_seqs=256)

prompts = [
    "You are a helpful assistant. User: What is the capital of France?",
    "You are a helpful assistant. User: What is the capital of Germany?",
    "You are a helpful assistant. User: What is the capital of Italy?",
]
sampling_params = SamplingParams(temperature=0.7, max_tokens=100)
outputs = llm.generate(prompts, sampling_params)
# The "You are a helpful assistant. User: " prefix is shared across all three.

Mixture-of-Experts models (Switch Transformer, GLaM, Mixtral, OLMoE, DeepSeek-V3) have many "expert" FFN layers per transformer block, but a small router network sends each token to only its top-1 or top-2 experts. The model has a lot of parameters but each forward pass uses a fraction of them. This is great for training compute (1.5T parameters at 1/5 the training cost of a dense 300B). For inference it gets weirder.

The first issue: at inference, you can't predict which experts will be used. Different tokens go to different experts, even within one sequence. The KV cache is fine (attention is dense), but the FFN computation hits a different subset of experts for each token. Two consequences:

1. All experts must be in memory. Even though only 2 of 32 experts are active per token, you have to load all 32 from HBM, because you don't know in advance whianother chapter. The memory cost is "as if dense", not "as if active params".
2. Routing is dynamic and uneven. Some experts get hit more often than others. A naive batched implementation processes all 32 experts on the same hardware and one becomes the bottleneck while others sit idle.

For models small enough to fit on one GPU (Mixtral-8x7B at INT4 ≈ 24GB), this is annoying but manageable. For very large MoEs (DeepSeek-V3, ~700B total params, ~37B active), you need expert parallelism: shard the experts across GPUs, route each token's compute to the GPU holding its expert.

This requires all-to-all communication every layer: each GPU sends each token to whatever GPU has its target expert(s). The amount of data shuffled is $O(T \cdot d_\text{model})$ per layer, and the latency is bounded by the slowest GPU-to-GPU link. On 8× H100s with NVLink (NVIDIA's high-speed interconnect between GPUs, far faster than PCIe), this is fine. On 8× A100s without NVLink, it's a disaster. On 64× distributed nodes, it's "you need a custom kernel and a lot of engineering".

DeepSeek's "DeepEP" kernel (2025) is the open-source state of the art: it overlaps the all-to-all comm with the FFN compute, hiding most of the latency. tutel, Megablocks, and FasterMoE are similar attempts. Production MoE serving is non-trivial systems work.

The practical advice as of 2026: if you can fit a dense model in your memory budget, use a dense model. Use MoE for the cases where you specifically need the parameter scaling — long context with many parallel queries, or training-compute-constrained scenarios. MoE inference is rich with sharp edges.

The model was trained with context length 4096. Your application needs context length 32768. Without changes, the model produces garbage past position 4096. RoPE (another chapter) creates the problem and is also the lever to fix it.

The math: RoPE rotates each Q/K dimension pair by an angle $p \cdot \theta$ where $p$ is the position and $\theta$ is a frequency that depends on the dimension index. The model has learned to interpret cos/sin patterns at positions $0..4096$. At position 5000, the angles are extrapolated; for high-frequency dimension pairs, this means many cycles of rotation that the model has never seen during training. The attention pattern breaks.

Position Interpolation (PI; Chen et al. 2023): rescale the positions to fit in the training range. If you want to extend from 4096 to 32768, scale all positions by 1/8 before applying RoPE. The model sees "positions" 0 through 4095, just at 1/8 spacing. Works, but loses resolution at short distances.

NTK-aware scaling (Reddit user bloc97, 2023): instead of rescaling positions uniformly, rescale the base frequency $\theta$ so that low-frequency dimensions (which carry long-range information) extend more, and high-frequency dimensions (short-range) extend less. The math: change $\theta = 10000^{-2d/D}$ to $\theta = (10000 \cdot \alpha)^{-2d/D}$ where $\alpha$ is the extension factor. Better than PI on most evals.

YaRN (Yet Another RoPE extensioN; Peng et al. 2023): a more careful NTK-aware approach that scales different frequency bands differently, with an attention temperature adjustment. Works well empirically and is what most open LLMs use for context extension.

LongRoPE (Ding et al. 2024): use evolutionary search to find a non-uniform rescaling that minimizes perplexity. Can extend Llama-2-7B from 4096 to 2 million tokens with reasonable retention (though the model still struggles at the extremes).

The practical effect: RoPE-scaling methods can extend the usable context of a model, usually with some combination of rescaling, continued training/fine-tuning, and evaluation at the target length. The result is model- and method-specific: changing a configuration field does not guarantee retained retrieval or reasoning quality.

Do not implement YaRN from a six-line frequency multiplier: the method uses wavelength-dependent interpolation/extrapolation and an attention-temperature correction. In practice, use a framework implementation whose version and model configuration you have pinned, then evaluate perplexity, retrieval, and task accuracy across positions up to the advertised limit. The framework's rope_scaling schema is versioned API, not a universal drop-in dictionary.

Knowledge distillation (KD) reuses an expensive teacher model to train a cheap student. For inference, the goal is to make the student fast enough to serve while inheriting the teacher's behavior. The KD loss in the simplest form:

Combining hard-label CE and soft-label KL with a temperature $T$ that smooths the teacher's distribution. The $T^2$ factor compensates for the gradient scaling that the temperature introduces. This is the formulation from Hinton et al. 2015.

For LLMs, the dominant recipe is response-level distillation: have the teacher generate responses to a corpus of prompts, then SFT-tune the student on those (prompt, response) pairs. This works because:

• You don't need access to the teacher's logits (just its outputs).
• It avoids the distribution-shift problem of fixed-corpus distillation.
• It composes with RLHF naturally. The student inherits aligned behaviors.

This is essentially what Vicuna, Alpaca, and Orca did with ChatGPT outputs in 2023. Now most "small open model" releases (Llama-3-8B-Instruct as distilled from Llama-3-70B, Mistral's smaller models, etc.) are some form of distilled-from-bigger-model.

A more aggressive recipe combines distillation with architecture pruning. The pipeline:

1. Identify which layers / heads / dimensions matter least, using calibration data.
2. Remove them. Now you have a smaller, broken model.
3. Distill the original (large, working) model into the pruned one.

NVIDIA's pruning-and-distillation paper (Muralidharan et al. 2024) shows this on Llama-3-8B compressed to ~3B parameters with the original model as teacher. Higher quality than training a 3B from scratch.

DistilBERT (Sanh et al. 2019) is the canonical case study. BERT-base distilled to half the parameters, 60% faster inference, 97% of the original accuracy on GLUE. The loss combined distillation, hard-label CE, and a cosine embedding loss between teacher and student hidden states. The hidden-state alignment turned out to matter. Pure logit distillation gave less retention.

Distillation makes a small model that emulates a big one. LoRA does something different: it keeps the big model, freezes it, and learns a tiny rank-r update that nudges its behavior. At inference, the update merges back into the base weights at zero cost.

The math is brief. A pre-trained linear layer has weight $W \in \mathbb{R}^{d_\text{out} \times d_\text{in}}$. LoRA freezes $W$ and learns two low-rank matrices $A \in \mathbb{R}^{r \times d_\text{in}}$ and $B \in \mathbb{R}^{d_\text{out} \times r}$ with $r \ll \min(d_\text{in}, d_\text{out})$. The new weight is:

where $\alpha$ is a scaling constant (typically 8–32). For Llama-3-8B with $d_\text{in} = d_\text{out} = 4096$ and $r = 16$, the per-layer LoRA adds $16 \cdot 4096 \cdot 2 = 131{,}072$ parameters versus the original layer's $4096^2 = 16.8\text{M}$. That is 0.78% of the original size, trainable.

You apply LoRA to specific layers. The Vaswani-block convention is to add it to $W_q$ and $W_v$ in attention (sometimes $W_k$ and $W_o$ too). Empirical: targeting attention projections gives most of the benefit; the FFN can be left untouched and the model still adapts.

Library path (HuggingFace PEFT):

from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM

base = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-8B", torch_dtype="bfloat16")
config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],   # attention only
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(base, config)
model.print_trainable_parameters()   # "trainable params: 6.3M || all params: 8B || trainable%: 0.08%"
# Train normally. Save just the adapter (a few MB) instead of the whole model.

From-scratch path (the math, in 20 lines):

import torch
import torch.nn as nn

class LoRALinear(nn.Module):
    def __init__(self, base: nn.Linear, r: int = 16, alpha: float = 32.0):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False
        self.A = nn.Parameter(torch.zeros(r, base.in_features))
        self.B = nn.Parameter(torch.zeros(base.out_features, r))
        nn.init.kaiming_uniform_(self.A, a=5**0.5)   # A non-zero
        # B stays zero at init -> output identical to base at step 0
        self.scaling = alpha / r

    def forward(self, x):
        return self.base(x) + (x @ self.A.T @ self.B.T) * self.scaling

    def merge(self):
        """Fold LoRA back into base weights. Zero inference overhead after this."""
        with torch.no_grad():
            self.base.weight.data += (self.B @ self.A) * self.scaling
        # After merge: forward() returns base(x) only; remove or zero A, B.

The B-zero-at-init detail matters. It guarantees that at training step 0 the model's output is identical to the frozen base. Training only changes the model gradually as A and B move off zero.

QLoRA (Dettmers et al. 2023) is the high-leverage combination: quantize the base model to NF4 (a 4-bit format optimized for normal-distributed weights), keep A and B in BF16, train. The result: you can fine-tune a 65B model on a single 48 GB GPU. The base stays quantized during training; only the LoRA adapters update. At inference, you keep the base quantized and store the adapter separately. Multiple adapters can share the same base in memory.

LoRA at inference. Three patterns:

1. Merged: W' = W + \alpha B A. Run the merged weights. Zero inference overhead. The cost is you cannot swap adapters quickly.
2. Hot-swap: keep base weights frozen in memory, swap adapter weights per request. Adds one matmul per LoRA-targeted layer per forward pass (cheap). vLLM and TGI both support this. Lets you serve N user-specific fine-tunes off one base model.
3. Multi-adapter batching (SLoRA, Punica): batch requests across adapters by reshuffling the per-adapter matmuls. This is the cutting edge as of 2026 and is what lets a single GPU serve thousands of user-specific fine-tunes simultaneously.

The reason this section sits inside an inference chapter (and not a training chapter) is that the inference patterns above are what make LoRA economically interesting. Training a LoRA is easy. Serving fine-tunes at scale is where the engineering lives.

Where do all these techniques live in practice? The serving-stack landscape as of 2026:

vLLM (UC Berkeley + community, open source). The reference implementation of paged attention, continuous batching, prefix caching. Supports most quantization formats (AWQ, GPTQ, FP8). The model coverage is wide. Most open LLMs are first-class supported within weeks of release. The user interface is OpenAI-compatible. This is the default open-source serving stack and what most labs deploy on their own hardware.

TensorRT-LLM (NVIDIA, semi-open). The fastest single-GPU implementation, especially on H100. Uses NVIDIA-proprietary kernels (TMA, async cudaMalloc, custom matrix multiplies) that vLLM doesn't access. Harder to use, smaller model coverage. The standard for cost-optimized commercial deployment on NVIDIA hardware.

TGI (HuggingFace Text Generation Inference, open). HuggingFace's serving stack, integrated with their model hub. Easier deployment if you already use the HF ecosystem. Slightly behind vLLM on raw throughput as of 2026.

SGLang (open source). Focuses on multi-step / structured generation workloads: KV cache reuse across multiple queries, RadixAttention (a trie-based KV cache for shared prefixes), efficient JSON/regex-constrained generation. Where vLLM and TGI optimize per-request, SGLang optimizes the cross-request structure.

Together AI, Fireworks, DeepInfra (commercial). Hosted versions of vLLM and TensorRT-LLM (plus their own proprietary improvements). Their value-add is reliability, multi-region, and aggressive performance optimization. They expose pretty much the same API and charge by token.

The architectural distinctions blur quickly. They all use paged attention, continuous batching, FlashAttention, speculative decoding, and the major quantization methods. What differs is the kernel-level engineering, the scheduler policies, and the integration polish.

The advice for new deployments:

• Prototyping: vLLM, FP16, batanother chapter. Get the model working.
• Cost-optimizing: AWQ or GPTQ quantization, prefix caching enabled, batanother chapter-256. Profile and tune.
• Latency-optimizing: smaller batch, speculative decoding with a draft model, fewer concurrent requests.
• Need exotic features (RAG with structured outputs, multi-turn with cache, etc.): SGLang or roll your own on top of vLLM.