Custom CUDA implementation of PagedAttention and a full LLM serving runtime stack inspired by vLLM.
This project rebuilds from scratch the core memory and scheduling mechanisms behind high-performance LLM serving engines.
- Custom CUDA PagedAttention kernel
- Block-based KV cache allocator
- Continuous (iteration-level) batching
- Priority-aware scheduling
- Hybrid preemption (swap + recompute fallback)
- CPU-side KV swap-out with pinned memory support
- Cancellation support
- Streaming API
- Async submission API
- Reproducible benchmark infrastructure
- 80+ tests covering full lifecycle
This runtime includes the core architectural ideas behind modern LLM serving systems:
- Block-based KV allocation (paged memory abstraction)
- Prefix sharing and reference counting
- Copy-on-write semantics
- Continuous dynamic batching
- Priority-aware admission
- Preemption under memory pressure
- Swap-out preemption to CPU
- Streaming token emission
- Async request handling
All implemented from scratch with explicit control over memory lifecycle and scheduling behavior.
Client Request
↓
Scheduler (Priority + Continuous Admission)
↓
Block Manager (KV Allocation + Ref Counting)
↓
PagedAttention CUDA Kernel
↓
Token Generation
↓
Swap / Preempt / Continue
- Fixed-size KV blocks
- O(1) allocation via free list
- Reference counting for prefix sharing
- Copy-on-write support
- Fragmentation tracking
- Iteration-level admission
- Priority-aware queue
- FIFO within same priority
- Controlled prefill throttling
- Detects memory pressure
- Selects victim via policy
- Swap-first strategy
- Recompute fallback
- CPU-side KV storage
- Optional pinned memory
- Swap-out on pressure
- Swap-in on re-admission
- Preserves generated tokens
Example:
for event in runtime.stream(generator): print(event)
Example:
future = runtime.submit(generator) result = future.result()
Benchmarks are fully reproducible:
PYTHONPATH=. python benchmarks/compare_serving_modes.py
PYTHONPATH=. python benchmarks/plot_compare.py
PYTHONPATH=. python benchmarks/aggregate_history.py
PYTHONPATH=. python benchmarks/generate_report.py
- Device: CUDA
- Heads: 4
- Head size: 8
- Block size: 16
- Requests: 32
- Max new tokens per request: 40
Log Scale Comparison
See: benchmarks/results/compare_serving_modes_log.png
Continuous vs Continuous + Swap
See: benchmarks/results/compare_continuous_only.png
- Maximum raw kernel throughput
- No scheduling overhead
- Not suitable for heterogeneous workloads
- Enables heterogeneous serving
- Introduces scheduler overhead
- Supports dynamic request arrival
- Discards KV cache
- Wastes compute under pressure
- Simple fallback mechanism
- Preserves KV cache in CPU memory
- Avoids recompute cost
- Improves throughput under pressure
- Adds PCIe transfer latency
- High priority requests admitted first
- Lower priority requests swapped/preempted first
- Fairness maintained without strict real-time guarantees
Includes:
- Stress test under memory pressure
- SLA workload fairness test
- Throughput benchmark
- Historical aggregation
- Auto-generated plots
- Auto-generated Markdown report
Example stress test:
PYTHONPATH=. python benchmarks/stress_runtime.py
Example SLA test:
PYTHONPATH=. python benchmarks/sla_priority_workload.py
To evaluate the impact of victim selection strategies under heterogeneous workloads and tight memory constraints, we ran an extreme stress test with:
- num_blocks = 8 (very constrained GPU memory)
- max_active_sequences = 6
- max_swapped_requests = 2
- Mixed workloads:
- Large KV requests (prompt_len=64, max_new_tokens=80)
- Small KV requests (prompt_len=8, max_new_tokens=10)
- Equal priority across requests
Under extreme heterogeneity:
| Policy | Tokens/sec | Elapsed (s) |
|---|---|---|
| FIFO | ~457 | ~0.13 |
| COST_BASED | ~3232 | ~0.019 |
Under extreme memory constraints and heterogeneous sequence lengths, COST_BASED significantly outperformed FIFO in this synthetic stress scenario.
Although both policies triggered a similar number of swap events, the effective scheduling dynamics differed, leading to a ~7× throughput gap.
This suggests that victim selection policy can materially affect decode parallelism and runtime behavior under tight memory limits.
Further instrumentation (kernel call counts, active batch size evolution, and recompute metrics) is required to fully attribute the performance gap.
Victim selection policy matters significantly under:
- Heterogeneous sequence lengths
- Constrained GPU memory
- Limited swap pool capacity
Simple heuristics (FIFO) are sufficient in homogeneous workloads, but cost-aware scheduling substantially improves performance under real serving pressure.
Run full test suite:
pytest tests/ -v
Covers:
- Block allocation
- Prefix sharing
- Copy-on-write
- Continuous batching
- Preemption
- Swap-out / swap-in
- Cancellation
- Kernel correctness
- Memory reclamation
- Priority scheduling
This repository implements a full serving runtime stack comparable to modern LLM serving engines.
It demonstrates:
- Systems-level understanding
- GPU memory management
- Scheduling algorithms
- Memory pressure handling
- Hybrid preemption strategies
- Real-world serving behavior
All built from scratch with explicit control over memory lifecycle and scheduling semantics.
MIT João Felipe De Souza