Depuração de vazamento de memória no vLLM: o mistério do UCX e do mmap além do heap

Resumo:

• Situação do problema: Em um ambiente de serving desagregado de Prefill/Decode do vLLM, ocorreu um vazamento de memória do sistema (RSS) de 400 MB por minuto, mas ele não foi detectado por profiladores Python comuns.
• Análise da causa: Com Heaptrack e pmap, foi confirmado que o vazamento ocorria em mapeamentos de memória anônimos (mmap), e não no heap, e a causa foi rastreada com BPFtrace e scripts automatizados de GDB.
• Identificação do culpado: A biblioteca de comunicação de alto desempenho UCX estava interceptando chamadas mmap/munmap para otimização, e a causa era o fato de a memória liberada não ser devolvida imediatamente, ficando acumulada indefinidamente em uma fila.
• Solução: O problema foi resolvido ao desativar o recurso de memory hooking do UCX definindo a variável de ambiente UCX_MEM_MMAP_HOOK_MODE=none.

Resumo detalhado:

1. Um vazamento de memória misterioso

A equipe da Mistral AI descobriu, em um ambiente de serving desagregado de Prefill/Decode com vLLM (baseado em NIXL), um fenômeno em que a memória do sistema aumentava linearmente em 400 MB por minuto.

• Sintoma: A memória heap do Python permanecia estável, mas o RSS (Resident Set Size) no nível do sistema operacional continuava crescendo até eventualmente levar a OOM (Out of Memory).
• Falha nas tentativas iniciais: Ferramentas Python como Memray e Guppy 3 indicavam normalidade, o GDB padrão fazia o processo crashar, e o Valgrind era lento demais para uso prático.

2. Análise profunda no nível do kernel

Percebendo que a causa do problema não estava no nível da aplicação (Python/C++), mas em uma camada mais baixa, a equipe recorreu a ferramentas de sistema.

• Heaptrack: Confirmou visualmente que as alocações de heap (malloc/free) estavam estáveis, enquanto o RSS aumentava. Isso sugeria que o vazamento ocorria em mapeamentos de memória anônimos (anonymous memory mappings), fora do gerenciamento de heap da glibc.
• pmap: Ao monitorar /proc/<pid>/maps, a equipe confirmou que uma determinada região de mapeamento anônimo continuava crescendo e mudando de endereço. Isso indicava a repetição de um ciclo de mremap ou munmap seguido de mmap.
• BPFtrace: Para rastrear chamadas de sistema que não eram capturadas nem com LD_PRELOAD (por contornarem a glibc), foi usado o BPFtrace. Como resultado, foi confirmado que as chamadas mmap estavam ocorrendo por meio de syscall direta.

3. Captura do culpado: scripting automatizado com GDB

Após identificar com BPFtrace o endereço da chamada de sistema problemática, a equipe escreveu um script para o GDB parar apenas naquele endereço (SYS_mmap).

Exemplo de script GDB usado:

# Define um breakpoint condicional para a syscall mmap (número 9)  
break syscall if $rdi == 9  
commands  
  silent  
  # Define um breakpoint temporário no ponto de retorno da syscall  
  tbreak *0x00007ffff7d9525d  
  commands  
    silent  
    # Exibe o stack trace e o endereço retornado  
    bt  
    printf "Syscall returned: rax = 0x%012lx\n", $rax  
    continue  
  end  
  continue  
end  
  

Por meio desse stack trace, a equipe encontrou a pista decisiva de que a biblioteca UCX (Unified Communication X) estava interceptando chamadas mmap/munmap do Python no meio do caminho.

4. A causa: otimização excessiva do UCX

O UCX faz hooking de alocação/liberação de memória para aumentar o desempenho de transmissão via InfiniBand.

• Mecanismo: Quando munmap é chamado, a memória não é devolvida imediatamente ao sistema operacional; em vez disso, ela é colocada em uma “fila de invalidação” para reutilização ou limpeza posterior.
• Bug: Devido à configuração padrão (UCX_RCACHE_MAX_UNRELEASED=inf), essa fila podia crescer indefinidamente e, em um padrão específico de uso do vLLM, a lógica de limpeza (ucp_worker_progress) não funcionava corretamente, fazendo a memória apenas se acumular.

5. Como resolver

No caso do vLLM, como basta registrar uma única grande região de memória do KVCache, o recurso complexo de memory hooking do UCX não era realmente necessário.

• Solução imediata: O vazamento foi interrompido ao definir a variável de ambiente UCX_MEM_MMAP_HOOK_MODE=none e desativar completamente o memory hooking do UCX.
• Alternativa: Também é possível limitar o tamanho da fila com algo como UCX_RCACHE_MAX_UNRELEASED=1024, forçando a limpeza.
• Ação tomada: A correção correspondente foi incorporada para a comunidade do vLLM, e o comportamento padrão também deverá ser melhorado em futuras versões do NIXL.

[jongyeans]

Que tipo de vida é preciso viver... para alcançar
esse nível

[ng0301]

Nem consigo imaginar o nível de domínio técnico que essas pessoas têm.