Memory Pool Management: Slab Allocators for GPU Inference

GPU memory allocation in LLM inference is a systems problem that most ML engineers encounter only when things break. A serving system allocates and frees memory thousands of times per second: KV cache blocks are created as requests arrive and released as requests complete, activation buffers resize with batch composition changes, and temporary buffers for attention computation are allocated and freed every forward pass. Under a naive allocator, this churn creates memory fragmentation that can render 30% or more of GPU memory unusable despite appearing free. The system reports 20 GB available, but no contiguous block larger than 512 MB exists, and the next KV cache allocation needs 2 GB.

Slab allocators solve this by pre-allocating all GPU memory upfront and dividing it into fixed-size blocks. Allocation becomes a pop from a free list. Deallocation becomes a push to a free list. Both operations are $O(1)$. Fragmentation is zero by construction because every block is the same size.

This post covers the fragmentation problem in detail, derives the slab allocator from first principles, extends it to multi-slab architectures, examines PyTorch’s built-in CUDA allocator, documents vLLM’s approach, and provides a complete, production-ready slab allocator implementation.

---

1. The Fragmentation Problem

1.1 How GPU Memory Allocation Works

CUDA provides two levels of memory allocation:

cudaMalloc / cudaFree: The driver-level allocator. Allocates contiguous virtual memory backed by physical GPU memory. Each call is a kernel-level operation that can take 100us-10ms depending on allocation size and memory pressure. This is far too slow for per-request allocation in an inference server handling thousands of requests per second.

cudaMallocAsync / Stream-ordered allocation (CUDA 11.2+): A pool-based allocator with stream-ordered semantics. Faster than raw cudaMalloc but still not designed for the allocation patterns of LLM inference.

The fundamental issue: cudaMalloc returns a pointer to a contiguous block of the requested size. After many allocations and frees of varying sizes, the free memory becomes fragmented into non-contiguous chunks.

1.2 Fragmentation Under Realistic Workloads

Consider an inference server handling variable-length requests. Each request needs a KV cache whose size depends on the sequence length. Requests arrive and complete at different times.

import random

def simulate_fragmentation(gpu_memory_mb, num_operations, max_alloc_mb=256):
    """
    Simulate GPU memory fragmentation under random alloc/free patterns.
    Returns fragmentation metrics.
    """
    allocations = {}  # id -> (offset, size)
    free_regions = [(0, gpu_memory_mb)]  # List of (offset, size) free regions
    next_id = 0
    active_ids = []
    peak_fragmentation = 0.0

    for op in range(num_operations):
        if len(active_ids) > 0 and random.random() < 0.4:
            # Free a random allocation
            free_id = random.choice(active_ids)
            active_ids.remove(free_id)
            offset, size = allocations.pop(free_id)
            _merge_free_region(free_regions, offset, size)
        else:
            # Allocate a random size
            size = random.randint(1, max_alloc_mb)
            offset = _find_free_block(free_regions, size)
            if offset is not None:
                allocations[next_id] = (offset, size)
                active_ids.append(next_id)
                next_id += 1

        # Measure fragmentation
        total_free = sum(s for _, s in free_regions)
        largest_free = max((s for _, s in free_regions), default=0)
        if total_free > 0:
            frag = 1.0 - (largest_free / total_free)
            peak_fragmentation = max(peak_fragmentation, frag)

    return peak_fragmentation

def _find_free_block(free_regions, size):
    """First-fit allocation."""
    for i, (offset, region_size) in enumerate(free_regions):
        if region_size >= size:
            free_regions.pop(i)
            if region_size > size:
                free_regions.append((offset + size, region_size - size))
                free_regions.sort()
            return offset
    return None  # Out of memory (fragmented)

def _merge_free_region(free_regions, offset, size):
    """Merge a freed block with adjacent free regions."""
    free_regions.append((offset, size))
    free_regions.sort()
    merged = [free_regions[0]]
    for curr_offset, curr_size in free_regions[1:]:
        prev_offset, prev_size = merged[-1]
        if curr_offset == prev_offset + prev_size:
            merged[-1] = (prev_offset, prev_size + curr_size)
        else:
            merged.append((curr_offset, curr_size))
    free_regions.clear()
    free_regions.extend(merged)

Running this simulation with realistic parameters:

