VRAM Fragmentation as a First-Class Failure: A Buddy Allocator for Inference Workloads

If you've ever run a production inference service, you've seen it: the GPU has 24GB of VRAM, but you can't allocate a 2GB tensor. The process doesn't crash—it just fails with CUDA out of memory. You restart the container, and it works for a while. Then it happens again. You blame the model, the framework, or the GPU. But the real culprit is fragmentation.

Fragmentation is not an edge case. It's a first-class failure mode for inference workloads. And unlike CPU RAM, VRAM fragmentation is harder to diagnose and fix because the allocation patterns are different and the memory is managed by drivers and libraries that you don't control.

In this article, I'll explain why VRAM fragmentation happens, why standard allocators fall short, and how to build a buddy allocator that keeps inference stable. I'll also show you how to integrate it into a real inference stack using PyTorch and CUDA.

Why VRAM Fragmentation Happens

Fragmentation occurs when memory is allocated and freed in non-contiguous chunks, leaving gaps that are too small to satisfy new requests. In inference, this happens constantly:

• Dynamic batch sizes: You allocate tensors for a batch of 8, then 16, then 4. The sizes vary, creating holes.
• KV cache growth: During generation, the KV cache grows token by token. Each growth step may allocate a new block, leaving the old one behind.
• Temporary tensors: Activations, intermediate results, and gradients are allocated and freed within each forward pass.

Over time, the GPU memory becomes a patchwork of allocated and free blocks. The total free memory might be 10GB, but the largest contiguous block is only 500MB. When you try to allocate a 1GB tensor, it fails.

Standard Allocators vs. Buddy Allocators

The default CUDA allocator in PyTorch (caching allocator) already tries to mitigate fragmentation by caching freed blocks. But it has limitations: it uses a best-fit strategy with a round-robin per-stream, which can still lead to fragmentation under varied allocation sizes.

A buddy allocator is a classic solution. It splits memory into power-of-two blocks. When you request a block of size S, it finds the smallest power-of-two block that fits, and splits it recursively until it reaches the exact size. When you free a block, it merges adjacent buddies back into larger blocks.

The key advantage is that merging is deterministic: blocks are only merged if they are buddies (i.e., they form a perfect pair). This guarantees that after a series of allocations and frees, the memory can coalesce into larger contiguous regions, reducing fragmentation.

Designing a Buddy Allocator for VRAM

Let's design a buddy allocator specifically for inference workloads. We'll assume a single GPU with, say, 24GB VRAM. The allocator manages a contiguous memory pool, which we'll reserve at startup.

Memory Layout

We'll divide the pool into blocks of size 2^k, for k from min_block_order to max_block_order. For example, min_block_size = 256 bytes (order 8), max_block_size = 16GB (order 34). We'll maintain a free list for each order.

Data Structures

• free_lists: an array of lists, one per order. Each list contains pointers to free blocks of that size.
• allocated_map: a dictionary mapping block pointers to their order (for O(1) free).
• buddy_map: a dictionary mapping a block pointer to its buddy's pointer (or we can compute it with XOR).

Allocation Algorithm

1. Given a request size, round up to the next power of two: block_size = 2^ceil(log2(size)).
2. Find the smallest order o such that 2^o >= block_size.
3. If the free list for order o is non-empty, pop a block and return it.
4. Else, find the smallest order o' > o with a free block. If none, return null (or trigger a compaction/eviction).
5. Split that block recursively: take a block of order o', split it into two buddies of order o'-1, add the buddy to the free list for o'-1, and repeat until we reach order o.
6. Mark the block as allocated and return its pointer.

Free Algorithm

1. Given a pointer and its order, mark it as free and add it to the free list for that order.
2. Check if its buddy is also free. If yes, remove both from the free list, merge them into a block of order o+1, and repeat the merge process.

Buddy Calculation

For a block of order o with base address addr, its buddy's address is addr ^ (1 << o) (assuming the pool base is aligned to the max block size). This XOR trick works because buddies are adjacent and differ in the bit at position o.

Handling Real-World Constraints

In practice, you can't just allocate a giant contiguous pool at startup because the GPU memory is shared with other processes and the driver. You need to allocate from the CUDA allocator in chunks, and manage those chunks with your buddy allocator.

One approach: reserve a large virtual memory region (using cudaMalloc with a large size, but not touching it), and then commit physical memory on demand. However, CUDA doesn't expose virtual memory easily until CUDA 11.2 with cuMemCreate and cuMemMap. If you're on older versions, you can use cudaMalloc for chunks of, say, 1GB, and manage them as a pool.

Let's implement a simple version using PyTorch's caching allocator as a back-end. We'll create a custom allocator that uses a buddy structure on top of PyTorch's tensors.

Implementation Sketch

Here's a simplified Python implementation using PyTorch tensors as the memory pool. This is for illustration; a production version would be in C++ for performance.

import torch
from math import log2, ceil

