LLM Cache Management — Production Code & Kernels

Ebook: https://shop.beacons.ai/aiengineeringinsider/31df9f7d-fc7d-43d1-857c-c720ed8815b5

Preview: https://drive.google.com/file/d/12cUngq6P0yBTVra7W60Df28kxttRa4no/view?usp=sharing

A practical, production-grade code repository accompanying the 15-chapter masterclass on Key-Value Cache Management in Large Language Models.

This repository contains 25 runnable Python implementations, C++ structures, and CUDA GPU kernels that demonstrate the core algorithms, memory allocators, quantization schemes, disaggregation transports, and latency profilers used in modern LLM serving engines like vLLM, SGLang, TensorRT-LLM, llama.cpp, NVIDIA Dynamo, and LMCache.

---

📁 Repository Code Structure

llm-cache-management-code/
├── README.md                     # This documentation file
├── .venv/                         # Isolated Python virtual environment with PyTorch & NumPy
│
├── # Part I: Foundations — Why the Cache Exists
├── causal_attention.py           # Ch 1: Reference MHA with per-layer append-only KV Cache
├── qk_scores.cu                  # Ch 1: Single-query score computation CUDA kernel
├── decode_benchmark.py           # Ch 2: Cached vs. uncached (recompute) decoding benchmark
├── kv_cache_layout.cpp           # Ch 2: Contiguous per-layer KV cache C++ structure
├── kv_memory_calculator.py       # Ch 3: KV cache size and GPU concurrency budget calculator
├── fp8_kv_read.cu                # Ch 3: FP8 (E4M3) dequantize-on-read CUDA kernel
│
├── # Part II: Making the Cache Smaller and Managing It
├── grouped_query_attention.py    # Ch 4: GQA implementation with shared KV heads
├── sliding_window_cache.py       # Ch 4: StreamingLLM bounded rolling cache with sinks
├── fragmentation_sim.py          # Ch 5: Memory fragmentation simulator (Contiguous vs. Paged)
├── hbm_bandwidth.cu              # Ch 5: STREAM-style HBM bandwidth measurement kernel
├── paged_allocator.py            # Ch 6: PagedAttention block allocator with Copy-on-Write
├── batching_sim.py               # Ch 6: Iteration-level scheduling (Static vs. Continuous)
├── paged_attention_gather.cu     # Ch 6: Block-table gather score CUDA kernel
│
├── # Part III: Reusing, Compressing, and Evicting the Cache
├── radix_prefix_cache.py         # Ch 7: RadixAttention shared-prefix tree (SGLang style)
├── kv_quantize.py                # Ch 8: Per-token vs. per-tensor INT8/INT4 quantization error
├── int4_kv_pack.cu               # Ch 8: INT4 2-element per byte packing & unpacking kernel
├── kv_eviction.py                # Ch 9: Attention-guided eviction (Heavy-Hitter Oracle H2O)
│
├── # Part IV: Scaling the Cache Across Memory and Machines
├── tiered_kv_cache.py            # Ch 10: Two-tier HBM -> CPU DRAM cache with async prefetch
├── kv_prefetch.cu                # Ch 10: Pinned-memory async copy CUDA stream prefetcher
├── kv_transport_cost.py          # Ch 11: Interconnect transport cost model (NVLink/PCIe/RDMA)
├── kv_rdma_transfer.cu           # Ch 11: GPUDirect RDMA device buffer registration
├── speculative_decode.py         # Ch 12: Speculative decoding draft/verify acceptance model
│
└── # Part V: Application Layer, Measurement, and Production Engines
    ├── app_cache.py              # Ch 13: Exact hash match & semantic similarity response cache
    ├── serving_benchmark.py      # Ch 14: Latency profiling harness (TTFT, TPOT, ITL, Goodput)
    └── engine_selector.py        # Ch 15: Production engine workload decision framework

---

🛠️ Prerequisites & Installation

Requirements

• Python: 3.10 or higher
• Python Libraries: torch, numpy (pre-installed inside the .venv directory)
• C++ Compiler: clang++ or g++ supporting C++17
• NVIDIA CUDA Toolkit (Optional for CUDA kernels): nvcc compiler with CUDA runtime support

Quick Setup

Activate the pre-configured virtual environment or set up a new environment:

# Clone or navigate to the directory
cd llm-cache-management-code

# Activate the virtual environment
source .venv/bin/activate

# Or install dependencies manually via pip:
pip install torch numpy

---

🚀 How to Run the Code

1. Running Python Implementations

All Python scripts are self-contained executable modules. You can run any script individually using Python:

# Run reference attention with KV cache (Chapter 1)
python causal_attention.py

# Run the cached vs. uncached decoding benchmark (Chapter 2)
python decode_benchmark.py

# Run the KV memory and concurrency budget calculator (Chapter 3)
python kv_memory_calculator.py

# Run Grouped-Query Attention (Chapter 4)
python grouped_query_attention.py

