Why Naive KV-Caches Waste >60% VRAM (And How Paged Memory Fixes It)

The Hidden Bottleneck in LLM Serving

When serving Large Language Models (LLMs), the primary memory bottleneck during multi-turn generation is rarely model weights—it is the Key-Value (KV) Cache.

In standard PyTorch transformer implementations (like Hugging Face generate()), memory for the KV-cache is allocated contiguously up to max_model_len (e.g., 2048 or 4096 tokens) per request:

[Request 1: Token 1..50] [-------- Contiguous Pre-allocated Waste (1998 slots) --------]
[Request 2: Token 1..120] [------- Contiguous Pre-allocated Waste (1928 slots) --------]

The Problem: Over 60% VRAM Waste

Because user prompt lengths and output generation counts are unknown ahead of time, contiguous allocators must reserve worst-case buffer pools. This causes two severe types of memory waste:
1. Internal Fragmentation: Reserved slots in the sequence buffer that are never filled by the user request.
2. External Fragmentation: Virtual memory address gaps between concurrent requests that cannot fit a new contiguous allocation, causing Out-Of-Memory (OOM) errors even when total free VRAM is plentiful.

On consumer laptop hardware (such as an NVIDIA RTX 3050 6GB Laptop GPU), contiguous pre-allocation caps your maximum concurrency to just 1 or 2 concurrent requests before crashing with CUDA OOM.

The Solution: Paged KV Memory Allocation in Pure PyTorch

Inspired by virtual memory paging in operating systems (and popularized by vLLM's CUDA kernels), mini-infer brings hardware-efficient Paged KV caching to pure PyTorch without requiring a C++ compiler or CUDA toolkit toolchain.

1. Physical Block Pool & Block Tables

Instead of allocating one giant tensor per sequence, mini-infer manages memory as a fixed pool of 16-token physical blocks:

Physical Block Pool (in VRAM):
┌──────────┬──────────┬──────────┬──────────┬──────────┐
│ Block 0  │ Block 1  │ Block 2  │ Block 3  │ Block 4  │ ...
│ (16 tok) │ (16 tok) │ (16 tok) │ (16 tok) │ (16 tok) │
└──────────┴──────────┴──────────┴──────────┴──────────┘

Each sequence maintains a Block Table that maps virtual token offsets to non-contiguous physical blocks:

# Virtual to Physical translation
# Sequence 1: 35 tokens -> requires ceil(35 / 16) = 3 blocks
sequence_block_table = [4, 12, 19]

# Physical allocation:
# Tokens 0..15  -> Physical Block 4
# Tokens 16..31 -> Physical Block 12
# Tokens 32..34 -> Physical Block 19 (only 3 slots used, 13 slots free)

2. Eliminating Memory Waste to Below 3.8%

Because memory is requested on-demand as tokens generate:
* Zero External Fragmentation: Physical blocks can be anywhere in VRAM.
* Minimal Internal Fragmentation: Waste is strictly limited to the tail block of an active sequence (at most $16 - 1 = 15$ token slots).
* Across a typical batch of 16 concurrent requests, memory utilization jumps from <40% up to >96.2%, dropping memory waste to <3.8%.

Implementation Walkthrough

In mini-infer, the BlockAllocator manages available block indices using a simple, O(1) free-list stack:

class BlockAllocator:
    def __init__(self, num_blocks: int, block_size: int = 16):
        self.block_size = block_size
        self.num_blocks = num_blocks
        self.free_blocks: list[int] = list(reversed(range(num_blocks)))

    def allocate(self) -> int:
        if not self.free_blocks:
            raise MemoryError("Out of physical KV blocks!")
        return self.free_blocks.pop()

    def free(self, block_id: int):
        self.free_blocks.append(block_id)

During attention computation, the KV cache tensor is sized strictly to the physical pool:

$$\text{Shape: } [N_{\text{blocks}}, \text{block_size}, N_{\text{heads}}, D_{\text{head}}]$$

Attention is gathered dynamically using the sequence's block table via PyTorch's native Scaled Dot-Product Attention (F.scaled_dot_product_attention), ensuring high CUDA kernel performance without custom C++ extensions.

Measured Results on Consumer Hardware

We benchmarked mini-infer against real hardware constraints on an RTX 3050 6GB Laptop GPU (maintaining a < 4GB total VRAM footprint) using servebench:

Try It Locally

You can install mini-infer directly from PyPI:

pip install mini-infer

Run a local OpenAI-compatible continuous-batching server:

mini-infer serve --model Qwen/Qwen2.5-0.5B-Instruct --port 8000

Read the full architecture docs and block table math at prajwal7-2002.github.io/mini-infer.