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 side），即透過 NIXL 和 UCX 發起 KVCache 傳輸的地方。

在 P/D 解耦式服務設定中，程序分為兩個階段：

Prefill phase: 一個路由器（router）將 prefill 請求發送至 vLLM 實例以計算 KVCache。
Decode phase: 路由器將 KVCache 中繼資料（metadata）和解碼請求傳送至解碼 vLLM 實例，在那裡使用傳輸的 KVCache 進行 token 生成。

Diagnostic Process: From Python to Kernel Tracing

Mistral AI 的工程團隊採用了系統化的方法來隔離洩漏，從軟體堆疊的各個層級向下深入：

High-Level Profiling

最初使用 Python 記憶體分析工具（如 Memray 和 Guppy 3）的嘗試均未發現洩漏。GDB 會導致程序崩潰，而 Valgrind 對於沉重的 vLLM 設定而言速度太慢。這使得團隊決定在 vLLM 儲存庫中開啟一個 GitHub issue，以確認其他人是否也遇到同樣的問題。

Heap Analysis with Heaptrack

使用 Heaptrack，團隊監控了 malloc 和 free 操作。雖然堆積記憶體（heap memory）保持穩定，但他們觀察到 Peak Resident Set Size (RSS) 的差異。這顯示洩漏發生在堆積之外，發生在透過 mmap 而非 glibc 的 malloc 分配的匿名記憶體映射（anonymous memory mappings）。

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 無法提供完整的用戶空間堆疊追蹤（stack traces），團隊因此使用 GDB 對 syscall 地址設置了條件斷點。透過僅在 SYS_mmap 時觸發並列印出完整的堆疊追蹤，他們發現 Python 是透過 UCX 呼叫 mmap 的。

Root Cause: UCX mmap Hooking

調查顯示，UCX 使用了一種 mmap 鉤子機制 來優化 InfiniBand 記憶體註冊（Registration Cache 或 RCache）。此機制會動態修補 Global Offset Table (GOT) 中 mmap 和 munmap 的條目，以預設攔截所有呼叫。

這種攔截造成了兩個主要問題：

Hook Bypassing: 它阻止了標準的除錯工具和 LD_PRELOAD 鉤子追蹤洩漏。
Memory Accumulation: UCX 在呼叫 munmap 時不會立即釋放記憶體；相反，它會將該區域移至失效隊列（invalidation queue）。在這種特定的邊緣案例中，管理該隊列的記憶體池會動態擴張，導致在 munmap 操作期間觸發 mmap 呼叫，進而導致 RSS 線性增長。

Resolution and Fixes

該洩漏已透過兩種可能的配置方式解決：

Disabling Hooks: 將環境變數 UCX_MEM_MMAP_HOOK_MODE=none 設置為 none 可以完全禁用鉤子機制。由於 vLLM 僅需在一次性註冊一個大型且連續的記憶體區域（KVCache Manager 記憶體）後，便不再需要進行額外操作，因此這對 vLLM 的效能沒有負面影響。
Limiting the Cache: 將 UCX_RCACHE_MAX_UNRELEASED=1024（而非預設的 inf）設置為 1024，會強制 UCX 在達到未釋放記憶體區域的閾值時啟動清理程序。

Mistral AI 已將修復程式併入 vLLM 儲存庫（PR #32181），並與 NIXL 和 UCX 團隊合作，在未來的 NIXL 版本中更改 UCX_RCACHE_MAX_UNRELEASED 的預設行為。