# 80 GB GPU, 1000 alloc/free cycles, allocations up to 256 MB
fragmentation = simulate_fragmentation(
    gpu_memory_mb=81920,
    num_operations=1000,
    max_alloc_mb=256
)
print(f"Peak fragmentation: {fragmentation:.1%}")
# Typical output: Peak fragmentation: 34.2%

After 1000 operations, 34% of free memory is fragmented — 27.9 GB free but the largest contiguous block might be only 2 GB. A request needing 4 GB of KV cache would fail even though total free memory is abundant. This is external fragmentation.

1.3 Types of Fragmentation

External fragmentation: Free memory exists but is scattered in non-contiguous regions too small to satisfy allocation requests. This is the dominant issue.

Internal fragmentation: Allocated blocks are larger than needed because the allocator rounds up to a fixed size. A 100 MB request gets a 128 MB block, wasting 28 MB. This is a tradeoff: more internal fragmentation means less external fragmentation.

The slab allocator eliminates external fragmentation entirely at the cost of controlled internal fragmentation.

---

2. The Slab Allocator: Design from First Principles

The slab allocator was introduced by Jeff Bonwick for the SunOS kernel in 1994. The core insight: if all allocations are the same size, fragmentation is impossible. Every freed block exactly fits the next allocation.

2.1 Data Structure

A slab is a contiguous region of memory divided into equal-sized blocks:

Slab (1 GB total, block_size = 2 MB, 512 blocks):

+--------+--------+--------+--------+-----+--------+
| Block 0| Block 1| Block 2| Block 3| ... |Block511|
+--------+--------+--------+--------+-----+--------+
  2 MB     2 MB     2 MB     2 MB           2 MB

Free list: [5, 12, 0, 3, ...]  (indices of free blocks)

The free list is a stack (LIFO). Allocation pops the top. Deallocation pushes to the top.

class SlabAllocator:
    """
    Fixed-size block allocator with O(1) alloc and free.
    Zero external fragmentation.
    """

    def __init__(self, total_size_bytes, block_size_bytes, base_ptr=0):
        """
        total_size_bytes: Total memory pool size.
        block_size_bytes: Size of each block.
        base_ptr: Base address of the memory pool (e.g., from cudaMalloc).
        """
        self.block_size = block_size_bytes
        self.num_blocks = total_size_bytes // block_size_bytes
        self.base_ptr = base_ptr

        # Free list: stack of available block indices
        self.free_stack = list(range(self.num_blocks))
        self.allocated = set()

        # Statistics
        self.total_allocs = 0
        self.total_frees = 0

    def alloc(self):
        """
        Allocate one block. Returns (ptr, block_id) or raises MemoryError.
        Time complexity: O(1).
        """
        if not self.free_stack:
            raise MemoryError(
                f"Slab exhausted: {self.num_blocks} blocks allocated, "
                f"0 free. Block size: {self.block_size} bytes."
            )

        block_id = self.free_stack.pop()
        self.allocated.add(block_id)
        self.total_allocs += 1

        ptr = self.base_ptr + block_id * self.block_size
        return ptr, block_id

    def free(self, block_id):
        """
        Free a block by its ID. Time complexity: O(1).
        """
        if block_id not in self.allocated:
            raise ValueError(f"Block {block_id} is not allocated (double free?)")

        self.allocated.remove(block_id)
        self.free_stack.append(block_id)
        self.total_frees += 1

    @property
    def num_free(self):
        return len(self.free_stack)

    @property
    def num_allocated(self):
        return len(self.allocated)

    @property
    def utilization(self):
        return self.num_allocated / self.num_blocks

    def __repr__(self):
        return (
            f"SlabAllocator(block_size={self.block_size}, "
            f"blocks={self.num_blocks}, "
            f"allocated={self.num_allocated}, "
            f"free={self.num_free}, "
            f"utilization={self.utilization:.1%})"
        )

2.2 Performance Characteristics

Both alloc() and free() are $O(1)$: a single list.pop() or list.append(). Compare with a general-purpose allocator:

2.3 Why LIFO Ordering Matters

The free stack uses LIFO (last-in, first-out) ordering. The most recently freed block is reallocated first. This is not arbitrary — it has a concrete benefit for cache performance:

1. A recently freed block is likely still in the GPU’s L2 cache.
2. Reallocating it means the next write hits warm cache lines.
3. FIFO ordering would allocate blocks that were freed longest ago, likely evicted from cache.

