A block-based paged memory allocator and attention kernel for LLM KV-caches on the GPU: OS-style virtual memory paging, implemented in CUDA.
This is a learning implementation of the idea behind vLLM's PagedAttention and TensorRT-LLM's paged KV-cache, not a production serving stack. The point is to measure why paging works (far less reserved GPU memory), what it costs (an extra gather in the attention kernel), and what goes wrong without a scheduler (thrashing).
| Naive (contiguous) | Paged | |
|---|---|---|
| KV memory reserved, 32 mixed-length sequences | 33.55 MB | 10.63 MB (3.16× less) |
| Utilization (stored / reserved) | 31.3% | 98.7% |
| Simulated throughput, no admission control | 1.00× | 0.59–0.67× (thrashing) |
| Simulated throughput, 5–30% watermark | 1.00× | ~2× |
| Actively decoding sequences (mean) | ~27 | ~58 |
Memory numbers are exact, computed from the workload. Throughput comes from the multithreaded CPU serving simulator (Stage 4), not from GPU serving, and varies about ±10% between runs. GPU kernel latency and accuracy are pending.
Autoregressive decoding keeps a growing KV-cache for every live sequence. The naive layout is one contiguous MAX_SEQ_LEN buffer per sequence. Real sequences finish at very different lengths, so most of that reservation is empty. You run out of GPU memory while a large fraction of the cache is unused.
| Operating system | PageFault |
|---|---|
| Virtual page | Logical KV block (BLOCK_SIZE tokens) |
| Physical frame | Slot in the GPU block pool |
| Page table | Per-sequence BlockTable |
| Page-table walk | resolveTokenAddress() / poolSlotIndex() in the attention kernel |
| Demand paging | A new page is allocated only when a sequence fills its last one |
| Low-memory watermark | Admission keeps a fraction of the pool free for running sequences to grow |
| Thrashing | Too many sequences admitted; they stall and preempt each other fighting for pages |
| Page fault with no free frame | Simulator: preempt the sequence (recompute). GPU allocator: throws |
GPU memory is split into fixed-size pages. Each sequence has a block table mapping logical token ranges to physical blocks anywhere in the pool. The paged attention kernel gathers K/V through that table instead of assuming contiguity.
logical tokens: [ 0..15 | 16..31 | 32..47 | ... ]
block table: [ 7 , 2 , 19 , ... ] physical frame ids
physical pool: block 2, block 7, block 19, ... scattered
- Naive baseline: contiguous per-sequence cache and scaled dot-product attention. This is the memory-waste and latency floor.
- Block allocator: one
cudaMallocfor the whole pool; a free list hands out frames. Sequences grow a page at a time. - Paged attention: same math as Stage 1, but every K/V load walks the block table. Checked against the naive kernel, a CPU reference (every sequence) and a host read-back through the page tables. Measured twice: with ordered frames (best-case locality) and shuffled frames (what a pool looks like after churn).
- Concurrent serving: many mixed-length sequences hold pages while they decode. Naive reserves
MAX_SEQ_LENper sequence; paged reserves the prompt and demand-pages the rest, with an admission watermark and recompute-style preemption.
All results below were produced on an Apple M-series MacBook Pro (libc++). Reproduce them with the commands in Calculate.
The workload is 50% short (32–128 tokens), 30% medium (256–512) and 20% long (1024–2048) sequences.
| Workload | Blocks used | Paged reserved | Naive reserved | Savings | Utilization (naive → paged) |
|---|---|---|---|---|---|
| seed 42 (8 long sequences) | 1298 / 4096 | 10.63 MB | 33.55 MB | 3.16× | 31.3% → 98.7% |
| seed 7 (5 long sequences) | 790 / 4096 | 6.47 MB | 33.55 MB | 5.18× | 18.9% → 98.1% |
The savings depend on the traffic mix: the more short requests, the more naive wastes. Measured internal fragmentation was 8.38 and 7.34 tokens per sequence, matching the expected (B − 1) / 2 = 7.5 for uniformly distributed lengths.
Block-size sweep for seed 42 (the page-size tradeoff):
| Block size | Reserved | Fragmentation | Utilization | Table entries per sequence |
|---|---|---|---|---|
| 1 | 10.50 MB | 0.00 MB | 100.0% | 2048 |
| 16 (default) | 10.63 MB | 0.14 MB | 98.7% | 128 |
| 64 | 11.01 MB | 0.51 MB | 95.3% | 32 |
| 256 | 13.24 MB | 2.74 MB | 79.3% | 8 |
Smaller pages waste less memory but mean longer block tables and more indirection per token in the kernel.
512 jobs, 128 concurrent workers, a pool of 4096 blocks. Each number is a single run.
| Watermark | Preemptions | Decoding (mean) | Admission wait | Throughput vs naive |
|---|---|---|---|---|
| naive | 0 | ~27 | ~105 ms | 1.00× |
| paged, 0% | 750–1000 | ~40 | < 1 ms | 0.59–0.67× |
| paged, 2% | 32 | 53.5 | 15 ms | 1.83× |
| paged, 5% | 14–19 | 55.2–55.9 | 15–17 ms | 1.92–1.98× |
| paged, 10% | 10–11 | 55.3–56.4 | 15–19 ms | 1.98–2.12× |
| paged, 30% | 0 | 58.3 | 18 ms | 2.33× |
| paged, 50% | 0 | 42.4 | 43 ms | 1.54× |
What this shows:
- Paging alone makes things worse. With no watermark, paging admits everyone, then about half the live sequences sit stalled waiting for a page, and preemptions throw away finished work. Throughput drops below naive.
- A small reserve fixes most of it. Even 2% cuts preemptions from about 1000 to 32.
- Too much reserve wastes memory. At 50%, new requests queue behind memory that running sequences never use, and throughput falls back to 1.54×.
- The peak here is high (around 30%) because the simulator preempts whichever sequence fails to get a page, often a long one that is nearly finished. Production schedulers choose victims deliberately (newest or lowest priority) and can swap pages to CPU memory, so they get away with much smaller reserves.
- The pool is the ceiling, not the workers. At 256 workers with a 5% watermark, decoding stayed at 57.4, essentially unchanged from 128 workers. Without a watermark, doubling demand made thrashing worse: 2947 preemptions and 0.48× naive.
Run-to-run noise: the naive and 0% columns do not depend on --watermark, yet they varied about ±10% in throughput across runs. Treat small differences (such as 5% vs 10%) as ties. The preemption and decoding counts are more stable than throughput.
Not yet run on an NVIDIA GPU. pagefault_bench will report:
| Metric | Result |
|---|---|
| GPU | TBD |
| Max |paged − naive| | TBD |
| Paged slowdown vs contiguous (shuffled frames) | TBD |
| Cost of scattering pages (shuffled / ordered) | TBD |
The host tests and calculator need only a C++17 compiler and CMake ≥ 3.18, so they build on macOS. The GPU benchmark additionally needs the CUDA Toolkit and is skipped automatically when no CUDA compiler is found.
cmake -S . -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build --config ReleaseIf your GPU is not in the default architecture list (75;80;86;89), pass it explicitly, for example -DCMAKE_CUDA_ARCHITECTURES=86.
Run these from the build directory. (The command blocks in this README have no inline # comments, because zsh on macOS treats them as arguments when pasted.)
ctest --output-on-failure
./pagefault_tests
./pagefault_tests alloc
ctest -L gpu --output-on-failure| Command | What it runs |
|---|---|
ctest --output-on-failure |
everything that was built |
./pagefault_tests |
the 22 host tests (any machine) |
./pagefault_tests alloc |
only tests whose name contains alloc |
ctest -L gpu --output-on-failure |
only the GPU benchmark |
| Layer | What proves it works | Needs GPU |
|---|---|---|
| Page-table math | hand-calculated addresses; no two tokens share a slot | no |
| Free list | exhaustion, double free, all-or-nothing, watermark, 8-thread stress | no |
| Attention math | uniform keys → mean of values; one-hot key → that value; paged == contiguous bitwise on a shuffled pool | no |
| Memory accounting | worked example {1,16,17,100} → 11 blocks, 42 tokens of fragmentation |
no |
| Scheduler | no leaks, naive ≤ 32 live, paged > naive, watermark cuts preemptions | no |
| CUDA kernels | naive vs CPU (all sequences), paged vs naive (ordered and shuffled), pool read-back | yes |
pagefault_bench exits non-zero if any GPU check fails. Useful flags: --seqs, --iters, --seed, --shuffle-seed 0 (ordered frames only), --tolerance (default 1e-3), --jobs, --workers, --watermark, --skip-sim.
pagefault_calc derives every memory number the benchmark prints, on the CPU, with the formulas shown. From the build directory:
./pagefault_calc
./pagefault_calc --seqs 32 --seed 7
./pagefault_calc --sim --watermark 0.30
./pagefault_calc --sim --workers 256
./pagefault_calc --model llama3-70b --budget-gb 80 --max-len 8192 --mean-len 2000
./pagefault_calc --model custom --layers 32 --kv-heads 8 --head-dim 128 --dtype-bytes 2| Command | What it shows |
|---|---|
./pagefault_calc |
config, workload, utilization, block-size sweep |
--seqs N --seed S |
a different workload |
--sim [--watermark F] [--workers N] |
Stage 4: naive vs paged vs paged with watermark |
--model llama2-7b|llama3-8b|llama3-70b|custom |
KV-cache sizing for a real model |
Watermark sweep, printing only the key rows (the last column is the watermark passed):
for wm in 0.05 0.10 0.30 0.50; do
echo "== watermark $wm"
./pagefault_calc --sim --watermark $wm | grep -E "preemptions|of which decoding|admission wait|throughput vs naive"
doneCore formulas (B = BLOCK_SIZE):
bytes/token (this repo) = 2 * HEAD_DIM * 4 = 512 B
bytes/token (real model) = 2 * layers * kv_heads * head_dim * dtype_bytes
naive reserved = num_seqs * MAX_SEQ_LEN * bytes/token
paged reserved = sum(ceil(len / B)) * B * bytes/token
utilization = sum(len) / reserved tokens
internal fragmentation = sum(ceil(len / B) * B - len) ~ (B-1)/2 tokens per seq
naive capacity = pool tokens / MAX_SEQ_LEN = 32 here
Example: Llama-3-70B stores 2 × 80 × 8 × 128 × 2 = 320 KB of KV-cache per token. An 80 GB budget holds 29 sequences reserved at 8192 tokens, or 122 paged sequences averaging 2000 tokens (4.2× more, assuming the average holds).
Random workloads come from std::uniform_int_distribution, which differs between libstdc++ (Linux), libc++ (macOS) and MSVC. The same seed gives different lengths on different platforms; the calculator and benchmark agree when built with the same standard library.
| Metric | What it means |
|---|---|
| Memory reserved | Naive: num_seqs * MAX_SEQ_LEN. Paged: pages actually mapped. |
| Utilization | Bytes of real tokens / bytes reserved |
| Attention latency | Median of CUDA-event timings (mean, sd, min also shown) |
| Cost of scattering pages | Shuffled-frame latency / ordered-frame latency |
| Max |paged − naive| | Largest output difference between the two kernels; must be within --tolerance |
| Mean decoding sequences | Live sequences not stalled waiting for a page; this is what becomes throughput |
pagefault/
├── CMakeLists.txt
├── README.md
├── src/
│ ├── config.hpp sizes, KVEntry, pure C++ (no CUDA)
│ ├── paging.hpp page-table walk shared by host and kernels
│ ├── frame_allocator.hpp host free list: watermark, peak, double-free checks
│ ├── workload.hpp deterministic lengths + memory math
│ ├── reference_attention.hpp CPU attention (contiguous and paged)
│ ├── serving_sim.hpp Stage 4 scheduler simulation
│ ├── thread_pool.hpp
│ ├── common.cuh CUDA_CHECK, deviceMalloc
│ ├── naive_cache.cuh/.cu Stage 1
│ ├── block_allocator.cuh/.cu Stage 2
│ ├── paged_attention.cuh/.cu Stage 3
│ └── benchmark.cpp
├── tests/test_host.cpp
├── tools/pagefault_calc.cpp
└── scripts/export_weights.py
- Single-head kernel.
NUM_HEADSis config for a later grid dimension (blockIdx.y = head), not a packed multi-head layout yet. - Preemption is recompute-style and only in the Stage 4 simulator. A preempted sequence frees everything and starts over; nothing is swapped to host memory. The victim is whichever sequence fails to get a page, not a deliberately chosen one.
- Internal fragmentation is only the unused tail of the last page (
(BLOCK_SIZE - seq_len % BLOCK_SIZE) % BLOCK_SIZEtokens). - Attention is not FlashAttention. Scores live in shared memory; V is gathered again for the weighted sum. Clear and comparable, not peak FLOPs.
- Allocator mutex. The free list is host-side and lock-protected. A lock-free list would be the next systems step.
- Run the GPU benchmark and fill in the Stage 3 results.
- Victim-selection policy for preemption (newest or lowest priority first), then re-run the watermark sweep to see the optimal reserve shrink.
- Swap preempted sequences' pages to host memory instead of recomputing.
- Multi-head attention using a second grid dimension.
- Copy-on-write sharing of common prompt prefixes.
nsys profile -o pagefault.nsys-rep ./build/pagefault_benchCompare naiveAttentionKernel and pagedAttentionKernel in the timeline: extra global-memory gathers on the paged path are expected.
Built to learn GPU memory management and inference systems, not to replace vLLM.