LLM Inference Optimization: KV Cache, FlashAttention, and Quantization

Serving cost depends on how a model processes prompts and generates responses. Small-batch decoding often spends much of its time moving weights and cached attention states through memory. Longer prompts, larger batches, and different kernels can shift the bottleneck toward arithmetic or communication. Inference optimization therefore starts with the workload: prompt lengths, response lengths, concurrency, hardware, and latency requirements.

Prefill and decode

Prefill computes representations for the prompt and fills the attention cache. Causal attention allows prompt positions to be processed together within each layer while blocking future positions. Larger matrix operations can use arithmetic units efficiently, but short prompts, long-context attention, and implementation overhead can prevent compute saturation. Servers may also split prefill into chunks.

Decode usually advances each active sequence by one token per step. At batch size one, dense projections resemble matrix–vector products and can require reading a large weight set for little computation. Counting a multiply-add as two FLOPs, an idealized dense projection with two-byte weights has about one FLOP per weight byte. This estimate omits KV reads, activations, cache residency, and communication; it is a starting point for a bandwidth model.

import math

def weight_read_bound(params_billions, bytes_per_param, bandwidth_gb_s):
    """Ideal single-stream bound if every weight byte is read once per step."""
    if not all(math.isfinite(v) and v > 0 for v in
               (params_billions, bytes_per_param, bandwidth_gb_s)):
        raise ValueError("All inputs must be positive and finite")
    seconds = params_billions * 1e9 * bytes_per_param / (bandwidth_gb_s * 1e9)
    return seconds * 1000, 1 / seconds

for bits, bpp in [(16, 2), (8, 1), (4, 0.5)]:
    ms, tps = weight_read_bound(7, bpp, 2000)
    print(f"{bits}-bit: ideal weight read {ms:.2f} ms, rate bound {tps:.0f} tokens/s")
# 16-bit: ideal weight read 7.00 ms, rate bound 143 tokens/s
# 8-bit: ideal weight read 3.50 ms, rate bound 286 tokens/s
# 4-bit: ideal weight read 1.75 ms, rate bound 571 tokens/s

For nominal 7B dense parameters at 2,000 decimal GB/s, reading 14 GB takes 7 ms. The table is an optimistic weight-read bound under its assumptions, not a measured decode speed. Packed lower-bit weights reduce payload size, but scales, dequantization, kernel support, KV traffic, and other operations affect end-to-end latency. Halving nominal precision does not guarantee doubling generation speed.

A small language model and generation loop are implemented in Building a Transformer Language Model in PyTorch. That example is not a production serving engine or an implementation of all the optimizations below. The Python examples here require PyTorch and should be run in order.

The KV cache

For fixed weights, positions, and an unchanged causal prefix in evaluation mode, earlier keys and values can be reused. Re-running full-prefix dense attention for a prefix of length \(T\) has a quadratic attention term; a new single-token query attends to \(T\) cached entries and has a linear attention term, with widths held fixed. Other projection and MLP costs remain. Reusing a cache after changing the prefix, model, adapter, or positional convention can give incorrect outputs.

The class below implements only multi-head attention, with no positional encoding or padding. Its mask uses the absolute position within the accumulated prefix: a two-token chunk after three cached tokens may attend through positions 3 and 4 respectively. A single new query may see every cached key, but a multi-token chunk needs this offset causal mask. In PyTorch’s scaled-dot-product API, boolean True means attention is allowed. An actual transformer must also preserve the correct position offsets, including when applying rotary embeddings.

import torch
import torch.nn.functional as F

class CachedAttention(torch.nn.Module):
    def __init__(self, width, heads):
        super().__init__()
        if heads <= 0 or width <= 0 or width % heads:
            raise ValueError("width must be positive and divisible by positive heads")
        self.h, self.d = heads, width // heads
        self.qkv = torch.nn.Linear(width, 3 * width)
        self.out = torch.nn.Linear(width, width)

    @torch.no_grad()
    def forward(self, x, past_kv=None):
        """Inference-only, unpadded (B, new_tokens, width) inputs."""
        if self.training:
            raise ValueError("Use eval mode for cached inference")
        if x.ndim != 3 or x.size(1) == 0 or x.size(2) != self.h * self.d:
            raise ValueError("Expected a nonempty (B, T, width) input")
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        q, k, v = (t.reshape(x.size(0), x.size(1), self.h, self.d).transpose(1, 2)
                   for t in (q, k, v))
        past = 0
        if past_kv is not None:
            pk, pv = past_kv
            if (pk.ndim != 4 or pk.shape != pv.shape or pk.shape[:2] != k.shape[:2]
                    or pk.size(-1) != self.d or pk.device != k.device or pv.device != k.device
                    or pk.dtype != k.dtype or pv.dtype != k.dtype):
                raise ValueError("Cache shape, device, or dtype mismatch")
            past = pk.size(2)
            k, v = torch.cat([pk, k], dim=2), torch.cat([pv, v], dim=2)
        query_positions = past + torch.arange(x.size(1), device=x.device)
        key_positions = torch.arange(k.size(2), device=x.device)
        allowed = key_positions[None, :] <= query_positions[:, None]
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed, dropout_p=0.0)
        y = y.transpose(1, 2).reshape(x.size(0), x.size(1), -1)
        return self.out(y), (k, v)