For GPU inference, L2 cache effects are significant. The H100’s 50 MB L2 cache can hold roughly 25 blocks of 2 MB each. With LIFO ordering, reallocated blocks have a high probability of L2 residency.

---

3. Multi-Slab Architecture

Real inference workloads need blocks of different sizes. KV cache blocks might be 2 MB each. Activation buffers might need 64 MB. Temporary attention score buffers might be 512 KB. A single block size forces everything to the largest required size, wasting memory (excessive internal fragmentation).

The solution: multiple slabs, each with a different block size.

class MultiSlabAllocator:
    """
    Multi-slab allocator with size classes.
    Requests are routed to the smallest slab that fits.
    """

    def __init__(self, total_memory_bytes, slab_configs):
        """
        slab_configs: list of (block_size, fraction_of_memory).
        Example: [(2*MB, 0.7), (64*MB, 0.2), (512*KB, 0.1)]
        """
        self.slabs = {}
        self.size_classes = []
        current_offset = 0

        for block_size, fraction in sorted(slab_configs, key=lambda x: x[0]):
            slab_memory = int(total_memory_bytes * fraction)
            slab = SlabAllocator(
                total_size_bytes=slab_memory,
                block_size_bytes=block_size,
                base_ptr=current_offset
            )
            self.slabs[block_size] = slab
            self.size_classes.append(block_size)
            current_offset += slab_memory

        # Sort size classes for binary search during allocation
        self.size_classes.sort()

    def alloc(self, requested_size):
        """
        Allocate a block of at least requested_size bytes.
        Routes to the smallest slab that fits.
        Returns: (ptr, block_id, block_size) or raises MemoryError.
        """
        for size_class in self.size_classes:
            if size_class >= requested_size:
                slab = self.slabs[size_class]
                try:
                    ptr, block_id = slab.alloc()
                    return ptr, block_id, size_class
                except MemoryError:
                    continue  # This slab is full, try next size class

        raise MemoryError(
            f"Cannot allocate {requested_size} bytes. "
            f"Status: {self.status()}"
        )

    def free(self, block_id, block_size):
        """Free a block, routing to the correct slab by block_size."""
        if block_size not in self.slabs:
            raise ValueError(f"Unknown block size: {block_size}")
        self.slabs[block_size].free(block_id)

    def status(self):
        """Return allocation status for all slabs."""
        return {
            size: {
                'allocated': slab.num_allocated,
                'free': slab.num_free,
                'total': slab.num_blocks,
                'utilization': f"{slab.utilization:.1%}"
            }
            for size, slab in self.slabs.items()
        }

Sizing the Slabs

The critical design decision: how much memory to assign to each slab. This depends on the workload.

For an LLM inference server:

MB = 1024 * 1024
GB = 1024 * MB

# H100 80 GB allocation plan
slab_configs = [
    # KV cache blocks: 2 MB each, 70% of memory (56 GB = 28,672 blocks)
    (2 * MB, 0.70),
    # Activation buffers: 32 MB each, 15% of memory (12 GB = 384 blocks)
    (32 * MB, 0.15),
    # Small buffers (logits, temp): 256 KB each, 10% of memory (8 GB = 32,768 blocks)
    (256 * 1024, 0.10),
    # Large buffers (weight loading, offload): 256 MB each, 5% of memory (4 GB = 16 blocks)
    (256 * MB, 0.05),
]

allocator = MultiSlabAllocator(80 * GB, slab_configs)

Internal Fragmentation Analysis

For a slab with block size $B$, a request of size $s$ where $s \leq B$ wastes $B - s$ bytes. The expected waste for uniformly distributed request sizes in $[1, B]$:

$$E[\text{waste}] = E[B - s] = B - \frac{B + 1}{2} \approx \frac{B}{2}$$

Expected internal fragmentation rate:

$$\text{frag}_\text{internal} = \frac{E[\text{waste}]}{B} \approx \frac{1}{2} = 50\%$$

This is worst-case for uniform distribution. In practice, KV cache allocations cluster around multiples of the block size because sequences are padded to block boundaries. Typical internal fragmentation is 5-15%.

---

4. PyTorch’s CUDACachingAllocator

PyTorch does not call cudaMalloc for every tensor allocation. It uses CUDACachingAllocator, a custom memory pool that caches freed blocks for reuse.