class BuddyAllocator:
    def __init__(self, total_size, min_block_size=256):
        self.total_size = total_size
        self.min_block_size = min_block_size
        self.min_order = int(log2(min_block_size))
        self.max_order = int(log2(total_size))
        self.free_lists = [[] for _ in range(self.max_order + 1)]
        self.allocated = {}  # addr -> (order, tensor)
        self.pool = torch.empty(total_size, dtype=torch.uint8, device='cuda')
        self.base_addr = self.pool.data_ptr()
        # Initially, the whole pool is one free block of max_order
        self.free_lists[self.max_order].append(self.base_addr)

    def _order_for_size(self, size):
        return max(self.min_order, int(ceil(log2(size))))

    def _split(self, addr, order):
        # Split a block of given order into two buddies of order-1
        if order <= self.min_order:
            return
        buddy_addr = addr + (1 << (order - 1))
        # Add the buddy to the free list for order-1
        self.free_lists[order - 1].append(buddy_addr)
        # The original block is now of order-1, but we'll handle it in the caller

    def allocate(self, size):
        order = self._order_for_size(size)
        # Find free block of order >= requested order
        for o in range(order, self.max_order + 1):
            if self.free_lists[o]:
                addr = self.free_lists[o].pop()
                # Split down to order
                while o > order:
                    self._split(addr, o)
                    o -= 1
                # Now we have a block of order 'order' at addr
                # Create a tensor view
                block_size = 1 << order
                tensor = torch.empty(block_size, dtype=torch.uint8, device='cuda')
                # But we need to point to the pool memory, not allocate new
                # We'll use from_blob
                tensor = torch.from_blob((addr - self.base_addr) + self.pool.data_ptr(), (block_size,), dtype=torch.uint8)
                # Actually, from_blob expects a CPU pointer, so this is tricky.
                # In practice, you'd use CUDA's virtual memory API.
                # We'll skip the actual tensor mapping for brevity.
                self.allocated[addr] = (order, tensor)
                return addr, tensor
        return None, None

    def free(self, addr):
        order, _ = self.allocated.pop(addr)
        self.free_lists[order].append(addr)
        # Try to merge with buddy
        buddy = addr ^ (1 << order)
        while order < self.max_order:
            if buddy in self.free_lists[order]:
                # Remove buddy from free list
                self.free_lists[order].remove(buddy)
                # Merge: the merged block is at min(addr, buddy)
                addr = min(addr, buddy)
                order += 1
                self.free_lists[order].append(addr)
                buddy = addr ^ (1 << order)
            else:
                break

This code is incomplete—mapping tensors to the pool is non-trivial. In a real system, you'd use CUDA's virtual memory management (cuMemCreate, cuMemMap, cuMemSetAccess) to reserve a large address range and commit physical memory on demand. That way, you can return pointers into that range and create tensors via torch.from_blob with the right device and data pointer.

Integration with Inference Workloads

How do you actually use this? You want your inference engine to allocate all its working memory through your buddy allocator. For a typical transformer, that includes:

• Model weights (static, allocate once)
• Activations (dynamic, per layer)
• KV cache (dynamic, grows with sequence length)
• Temporary buffers for attention, softmax, etc.

You can wrap your allocator in a PyTorch custom allocator via torch.cuda.memory.set_allocator (only in C++ extensions). Alternatively, you can pre-allocate a pool and then use torch.from_blob to create tensors that point into the pool, and manage them manually. This is more work but gives you full control.

A simpler approach: use your buddy allocator to manage the KV cache. The KV cache is the biggest source of fragmentation because it grows unpredictably with batch size and sequence length. By using a buddy allocator, you can allocate blocks of power-of-two sizes and merge them when sequences end.

Performance Considerations

Buddy allocators have a small overhead: O(log N) for allocation and free, which is negligible compared to the GPU kernel launch overhead. However, they can cause internal fragmentation: if you request 1.5GB, you get a 2GB block, wasting 0.5GB. To mitigate this, you can use size classes: for requests close to a power of two, you might use a different strategy.

In inference, you often know the typical sizes: KV cache blocks are often fixed-size (e.g., 512 tokens). You can allocate them in power-of-two sizes to minimize waste. For activations, you can use a slab allocator for fixed-size tensors.

Real-World Experience

At my previous company, we ran a large-scale LLM inference service on A100s. We had constant OOM failures, especially during peak traffic. We switched to a custom buddy allocator for the KV cache, and the failure rate dropped to zero. The key was that we also added monitoring: we tracked the number of free blocks per order, and when fragmentation reached a threshold, we triggered a memory compaction (by reordering sequences).

One lesson: don't rely on the CUDA caching allocator alone. It's optimized for speed, not fragmentation. For inference, you need predictability.

Alternatives and Trade-offs

There are other approaches:

• Memory pooling with fixed-size blocks: Simple, but wastes memory if sizes vary.
• Segregated free lists: Similar to buddy, but with multiple size classes.
• Compaction: Move allocated blocks to consolidate free space. This is expensive on GPU because you have to copy data.

Buddy allocator is a good middle ground: it's simple, fast, and provides good coalescing.

Conclusion

VRAM fragmentation is a real problem that can bring down inference services. A buddy allocator is a proven solution that you can implement with moderate effort. Start by applying it to the KV cache, then expand to other dynamic allocations.

Remember: the goal is not to eliminate fragmentation entirely—that's impossible—but to make it predictable and manageable. With a buddy allocator, you can guarantee that if there's enough total free memory, you can find a contiguous block for any request up to the largest free block.

If you're building a self-hosted inference stack, don't ignore memory management. It's not glamorous, but it's the difference between a service that runs for months and one that crashes every hour.

Now go allocate some memory.