Author: João Felipe De Souza
Year: 2026
Implement PagedAttention from scratch, including:
- Block-based KV cache allocation
- Custom CUDA attention kernel with advanced optimizations
- Prefix sharing between sequences
- Copy-on-write block semantics
- Concurrent sequence batching
- Integration into real LLM (Qwen2)
This project focuses on correctness, memory efficiency, kernel performance, and architectural clarity.
System component flow:
- Client sends inference request
- Runtime Layer receives and orchestrates
- BlockManager handles block allocation, freeing, prefix sharing, and copy-on-write
- Physical KV Cache stores K and V tensors indexed by physical block ID
- Block Table maps logical sequence positions to physical blocks
- CUDA PagedAttention Kernel iterates blocks with indirection and computes attention
Kernel v3 features:
- Block-wise iteration with indirection
- Shared memory KV tiling
- Online softmax (streaming, single-pass)
- Warp-level reduction
- half2 vectorization
Each block represents a fixed number of tokens (16 by default).
Key properties:
- total_blocks: total physical blocks
- block_size: tokens per block (16)
- free_blocks: free-list
- ref_count: reference count per block
- sequence_blocks: mapping seq_id to list of block_ids
Operations:
- allocate_blocks(seq_id, n)
- free_sequence(seq_id)
- share_prefix(source, target, n_blocks)
- cow_block(seq_id, block_index)
Prefix sharing increases ref_count. Copy-on-write allocates a new block when modifying shared memory. Blocks return to free-list when ref_count reaches zero.
- gridDim.x = batch_size
- gridDim.y = num_heads
- blockDim.x = head_size (typically 64 or 128)
For each sequence and head:
-
Initialize online softmax accumulators in registers
- m_i = -INFINITY (running max)
- l_i = 0 (running sum)
- acc[head_size] = 0 (running output)
-
For each logical block in the sequence:
- Resolve physical block via block_table indirection
- Cooperatively load K, V for this block into shared memory
- Compute Q dot K.T for all tokens in block
- Update online softmax:
- New max: m_new = max(m_i, max(scores))
- Rescale accumulator: acc *= exp(m_i - m_new)
- Compute weights: p = exp(scores - m_new)
- Update sum: l_i = l_i * exp(m_i - m_new) + sum(p)
- Accumulate output: acc += p dot V
- m_i = m_new
-
Normalize final output: output = acc / l_i
Online Softmax (Milakov & Gimelshein 2018)
- Single-pass computation
- Numerical stability via running max
- No need to materialize full scores tensor
- FP32 accumulation for stability
Warp-Level Reduction
- Dot products use __shfl_down_sync for warp-wide sum
- Eliminates shared memory round-trips for partial sums
- Warp = 32 threads, matches typical head_size / 2
Shared Memory KV Tiling
- Each block loaded once, reused across all query computations
- 4 KB per K tile + 4 KB per V tile (block_size x head_size x FP16)
- Fits comfortably in 64 KB shared memory budget
half2 Vectorization
- 128-bit loads via half2 pairs
- Fused multiply-add on 2 FP16 values per instruction
- Doubles effective memory bandwidth
Multiple sequences are processed simultaneously:
- Each sequence has independent block table
- Shared physical KV cache pool
- Concurrent grid execution (batch dim in gridDim.x)
- Heterogeneous sequence lengths supported
Throughput increases with batch size until memory-bound. Measured 5.32x scaling from 1 to 8 sequences.
Qwen2 uses Grouped Query Attention:
- num_heads = 14
- num_kv_heads = 2
- group_size = num_heads / num_kv_heads = 7
Each KV head is shared across 7 query heads. Kernel expands KV heads to match query heads via broadcast-style indexing.
Correctness note: GQA expansion is the subtlest bug source. Silent numerical drift can occur if expansion index is off-by-one. Validated with bit-close comparison against Python reference.
Validation performed at three layers:
-
Python reference vs Paged Python:
- Ensures block-based algorithm matches naive attention
- Test file: tests/test_paged_attention_python.py
-
CUDA kernel vs Python reference:
- Ensures GPU implementation matches Python
- Test file: tests/test_paged_attention_cuda.py
- Uses realistic shapes (Qwen2 head_size=64, num_heads=14)
-
Runtime layer vs raw kernel:
- Ensures orchestration layer preserves correctness
- Test files: tests/test_paged_attention_runtime.py, tests/test_concurrent_runtime.py
All validations pass with output equivalence up to FP16 precision noise.
Throughput on RTX 2070 (SM75):
| Concurrent Sequences | Tokens/sec | Speedup |
|---|---|---|
| 1 | 542,981 | 1.00x |
| 2 | 844,423 | 1.55x |
| 4 | 1,239,648 | 2.28x |
| 8 | 2,890,839 | 5.32x |
For attention-layer throughput on Qwen2-0.5B, ctx=512:
| Batch | HuggingFace | Paged Kernel | Speedup |
|---|---|---|---|
| 1 | 4,687 tok/s | 963,762 tok/s | 205x |
| 4 | 7,173 tok/s | 2,852,660 tok/s | 397x |
HF baseline includes full model forward pass; paged kernel measures attention only. Direct model-level comparison requires E2E integration.
| Version | Optimization | 8-seq throughput |
|---|---|---|
| v1 (naive) | Baseline block iteration | ~100k tok/s |
| v2 | Shared memory tiling | ~800k tok/s |
| v3 | + Warp reduction + Online softmax + half2 | 2.89M tok/s |
Total gain: ~29x through targeted optimizations.
Planned improvements (bounded by SM75 architectural limits):
Software-level:
- Iteration-level batching (requests enter/leave mid-generation)
- Preemption policy for memory pressure
- Async request scheduler
- Full E2E Qwen2 generation with logit validation
Kernel-level (limited by SM75):
- FlashAttention-style full Q tiling (limited gain expected without cp.async)
- Tensor Core integration for prefill (M >= 16)
- Multi-warp cooperation for larger head_size
Not viable on SM75:
- cp.async for load/compute overlap (Ampere+ only)
- Native INT4 MMA operands (Ampere+ only)
- Warp specialization (Hopper+ only)
- Not a production-grade vLLM replacement
- No KV eviction strategy for OOM scenarios
- No async scheduler
- No dynamic memory compaction
- No iteration-level batching
- Single-GPU only
- E2E Qwen2 generation in progress (attention-layer equivalence validated)
This is an architectural and systems-level implementation focused on correctness, kernel optimization discipline, and understanding of core algorithm design.
For production LLM serving, use vLLM directly. For understanding how vLLM works internally and building intuition for LLM serving systems, this project is designed to be readable and educational.
- Balances internal fragmentation vs allocation overhead
- Aligns with common head sizes (64, 128 divisible)
- Matches vLLM default (proven in production)
- 4 KB per K tile fits comfortably in shared memory
- Single kernel pass instead of traditional two-pass
- No need to materialize full scores tensor
- Better shared memory utilization
- Numerical stability preserved via FP32 accumulation
- SM75 lacks cp.async for load/compute overlap
- FlashAttention-style full tiling has marginal gain in Turing
- Streaming approach is simpler and correct
- Reserves complexity budget for portability to Ampere+
- FP16 accumulation loses precision after ~1000 accumulations
- FP32 accumulator preserves ~7 decimal digits
- Cost is minor (accumulator only, not weights)
- Standard practice in production kernels