El asignador de caché de CUDA en PyTorch organiza la memoria de la GPU en dos niveles: segmentos (obtenidos mediante cudaMalloc) y bloques (subdivisiones dentro de los segmentos). Aunque la abstracción promete un manejo simple de memoria (asignar reduce la disponible, liberar la aumenta), la implementación interna puede provocar fragmentación: el asignador no puede servir una solicitud aunque haya suficiente memoria libre total. Esto sucede porque los bloques libres solo se fusionan con vecinos adyacentes dentro del mismo segmento, y los bloques en segmentos distintos nunca pueden unirse. El orden de asignación es crítico. Si se asignan primero bloques pequeños (por ejemplo, 8 de 16 MiB), cada uno crea un segmento independiente de 16 MiB. Al liberarlos, cada segmento contiene un bloque libre de 16 MiB, pero permanecen separados. Una solicitud posterior de 32 MiB no puede servirse desde ningún bloque existente, lo que obliga a nuevas llamadas cudaMalloc y duplica la memoria reservada (de 128 a 256 MiB). En cambio, si se asignan primero los grandes (4x32 MiB), luego al liberarlos y solicitar los pequeños, el asignador divide los segmentos de 32 MiB para servir las peticiones de 16 MiB, reutilizando la misma memoria y manteniendo la reserva en 128 MiB. El post también presenta los segmentos expandibles, basados en la API de gestión de memoria virtual de CUDA (cuMemMap). Estos segmentos pueden crecer bajo demanda y permiten fusionar bloques de segmentos distintos, eliminando la rigidez impuesta por cudaMalloc. Esto resuelve la fragmentación en escenarios como la grabación de gráficos CUDA, donde el orden de asignación suele estar restringido. Se incluyen ejemplos de código que ilustran tanto los casos de fragmentación como la mitigación ofrecida por los segmentos expandibles, un avance importante para aplicaciones de inferencia de LLM que requieren exprimir al máximo la memoria de la GPU.
Cuándo ocurre fragmentación en el asignador de caché de CUDA