4.1 Architecture

The allocator maintains two pools:

• Small pool: for allocations up to 1 MB (default threshold). Uses 512-byte and 1 MB block granularity.
• Large pool: for allocations above 1 MB. Uses 2 MB block granularity.

Each pool is a set of memory segments obtained from cudaMalloc. Within each segment, blocks are managed as a doubly-linked list with splitting and coalescing.

# Pseudocode of PyTorch's CUDACachingAllocator logic
class CUDACachingAllocator:
    def __init__(self):
        self.small_pool = BlockPool(block_size_threshold=1 * MB)
        self.large_pool = BlockPool(block_size_threshold=float('inf'))

    def malloc(self, size):
        pool = self.small_pool if size <= 1 * MB else self.large_pool

        # 1. Try to find a free block in the pool
        block = pool.find_free_block(size)
        if block:
            if block.size > size + MIN_SPLIT_SIZE:
                # Split the block: use the first part, return the rest
                remainder = block.split(size)
                pool.add_free_block(remainder)
            block.allocated = True
            return block.ptr

        # 2. No suitable free block: allocate from CUDA
        try:
            segment = cuda_malloc(round_up(size, 2 * MB))
        except CudaOOMError:
            # 3. Try garbage collection first
            self.garbage_collect()
            segment = cuda_malloc(round_up(size, 2 * MB))

        pool.add_segment(segment)
        block = segment.create_block(size)
        return block.ptr

    def free(self, ptr):
        block = self.find_block(ptr)
        block.allocated = False

        # Coalesce with adjacent free blocks
        if block.prev and not block.prev.allocated:
            block = self.merge(block.prev, block)
        if block.next and not block.next.allocated:
            block = self.merge(block, block.next)

        block.pool.add_free_block(block)

    def garbage_collect(self):
        """Release all cached blocks back to CUDA driver."""
        for pool in [self.small_pool, self.large_pool]:
            for segment in pool.segments:
                if all(not b.allocated for b in segment.blocks):
                    cuda_free(segment.ptr)
                    pool.remove_segment(segment)

4.2 Block Splitting

When the allocator finds a free block of 8 MB but only needs 2 MB, it splits the block:

Before split:
[  Free block: 8 MB  ]

After split (allocating 2 MB):
[ Allocated: 2 MB ][ Free: 6 MB ]

Splitting creates smaller blocks that may not be reusable for future larger allocations. Over time, this leads to fragmentation.

4.3 Block Coalescing

When a block is freed, the allocator checks its neighbors. If either neighbor is free, they are merged:

Before coalesce (freeing middle block):
[ Free: 2 MB ][ Freeing: 3 MB ][ Free: 1 MB ]

After coalesce:
[         Free: 6 MB              ]

Coalescing helps but cannot eliminate fragmentation because it only merges adjacent blocks. If a 1 MB allocated block sits between two free regions, those regions cannot be merged even though they are collectively large.

4.4 Garbage Collection

When cudaMalloc fails, PyTorch triggers garbage collection: all cached free blocks are released back to the CUDA driver via cudaFree. This is expensive (milliseconds) and disrupts serving latency. Production systems try to avoid triggering GC by reserving memory upfront.

# PyTorch memory management controls
import torch

# Set maximum split size to reduce fragmentation
# (prevents splitting blocks larger than this threshold)
torch.cuda.memory.set_per_process_memory_fraction(0.9)  # Use 90% of GPU

# Aggressive GC when running low
torch.cuda.memory.set_allocator_settings(
    "expandable_segments:True,"
    "max_split_size_mb:512"  # Don't split blocks larger than 512 MB
)

# Manual GC
torch.cuda.empty_cache()  # Release cached blocks to CUDA driver

# Memory snapshot for debugging
torch.cuda.memory._record_memory_history()
# ... run workload ...
snapshot = torch.cuda.memory._snapshot()

---

5. vLLM’s Memory Management Architecture

vLLM takes a radically different approach from PyTorch’s general-purpose allocator. It pre-allocates virtually all GPU memory at startup and manages it with purpose-built allocators for each memory type.

5.1 Two-Pool Architecture

vLLM divides GPU memory into two regions:

1. KV cache pool: Managed by a dedicated slab allocator. Fixed block size (typically 16 tokens per block). This pool gets the majority of GPU memory (60-85% depending on model size).

