Debugging a Memory Leak in vLLM

Mistral AI identified and resolved a system memory leak in vLLM that occurred during pre-production testing of disaggregated serving. The leak, which caused system memory to increase by 400 MB per minute, was traced to the UCX (Unified Communication X) library's memory hooking mechanism used for InfiniBand optimizations.

The Leak: Conditions and Symptoms

The memory leak only manifested under a specific set of conditions: using vLLM with the Mistral Medium 3.1 model, graph compilation enabled, and a Prefill/Decode (P/D) disaggregated serving setup using NIXL. The leak occurred specifically on the decode side of the setup, where KVCache transfer is initiated through NIXL and UCX.

In a P/D disaggregated setup, the process is split into two phases:

1. Prefill phase: A router sends a prefill request to a vLLM instance to compute the KVCache.
2. Decode phase: The router transfers KVCache metadata and a decode request to a decode vLLM instance, where token generation occurs using the transferred KVCache.

Diagnostic Process: From Python to Kernel Tracing

Mistral AI's engineering team used a methodical approach to isolate the leak, descending through several layers of the software stack:

High-Level Profiling

Initial attempts using Python memory profiling tools like Memray and Guppy 3 showed no leaks. GDB caused the process to crash, and Valgrind was too slow for the heavy vLLM setup. This led the team to open a GitHub issue in the vLLM repository to confirm if others were experiencing the same issue.

Heap Analysis with Heaptrack

Using Heaptrack, the team monitored malloc and free operations. While heap memory remained stable, they observed a discrepancy in the Peak Resident Set Size (RSS). This indicated that the leak was occurring outside the heap, in anonymous memory mappings allocated via mmap rather than glibc's malloc.

System-Level Memory Mapping

Using the pmap command to read /proc/<pid>/maps, the team identified that certain anonymous memory regions were growing and their start addresses were changing. This behavior suggested the use of mremap or repeated cycles of mmap and munmap without proper release.

Kernel Tracing with BPFtrace

To confirm the source of the allocations, the team used BPFtrace to log every mmap, munmap, and mremap system call. They discovered that the leaking addresses were obtained through mmap calls originating from glibc's raw syscall wrapper (syscall+29), bypassing standard glibc wrappers and LD_PRELOAD hooks.

Targeted GDB Automation

Because BPFtrace could not provide full user-space stack traces due to disabled frame pointers in some dependencies, the team used GDB to set conditional breakpoints on the syscall address. By triggering only on SYS_mmap and printing the full stack trace, they discovered that Python was invoking mmap through UCX.

Root Cause: UCX mmap Hooking

The investigation revealed that UCX employs a mmap hooking mechanism to optimize InfiniBand memory registration (Registration Cache or RCache). This mechanism dynamically patches the Global Offset Table (GOT) entries for mmap and munmap to intercept all calls by default.

This interception caused two primary issues:

1. Hook Bypassing: It prevented standard debugging tools and LD_PRELOAD hooks from tracking the leak.
2. Memory Accumulation: UCX does not immediately free memory when munmap is called; instead, it moves the region to an invalidation queue. In this specific edge case, the memory pool managing this queue expanded dynamically, leading to mmap calls being triggered during munmap operations, causing the RSS to grow linearly.

Resolution and Fixes

The leak was resolved using two possible configurations:

• Disabling Hooks: Setting the environment variable UCX_MEM_MMAP_HOOK_MODE=none disables the hooking mechanism entirely. This had no negative impact on vLLM performance because vLLM only needs to register one large, contiguous memory region (the KVCache Manager memory) once.
• Limiting the Cache: Setting UCX_RCACHE_MAX_UNRELEASED=1024 (instead of the default inf) forces UCX to initiate cleanup once the threshold of unreleased memory regions is reached.

Mistral AI merged a fix into the vLLM repository (PR #32181) and collaborated with the NIXL and UCX teams to change the default behavior of UCX_RCACHE_MAX_UNRELEASED in future NIXL releases.