# Run PagedAttention block allocator with Copy-on-Write (Chapter 6)
python paged_allocator.py

# Run RadixAttention prefix caching simulation (Chapter 7)
python radix_prefix_cache.py

# Run INT8/INT4 quantization error benchmark (Chapter 8)
python kv_quantize.py

# Run the full serving benchmark harness (Chapter 14)
python serving_benchmark.py

# Run the engine selector decision framework (Chapter 15)
python engine_selector.py

Run All Python Scripts at Once:

for file in *.py; do
    echo "=== Running $file ==="
    python "$file"
    echo ""
done

---

2. Compiling and Running C++ Internals

To compile and run host-side C++ inference engine structures (e.g., kv_cache_layout.cpp):

# Compile using clang++ or g++ with C++17 standard
clang++ -std=c++17 kv_cache_layout.cpp -o kv_cache_layout

# Run the compiled binary
./kv_cache_layout

Expected Output:

LayerKVCache created successfully.
Stored tokens: 10
Head 0 key offset for token 0: 0

---

3. Compiling CUDA GPU Kernels (NVIDIA GPU Required)

The CUDA kernels (.cu files) demonstrate GPU-level HBM bandwidth optimization, FP8 dequantize-on-read, INT4 nibble packing, PagedAttention table gather, and GPUDirect RDMA zero-copy transfers.

To compile a CUDA kernel with nvcc:

# Compile HBM bandwidth microbenchmark
nvcc -O3 hbm_bandwidth.cu -o hbm_bandwidth
./hbm_bandwidth

# Compile single-query QK score kernel
nvcc -c qk_scores.cu -o qk_scores.o

# Compile FP8 dequantize-on-read kernel
nvcc -c fp8_kv_read.cu -o fp8_kv_read.o

# Compile INT4 packing kernel
nvcc -c int4_kv_pack.cu -o int4_kv_pack.o

---

📊 Summary of What Each Script Demonstrates

Script / Kernel	Chapter	Core System Concept Demonstrated
causal_attention.py	Ch 1	Appends $K$ and $V$ tensors along sequence axis dim=2; Query is never cached.
decode_benchmark.py	Ch 2	Quantifies speedup of KV caching vs. $O(N^2)$ full prompt recomputation.
kv_memory_calculator.py	Ch 3	Calculates per-token KV bytes and max GPU batch concurrency under memory limits.
grouped_query_attention.py	Ch 4	Demonstrates GQA head-sharing ($4\times$ to $8\times$ memory footprint reduction).
sliding_window_cache.py	Ch 4	Implements StreamingLLM rolling cache with pinned initial attention sinks.
fragmentation_sim.py	Ch 5	Measures $90%+$ memory waste in contiguous allocation vs. $5%$ in paged allocation.
hbm_bandwidth.cu	Ch 5	Measures effective GPU memory bandwidth (GB/s) that bounds decode speed.
paged_allocator.py	Ch 6	Implements vLLM-style virtual block table with Copy-on-Write reference counting.
batching_sim.py	Ch 6	Compares static batching vs. continuous iteration-level scheduling.
paged_attention_gather.cu	Ch 6	Implements one-indirection logical-to-physical block table lookup in CUDA.
radix_prefix_cache.py	Ch 7	Implements SGLang RadixTree token-level prefix sharing and measures cache hit rate.
kv_quantize.py	Ch 8	Benchmarks FP16 vs. INT8/INT4 per-token vs. per-tensor quantization error.
int4_kv_pack.cu	Ch 8	Packs two 4-bit signed values into 1 byte to double effective HBM throughput.
kv_eviction.py	Ch 9	Compares Heavy-Hitter Oracle ($H_2O$) attention-guided eviction vs. recency window.
tiered_kv_cache.py	Ch 10	Simulates HBM $\to$ CPU DRAM offloading and prefetching to eliminate PCIe stalls.
kv_prefetch.cu	Ch 10	Uses CUDA streams and cudaHostAlloc pinned memory for asynchronous DMA copy.
kv_transport_cost.py	Ch 11	Models KV transfer latency over PCIe Gen5, NVLink (900 GB/s), and InfiniBand RDMA.
kv_rdma_transfer.cu	Ch 11	Registers GPU device memory pointers for GPUDirect RDMA zero-copy transfers.
speculative_decode.py	Ch 12	Models speculative tokens per target forward pass as a function of draft acceptance rate.
app_cache.py	Ch 13	Implements exact SHA256 prompt hash matching and semantic vector distance search.
serving_benchmark.py	Ch 14	Measures percentiles ($P_{50}, P_{95}, P_{99}$) for TTFT, TPOT, ITL, and SLO Goodput.
engine_selector.py	Ch 15	Programmatically maps workload requirements to vLLM, SGLang, TensorRT-LLM, llama.cpp, or Dynamo.

---

📜 License

This repository is published by AI Engineering Insider for educational and reference purposes as part of the Cache Management in Large Language Models ebook series.

llm-cache-management-code