Mistral arregla fuga de memoria en vLLM causada por UCX | Keryc
Hace poco el equipo de Mistral detectó una fuga de memoria que solo aparecía en un escenario concreto: vLLM con disgregación Prefill/Decode, usando el modelo Mistral Medium 3.1 y la compilación de grafo activada. No había errores ni crashes, solo un aumento lineal en memoria de aproximadamente 400 MB por minuto hasta que el proceso se quedaba sin memoria.
¿Suena aterrador? Sí. ¿Imposible de investigar? Para nada. Esta historia muestra paso a paso cómo llegaron al fondo del problema: desde perfiles en Python hasta trazas a nivel de kernel, y cómo una dependencia de bajo nivel (UCX) terminaba siendo la culpable.
¿Qué ocurrió y por qué importó tanto?
El patrón era raro: la fuga se presentaba únicamente en la parte decode de la arquitectura disgregada (Prefill/Decode) y solo cuando el flujo de transferencia KVCache pasaba por NIXL. Eso apuntó a la ruta de transferencia de memoria como origen.
Prefill/Decode funciona así, a grandes rasgos:
El router envía una petición de prefill para calcular la KVCache.
Luego se transfiere esa KVCache a un worker decode que genera tokens extendiendo esa KVCache.
La transferencia de KVCache se hace con NIXL, que depende de UCX para comunicaciones rápidas (RDMA, Infiniband, etc.).
Si todo esto suena técnico, la lección práctica es clara: cuando mueves grandes bloques de memoria entre procesos y usas bibliotecas de alto rendimiento, hay más puntos donde algo puede fallar.
Cómo lo investigaron (sin perderse en jerga)
Primero usaron herramientas de perfilado en Python: Memray, Guppy 3 y luego Heaptrack. Curiosamente, Heaptrack mostró que el heap estaba estable, pero el RSS (memoria residente) crecía. ¿Cómo es eso posible? Porque RSS incluye más que el heap: también mapas anónimos creados por mmap, hojas de memoria grandes y memoria gestionada fuera de glibc.
Para ver los mapas en tiempo real usaron pmap en modo repetido (con watch) y observaron que ciertas regiones anónimas crecían y cambiaban de dirección: un signo claro de que había llamadas a mmap/mremap o ciclos de munmap + mmap que no liberaban realmente la memoria.
Heaptrack solo atrapa malloc/free de glibc, así que había que bajar un nivel.
BPFtrace y por qué fue útil
Con un script de BPFtrace trazaron mmap, munmap y mremap a nivel de syscall. Esa traza mostró que las llamadas venían desde syscall+29, es decir, desde un wrapper que hace llamadas directas al kernel (raw syscall). Eso indicaba que la biblioteca que hacía el trabajo hacía llamadas directas al kernel, eludiendo las rutas habituales y nuestros hooks.
Aun así, BPFtrace no devolvía la pila completa de usuarios que necesitaban para identificar el responsable final, solo la traza hasta la instrucción syscall.
GDB con breakpoints condicionales: el truco final
Como solución pragmática montaron un breakpoint muy dirigido en la instrucción syscall que solo saltara cuando el número de syscall fuera SYS_mmap. Al capturar la salida y el stack completo en esos instantes pudieron correlacionar direcciones devueltas por mmap con las regiones crecientes observadas en pmap.
Esa traza completa mostró al culpable: UCX (a través de UCM/ucs) estaba llamando a mmap desde sus hooks internos. En algunos momentos, munmap dentro de UCX disparaba operaciones que a su vez provocaban más mmap. En resumen: la gestión interna de la memoria de UCX (su registration cache o RCache) acumulaba regiones y no las liberaba inmediatamente.
¿Por qué UCX hacía esto?
UCX optimiza transferencias RDMA registrando (pinning) páginas de memoria para que las tarjetas de red puedan acceder a ellas sin CPU. Para eso, UCX parchea en tiempo de ejecución las entradas de la GOT y añade hooks a mmap/munmap, con el objetivo de controlar y acelerar el registro de memoria. Es potente, pero rompe las suposiciones de herramientas de depuración y puede interferir con hooks simples.
Además, UCX no libera regiones inmediatamente: las pone en una cola de invalidación gestionada por su memory pool. Si esa cola crece sin límites (valor por defecto inf), el proceso seguirá pidiendo más mmap y nunca recuperará RSS.
Cómo lo arreglaron (soluciones inmediatas y de fondo)
La buena noticia: había soluciones claras y seguras para este caso de uso concreto.
Solución inmediata: desactivar el hook de mmap de UCX con la variable de entorno UCX_MEM_MMAP_HOOK_MODE=none. Esto eliminó la fuga sin afectar el rendimiento en el escenario vLLM/NIXL, porque en este flujo solo se necesita registrar una gran región continua (KVCache) una vez.
Otra medida útil: limitar la cola de regiones no liberadas con UCX_RCACHE_MAX_UNRELEASED=1024 (el valor por defecto era infinito). Con esto UCX fuerza limpiezas periódicas y evita acumulaciones sin control.
A mediano plazo, vLLM y los mantenedores de NIXL/UCX acordaron ajustar valores por defecto y cambiar comportamientos para evitar que este patrón vuelva a aparecer.
Lecciones prácticas para equipos que corren inferencia a escala
No confíes solo en herramientas de perfilado de alto nivel: Heaptrack, Memray o Guppy son útiles, pero es vital mirar RSS, pmap y syscalls cuando la memoria que crece no está en el heap.
Las bibliotecas de alto rendimiento (UCX, RDMA, gestores de GPU) pueden parchear funciones en tiempo de ejecución. Eso es potente para rendimiento, pero complica debugging. ¿Cómo lo afrontas? Ten siempre trazas a nivel de kernel (BPFtrace/strace) y estrategias para obtener stacks completos cuando BPFtrace no basta (GDB dirigido).
Si ves mmap creciendo en RSS, piensa en registradores/pools de memoria y en colas internas de bibliotecas: el problema puede estar fuera de tu código Python o C principal.
Colabora con mantenedores de las librerías. En este caso la investigación y la coordinación con vLLM, NIXL y UCX fue clave para producir parches y mejores defaults.
Reflexión final
Esta investigación es un recordatorio de que la infraestructura de inferencia moderna es una pila de capas: cada dependencia trae optimizaciones que ayudan (y a veces generan nuevas fuentes de errores). Darte el tiempo de bajar al kernel, correlacionar pmap + BPFtrace + GDB y entender la intención de una librería es lo que separa una solución temporal de una solución robusta.
Si te toca gestionar despliegues de modelos grandes y disgregados, guarda estas herramientas y patrones para la próxima vez: puede ahorrarte días de frustración.