torch.manual_seed(0)
attention = CachedAttention(16, 4).eval()
x = torch.randn(2, 6, 16)
full, _ = attention(x)
cache, pieces = None, []
for start, stop in [(0, 3), (3, 5), (5, 6)]:
    part, cache = attention(x[:, start:stop], cache)
    pieces.append(part)
print("full versus cached", torch.allclose(full, torch.cat(pieces, dim=1), atol=1e-6, rtol=1e-5))
print("cache shape", list(cache[0].shape))
# full versus cached True
# cache shape [2, 4, 6, 4]

The comparison checks full causal attention against a three-token prefill, a two-token continuation, and a one-token continuation. It does not benchmark a GPU kernel. Appending with torch.cat copies the growing cache; a serving implementation normally uses preallocated storage or blocks. Dropping old entries is a separate windowing policy and need not match recomputing a truncated context, because retained states may already encode earlier tokens. For an equal-length dense cache, storage is:

\[\text{KV bytes}=2L H_{KV}d_{\text{head}}T B\,b.\]

def kv_cache_gib(layers, kv_heads, head_dim, seq_len, batch, bytes_per=2):
    return 2 * layers * kv_heads * head_dim * seq_len * batch * bytes_per / 2**30

for batch in [1, 8, 32]:
    print(batch, round(kv_cache_gib(32, 32, 128, 4096, batch), 2), "GiB")
# 1 2.0 GiB
# 8 16.0 GiB
# 32 64.0 GiB

Here \(L\) is the number of layers, \(H_{KV}\) the number of key/value heads, \(d_{\text{head}}\) their width, \(T\) the stored prompt-plus-generated length, \(B\) the batch size, and \(b\) bytes per entry. The factor two stores both keys and values. The example’s 64 GiB cache at batch 32 is about 4.9 times the 13.04 GiB of nominal 7B fp16 weights. Both must fit, along with working buffers. Unequal lengths, padding, allocation granularity, sharding, and quantization metadata change actual allocation.

Shrinking the cache: MQA, GQA, and paging

Multi-query attention uses one KV head for all query heads. Grouped-query attention uses one KV head per query-head group. Replacing 32 KV heads with 8 reduces the cache payload by four at fixed dimensions; one KV head reduces it by 32. These are architectural choices, not universally fixed configurations or predictable quality penalties. Converting a multi-head checkpoint by pooling KV projections is possible, but it changes the function; the GQA paper studied additional training to recover quality.

PagedAttention maps logical cache positions to physical fixed-size blocks. Blocks can be allocated as a request grows instead of reserving a full contiguous maximum-length buffer. This reduces allocation waste without reducing the bytes required for live KV entries. Compatible shared prefixes can reuse blocks; diverging requests need appropriate ownership or copy-on-write behavior. Prefix identity includes the model configuration and token history, not just a similar-looking prompt.

FlashAttention

A straightforward dense attention implementation materializes score and probability arrays over query–key pairs in high-bandwidth memory. For \(T=8192\), one fp16 \(T\times T\) array occupies 128 MiB per head, or about 134 decimal MB. This illustrates a large intermediate allocation; it does not identify the bottleneck of every attention workload.

FlashAttention tiles the computation and updates softmax statistics incrementally, avoiding the full score/probability array in high-bandwidth memory. For a row of scores, a new block with maximum \(m_b\) changes the running maximum to \(m_{\mathrm{new}}=\max(m,m_b)\). Old exponential sums and weighted-value sums are rescaled by \(e^{m-m_{\mathrm{new}}}\) before adding the new block’s contributions. Dividing the final weighted sum by the final exponential sum gives the same softmax-weighted result. This reduces auxiliary attention storage from quadratic to linear in sequence length at fixed widths; dense attention arithmetic remains quadratic.

“Exact” means computing the same attention function without a sparsity or low-rank approximation; rounding and operation order can still change floating-point outputs. Speed depends on shapes, hardware, dtype, and the baseline. During one-token decode the score shape is 1×T, so avoiding a T×T intermediate is primarily a prefill or training argument. FlashAttention also does not remove the persistent KV cache. Calling PyTorch’s scaled-dot-product attention API does not guarantee that a FlashAttention kernel was selected.