2. Everything else: Model weights (static, allocated once), activation buffers (managed by PyTorch’s allocator within a fixed reservation).

class VLLMMemoryManager:
    """
    Simplified model of vLLM's two-pool memory architecture.
    """

    def __init__(self, total_gpu_memory, model_weight_size, kv_block_size_bytes):
        self.total_memory = total_gpu_memory

        # Static allocations (model weights + activation buffer reservation)
        self.weight_memory = model_weight_size
        activation_reserve = total_gpu_memory * 0.05  # 5% for activations
        self.activation_memory = activation_reserve

        # KV cache pool gets everything else
        kv_pool_size = total_gpu_memory - model_weight_size - activation_reserve
        self.kv_allocator = SlabAllocator(
            total_size_bytes=int(kv_pool_size),
            block_size_bytes=kv_block_size_bytes,
        )

        print(f"Total GPU memory: {total_gpu_memory / GB:.1f} GB")
        print(f"Model weights: {model_weight_size / GB:.1f} GB")
        print(f"Activation reserve: {activation_reserve / GB:.1f} GB")
        print(f"KV cache pool: {kv_pool_size / GB:.1f} GB")
        print(f"KV blocks available: {self.kv_allocator.num_blocks}")

    def allocate_kv_blocks(self, num_blocks):
        """Allocate num_blocks KV cache blocks for a request."""
        blocks = []
        for _ in range(num_blocks):
            ptr, block_id = self.kv_allocator.alloc()
            blocks.append(block_id)
        return blocks

    def free_kv_blocks(self, block_ids):
        """Free KV cache blocks when a request completes."""
        for block_id in block_ids:
            self.kv_allocator.free(block_id)

5.2 KV Cache Block Size

vLLM’s KV cache block stores $B_t$ tokens worth of K and V tensors for a single layer, for all KV heads:

$$\text{block\_size} = 2 \times n_\text{kv\_heads} \times d_\text{head} \times B_t \times d_\text{type}$$

For Llama 70B GQA with $n_\text{kv\_heads} = 8$, $d_\text{head} = 128$, $B_t = 16$ tokens, BF16:

$$\text{block\_size} = 2 \times 8 \times 128 \times 16 \times 2 = 65{,}536 \text{ bytes} = 64 \text{ KB per layer}$$

Across all 80 layers:

$$64 \text{ KB} \times 80 = 5{,}120 \text{ KB} = 5 \text{ MB per block}$$

On an 80 GB H100 with 28 GB for model weights (Llama 70B in INT4):

$$\text{KV pool} = 80 - 28 - 4 \text{ (activations)} = 48 \text{ GB}$$ $$\text{blocks} = \frac{48 \times 1024}{5} = 9{,}830 \text{ blocks}$$ $$\text{total tokens} = 9{,}830 \times 16 = 157{,}280 \text{ tokens}$$

At 2K tokens per request: $157{,}280 / 2{,}048 \approx 76$ concurrent requests.

The total token capacity is the same regardless of block size (it is determined by total pool memory divided by per-token KV size). The block size affects internal fragmentation: with 16 tokens per block, the last block of each request wastes $16/2 = 8$ token slots on average. With 64 tokens per block, the waste is $64/2 = 32$ token slots. Smaller blocks reduce internal fragmentation but increase the number of blocks the scheduler must track.

5.3 PagedAttention: Why Fixed-Size Blocks Enable Efficient Attention

vLLM’s slab allocator is not just a memory optimization — it is tightly coupled with the PagedAttention kernel. PagedAttention computes attention over non-contiguous KV cache blocks by using a block table that maps logical sequence positions to physical block IDs:

