The KV pool, the page table and the token pool¶
At startup the engine allocates three blocks of memory once and for all: the KV pool that holds every request's K and V, the page table that records where in the pool each token's KV lives, and the token pool that records which token sits at each position of each request. The attention backend, the scheduler and overlap scheduling all work around these three blocks. This chapter builds them and explains one design that is particular to mini-sglang: the page table always stores per-token locations, whatever the page size.
Self-test: if you can answer these, skip the chapter
- How many bytes does one token's KV take for Qwen3-0.6B in bf16? How many tokens fit in 40 GiB?
- What problem does paging (page size > 1) solve? Is a bigger page always better?
- Why is a copy of the input tokens kept on the GPU (the token pool) instead of sending them from the CPU each step?
- Why does the KV pool allocate one extra page?
Answers (try it yourself first, then expand)
- 2 × 28 layers × 8 KV heads × 128 dims × 2 bytes = 114688 bytes (112 KiB); 40 GiB holds about 374,000 tokens (40 × 1024 × 1024 ÷ 112).
- Allocating and managing by the page makes the block table shorter and the number of allocations and frees smaller, and it also makes sharing a prefix or transferring data page by page easy. Bigger is not always better: the last page leaves more internal fragmentation, and the prefix cache can only cache whole pages, which coarsens the hit granularity.
- The sampled token is written straight back into the token pool on the GPU, and the next step's input is read from there, which keeps the CPU off the critical path. That is also what makes overlap scheduling possible.
- For the dummy request that pads the batch for CUDA Graph: it writes KV too, and the extra last page is where it writes, so it never steps on a real request's data.
Files you will write: kvcache/base.py, kvcache/mha_pool.py, kvcache/naive_cache.py, kvcache/__init__.py, scheduler/table.py.
minisgl/
├── __init__.py
├── core.py
├── distributed/ (3 files, unchanged)
├── kernel/ (2 files, unchanged)
├── kvcache/
│ ├── __init__.py ← new in this chapter
│ ├── base.py ← new in this chapter
│ ├── mha_pool.py ← new in this chapter
│ ├── naive_cache.py ← new in this chapter
├── layers/ (8 files, unchanged)
├── models/ (7 files, unchanged)
├── scheduler/
│ ├── table.py ← new in this chapter
│ ├── __init__.py ← brought in early (explained in chapter 7)
├── utils/ (6 files, unchanged)
examples/ch04_kvcache.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.
The KV pool: one preallocated block of memory¶
During inference each request's KV cache grows as it generates, and requests come and go. Letting each request allocate its own memory fragments it quickly. Every modern inference engine does the same thing: allocate the remaining memory into one big pool at startup and manage it in fixed-size pages, handing out pages on demand and taking them back when done. That is the idea behind PagedAttention (the concept is in the Inference Systems handbook's paged KV cache chapter).
class MHAKVCache(BaseKVCachePool):
"""One preallocated block of KV memory: [2, layers, pages, page_size, this rank's KV heads, head_dim]."""
def __init__(self, num_kv_heads: int, num_layers: int, head_dim: int, num_pages: int,
page_size: int, dtype: torch.dtype, device: torch.device) -> None:
local_kv_heads = div_even(num_kv_heads, get_tp_info().size, allow_replicate=True)
self._kv_buffer = torch.empty(
(2, num_layers, num_pages, page_size, local_kv_heads, head_dim),
device=device, dtype=dtype,
)
self._k_buffer = self._kv_buffer[0]
self._v_buffer = self._kv_buffer[1]
self._num_layers = num_layers
self._device = torch.device(device)
# the shape once flattened per token: row i is the location of the pool's i-th token
self._storage_shape = (num_pages * page_size, local_kv_heads, head_dim)
def k_cache(self, index: int) -> torch.Tensor:
return self._k_buffer[index]
def v_cache(self, index: int) -> torch.Tensor:
return self._v_buffer[index]
def store_kv(self, k: torch.Tensor, v: torch.Tensor, out_loc: torch.Tensor,
layer_id: int) -> None:
store_cache(
k_cache=self._k_buffer[layer_id].view(self._storage_shape),
v_cache=self._v_buffer[layer_id].view(self._storage_shape),
indices=out_loc, k=k, v=v,
)
@property
def device(self) -> torch.device:
return self._device
@property
def dtype(self) -> torch.dtype:
return self._kv_buffer.dtype
@property
def num_layers(self) -> int:
return self._num_layers
[2, layers, pages, page_size, this rank's KV heads, head_dim]:
- dimension 0 separates K from V, and
k_cache(i)andv_cache(i)give layer i's[pages, page_size, heads, head_dim]; - flatten "pages × page_size" and row j is the location of the pool's j-th token, which is what
_storage_shapedescribes; store_kvwrites this step's new K and V byout_loc, each new token's location in the pool. The attention backend calls it before computing attention.
def store_cache(
k_cache: torch.Tensor,
v_cache: torch.Tensor,
indices: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
) -> None:
"""Write this step's new K and V into the KV pool at indices. k_cache: [num_tokens, H, D]; k: [n, H*D] or [n, H, D]."""
idx = indices.long()
k_cache[idx] = k.view(k.shape[0], *k_cache.shape[1:])
v_cache[idx] = v.view(v.shape[0], *v_cache.shape[1:])
"""Chapter 4: the KV pool's layout, its capacity, and what the page table and token pool look like."""
import torch
from minisgl.distributed import set_tp_info
from minisgl.kvcache.mha_pool import MHAKVCache
from minisgl.models import ModelConfig
from minisgl.utils import cached_load_hf_config
set_tp_info(0, 1)
c = ModelConfig.from_hf(cached_load_hf_config("models/Qwen3-0.6B"))
per_token = 2 * c.num_layers * c.num_kv_heads * c.head_dim * 2 # K+V，bf16
print(f"Qwen3-0.6B KV per token: {per_token} bytes = {per_token / 1024:.0f} KiB")
for gib in (4, 40):
print(f" {gib:>2} GiB of memory holds {gib * 2**30 // per_token:>8} tokens")
pool = MHAKVCache(num_kv_heads=c.num_kv_heads, num_layers=c.num_layers, head_dim=c.head_dim,
num_pages=64 + 1, page_size=16, dtype=torch.bfloat16, device=torch.device("cpu"))
print("KV pool shape [K/V, layer, page, offset, head, dim]:", tuple(pool._kv_buffer.shape))
print("layer 3's K:", tuple(pool.k_cache(3).shape), " flattened per token:", pool._storage_shape)
# two requests' page table (page_size = 16, locations per token)
page_table = torch.zeros((3, 64), dtype=torch.int32)
page_table[0, :20] = torch.cat([torch.arange(32, 48), torch.arange(80, 84)]) # page 2 plus the first 4 of page 5
page_table[1, :5] = torch.arange(160, 165) # the first 5 of page 10
print("request 0, first 20 locations:", page_table[0, :20].tolist())
print("request 1, first 5 locations: ", page_table[1, :5].tolist())
Qwen3-0.6B KV per token: 114688 bytes = 112 KiB
4 GiB of memory holds 37449 tokens
40 GiB of memory holds 374491 tokens
KV pool shape [K/V, layer, page, offset, head, dim]: (2, 28, 65, 16, 8, 128)
layer 3's K: (65, 16, 8, 128) flattened per token: (1040, 8, 128)
request 0, first 20 locations: [32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 80, 81, 82, 83]
request 1, first 5 locations: [160, 161, 162, 163, 164]
Why the extra page? The engine keeps a dummy request for CUDA Graph (chapter 18), and it writes KV too when it pads the batch. The extra last page is where it writes, so it never steps on a real request's data.
The page table: locations per token¶
The page table is a two-dimensional int32 tensor of shape [max_running_req + 1, max_seq_len rounded up to 32], one row per running request (the last row for the dummy), where page_table[r, j] is the location in the pool of that request's j-th token's KV.
One thing is particular to mini-sglang: whatever the page size, the page table stores token locations rather than page numbers. In the example above the page size is 16 and request 0 occupies pages 2 and 5, so its page-table row is [32, 33, …, 47, 80, 81, 82, 83], not [2, 5].
The benefit is simplicity:
- when writing KV,
out_loc = page_table[row, cached_len:device_len]gives each new token's location directly, with no "page number × page size + offset" arithmetic; - the FlashInfer backend treats the KV pool as if page size were 1, so one row of the page table is exactly the
indicesit wants; - FlashAttention wants page numbers, so it takes every page_size-th location and divides:
page_table[r, ::page_size] // page_size(chapter 17).
The price is a page table page_size times larger, but it is only int32, at most tens of MB. Rounding to 32 int32 values (128 bytes) aligns the start of each row, which the GPU reads more efficiently.
A page size above 1 matters for the granularity of prefix caching and allocation: the radix cache only caches whole pages (chapter 9) and the allocator hands out a whole page at a time (chapter 8), so larger pages mean less management overhead, at the cost of half a page wasted on average at the end and prefix matching that can only align to pages. Some kernels, such as TensorRT-LLM's attention, only support particular page sizes.
The token pool: inputs live on the device too¶
class TableManager:
"""Manages the rows of the page table / token pool: one row per running request."""
def __init__(self, max_running_reqs: int, page_table: torch.Tensor) -> None:
self._free_slots = list(range(max_running_reqs))
self.page_table = page_table
# token_pool[i, j]: the token id at position j of the request in row i. Same shape as page_table, on the device,
# so the model input can be gathered by index on the device and the sample written straight back (the key to overlap scheduling, chapter 11).
# the dummy request reads its input from here too, so this is initialized to a valid token id (0).
self.token_pool = torch.zeros_like(page_table, dtype=torch.int32)
@property
def available_size(self) -> int:
return len(self._free_slots)
def allocate(self) -> int:
return self._free_slots.pop()
def free(self, slot: int) -> None:
self._free_slots.append(slot)
TableManager manages the page table's rows: a request gets a row when it is admitted and returns it when it finishes. It also owns the token pool, an int32 tensor of the same shape as the page table, where token_pool[r, j] is the token id at that request's j-th position.
Why keep another copy of the tokens on the device? Because this step's input is the previous step's sample. Once the engine has next_tokens, the scheduler runs token_pool[row, device_len] = next_tokens right on the device, and when the next batch is assembled input_ids = token_pool[row, positions] also happens on the device. The CPU never needs to know what the new token is. This is exactly what overlap scheduling (chapter 11) relies on: while the GPU computes step N+1, the CPU may not yet have step N's results.
The dummy request also reads its input from the token pool, so the pool is initialized to a valid token id (0) rather than random values.
The prefix-cache interface¶
The KV pool only stores; it has no idea which location belongs to whom. What decides which locations can be reused is the prefix cache, whose interface is in kvcache/base.py:
class BasePrefixCache(ABC):
@abstractmethod
def lock_handle(self, handle: BaseCacheHandle, unlock: bool = False) -> None:
"""Lock (or unlock) a handle: once locked, that prefix cannot be evicted. It only changes counts, never the cached content."""
@abstractmethod
def match_prefix(self, input_ids: torch.Tensor) -> MatchResult:
"""Find the longest cached prefix. The returned handle can be evicted at any moment until it is locked."""
@abstractmethod
def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> InsertResult:
"""Insert a stretch of tokens and their KV locations into the cache."""
@abstractmethod
def evict(self, size: int) -> torch.Tensor:
"""Evict at least size tokens and return the KV locations freed. evict(0) is always a safe no-op."""
@abstractmethod
def reset(self) -> None: ...
@property
@abstractmethod
def size_info(self) -> SizeInfo: ...
@abstractmethod
def check_integrity(self) -> None: ...
NaivePrefixCache: it never hits and keeps nothing on insert, so a request's pages are all freed when it ends.
class NaivePrefixCache(BasePrefixCache):
"""No prefix reuse at all: it never hits and keeps nothing on insert (the caller then frees every page)."""
def __init__(self, device: torch.device):
self.empty_tensor = torch.empty(0, dtype=torch.int32, device=device)
def lock_handle(self, handle: BaseCacheHandle, unlock: bool = False) -> None:
pass
def match_prefix(self, input_ids: torch.Tensor) -> MatchResult:
return MatchResult(NaiveCacheHandle(self.empty_tensor))
def insert_prefix(self, input_ids: torch.Tensor, indices: torch.Tensor) -> InsertResult:
return InsertResult(0, NaiveCacheHandle(self.empty_tensor))
def evict(self, size: int) -> torch.Tensor:
if size == 0:
return self.empty_tensor
raise RuntimeError("NaivePrefixCache has nothing to evict")
def reset(self) -> None:
pass
@property
def size_info(self) -> SizeInfo:
return SizeInfo(evictable_size=0, protected_size=0)
def check_integrity(self) -> None:
pass
naive, radix) through the SUPPORTED_CACHE_MANAGER registry.
The official implementation
- the KV pool:
kvcache/mha_pool.py · MHAKVCache (lines 10-68), with the KV writes done by a custom CUDA kernel (kernel/csrc/jit/store.cu) - creating and aligning the page table: the "Page table initialization" section in
engine/engine.py · Engine.__init__ (lines 30-110), whose comment reads "aligned to 128 bytes; store raw locations instead of pages"; theContextincore.pyalso notes "this table always treat page_size = 1" - the token pool:
scheduler/table.py · TableManager (lines 4-21)
Tests¶
def test_kv_pool_layout_and_store():
set_tp_info(0, 1)
pool = MHAKVCache(num_kv_heads=2, num_layers=3, head_dim=4, num_pages=5, page_size=2,
dtype=torch.float32, device=torch.device("cpu"))
assert pool._kv_buffer.shape == (2, 3, 5, 2, 2, 4) # K/V, layer, page, offset in page, head, dim
k = torch.arange(3 * 8, dtype=torch.float32).view(3, 8) # 3 tokens, 2 heads x 4 dims each
pool.store_kv(k, -k, out_loc=torch.tensor([9, 0, 4], dtype=torch.int32), layer_id=1)
flat_k = pool.k_cache(1).reshape(10, 2, 4) # flattened per token: token 9 sits at offset 2 of page 4
assert torch.equal(flat_k[9].flatten(), k[0]) and torch.equal(flat_k[4].flatten(), k[2])
assert torch.equal(pool.v_cache(1)[4, 1].flatten(), -k[0])
How to explain it
On the KV pool: allocate one block of memory at startup, [2, layers, pages, page_size, KV heads, head_dim], and hand it out in pages, which avoids runtime allocation and fragmentation; one token's KV is 2 × layers × KV heads × head_dim × 2 bytes, or 112 KB for Qwen3-0.6B. The page table has one row per request and stores KV locations per token, and an attention backend that wants page numbers converts them itself. Paging lets memory grow on demand and makes prefix sharing possible; bigger pages mean less metadata and more efficient kernels but more waste at the tail and coarser prefix reuse. The token pool keeps the input tokens on the device too, so the sample is written back and read as the next input without the CPU on the critical path. One extra page is allocated for the dummy request used for padding.
Exercises¶
- How large is one token's KV for Llama-3.1-8B (32 layers, 8 KV heads, head_dim 128) in bf16? After 16 GB of weights on an 80 GB card, how many tokens fit?
- If the page table stored page numbers instead, how would
out_locbe computed when writing KV? Write the PyTorch expression. - With a page size of 16, how many pages does a 17-token request take? How much space does an average request waste?
Answers
- 2 × 32 × 8 × 128 × 2 = 131072 bytes (128 KiB); with about 64 GiB usable that is roughly 520,000 tokens (in practice activations and CUDA Graph memory come off the top too).
page_ids[pos // page_size] * page_size + pos % page_size, whereposis[cached_len, device_len).- 2 pages; the last page is half used on average, so the waste averages about page_size / 2 = 8 tokens.
Summary¶
- The KV pool is one big tensor
[2, layers, pages, page_size, heads, dim]allocated at startup and handed out in pages, with one extra page for the dummy request. - The page table has one row per request and stores KV locations per token regardless of page size; a backend that wants page numbers converts them itself.
- The token pool keeps each request's token at each position on the device, so the sample is written back directly and the next step's input read directly, keeping the CPU off the critical path.
- The prefix cache decides which locations can be reused; for now it is the naive version that never hits.
Exercises for this chapter (1)
Write the code and run the tests in your browser, or judge locally with python practice/judge.py. The exercises are in Chinese for now.
- 188. 按页分配 KV 并写进 page table (medium)
Found a mistake or a better way to say it? Edit this page on GitHubReport an issue