Quantization

A quantizer replaces a value with an integer code and a scale: approximately w = scale × code for a symmetric scheme. A large value in a quantization group can make the step size coarse for smaller values. Activation outliers are particularly relevant to LLM.int8()’s mixed-precision path; this is not a complete explanation of every weight-only quantization error. AWQ uses activation information to choose protective weight scaling, while GPTQ uses layer-input statistics to reduce reconstruction error. Neither guarantees a fixed quality loss.

In this numerical illustration, a shared scale maps [0, 0.1, 0.2, 8] onto codes from −7 to 7. The two small nonzero weights round to zero. Smaller groups isolate some values from the outlier. This uses 15 signed levels, which fit in four bits, but stores the codes in int8 tensors for readability; it is not a packed 4-bit kernel or a memory benchmark.

def round_group(values):
    scale = values.abs().max() / 7
    if scale == 0:
        scale = values.new_tensor(1.0)
    codes = (values / scale).round().clamp(-7, 7).to(torch.int8)
    return codes, codes.to(values.dtype) * scale

weights = torch.tensor([0., 0.1, 0.2, 8.])
codes, restored = round_group(weights)
grouped = torch.cat([round_group(group)[1] for group in weights.split(2)])
print("integer codes", codes.tolist())
print("one group", [round(v, 3) for v in restored.tolist()])
print("two groups", [round(v, 3) for v in grouped.tolist()])
# integer codes [0, 0, 0, 7]
# one group [0.0, 0.0, 0.0, 8.0]
# two groups [0.0, 0.1, 0.0, 8.0]

Weight-only quantization reduces stored weight payload, while kernels commonly dequantize values for higher-precision arithmetic. Actual compression includes scales and packing overhead. Weight quantization does not automatically quantize the KV cache or activations. Lower-precision activation arithmetic can help on supported hardware, but calibration, kernels, and workload determine both speed and error; unsupported paths can be slower.

Compare task performance as well as held-out language-model loss before and after quantization. Also measure latency and peak memory with the intended context lengths and concurrency. An acceptable average loss change does not establish that every task is unaffected, and a smaller checkpoint file does not prove faster serving.

Batching and speculative decoding

Batching reuses weight reads across active sequences and can improve arithmetic intensity. It also increases computation and per-request KV traffic, so additional sequences are not free. Throughput can saturate before memory is full, and per-request latency can rise. The useful batch size depends on both the latency target and the memory budget.

A fixed batch may retain unused slots or padding work after some requests finish. Continuous batching releases completed requests and admits new ones between iterations, improving occupancy under variable request lengths. It does not guarantee the same latency or a particular speedup. Scheduling long prefills, sometimes in chunks, alongside decode requests helps manage interference and queueing.

Speculative decoding uses a cheaper draft model to propose \(k\) tokens and a target model to score the proposed continuation in parallel. For exact stochastic sampling, with target distribution \(p\) and draft distribution \(q\) at the same accepted prefix, accept a proposed token with probability \(\min(1,p/q)\). On the first rejection, sample from normalized \((p-q)_+\) and discard the later drafts. If every draft is accepted, sample an extra token from the target. This correction preserves the target distribution mathematically; simply keeping tokens that look “correct” does not.

For each token, accepted probability mass is min(p, q), and the rejection branch contributes (p − q)₊, whose sum is p. Exactness requires consistent conditional distributions, including any temperature or sampling filters, and correct cache rollback after rejected drafts. It means the same distribution, not necessarily the same text with the same random seed. Draft time, verification cost, and conditional acceptance rates determine speed. Verification is one batched target call, but that call is not guaranteed to cost the same as a one-token decode.

How these models are pretrained at scale is covered in LLM Pretraining: Objectives, Data, and Scaling Laws.

What to measure

• Time to first token — request arrival to first output token; includes queueing and preprocessing as well as prefill.
• Time per output token — streaming time after the first token divided by subsequent tokens; also inspect individual inter-token delays.
• Throughput — aggregate output tokens or completed requests per second; state the definition and quality/latency constraints.
• Latency percentiles — report the median (p50) and slower-request percentiles such as p95 and p99 under a specified request load.

Report latency and throughput together with input/output length distributions, concurrency or arrival rate, batch policy, model and quantization settings, hardware, and warm-up procedure. Distinguish isolated kernel timing from end-to-end service measurements. For asynchronous GPU work, use synchronization or device events appropriately; a host timer around a launch can undercount the work. The small calculations here estimate storage and idealized costs, not production throughput.

Compressing and serving the finished model is covered in Model Compression and Deployment: Distillation, Pruning, and Serving.

