Debugging eines vLLM-Speicherlecks: Das Rätsel von UCX und `mmap` jenseits des Heaps

Zusammenfassung:

• Problemsituation: In einer disaggregierten Prefill/Decode-Serving-Umgebung von vLLM trat ein Leak des Systemarbeitsspeichers (RSS) von 400 MB pro Minute auf, das mit gewöhnlichen Python-Profilern jedoch nicht erkannt wurde.
• Ursachenanalyse: Mit Heaptrack und pmap wurde bestätigt, dass das Leak nicht im Heap, sondern in anonymen Speicherabbildungen (mmap) lag; mit BPFtrace und automatisierten GDB-Skripten wurde die Ursache weiter zurückverfolgt.
• Täter ermittelt: Die High-Performance-Kommunikationsbibliothek UCX fing zur Optimierung mmap/munmap-Aufrufe ab und gab freigegebenen Speicher nicht sofort zurück, sondern sammelte ihn unbegrenzt in einer Queue.
• Lösung: Das Problem wurde behoben, indem die Memory-Hooking-Funktion von UCX über die Umgebungsvariable UCX_MEM_MMAP_HOOK_MODE=none deaktiviert wurde.

Ausführliche Zusammenfassung:

1. Das rätselhafte Speicherleck

Das Team von Mistral AI entdeckte in einer disaggregierten Prefill/Decode-Serving-Umgebung (auf Basis von NIXL) mit vLLM ein Phänomen, bei dem der Systemspeicher linear um 400 MB pro Minute anwuchs.

• Symptom: Der Python-Heap-Speicher blieb stabil, aber die RSS (Resident Set Size) auf Betriebssystemebene stieg weiter an und führte schließlich zu OOM (Out of Memory).
• Erste Versuche scheiterten: Python-Tools wie Memray und Guppy 3 meldeten keine Auffälligkeiten, das Standard-GDB ließ den Prozess abstürzen und Valgrind war zu langsam für den praktischen Einsatz.

2. Tiefenanalyse auf Kernel-Ebene

Da klar wurde, dass die Ursache nicht auf Anwendungsebene (Python/C++) lag, sondern tiefer im Stack, kamen Systemwerkzeuge zum Einsatz.

• Heaptrack: Visuell wurde bestätigt, dass Heap-Allokationen (malloc/free) stabil waren, während die RSS anstieg. Das deutete darauf hin, dass das Leak außerhalb der Heap-Verwaltung von glibc, nämlich in anonymen Speicherabbildungen (anonymous memory mappings), auftrat.
• pmap: Durch das Beobachten von /proc/<pid>/maps wurde festgestellt, dass bestimmte anonyme Mapping-Bereiche weiter anwuchsen und ihre Adressen wechselten. Das deutete auf wiederholte Zyklen aus mremap oder munmap gefolgt von mmap hin.
• BPFtrace: Um Systemaufrufe zu verfolgen, die selbst mit LD_PRELOAD nicht sichtbar waren (weil sie glibc umgingen), wurde BPFtrace eingesetzt. Das Ergebnis zeigte, dass mmap-Aufrufe über direkte syscalls erfolgten.

3. Täter gefunden: automatisiertes GDB-Scripting

Nachdem mit BPFtrace die Adresse des problematischen Systemaufrufs ermittelt worden war, wurde ein GDB-Skript geschrieben, das nur an dieser Adresse (SYS_mmap) anhält.

Beispiel des verwendeten GDB-Skripts:

# Bedingten Breakpoint für den mmap-Systemaufruf (Nummer 9) setzen  
break syscall if $rdi == 9  
commands  
  silent  
  # Temporären Breakpoint am Rückkehrpunkt des Systemaufrufs setzen  
  tbreak *0x00007ffff7d9525d  
  commands  
    silent  
    # Stack-Trace und zurückgegebene Adresse ausgeben  
    bt  
    printf "Syscall returned: rax = 0x%012lx\n", $rax  
    continue  
  end  
  continue  
end  
  

Über diesen Stack-Trace wurde der entscheidende Hinweis gefunden, dass die Bibliothek UCX (Unified Communication X) die Python-Aufrufe mmap/munmap unterwegs abfängt (intercept).

4. Ursache: übermäßige Optimierung in UCX

UCX hookt Speicherallokation und -freigabe, um die InfiniBand-Übertragungsleistung zu steigern.

• Mechanismus: Wenn munmap aufgerufen wird, wird der Speicher nicht sofort an das Betriebssystem zurückgegeben, sondern zur späteren Wiederverwendung oder Bereinigung in eine „Invalidierungs-Queue“ gelegt.
• Bug: Durch die Standardeinstellung (UCX_RCACHE_MAX_UNRELEASED=inf) konnte diese Queue unbegrenzt wachsen, und bei einem bestimmten Nutzungsmuster von vLLM funktionierte die Bereinigungslogik (ucp_worker_progress) nicht korrekt, sodass sich immer mehr Speicher ansammelte.

5. Lösung

Im Fall von vLLM muss lediglich ein großer KVCache-Speicherbereich registriert werden, sodass die komplexe Memory-Hooking-Funktion von UCX nicht unbedingt nötig war.

• Sofortige Lösung: Das Leak wurde gestoppt, indem das Memory Hooking von UCX mit der Umgebungsvariable UCX_MEM_MMAP_HOOK_MODE=none vollständig deaktiviert wurde.
• Alternative: Man kann die Größe der Queue auch begrenzen, etwa mit UCX_RCACHE_MAX_UNRELEASED=1024, um Bereinigung zu erzwingen.
• Maßnahme: Die entsprechende Änderung wurde für die vLLM-Community bereits gemergt; außerdem soll das Standardverhalten in künftigen NIXL-Releases verbessert werden.