The 8-Hour Silent CUDA Leak: Tracing Attention-Mask Fragmentation in a Swarm
How a subtle fragmentation bug in attention masks ate 14GB of VRAM across 6 GPUs—and how we caught it.
The 8-Hour Silent CUDA Leak: Tracing Attention-Mask Fragmentation in a Swarm
We run a swarm of autonomous agents that process documents, summarize, and answer queries. Each agent uses a small LLM (Llama-3.2-3B) with a custom attention mask to handle variable-length sequences. The swarm runs on a single node with 6× RTX 4090 GPUs (24GB each). It's a self-hosted setup—no cloud, no managed inference. We control everything, which is both a blessing and a curse.
Last week, we hit a nasty bug: the swarm would run fine for about 8 hours, then start throwing CUDA OOM errors. At first, we thought it was a slow leak in our inference server (vLLM). But vLLM's memory usage was flat. The leak was in our own code—specifically, in the attention masks we build for each request.
This post is a case study in debugging a silent memory leak that took us a full day to trace. It's a story about fragmentation, PyTorch's caching allocator, and the subtle way that a few extra tensors can accumulate into gigabytes of VRAM.
The Setup
Our swarm is a set of Python processes (one per agent) that communicate via Redis. Each agent receives a task, fetches context from a Postgres database (with pgvector for embeddings), and then calls a local vLLM server to generate a response. The vLLM server is the only process that touches the GPUs directly. The agents just send HTTP requests with a prompt and a set of attention masks.
Why do we need custom attention masks? Our agents work with documents that are split into chunks. For each request, we concatenate the relevant chunks with a system prompt and a user query. The model is trained on sequences with a special token that separates instructions from context. To prevent the model from attending to certain parts (e.g., the system prompt when generating the final answer), we build a mask that zeroes out those positions.
We build the mask as a list of Python lists: mask = [[0, 1, 1, ...], ...]. Then we convert it to a tensor on the CPU and move it to the GPU: mask_tensor = torch.tensor(mask, device='cuda'). That's the first mistake.
The Leak
vLLM has a feature that allows you to pass a custom attention mask as part of the request. The mask is used to override the default causal mask. It's a 2D tensor of shape [batch_size, seq_len], where 1 means attend and 0 means not attend.
We built that mask for every request. The mask size is proportional to the sequence length, which can be up to 4096 tokens. That's 4096 floats per row, or 16KB per row. For a batch of 8 sequences, that's 128KB. Not huge.
But here's the thing: vLLM copies the mask to the GPU for every forward pass. And because we're using PyTorch's caching allocator, that memory isn't freed immediately. It's cached for future allocations. That's fine if the allocations are of similar sizes. But our masks vary in size depending on the number of chunks and the length of the query. So the allocator ends up with a fragmented set of blocks.
Over time, the fragmentation grows. Each request allocates a new mask tensor, and the old one is released, but the allocator doesn't coalesce contiguous free blocks. So after thousands of requests, the GPU memory is a mosaic of small free blocks that can't be used for large allocations. Eventually, a request needs a 200MB contiguous block, and the allocator fails, even though total free memory is 5GB.
That's the classic fragmentation leak. But it took us hours to figure out, because the memory usage graph looked like a sawtooth—up and down—but the baseline slowly crept upward. We initially suspected a reference cycle in our agent code, or a bug in vLLM's handling of masks.
The Investigation
We started by monitoring GPU memory with nvidia-smi every 10 seconds. We saw the sawtooth pattern, but the troughs were rising. That pointed to a leak, not just fragmentation. But fragmentation can also cause that if the allocator keeps expanding the pool.
We then used torch.cuda.memory_summary() to see what was allocated. The summary showed a large number of small allocations from a single source: aten::zeros. That was suspicious. We weren't calling zeros explicitly, but vLLM might be. We looked at the vLLM source code (v0.6.3) and found that it uses torch.zeros to create the mask tensor when a custom mask is provided. The mask is passed as a list, and vLLM converts it to a tensor on the GPU.
That's the culprit. Every request creates a new tensor on the GPU, and the previous one is released. But the allocator keeps the memory around. Over 8 hours, with an average request rate of 5 per second, that's 144,000 allocations. Each allocation is at least 16KB, but many are larger. The allocator's block size granularity is 512 bytes, so the fragmentation overhead is significant.
We confirmed by adding a custom allocator hook: torch.cuda.memory._set_allocator_settings to log all allocations. The log showed thousands of allocations of sizes like 16KB, 32KB, 48KB—all from aten::zeros. The total allocated was 14GB, but the actual live tensors were only 2GB. The rest was fragmentation.
The Fix
We fixed it in two ways.
First, we stopped building masks on the fly. Instead, we precompute a set of masks for the most common sequence lengths (256, 512, 1024, 2048, 4096) and cache them on the GPU. When a request comes in, we look up the closest mask and slice it to the required length. This avoids new allocations entirely.
Second, we changed the way we pass the mask to vLLM. Instead of a list of lists, we now pass a pre-built tensor. vLLM still copies it, but because we reuse the same tensor, the copy is cheap and the allocator can reuse the same block.
We also added a watchdog script that monitors GPU memory and alerts us if the trough of the sawtooth rises above a threshold.
The Deeper Lesson
This bug was silent because it didn't crash immediately. It degraded performance over hours. In a swarm, where agents run indefinitely, such leaks are especially dangerous because they accumulate across restarts. The fix was simple, but the diagnosis required understanding the interaction between PyTorch's caching allocator and vLLM's request handling.
If you're building self-hosted inference pipelines, here's what I recommend:
- Preallocate and reuse. Any tensor that is created per-request should be pooled. This includes masks, position IDs, and even input token IDs if they're fixed.
- Monitor allocation patterns. Use
torch.cuda.memory_summary()andnvidia-smito see if the number of allocations is growing. A high allocation count with low live memory is a sign of fragmentation. - Set
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:Truefor PyTorch 2.0+. This makes the allocator use expandable segments, which reduces fragmentation. We didn't have that set. - Use
torch.cuda.memory._dump_snapshot()to get a detailed breakdown of allocations. It's a bit hidden, but it's invaluable.
Conclusion
Our swarm now runs for days without memory issues. The fix saved us from having to restart every 8 hours, which was disrupting our agents' workflows. We also learned to treat every per-request tensor as a potential leak source.
Self-hosting means you own the whole stack, which also means you own the bugs. This one was subtle, but with the right tools and a methodical approach, we traced it to a single line in vLLM's code. The lesson: always check the allocator's behavior before blaming the model.
If you're running a similar swarm, keep an eye on your GPU memory troughs. And if you see a slow rise, check for fragmentation before hunting ghosts.