Retrieval and tool use built on top of a served model are covered in Retrieval-Augmented Generation and Tool-Using Agents.

Exercises

1. Cache payload and total memory. Compute fp16 KV payload for 32 layers, 32 KV heads, and head width 128, for one 4096-token sequence and a batch of 64. Repeat the batch calculation with eight KV heads and compare with nominal 7B fp16 weights.

You should get: a cache that exceeds the weights themselves at moderate batch size.

L, H, hd, bytes_per = 32, 32, 128, 2

def kv_bytes(batch, seq, kv_heads):
    return 2 * L * batch * seq * kv_heads * hd * bytes_per   # 2 for K and V

for label, b, kvh in (("1 seq,  MHA", 1, 32), ("64 seqs, MHA", 64, 32),
                      ("64 seqs, GQA", 64, 8)):
    gb = kv_bytes(b, 4096, kvh) / 2**30
    print(f"{label}  {gb:8.2f} GiB")
print(f"weights in fp16 {7e9 * 2 / 2**30:.2f} GiB")
# 1 seq,  MHA      2.00 GiB
# 64 seqs, MHA    128.00 GiB
# 64 seqs, GQA     32.00 GiB
# weights in fp16 13.04 GiB

2. A weight-only roofline estimate. Compare 2048-position prefill and one-position decode using 2 FLOPs per parameter per position and a single fp16 weight read. Use a hypothetical device with 312 TFLOPs/s and 2,039 GB/s bandwidth. Which resource limits this simplified model?

You should get: two intensities on opposite sides of the hardware’s break-even point.

params = 7e9
bytes_read = params * 2            # fp16 weights, read once per forward

for label, tokens in (("prefill 2048", 2048), ("decode 1", 1)):
    flops = 2 * params * tokens
    print(f"{label:14s} {flops:10.3e} FLOPs / {bytes_read:.3e} B"
          f"  = {flops / bytes_read:8.1f} FLOP/byte")

print(f"hardware break-even {312e12 / 2039e9:.1f} FLOP/byte")
# prefill 2048    2.867e+13 FLOPs / 1.400e+10 B  =   2048.0 FLOP/byte
# decode 1        1.400e+10 FLOPs / 1.400e+10 B  =      1.0 FLOP/byte
# hardware break-even 153.0 FLOP/byte

3. Speculative tokens per round versus speedup. For four draft tokens and independent acceptance probability 0.9, 0.7, or 0.5 at each position, compute expected emitted tokens, including a correction or bonus token. Then estimate speedup if drafting costs 0.10 baseline target-step times per token and verification costs 1.20 baseline steps.

You should get: expected token counts and a separate latency estimate under stated cost assumptions.

def expected_emitted(k, acceptance):
    if not isinstance(k, int) or k < 0 or not 0 <= acceptance <= 1:
        raise ValueError("Use a nonnegative draft count and acceptance in [0, 1]")
    return sum(acceptance ** i for i in range(k + 1))

k, draft_cost, verification_cost = 4, 0.10, 1.20
for acceptance in (0.9, 0.7, 0.5):
    emitted = expected_emitted(k, acceptance)
    speedup = emitted / (k * draft_cost + verification_cost)
    print(f"acceptance {acceptance}  emitted {emitted:.3f}  estimated speedup {speedup:.3f}x")

p = torch.tensor([0.6, 0.3, 0.1], dtype=torch.float64)
q = torch.tensor([0.2, 0.5, 0.3], dtype=torch.float64)
accepted_mass = torch.minimum(p, q)
rejected_mass = (p - q).clamp_min(0)
print("corrected distribution", [round(v, 3) for v in (accepted_mass + rejected_mass).tolist()])
# acceptance 0.9  emitted 4.095  estimated speedup 2.559x
# acceptance 0.7  emitted 2.773  estimated speedup 1.733x
# acceptance 0.5  emitted 1.938  estimated speedup 1.211x
# corrected distribution [0.6, 0.3, 0.1]

References

• Shazeer (2019). Fast Transformer Decoding: One Write-Head Is All You Need. arXiv preprint.
• Dettmers, Lewis, Belkada, and Zettlemoyer (2022). LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale.
• Dao, Fu, Ermon, Rudra, and Re (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS.
• Yu et al. (2022). Orca: A Distributed Serving System for Transformer-Based Generative Models. OSDI.
• Frantar, Ashkboos, Hoefler, and Alistarh (2023). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers.
• Leviathan, Kalman, and Matias (2023). Fast Inference from Transformers via Speculative Decoding. ICML.
• Kwon et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP.
• Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. EMNLP.
• Dao (2024). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR.
• Lin et al. (2024). AWQ: Activation-Aware Weight Quantization for On-Device LLM Compression and Acceleration. MLSys.