class PagedKVCache:
    """
    Paged KV cache with block table for non-contiguous storage.
    """

    def __init__(self, slab_allocator, num_layers, kv_heads, head_dim, tokens_per_block):
        self.slab = slab_allocator
        self.num_layers = num_layers
        self.kv_heads = kv_heads
        self.head_dim = head_dim
        self.tokens_per_block = tokens_per_block

    def allocate_sequence(self, seq_len):
        """
        Allocate blocks for a sequence. Returns block table.
        """
        num_blocks = (seq_len + self.tokens_per_block - 1) // self.tokens_per_block
        block_table = []
        for _ in range(num_blocks):
            _, block_id = self.slab.alloc()
            block_table.append(block_id)
        return block_table

    def extend_sequence(self, block_table, new_tokens):
        """
        Extend a sequence by new_tokens. Allocate additional blocks if needed.
        """
        current_tokens = len(block_table) * self.tokens_per_block
        # Check if current last block has space
        used_in_last = current_tokens % self.tokens_per_block
        remaining_in_last = self.tokens_per_block - used_in_last if used_in_last > 0 else 0

        extra_tokens = new_tokens - remaining_in_last
        if extra_tokens > 0:
            extra_blocks = (extra_tokens + self.tokens_per_block - 1) // self.tokens_per_block
            for _ in range(extra_blocks):
                _, block_id = self.slab.alloc()
                block_table.append(block_id)

        return block_table

    def free_sequence(self, block_table):
        """Free all blocks for a completed sequence."""
        for block_id in block_table:
            self.slab.free(block_id)

The block table indirection means sequences do not need contiguous memory. A sequence’s KV cache can be scattered across the physical memory pool. This is why the slab allocator achieves zero fragmentation: blocks are interchangeable, and non-contiguity is handled by the attention kernel.

---

6. Complete Production Slab Allocator

Here is a complete slab allocator implementation with GPU integration, reference counting (for copy-on-write), and monitoring:

import torch
from dataclasses import dataclass, field
from collections import defaultdict
from typing import Optional

@dataclass
class BlockMetadata:
    """Metadata for a single block in the slab."""
    block_id: int
    ref_count: int = 0        # Reference count for copy-on-write
    last_access: int = 0      # Timestamp for LRU eviction
    sequence_id: Optional[int] = None  # Owning sequence (None if free)

