vLLMにおけるメモリリークのデバッグ

Mistral AIは、disaggregated servingのプレプロダクション・テスト中に発生したvLLMのシステムメモリリークを特定し、解決しました。このリークは、システムメモリが毎分400 MB増加させる原因となり、InfiniBandの最適化に使用されるUCX (Unified Communication X) ライブラリのメモリフッキングメカニズムに起因することが判明しました。

リーク：条件と症状

メモリリークは、特定の条件下でのみ発生しました：vLLMをMistral Medium 3.1モデルで使用し、graph compilationを有効にし、NIXLを使用したPrefill/Decode (P/D) disaggregated servingのセットアップを使用する場合です。リークは、セットアップのdecode側で具体的に発生しており、そこではNIXLとUCXを通じてKVCacheの転送が開始されます。

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

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

診断プロセス：Pythonからカーネル・トレーシングまで

Mistral AIのエンジニアリング・チームは、リークを分離するために、ソフトウェア・スタックのいくつかのレイヤーを下降していく体系的なアプローチを採用しました：

高レベル・プロファイリング

MemrayやGuppy 3のようなPythonメモリ・プロファイリング・ツールを使用した初期の試みでは、リークは検出されませんでした。GDBはプロセスをクラッシュさせ、Valgrindは重いvLLMのセットアップに対しては遅すぎました。これにより、チームは他のユーザーも同じ問題に遭遇しているかを確認するために、vLLMリポジトリにGitHub issueを開きました。

Heaptrackによるヒープ解析

Heaptrackを使用して、mallocとfree操作を監視しました。ヒープメモリは安定していましたが、**Peak Resident Set Size (RSS)**に不一致が観察されました。これは、リークがヒープの外側、つまりglibcのmallocではなく、mmapを介して割り当てられた匿名メモリ・マッピングに発生していることを示していました。

システム・レベルのメモリ・マッピング

pmapコマンドを使用して/proc/<pid>/mapsを読み取ったところ、チームは特定の匿名メモリ領域が拡大し、その開始アドレスが変化していることを特定しました。この挙動は、mremapの使用、または適切なリリースなしにmmapとmunmapのサイクルが繰り返されていることを示唆していました。

BPFtraceによるカーネル・トレーシング

割り当てのソースを特定するために、チームはBPFtraceを使用して、すべてのmmap、munmap、およびmremapシステムコールをログに記録しました。その結果、リークしているアドレスは、標準的なglibcのラッパーやLD_PRELOADフックをバイパスして、glibcの生のsyscallラッパー（syscall+29）から発生するmmapコールによって取得されていることが判明しました。

ターゲットを絞ったGDB自動化

一部の依存関係でフレーム・ポインタが_disabled_されているため、BPFtraceでは完全なユーザー空間スタック・トレースを提供できないため、チームはGDBを使用してsyscallアドレスに条件付きブレークポイントを設定しました。SYS_mmapのみをトリガーし、完全なスタック・トレースを表示させることで、PythonがUCXを通じてmmapを呼び出していることを発見しました。

根本原因：UCX mmapフッキング

調査の結果、UCXはInfiniBandメモリ登録（Registration CacheまたはRCache）を最適化するために、mmapフッキングメカニズムを採用していることが判明しました。このメカニズムは、デフォルトですべてのコールをインターセプトするため、mmapとmunmapのGlobal Offset Table (GOT)エントリを動的にパッチします。

このインターセプトは、主に2つの問題を引き起こしました：

フックのバイパス: 標準的なデバッグ・ツールやLD_PRELOADフックがリークを追跡することを妨げました。
メモリの蓄積: UCXはmunmapが呼び出されたときにメモリを即座に自由にするわけではなく、代わりにその領域を無効化キュー（invalidation queue）に移動します。この特定の境界条件において、このキューを管理するメモリ・プールが動的に拡大し、munmap操作中にmmap callsがトリガーされるようになり、RSSが線形に増加しました。

解決策と修正

リークは、以下の2つの設定を使用して解決されました：

フックの無効化: 環境変数UCX_MEM_MMAP_HOOK_MODE=noneを設定することで、フッキングメカニズムを全く無効化できます。vLLMは、一度だけ大きな連続したメモリ領域（KVCache Managerのメモリ）を登録する必要があるため、これによってvLLMのパフォーマンスに悪的な影響はありません。
キャッシュの制限: UCX_RCACHE_MAX_UNRELEASED=1024を設定する（デフォルトのinfの代わりに）ことで、UCXに、未リリースのメモリ領域のしきい値に達したときにクリーンアップを開始するように強制します。

Mistral AIは、vLLMリポジトリに修正をマックージ（PR #32181）にマージし、今後のNIXLリリースにおいてUCX_RCACHE_MAX_UNRELEASEDのデフォルトの挙動動向を変更するために、NIXLおよびUCXチームと協力しました。