Skip to content

Latest commit

 

History

9 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

PageFault

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).

Results at a glance

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.

The problem

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.

The OS analogy

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

Four stages (each one is measured)

  1. Naive baseline: contiguous per-sequence cache and scaled dot-product attention. This is the memory-waste and latency floor.
  2. Block allocator: one cudaMalloc for the whole pool; a free list hands out frames. Sequences grow a page at a time.
  3. 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).
  4. Concurrent serving: many mixed-length sequences hold pages while they decode. Naive reserves MAX_SEQ_LEN per sequence; paged reserves the prompt and demand-pages the rest, with an admission watermark and recompute-style preemption.

Results

All results below were produced on an Apple M-series MacBook Pro (libc++). Reproduce them with the commands in Calculate.

Stage 2: memory

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.

Stage 4: serving under contention

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.

Stage 3: GPU kernel (pending)

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

Build

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 Release

If your GPU is not in the default architecture list (75;80;86;89), pass it explicitly, for example -DCMAKE_CUDA_ARCHITECTURES=86.

Test

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.

Calculate

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"
done

Core 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.

What the benchmark prints

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

Layout

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

Design notes (intentional simplifications)

  • Single-head kernel. NUM_HEADS is 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_SIZE tokens).
  • 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.

Roadmap

  • 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.

Profiling

nsys profile -o pagefault.nsys-rep ./build/pagefault_bench

Compare naiveAttentionKernel and pagedAttentionKernel in the timeline: extra global-memory gathers on the paged path are expected.

License / intent

Built to learn GPU memory management and inference systems, not to replace vLLM.

About

A block-based paged memory allocator and attention kernel for LLM KV-caches on the GPU; OS-style virtual memory paging, reimplemented in CUDA

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages