Débogage d’une fuite mémoire dans vLLM : le mystère d’UCX et de `mmap` au-delà du heap

Résumé :

• Situation du problème : dans un environnement de serving vLLM Prefill/Decode désagrégé, une fuite de mémoire système (RSS) de 400 Mo par minute se produisait, sans être détectée par les profileurs Python classiques.
• Analyse de la cause : avec Heaptrack et pmap, l’équipe a confirmé que la fuite ne venait pas du heap mais de mappages mémoire anonymes (mmap), puis a remonté la piste grâce à BPFtrace et à des scripts GDB automatisés.
• Identification du coupable : UCX, une bibliothèque de communication hautes performances, interceptait les appels mmap/munmap pour optimiser les performances, et plaçait la mémoire libérée dans une file d’attente sans jamais la restituer immédiatement.
• Solution : le problème a été résolu en désactivant la fonctionnalité de hook mémoire d’UCX via la variable d’environnement UCX_MEM_MMAP_HOOK_MODE=none.

Résumé détaillé :

1. Une fuite mémoire mystérieuse

L’équipe de Mistral AI a découvert qu’en environnement de serving Prefill/Decode désagrégé avec vLLM (basé sur NIXL), la mémoire système augmentait de façon linéaire de 400 Mo par minute.

• Symptôme : la mémoire du heap Python restait stable, mais la RSS (Resident Set Size) au niveau du système d’exploitation continuait d’augmenter jusqu’à provoquer un OOM (Out of Memory).
• Échec des premières tentatives : des outils Python comme Memray et Guppy 3 n’indiquaient rien d’anormal, GDB standard faisait planter le processus, et Valgrind était trop lent pour être exploitable.

2. Analyse approfondie au niveau noyau

En pressentant que l’origine du problème se situait à un niveau plus bas que l’application (Python/C++), l’équipe s’est tournée vers des outils système.

• Heaptrack : a permis de visualiser que les allocations heap (malloc/free) restaient stables alors que la RSS augmentait. Cela suggérait une fuite dans des mappages mémoire anonymes (anonymous memory mappings), hors du contrôle du gestionnaire de heap de glibc.
• pmap : la surveillance de /proc/<pid>/maps a montré qu’une certaine zone de mappage anonyme continuait à grossir et à changer d’adresse. Cela indiquait des cycles répétés de mremap, ou de munmap suivi de mmap.
• BPFtrace : pour suivre des appels système invisibles même avec LD_PRELOAD (car contournant glibc), l’équipe a utilisé BPFtrace. Le résultat a montré que les appels mmap étaient effectués via des syscall directs.

3. Arrestation du coupable : scripts GDB automatisés

Après avoir identifié avec BPFtrace l’adresse des appels système problématiques, l’équipe a écrit un script GDB pour n’arrêter l’exécution qu’à cette adresse (SYS_mmap).

Exemple de script GDB utilisé :

# Définir un point d’arrêt conditionnel sur l’appel système mmap (numéro 9)  
break syscall if $rdi == 9  
commands  
  silent  
  # Définir un point d’arrêt temporaire sur le point de retour de l’appel système  
  tbreak *0x00007ffff7d9525d  
  commands  
    silent  
    # Afficher la pile d’appels et l’adresse retournée  
    bt  
    printf "Syscall returned: rax = 0x%012lx\n", $rax  
    continue  
  end  
  continue  
end  
  

Cette stack trace a fourni l’indice décisif : la bibliothèque UCX (Unified Communication X) interceptait les appels mmap/munmap de Python.

4. Cause : une optimisation excessive d’UCX

UCX hooke les allocations/libérations mémoire pour améliorer les performances de transport InfiniBand.

• Mécanisme : lorsqu’un munmap est appelé, UCX ne rend pas immédiatement la mémoire au système d’exploitation, mais la place dans une « file d’invalidation » pour une réutilisation ou un nettoyage ultérieurs.
• Bug : avec la configuration par défaut (UCX_RCACHE_MAX_UNRELEASED=inf), cette file pouvait grossir sans limite. Dans certains schémas d’utilisation de vLLM, la logique de nettoyage (ucp_worker_progress) ne fonctionnait pas correctement, de sorte que la mémoire ne faisait que s’accumuler.

5. Méthode de résolution

Dans le cas de vLLM, comme un seul grand espace mémoire de KVCache devait être enregistré, la fonctionnalité complexe de hook mémoire d’UCX n’était pas vraiment nécessaire.

• Correctif immédiat : la fuite a été stoppée en désactivant complètement le hook mémoire d’UCX via la variable d’environnement UCX_MEM_MMAP_HOOK_MODE=none.
• Alternative : il est aussi possible de limiter la taille de la file avec UCX_RCACHE_MAX_UNRELEASED=1024 afin de forcer le nettoyage.
• Mesure prise : le correctif a été fusionné pour la communauté vLLM, et le comportement par défaut devrait être amélioré dans de futures versions de NIXL.