class GPUSlabAllocator:
    """
    Production GPU slab allocator with:
    - O(1) alloc/free
    - Reference counting for copy-on-write
    - Memory-mapped GPU tensors
    - Utilization monitoring
    """

    def __init__(
        self,
        num_blocks,
        block_shape,
        dtype=torch.float16,
        device='cuda',
    ):
        """
        num_blocks: Number of blocks in the slab.
        block_shape: Shape of each block, e.g., (2, kv_heads, tokens_per_block, head_dim)
                     where 2 is for K and V.
        """
        self.num_blocks = num_blocks
        self.block_shape = block_shape
        self.dtype = dtype
        self.device = device

        # Allocate the entire slab as a single contiguous GPU tensor
        full_shape = (num_blocks, *block_shape)
        self.gpu_memory = torch.zeros(full_shape, dtype=dtype, device=device)

        # Block metadata
        self.metadata = {
            i: BlockMetadata(block_id=i) for i in range(num_blocks)
        }

        # Free list (LIFO stack)
        self.free_stack = list(range(num_blocks - 1, -1, -1))

        # Statistics
        self.stats = {
            'total_allocs': 0,
            'total_frees': 0,
            'peak_utilization': 0.0,
            'oom_events': 0,
            'cow_copies': 0,
        }

        # Global timestamp for LRU
        self.timestamp = 0

    def alloc(self, sequence_id=None):
        """
        Allocate a single block. Returns block_id.
        Raises MemoryError if no blocks available.
        """
        if not self.free_stack:
            self.stats['oom_events'] += 1
            raise MemoryError(
                f"Slab OOM: {self.num_blocks} blocks allocated, 0 free. "
                f"Total allocs: {self.stats['total_allocs']}, "
                f"Total frees: {self.stats['total_frees']}"
            )

        block_id = self.free_stack.pop()
        meta = self.metadata[block_id]
        meta.ref_count = 1
        meta.sequence_id = sequence_id
        meta.last_access = self.timestamp
        self.timestamp += 1

        self.stats['total_allocs'] += 1
        utilization = 1.0 - len(self.free_stack) / self.num_blocks
        self.stats['peak_utilization'] = max(
            self.stats['peak_utilization'], utilization
        )

        return block_id

    def free(self, block_id):
        """
        Decrement reference count. Actually free if ref_count reaches 0.
        """
        meta = self.metadata[block_id]
        if meta.ref_count <= 0:
            raise ValueError(f"Block {block_id}: ref_count is {meta.ref_count} (double free?)")

        meta.ref_count -= 1
        if meta.ref_count == 0:
            meta.sequence_id = None
            self.free_stack.append(block_id)
            self.stats['total_frees'] += 1

    def add_ref(self, block_id):
        """Increment reference count (for copy-on-write sharing)."""
        meta = self.metadata[block_id]
        if meta.ref_count <= 0:
            raise ValueError(f"Block {block_id}: cannot add ref to free block")
        meta.ref_count += 1

    def copy_on_write(self, block_id):
        """
        If block has ref_count > 1, allocate a new block, copy data,
        decrement old ref_count, return new block_id.
        If ref_count == 1, return same block_id (no copy needed).
        """
        meta = self.metadata[block_id]
        if meta.ref_count == 1:
            return block_id  # Exclusive owner, no copy needed

        # Need to copy
        new_block_id = self.alloc(sequence_id=meta.sequence_id)
        self.gpu_memory[new_block_id].copy_(self.gpu_memory[block_id])
        self.free(block_id)  # Decrement ref on old block
        self.stats['cow_copies'] += 1
        return new_block_id

    def get_block_tensor(self, block_id):
        """
        Get a tensor view of a specific block.
        Shape: block_shape (e.g., (2, kv_heads, tokens_per_block, head_dim))
        """
        meta = self.metadata[block_id]
        meta.last_access = self.timestamp
        self.timestamp += 1
        return self.gpu_memory[block_id]

    def get_blocks_tensor(self, block_ids):
        """
        Get a batched tensor view of multiple blocks.
        Shape: (len(block_ids), *block_shape)
        Uses advanced indexing (creates a copy, not a view).
        """
        indices = torch.tensor(block_ids, dtype=torch.long, device=self.device)
        return self.gpu_memory[indices]

    @property
    def num_free(self):
        return len(self.free_stack)

    @property
    def num_allocated(self):
        return self.num_blocks - len(self.free_stack)

    @property
    def utilization(self):
        return self.num_allocated / self.num_blocks

    @property
    def memory_bytes(self):
        """Total GPU memory used by this slab."""
        return self.gpu_memory.nelement() * self.gpu_memory.element_size()

    def report(self):
        """Print allocation report."""
        print(f"=== Slab Allocator Report ===")
        print(f"Total blocks:      {self.num_blocks}")
        print(f"Allocated:         {self.num_allocated}")
        print(f"Free:              {self.num_free}")
        print(f"Utilization:       {self.utilization:.1%}")
        print(f"Peak utilization:  {self.stats['peak_utilization']:.1%}")
        print(f"Total allocs:      {self.stats['total_allocs']}")
        print(f"Total frees:       {self.stats['total_frees']}")
        print(f"OOM events:        {self.stats['oom_events']}")
        print(f"CoW copies:        {self.stats['cow_copies']}")
        print(f"GPU memory:        {self.memory_bytes / (1024**3):.2f} GB")

Usage Example: KV Cache for Llama 70B

# Llama 70B GQA: 8 KV heads, 128 head_dim, 16 tokens per block
# Per-layer block shape: (2, 8, 16, 128) = K and V
# 2 * 8 * 16 * 128 * 2 bytes (BF16) = 65,536 bytes = 64 KB per layer per block

# For all 80 layers, we create one slab per layer (simpler)
# or one slab for all layers (more efficient, block shape includes layers)

# Option 1: Single slab, blocks store all layers
# Block shape: (num_layers, 2, kv_heads, tokens_per_block, head_dim)
# = (80, 2, 8, 16, 128) = 5 MB per block

kv_slab = GPUSlabAllocator(
    num_blocks=9830,  # ~48 GB / 5 MB per block
    block_shape=(80, 2, 8, 16, 128),
    dtype=torch.bfloat16,
    device='cuda',
)

# Allocate blocks for a request with 1000 tokens
# Need ceil(1000 / 16) = 63 blocks
block_table = []
for _ in range(63):
    block_id = kv_slab.alloc(sequence_id=42)
    block_table.append(block_id)

print(f"Allocated {len(block_table)} blocks for sequence 42")
print(f"Slab utilization: {kv_slab.utilization:.2%}")

# During attention, access blocks by index
for layer_idx in range(80):
    for block_id in block_table:
        block_tensor = kv_slab.get_block_tensor(block_id)
        # block_tensor shape: (80, 2, 8, 16, 128)
        k = block_tensor[layer_idx, 0]  # (8, 16, 128)
        v = block_tensor[layer_idx, 1]  # (8, 16, 128)
        # ... use k, v in attention computation ...

