Debugging a Memory Leak in vLLM

Mistral AI 识别并解决了 vLLM 中的一个系统内存泄漏问题，该问题发生在解耦式服务（disaggregated serving）的预生产测试期间。该泄漏导致系统内存每分钟增加 400 MB，经追溯发现是由于用于 InfiniBand 优化的 UCX (Unified Communication X) 库的内存钩子（hooking）机制引起的。

The Leak: Conditions and Symptoms

内存泄漏仅在特定条件下显现：使用 vLLM 配合 Mistral Medium 3.1 模型、启用图编译（graph compilation）、并使用 NIXL 进行 Prefill/Decode (P/D) 解耦式服务设置。泄漏特别发生在设置的 decode 端，即通过 NIXL 和 UCX 发起 KVCache 传输的地方。

在 P/D 解耦式设置中，流程分为两个阶段：

Prefill phase: 一个路由器将 prefill 请求发送到 vLLM 实例以计算 KVCache。
Decode phase: 路由器将 KVCache 元数据和 decode 请求传输到 decode vLLM 实例，在那里使用传输的 KVCache 进行 token 生成。

Diagnostic Process: From Python to Kernel Tracing

Mistral AI 的工程团队采用了一种系统化的方法来隔离泄漏，通过软件栈的多个层级进行深入排查：

High-Level Profiling

最初使用 Memray 和 Guppy 3 等 Python 内存分析工具的尝试均未发现泄漏。GDB 会导致进程崩溃，而 Valgrind 对于沉重的 vLLM 设置来说速度太慢。这促使团队在 vLLM 仓库中提交了一个 GitHub issue，以确认其他人是否也遇到了同样的问题。

Heap Analysis with Heaptrack

使用 Heaptrack，团队监控了 malloc 和 free 操作。虽然堆内存保持稳定，但他们观察到了 Peak Resident Set Size (RSS) 的差异。这表明泄漏发生在堆之外，即通过 mmap 分配的匿名内存映射，而不是通过 glibc 的 malloc。

System-Level Memory Mapping

使用 pmap 命令读取 /proc/<pid>/maps，团队发现某些匿名内存区域正在增长，且其起始地址正在发生变化。这种行为表明使用了 mremap 或在没有正确释放的情况下重复进行 mmap 和 munmap 循环。

Kernel Tracing with BPFtrace

使用 BPFtrace 记录了每一次 mmap、munmap 和 mremap 系统调用。他们发现泄漏的地址是通过 glibc 的原始系统调用封装器 (syscall+29) 获取的，这绕过了标准的 glibc 封装器和 LD_PRELOAD 钩子。

Targeted GDB Automation

由于某些依赖项中禁用了帧指针（frame pointers），导致 BPFtrace 无法提供完整的用户空间堆栈轨迹，团队使用 GDB 对 syscall 地址设置了条件断点。通过仅在 SYS_mmap 时触发并打印完整的堆栈轨迹，他们发现 Python 正在通过 UCX 调用 mmap。

Root Cause: UCX mmap Hooking

调查显示，UCX 采用了 mmap 钩子机制 来优化 InfiniBand 内存注册（Registration Cache 或 RCache）。该机制通过动态修补 Global Offset Table (GOT) 中 mmap 和 munmap 的 entries，从而默认拦截所有调用。

这种拦截导致了两个主要问题：

Hook Bypassing: 它阻止了标准的调试工具和 LD_PRELOAD 钩子来追踪泄漏。
Memory Accumulation: UCX 在调用 munmap 时不会立即释放内存；相反，它会将该区域移动到失效队列（invalidation queue）中。在这种特定的边缘情况（edge case）下，管理该队列的内存池会动态扩展，导致在 munmap 操作期间触发 mmap 调用，从而导致 RSS 线性增长。

Resolution and Fixes

该泄漏通过两种可能的配置进行解决：

Disabling Hooks: 设置环境变量 UCX_MEM_MMAP_HOOK_MODE=none 可以完全禁用钩子机制。由于 vLLM 仅需在一次性注册一个大型且连续的内存区域（KVCache Manager 内存）后，因此这不会对 vLLM 性能产生负面影响。
Limiting the Cache: 设置 UCX_RCACHE_MAX_UNRELEASED=1024（而不是默认的 inf）会强制 UCX 在达到未释放内存区域的阈值时启动清理工作。

Mistral AI 已将修复方案合并到了 vLLM 仓库（PR #32181），并与 NIXL 和 UCX 团队协作，在未来的 NIXL 版本中更改 UCX_RCACHE_MAX_UNRELEASED 的默认行为。