AI Computer Institute
Expert-curated CS & AI curriculum aligned to CBSE standards. A bharath.ai initiative. About Us

KV Cache Optimization and Management

📚 Model Inference⏱️ 20 min read🎓 Grade 12
✍️ AI Computer Institute Editorial Team Updated: September 2026 CBSE-aligned · Peer-reviewed · 20 min read
Content curated by subject matter experts with IIT/NIT backgrounds. All chapters are fact-checked against official CBSE/NCERT syllabi.

A Bengaluru food-delivery platform runs a support chatbot on a cluster of GPUs. At any moment it is holding open conversations with thousands of users. Most send something short — "where is my order" — and get an answer in six or seven tokens. A few paste their entire order history and a three-paragraph complaint, and the model has to generate a long, careful reply. The engineering team cannot know, when a conversation starts, whether it will end in 40 tokens or 4,000. Yet every one of those conversations needs a slice of GPU memory reserved for its key-value (KV) cache the moment it begins, and that memory has to be found, tracked, and eventually reclaimed without ever colliding with another user's data. The question this chapter answers is not "how big is one request's cache" — that arithmetic (bytes per layer, per head, per token) is a different problem you have likely already worked through. The question here is: once you know how big each cache could get, how does a serving system actually manage thousands of these caches, of unpredictable and constantly changing size, on a fixed pool of GPU memory, without wasting most of it?

A one-paragraph recap, and then the real problem

During autoregressive decoding, a transformer computes a key vector and a value vector for every token, at every layer, from the input embeddings. Instead of recomputing the keys and values for all previous tokens at every new decoding step, an inference engine caches them and computes only the new token's K and V each step, reusing the rest. That cache grows by one token's worth of K/V pairs per layer for every token generated, and it must stay resident in GPU memory for as long as that conversation is active. So far, so familiar. The problem this chapter is actually about begins one level up: a production server is not running one conversation, it is running a batch of hundreds running concurrently, each cache growing at its own unpredictable rate, and the GPU has a fixed number of gigabytes to hand out. How you allocate and re-allocate that memory, live, is a systems problem in its own right — and it is the problem that determined how fast and how cheap LLM serving became between 2022 and 2024.

Internal fragmentation: the cost of guessing wrong on purpose

The naive approach, and the one every first implementation of batched generation used, is to reserve a contiguous block of GPU memory for each request sized to the model's maximum context length — say 2,048 tokens — the moment the request arrives, because you don't know in advance how long the reply will be and a contiguous array can't be resized without copying it. This is exactly the strategy a beginner would reach for, and it is exactly what breaks at scale. If a request finishes after 47 tokens, the other 2,001 reserved token-slots in its buffer sat idle for the whole conversation, unusable by any other request, because contiguous memory can't be handed out in pieces to two different owners. Multiply this by a batch of hundreds of concurrent chats with wildly different actual lengths, and the majority of a GPU's expensive high-bandwidth memory is reserved but empty. This wasted, unusable-but-allocated memory is called internal fragmentation, and it is the single largest reason early LLM-serving stacks needed far more GPUs than the raw compute would suggest.

PagedAttention: borrowing virtual memory from the operating system