# Free when request completes
for block_id in block_table:
    kv_slab.free(block_id)

Copy-on-Write for Prefix Sharing

# Two requests share the same system prompt (500 tokens = 32 blocks)
system_prompt_blocks = []
for _ in range(32):
    block_id = kv_slab.alloc(sequence_id=100)
    system_prompt_blocks.append(block_id)
    # ... fill with KV cache for system prompt ...

# Request A: shares system prompt, then continues
request_a_blocks = list(system_prompt_blocks)  # Copy block table (not data)
for block_id in system_prompt_blocks:
    kv_slab.add_ref(block_id)  # Increment ref count
# Request A's unique continuation
for _ in range(10):
    block_id = kv_slab.alloc(sequence_id=200)
    request_a_blocks.append(block_id)

# Request B: also shares system prompt
request_b_blocks = list(system_prompt_blocks)
for block_id in system_prompt_blocks:
    kv_slab.add_ref(block_id)

# Memory used: 32 shared blocks + 10 unique blocks = 42 blocks
# Without sharing: 32 + 32 + 10 = 74 blocks
# Savings: 43%

# When request A completes:
for block_id in request_a_blocks:
    kv_slab.free(block_id)  # Shared blocks: ref_count 2->1 (not freed)
                             # Unique blocks: ref_count 1->0 (freed)

---

7. Benchmarking: Slab vs General-Purpose Allocation

import time

def benchmark_allocator(allocator_type, num_ops=10000):
    """
    Benchmark allocation throughput.
    """
    if allocator_type == 'slab':
        alloc = SlabAllocator(
            total_size_bytes=10 * GB,
            block_size_bytes=5 * MB,
        )
        allocated = []

        start = time.perf_counter()
        for i in range(num_ops):
            if len(allocated) > 100 and i % 3 == 0:
                # Free a random block
                idx = i % len(allocated)
                alloc.free(allocated[idx])
                allocated[idx] = allocated[-1]
                allocated.pop()
            else:
                try:
                    _, block_id = alloc.alloc()
                    allocated.append(block_id)
                except MemoryError:
                    break
        elapsed = time.perf_counter() - start

    elif allocator_type == 'pytorch':
        allocated = []
        torch.cuda.empty_cache()

        start = time.perf_counter()
        for i in range(num_ops):
            if len(allocated) > 100 and i % 3 == 0:
                idx = i % len(allocated)
                del allocated[idx]
            else:
                t = torch.empty(5 * 1024 * 1024 // 2, dtype=torch.float16, device='cuda')
                allocated.append(t)
        elapsed = time.perf_counter() - start

    return elapsed, num_ops / elapsed

The Python slab allocator is 7x faster than PyTorch’s C++ allocator because it avoids all the splitting, coalescing, and pool management overhead. A C++ slab allocator would be even faster. The raw cudaMalloc path is 267x slower and fragments heavily.

---

8. Advanced: Defragmentation Without Slab Allocators

For systems that cannot use fixed-size blocks (e.g., variable-size activation buffers), defragmentation is sometimes necessary. The approach is to compact allocated blocks by moving them to be contiguous:

def defragment(allocations, total_memory):
    """
    Compact allocated blocks to eliminate external fragmentation.
    Requires copying data, which is expensive.

    allocations: dict of id -> (offset, size, gpu_tensor)
    Returns: new allocations dict with compacted offsets.
    """
    # Sort by current offset
    sorted_allocs = sorted(allocations.items(), key=lambda x: x[1][0])

    new_offset = 0
    new_allocations = {}

    for alloc_id, (old_offset, size, tensor) in sorted_allocs:
        if new_offset != old_offset:
            # Need to move this block
            # This requires a GPU memcpy (expensive but necessary)
            new_tensor = torch.empty_like(tensor)
            new_tensor.copy_(tensor)
            new_allocations[alloc_id] = (new_offset, size, new_tensor)
        else:
            new_allocations[alloc_id] = (new_offset, size, tensor)

        new_offset += size

    return new_allocations

Defragmentation requires $O(n)$ copies where $n$ is the number of allocated blocks, and each copy is a GPU memcpy (potentially gigabytes). This is far too expensive for real-time serving. The slab allocator avoids the problem entirely.