vLLMのメモリリークをデバッグする：ヒープを超えたUCXとmmapのミステリー

要約:

• 問題の状況: vLLMのPrefill/Decode分離（disaggregated）サービング環境で、毎分400MBのシステムメモリ（RSS）リークが発生したが、一般的なPythonプロファイラでは検出できなかった。
• 原因分析: Heaptrackとpmapで、ヒープではなく匿名メモリマッピング（mmap）でリークが起きていることを確認し、BPFtraceと自動化されたGDBスクリプトで原因を追跡した。
• 犯人の特定: 高性能通信ライブラリのUCXが最適化のためにmmap/munmap呼び出しをフックしており、解放されたメモリを即座に返却せず無制限にキューへ溜め込んでいたことが原因だった。
• 解決策: 環境変数UCX_MEM_MMAP_HOOK_MODE=noneを設定してUCXのメモリフック機能を無効化することで問題を解決した。

詳細要約:

1. 謎のメモリリーク

Mistral AIチームは、vLLMを使ったPrefill/Decode分離サービング（NIXLベース）環境で、毎分400MBずつシステムメモリが線形に増加する現象を発見しました。

• 症状: Pythonヒープメモリは安定していた一方で、OSレベルのRSS（Resident Set Size）は増え続け、最終的にOOM（Out of Memory）に至った。
• 初期の試みは失敗: Memray、Guppy 3といったPythonツールでは正常に見え、標準のGDBはプロセスをクラッシュさせ、Valgrindは遅すぎて使えなかった。

2. カーネルレベルへの深掘り分析

原因がアプリケーションレベル（Python/C++）ではなく、より低いレベルにあると直感し、システムツールを活用しました。

• Heaptrack: ヒープ割り当て（malloc/free）は安定しているのにRSSが増加していることを可視化して確認。これはリークがglibcのヒープ管理の外側にある**匿名メモリマッピング（anonymous memory mappings）**で発生していることを示唆していた。
• pmap: /proc/<pid>/mapsを監視し、特定の匿名マッピング領域が増え続け、アドレスも変化していることを確認した。これはmremapまたはmunmap後のmmapサイクルが繰り返されていることを意味していた。
• BPFtrace: LD_PRELOADでも捕捉できない（glibcをバイパスする）システムコールを追跡するため、BPFtraceを使用した。その結果、mmap呼び出しが直接syscallによって発生していることを確認した。

3. 犯人の特定：自動化されたGDBスクリプティング

BPFtraceで問題のシステムコールアドレスを確認した後、GDBを使ってそのアドレス（SYS_mmap）でのみ停止するようスクリプトを書きました。

使用したGDBスクリプト例:

# mmapシステムコール（番号9）に条件付きブレークポイントを設定  
break syscall if $rdi == 9  
commands  
  silent  
  # システムコールのリターン地点に一時ブレークポイントを設定  
  tbreak *0x00007ffff7d9525d  
  commands  
    silent  
    # スタックトレースと返されたアドレスを出力  
    bt  
    printf "Syscall returned: rax = 0x%012lx\n", $rax  
    continue  
  end  
  continue  
end  
  

このスタックトレースにより、**UCX（Unified Communication X）**ライブラリがPythonのmmap/munmap呼び出しを途中でフック（intercept）しているという決定的な手がかりを得ました。

4. 原因：UCXの過剰な最適化

UCXはInfiniBand転送性能を高めるため、メモリ割り当て/解放をフックします。

• メカニズム: munmapが呼ばれると、メモリをすぐOSに返却せず、後で再利用または整理するために「無効化キュー（invalidation queue）」に入れておく。
• バグ: デフォルト設定（UCX_RCACHE_MAX_UNRELEASED=inf）のためこのキューは無制限に大きくなり得て、vLLMの特定の利用パターンではクリーンアップロジック（ucp_worker_progress）が正しく動作せず、メモリが蓄積し続けていた。

5. 解決方法

vLLMでは巨大なKVCacheメモリ領域を1つ登録すればよいため、UCXの複雑なメモリフック機能は必ずしも必要ではありませんでした。

• 即時の解決策: 環境変数**UCX_MEM_MMAP_HOOK_MODE=none**を設定してUCXのメモリフックを完全に無効化することで、リークを防いだ。
• 代替案: UCX_RCACHE_MAX_UNRELEASED=1024のようにキューのサイズを制限し、強制的にクリーンアップを発生させることもできる。
• 対応: この修正はvLLMコミュニティ向けにマージされており、今後のNIXLリリースでもデフォルト動作の改善が予定されている。

[jongyeans]

いったいどんな人生を送れば…このレベルの
境地に至れるのか

[ng0301]

こういう人たちの実力がどれほどのものなのか、まったく見当がつきません