The fix, introduced by Kwon et al. in "Efficient Memory Management for Large Language Model Serving with PagedAttention" (SOSP 2023, the UC Berkeley team behind vLLM), is a direct import of an idea operating systems have used for KV caches' exact opposite problem — process memory — since the 1960s: don't allocate contiguous memory at all. Instead, chop GPU memory into small fixed-size physical blocks (a common choice is 16 tokens' worth of K/V per block), and give every request a logical view of its cache as a numbered sequence of blocks — block 0, block 1, block 2 — that is mapped, through a small per-request block table, to whichever physical blocks happen to be free anywhere in the pool. A new block is only allocated from the pool when the sequence actually generates enough tokens to need it, one block at a time, not all 2,048 slots up front. The attention kernel is rewritten to walk this block table and gather the scattered physical blocks at each layer instead of assuming one contiguous array — this is the one piece of extra work PagedAttention asks of the compute kernel, and it is a memory-indexing cost, not an increase in the number of attention computations performed.

Because physical blocks are now just entries a table can point to, the same physical block can be pointed to by two different requests' block tables simultaneously. If two conversations happen to share an identical run of tokens — most commonly a system prompt every user gets prepended to their query — the block holding that shared prefix's K/V vectors can be computed once and referenced by every sequence that needs it, with a reference count tracking how many block tables currently point at it. A block is only returned to the free pool once its reference count drops to zero. If one of the sharing sequences ever needs to write new, sequence-specific tokens into what was a shared block, the engine gives it a private copy first — copy-on-write, the same mechanism operating systems use when a forked process modifies memory it initially shared with its parent.

PagedAttention: logical blocks mapped to a physical block pool Two sequences, one shared prefix block (P0) held with reference count 2 Sequence A — block table [P0, P2, P5] L0 (prefix, shared) L1 L2 Sequence B — block table [P0, P4] L0 (prefix, shared) L1 Physical GPU memory pool (8 blocks) P0 ref=2 P1 free P2 (A.L1) P3 free P4 (B.L1) P5 (A.L2) P6 free P7 free Allocated, single owner Shared, ref count > 1 Free, in pool

Worked Example 1: how much memory fragmentation actually costs

Take a batch of four concurrent requests on a server with maximum context length 2,048 tokens, and suppose they generate 47, 612, 130, and 1,900 tokens respectively before finishing. Under the naive contiguous-allocation strategy, each request reserves the full 2,048-token buffer the moment it starts, whether or not it ends up using it. Under PagedAttention with a block size of 16 tokens, each request is allocated blocks one at a time, rounding its actual length up to the nearest multiple of 16.

def paged_vs_naive(seq_lengths, block_size, max_seq_len):
    total_blocks = 0
    total_actual = 0
    for length in seq_lengths:
        blocks = -(-length // block_size)  # ceiling division
        total_blocks += blocks
        total_actual += length
    paged_slots = total_blocks * block_size
    naive_slots = max_seq_len * len(seq_lengths)
    paged_frag = (paged_slots - total_actual) / paged_slots * 100
    naive_frag = (naive_slots - total_actual) / naive_slots * 100
    return {
        "total_actual_tokens": total_actual,
        "paged_slots": paged_slots,
        "naive_slots": naive_slots,
        "paged_fragmentation_pct": round(paged_frag, 2),
        "naive_fragmentation_pct": round(naive_frag, 2),
        "memory_reduction_factor": round(naive_slots / paged_slots, 2),
    }

result = paged_vs_naive([47, 612, 130, 1900], block_size=16, max_seq_len=2048)
print(result)
# {'total_actual_tokens': 2689, 'paged_slots': 2720, 'naive_slots': 8192,
#  'paged_fragmentation_pct': 1.14, 'naive_fragmentation_pct': 67.18,
#  'memory_reduction_factor': 3.01}

Tracing it by hand confirms the printed values. Blocks needed per request: ⌈47/16⌉ = 3, ⌈612/16⌉ = 39, ⌈130/16⌉ = 9, ⌈1900/16⌉ = 119, totalling 170 blocks, or 2,720 token-slots. The four requests actually used 2,689 tokens combined, so only 2,720 − 2,689 = 31 slots are wasted to rounding up to the nearest block — 1.14% fragmentation. The naive strategy reserved 4 × 2,048 = 8,192 slots for the same requests, wasting 8,192 − 2,689 = 5,503 of them — 67.18% fragmentation, meaning two-thirds of the reserved memory sat empty. The ratio 8,192 / 2,720 ≈ 3.01 is the practical payoff: the same GPU memory that held this naive batch of four requests can hold roughly three times as many requests once it is block-managed instead — directly more concurrent users served per GPU, at zero cost to model quality, because nothing about the attention computation itself changed.

RadixAttention: caching, not just packing, the prefix

PagedAttention's block indirection makes sharing possible, but it doesn't by itself decide when two requests' prefixes are identical and should share. That is a separate, orthogonal problem — recognizing reusable prefixes across a stream of arriving requests — solved by SGLang's RadixAttention (Zheng et al., 2024). The idea is to organize all currently cached token sequences into a radix tree: a tree where each edge is labeled with a run of tokens, and any two sequences that share a prefix share the same path down from the root before their edges diverge. When a new request arrives, the engine walks the tree matching its tokens against existing edges; every token it matches is a token whose K/V blocks are already computed and sitting in GPU memory, ready to be referenced with no recomputation. Only the tokens past the point of divergence — the truly new, request-specific part — need a fresh prefill pass. When GPU memory runs low, the tree is trimmed using an LRU policy over its leaves, evicting the least recently used cached branches first, while reference counting (the same mechanism PagedAttention uses) protects any branch still actively referenced by a running request from being evicted out from under it.

This matters most exactly where you'd expect: chat applications that prepend a long, fixed system prompt — persona instructions, safety rules, a few worked examples for few-shot prompting — to every single user turn. Without prefix caching, that fixed prompt is re-encoded from scratch, token by token, for every request that arrives. With it, the fixed prompt is encoded once, and every subsequent request that starts with the identical token sequence skips straight to its own unique suffix.

Worked Example 2: what a shared system prompt is actually worth

Suppose a tutoring chatbot prepends an identical 380-token system prompt to every student question, average unique question length 90 tokens, block size 16, and consider 500 requests arriving over an hour. Without prefix sharing, every request needs its own full copy: ⌈380/16⌉ = 24 blocks for the prompt plus ⌈90/16⌉ = 6 blocks for the question, 30 blocks per request (this also equals ⌈470/16⌉ = 30 computed directly, since the prompt happens to leave the boundary between prompt and question aligned to a fresh block). Across 500 requests that is 30 × 500 = 15,000 blocks. With RadixAttention-style sharing, the 24 prompt blocks are computed and stored exactly once, and only the 6 question-specific blocks are allocated per request: 24 + 6 × 500 = 24 + 3,000 = 3,024 blocks. The reduction factor is 15,000 / 3,024 ≈ 4.96 — essentially five times less GPU memory consumed by the batch, and, just as valuable in a live chat product, every request after the first skips the prefill compute for those 380 shared tokens entirely, cutting time-to-first-token for a real user waiting on a reply.

Common misconception: "PagedAttention makes attention faster, like FlashAttention"

Because both names circulate together in vLLM's documentation and both are pitched as inference speedups, students routinely conflate PagedAttention with a numerical optimization of the attention computation itself, the way FlashAttention reduces memory reads by fusing the softmax and the matrix multiplies into one tiled kernel. That's incorrect, and the distinction matters for reasoning about serving systems correctly. PagedAttention performs exactly the same number of floating-point operations as before — every query still attends to every key in its cache, with the identical scores and the identical softmax. What changes is purely where those keys and values physically sit in memory and how the kernel locates them: a gather over scattered physical blocks via a block table, instead of a read from one contiguous array. It is a memory-allocation and -sharing innovation, not an algorithmic one, which is exactly why the two compose rather than compete — vLLM's actual serving kernels run FlashAttention-style tiled compute on top of PagedAttention-style paged memory layout, taking the speed benefit of one and the memory-efficiency benefit of the other simultaneously.

Active recall

Attempt each question before reading the worked answer beneath it.

Q1. Why can't a serving system just reserve a contiguous buffer of size max_seq_len per request, coded the straightforward way? What specifically breaks as the batch grows?

Q2. A server uses block size 16 and max_seq_len 512. Three requests generate 25, 200, and 33 tokens. Compute the blocks needed, total paged slots, paged fragmentation %, naive fragmentation %, and the memory reduction factor.

Q3. The platform now switches block size from 16 to 64 tokens (fewer block-table entries to track, at the cost of coarser granularity), for the exact same three requests from Q2. Recompute the fragmentation. Then trace the full ripple: how does this same block-size change affect the number of block-table entries, and how does it affect how much of a 380-token shared system prompt (as in Worked Example 2) two requests can actually reuse?

Q4. A colleague says "we added PagedAttention, so our attention layers now do fewer FLOPs." Is this correct? What did actually change?

Q5. A different deployment shares a 512-token system prompt across 1,000 requests, average unique suffix 60 tokens, block size 16. Compute the blocks needed with and without prefix sharing, and the reduction factor.

Q6. Sequence A and Sequence B both reference physical block P0 (ref count 2, holding their shared prefix). Sequence A finishes and its block table is torn down. Can P0 be returned to the free pool immediately? Why or why not?

Answers.

A1. Reserving max_seq_len per request wastes GPU memory as internal fragmentation on every request that finishes short of the maximum — which is most of them in a real workload. Since that reserved-but-unused memory can't be handed to any other request, the number of concurrent sequences a GPU can hold is capped by worst-case length, not actual usage, so far fewer real users get served per GPU than the hardware's raw memory would suggest.

A2. Blocks: ⌈25/16⌉=2 (32 slots, waste 7), ⌈200/16⌉=13 (208 slots, waste 8), ⌈33/16⌉=3 (48 slots, waste 15). Total 18 blocks = 288 slots, actual tokens 258, waste 30, paged fragmentation = 30/288 ≈ 10.42%. Naive: 3×512 = 1,536 slots, waste 1,536−258=1,278, naive fragmentation = 1,278/1,536 ≈ 83.20%. Reduction factor = 1,536/288 ≈ 5.33×.

A3. At block size 64: ⌈25/64⌉=1 (64 slots, waste 39), ⌈200/64⌉=4 (256 slots, waste 56), ⌈33/64⌉=1 (64 slots, waste 31). Total 6 blocks = 384 slots, waste 126, fragmentation = 126/384 ≈ 32.81% — roughly three times worse than the 10.42% at block size 16, because each request now rounds up to a much coarser unit. Block-table entries drop from 18 (at size 16) to 6 (at size 64), a clean 3× reduction in per-sequence bookkeeping the GPU has to index through on every kernel call — the direct benefit the larger block size was chosen for. The ripple reaches prefix sharing too: with a 380-token shared prompt, only whole blocks that are entirely prefix (no divergence inside them) can be shared. At block size 16, ⌊380/16⌋=23 full blocks are shareable — 368 of the 380 prefix tokens (96.8%). At block size 64, only ⌊380/64⌋=5 full blocks are shareable — 320 of 380 tokens (84.2%); the block straddling the divergence point must be duplicated per request rather than shared, exactly like copy-on-write triggering earlier. So the same 16→64 change that cut block-table overhead 3× also tripled fragmentation waste and shrank reusable prefix by 48 tokens (368→320) — a genuine three-way tradeoff, not a free win.

A4. Incorrect. Every query still attends to every cached key with the same scores and the same softmax — the FLOP count of attention is unchanged. What PagedAttention changes is memory layout: keys and values live in scattered fixed-size physical blocks instead of one contiguous per-request array, located via a block table. It is a memory-management and sharing mechanism, orthogonal to and combinable with compute-kernel speedups like FlashAttention.

A5. Prefix blocks: ⌈512/16⌉=32 (512 slots, no waste, since 512 is an exact multiple of 16). Suffix blocks per request: ⌈60/16⌉=4 (64 slots). Without sharing: (32+4)×1,000 = 36,000 blocks. With sharing: 32 + 4×1,000 = 4,032 blocks. Reduction factor = 36,000/4,032 ≈ 8.93×.

A6. No. P0's reference count drops from 2 to 1 when Sequence A's table is torn down, but Sequence B's block table still points at it, so the block is still live. It can only be returned to the free pool once its reference count reaches zero — i.e., after every sequence referencing it has either finished or diverged away from it via copy-on-write. Freeing it while a reference remains would corrupt Sequence B's cache.

Think About It

Think about this: How would you explain kv cache optimization and management to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

  • Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
  • Exercise 2: Find a real-world example where kv cache optimization and management is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
  • Exercise 3: Create a mind-map connecting kv cache optimization and management to at least 3 other topics you have studied.
← Flash Attention and Memory-Efficient Attention VariantsSpeculative Decoding: Speeding Sequential Generation →

Found this useful? Share it!

📱 WhatsApp 🐦 Twitter 💼 LinkedIn