Depuración de una fuga de memoria en vLLM: el misterio de UCX y `mmap` más allá del heap

Resumen:

• Situación del problema: en un entorno de serving desagregado de Prefill/Decode de vLLM, se producía una fuga de memoria del sistema (RSS) de 400 MB por minuto, pero los perfiladores habituales de Python no la detectaban.
• Análisis de la causa: con Heaptrack y pmap se confirmó que la fuga no estaba en el heap sino en mapeos de memoria anónimos (mmap), y se rastreó el origen mediante BPFtrace y scripts automatizados de GDB.
• Identificación del culpable: UCX, una biblioteca de comunicación de alto rendimiento, estaba interceptando llamadas a mmap/munmap para optimización, y la causa era que no devolvía de inmediato la memoria liberada sino que la acumulaba indefinidamente en una cola.
• Solución: el problema se resolvió desactivando la función de hooking de memoria de UCX mediante la variable de entorno UCX_MEM_MMAP_HOOK_MODE=none.

Resumen detallado:

1. Una fuga de memoria misteriosa

El equipo de Mistral AI descubrió que, en un entorno de serving desagregado de Prefill/Decode con vLLM (basado en NIXL), la memoria del sistema aumentaba linealmente en 400 MB por minuto.

• Síntoma: la memoria del heap de Python se mantenía estable, pero el RSS (Resident Set Size) a nivel del sistema operativo seguía creciendo hasta terminar en OOM (Out of Memory).
• Fracaso de los intentos iniciales: herramientas de Python como Memray y Guppy 3 mostraban todo como normal, GDB estándar hacía colapsar el proceso y Valgrind era demasiado lento para resultar útil.

2. Análisis profundo a nivel de kernel

Al intuir que la causa no estaba en el nivel de la aplicación (Python/C++) sino más abajo, se recurrió a herramientas del sistema.

• Heaptrack: confirmó visualmente que las asignaciones del heap (malloc/free) eran estables mientras el RSS seguía creciendo. Esto sugería que la fuga se producía en mapeos de memoria anónimos (anonymous memory mappings) fuera de la gestión del heap de glibc.
• pmap: al monitorear /proc/<pid>/maps, se observó que ciertas regiones de mapeo anónimo seguían creciendo y cambiando de dirección. Esto implicaba un ciclo repetido de mremap o de munmap seguido de mmap.
• BPFtrace: se utilizó BPFtrace para rastrear llamadas al sistema que no eran visibles ni siquiera con LD_PRELOAD (porque evitaban glibc). El resultado mostró que las llamadas a mmap se estaban haciendo mediante syscall directa.

3. Captura del culpable: automatización con scripts de GDB

Tras identificar con BPFtrace la dirección de la llamada al sistema problemática, se escribió un script de GDB para detenerse solo en esa dirección (SYS_mmap).

Ejemplo del script de GDB utilizado:

# Configurar un breakpoint condicional en la syscall mmap (número 9)  
break syscall if $rdi == 9  
commands  
  silent  
  # Configurar un breakpoint temporal en el punto de retorno de la syscall  
  tbreak *0x00007ffff7d9525d  
  commands  
    silent  
    # Imprimir el stack trace y la dirección devuelta  
    bt  
    printf "Syscall returned: rax = 0x%012lx\n", $rax  
    continue  
  end  
  continue  
end  
  

Ese stack trace dio la pista decisiva de que la biblioteca UCX (Unified Communication X) estaba interceptando las llamadas mmap/munmap de Python.

4. La causa: una optimización excesiva de UCX

UCX hace hooking de asignaciones y liberaciones de memoria para mejorar el rendimiento de transporte sobre InfiniBand.

• Mecanismo: cuando se llama a munmap, en lugar de devolver la memoria inmediatamente al sistema operativo, la coloca en una “cola de invalidación” para reutilizarla o limpiarla más adelante.
• Bug: con la configuración predeterminada (UCX_RCACHE_MAX_UNRELEASED=inf), esa cola podía crecer sin límite, y en un patrón de uso específico de vLLM la lógica de limpieza (ucp_worker_progress) no funcionaba correctamente, de modo que la memoria solo seguía acumulándose.

5. Cómo se resolvió

En el caso de vLLM, solo hacía falta registrar una única gran región de memoria de KVCache, por lo que la compleja función de hooking de memoria de UCX no era realmente necesaria.

• Solución inmediata: se evitó la fuga configurando la variable de entorno UCX_MEM_MMAP_HOOK_MODE=none para desactivar por completo el hooking de memoria de UCX.
• Alternativa: también se puede forzar la limpieza limitando el tamaño de la cola con algo como UCX_RCACHE_MAX_UNRELEASED=1024.
• Acción tomada: la corrección ya fue integrada para la comunidad de vLLM y se espera que el comportamiento predeterminado mejore en futuras versiones de NIXL.

[jongyeans]

¿Qué clase de vida hay que vivir para llegar… a semejante
nivel?

[ng0301]

No me hago una idea del nivel de experiencia que tienen estas personas.