JuliaPagedAttn

A pure-Julia CPU implementation of PagedAttention for efficient KV cache management during LLM inference, with full GPT-2 transformer integration and interactive chat.

Features

• PagedAttention kernel — SIMD-optimized (@turbo) with online softmax, software prefetch, and configurable double buffering
• Unified K/V blocks — keys and values for the same token in one block, single block table, single pool, single lock
• GPT-2 transformer — complete PagedGPT2 model with prefill, autoregressive decode, and streaming output
• Real GPT-2 weights — load HuggingFace GPT-2 small (124M params) via safetensors
• BPE tokenizer — pure-Julia GPT-2 BPE tokenizer with priority-based merge selection
• Interactive chat — streaming token-by-token chat with top-p/top-k sampling and repetition penalty
• Copy-on-Write — fork shared prefix blocks between sequences without copying
• Thread-safe — ReentrantLock on the block allocator for concurrent access
• Batched prefill — all prompt positions processed in a single kernel call
• Online softmax — single-pass attention, O(1) memory per sequence

Quick Start

# Clone
git clone https://github.com/yourname/JuliaPagedAttn.jl
cd JuliaPagedAttn.jl

# Run tests
julia --project=. test/runtests.jl

# Run demo (random weights)
julia --project=. demo.jl

Interactive Chat with GPT-2

# Download GPT-2 weights (requires Python + transformers)
pip install transformers safetensors torch
python tools/convert_gpt2.py

# Copy weight files to project root
# (gpt2_small.safetensors, gpt2_vocab.json, gpt2_merges.txt)

# Start chat
julia --project=. chat.jl

The chat loads GPT-2 small (124M params) and runs inference through the paged attention engine. Uses top-p sampling, repetition penalty, and streaming output.

Without GPT-2 weights, demo.jl and chat.jl run with a small random-weight model.

Architecture

User Prompt
    |
    v
+-------------------------------+
|   Tokenizer (GPT-2 BPE)      |
|   or CharTokenizer (demo)     |
+-------------------------------+
    |
    v
+-------------------------------+
|   PagedGPT2 Transformer       |
|   - Embedding + Position      |
|   - N x TransformerBlock      |
|     - LayerNorm + Attention   |
|     - LayerNorm + FFN (GELU)  |
|   - Final LayerNorm + LM Head |
+-------------------------------+
    |
    v
+-------------------------------+
|   PagedEngine                 |
|   - UnifiedCacheManager       |
|   - K_cache / V_cache (5D)    |
|   - Block table indirection   |
+-------------------------------+
    |
    v
+-------------------------------+
|   paged_attention_cpu!        |
|   - @turbo SIMD dot products  |
|   - @batch sequence parallel  |
|   - Online softmax            |
|   - Software prefetch         |
+-------------------------------+
    |
    v
  Output

Block Layout

Each physical block stores K and V for block_size tokens across all KV heads:

K: [head_dim, block_size, num_kv_heads]
V: [head_dim, block_size, num_kv_heads]

Online Softmax

Single-pass attention with running maximum and sum — no full attention matrix materialization:

running_max = -Inf; running_sum = 0; acc = 0
for each token t:
    score = dot(Q, K[t]) * scale
    update running_max, running_sum, acc with correction
    acc += exp(score - running_max) * V[t]
output = acc / running_sum

Copy-on-Write

When multiple sequences share a prefix (e.g., system prompt), their block tables point to the same physical blocks with refcount > 1. On write: allocate new block, copy data, update block table, decrement refcount.

API Reference

Core Types

PagedEngine(head_dim, num_heads, num_kv_heads, max_seq_len, num_layers,
            block_size, num_blocks, max_sequences)

PagedGPT2(vocab_size; embed_dim=128, num_heads=8, num_kv_heads=2,
          max_seq_len=1024, num_layers=4)

Block Management

allocate_seq_blocks!(engine, seq_idx, context_len)
free_seq_blocks!(engine, seq_idx)
free_all_blocks!(engine)
cow_fork!(manager, block_id)

Cache Operations

write_kv_cache!(engine, K_new, V_new, seq_idx, start_pos, layer)

Attention

paged_attention_cpu!(O, Q, K_cache, V_cache, block_table, context_lengths, layer, scale;
                    config=DEFAULT_CONFIG)

Transformer

prefill!(engine, model, tokens; seq_idx=1)
next_token = decode_step!(engine, model, history_tokens;
    temperature=0.7, top_k=50, top_p=0.9, repetition_penalty=1.2)
tokens = generate(model, engine, prompt_tokens; max_new_tokens=50)

Weight Loading

weights = load_gpt2_weights("gpt2_small.safetensors")
load_gpt2!(model, weights)
tokenizer = load_gpt2_tokenizer("gpt2_vocab.json", "gpt2_merges.txt")

Running Tests

julia --project=. test/runtests.jl

22 correctness tests covering:

• Single/multi block attention
• Variable-length sequences
• GQA head grouping
• Partial last blocks
• CoW fork correctness
• Decode step (context_len=1)
• Block manager allocation/deallocation
• All kernel configs (naive, default, aggressive)

Running Benchmarks

julia --project=. bench/benchmark.jl

Sweeps context lengths (128–2048) and batch sizes (1–8), comparing reference, naive, default, and aggressive kernel configs.

Project Structure

JuliaPagedAttn.jl/
├── src/
│   ├── JuliaPagedAttn.jl      # Module definition
│   ├── block_manager.jl       # UnifiedCacheManager, CoW, locks
│   ├── cache_layout.jl        # PagedEngine, slot mapping, write_kv_cache!
│   ├── kernels.jl             # paged_attention_cpu!, online softmax
│   ├── prefetch.jl            # DoubleBuffer, prefetch_block!
│   ├── scheduler.jl           # RequestScheduler, continuous batching
│   ├── model.jl               # PagedGPT2, LayerNorm, Linear, attention
│   ├── safetensors.jl         # Pure-Julia safetensors reader
│   ├── weight_loader.jl       # load_gpt2!, load_fused_qkv!
│   ├── tokenizer.jl           # CharTokenizer (demo)
│   ├── tokenizer_gpt2.jl      # GPT2Tokenizer struct
│   └── tokenizer_gpt2_fns.jl  # BPE encode/decode
├── test/
│   └── runtests.jl
├── bench/
│   └── benchmark.jl
├── tools/
│   └── convert_gpt2.py        # Download GPT-2 weights
├── demo.jl                     # End-to-end demo
├── chat.jl                     # Interactive chat
├── Project.toml
└── .gitignore

Dependencies

• Julia 1.9+
• LoopVectorization.jl — @turbo SIMD
• Polyester.jl — @batch parallelism
• JSON.jl — tokenizer loading

References

• Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention," SOSP 2023. arXiv:2309.06180
• vLLM — PagedAttention reference implementation

License

MIT