mini-vllm: Implementation of PagedAttention-style KV Cache Management Based on NanoGPT

Abstract: A minimal LLM inference engine that implements a PagedAttention-style KV cache management mechanism on NanoGPT, significantly improving memory utilization efficiency and inference speed. Keywords: LLM, PagedAttention, KV Cache, Inference Optimization, NanoGPT, Memory Management, vLLM

This post will detail the background, core technologies, architectural design, performance, and future plans of the mini-vllm project, helping everyone understand the implementation and value of PagedAttention-style KV cache optimization.

In large language model (LLM) inference, traditional Transformer implementations need to recalculate the KV tensors of all historical tokens every time a new token is generated, leading to a quadratic increase in complexity with sequence length. Although KV cache technology solves this problem, traditional implementations use continuous memory space and pre-reserve memory for the maximum length, resulting in severe memory fragmentation issues, which become a bottleneck for inference tasks.

PagedAttention draws on the virtual memory paging mechanism of operating systems, storing KV tensors in dynamically allocated non-continuous memory blocks that can be shared between different requests, greatly reducing memory waste. The original vLLM paper shows that compared to continuous memory allocation, PagedAttention achieves a 2-4x improvement in memory efficiency, which is of great significance for deploying large models in resource-constrained environments.

mini-vllm adopts a modular design, with core components including:

1. BlockAllocator: Maintains the state of memory blocks, allocates and releases memory for the pre-filling/decoding phases, similar to operating system memory management strategies.
2. KVCache: Uses the block table of BlockAllocator to determine the read/write positions of KV tensors in GPU tensors, enabling non-continuous storage and logical continuity.
3. InferenceEngine: A coordinator that receives prompts, runs pre-filling and decoding loops, manages memory block allocation and release, and returns results.

Process: prompt → InferenceEngine → BlockAllocator (manages physical blocks) + KVCache (stores/retrieves KV tensors by block).

The project provides benchmark.py to evaluate performance, comparing two generation methods:

• KV cache optimized version: Uses InferenceEngine, custom KVCache/BlockAllocator, and only calculates new tokens during the decoding phase.
• Baseline version without KV cache: Standard NanoGPT, which performs a full forward pass every time a new token is generated, leading to a sharp increase in cost as the sequence grows.

Test metrics: Average generation time, TPS, speedup ratio; supports configuration of prompt length, number of tokens, and memory parameters; results are saved in CSV (benchmark_results.csv), and benchmark_results_viz.ipynb is provided for visualization (runtime-length relationship, speedup ratio comparison, throughput comparison).

Tests in a CPU environment show that the KV cache implementation provides a 2-10x speedup for medium sequence lengths (200-500 tokens), and latency is significantly reduced as the sequence grows. Even in resource-constrained environments, PagedAttention-style memory management brings substantial performance improvements.

Future Plans:

1. Rewrite BlockAllocator in C++ to improve memory allocation performance;
2. Prefix-sharing to support block sharing between different requests, further improving memory efficiency.

Related Resources:

• PagedAttention paper: https://arxiv.org/abs/2309.06180
• vLLM GitHub: https://github.com/vllm-project/vllm
• NanoGPT project: https://github.com/karpathy/nanoGPT

mini-vllm provides a concise and complete reference implementation for understanding and implementing PagedAttention technology. Through its modular architecture and detailed benchmark tests, developers can deeply understand the principles of KV cache optimization and apply them to their own LLM inference systems. It is an extremely valuable learning resource for developers deploying large models efficiently in resource-constrained environments.