Asincronía en continuous batching optimiza inferencia LLM | Keryc
La CPU y la GPU a menudo se quedan esperando la una a la otra, y eso se come tiempo y dinero. ¿Te imaginas pagar 140 USD al día por una H200 y ver que una cuarta parte del tiempo la GPU está idle esperando a la CPU? Aquí te explico cómo separar las cargas de CPU y GPU para que ambas trabajen en paralelo y así exprimir la inferencia de modelos grandes.
Por qué esto importa
Si ejecutas inference a escala —por ejemplo en endpoints con H200— cada minuto cuenta. Continuous batching ya mejora la utilización porque reduce padding y agrupa requests de forma eficiente. Pero el siguiente cuello de botella es la sincronía: CPU y GPU se turnan y en loops con cientos de pasos por segundo esos huecos suman.
En un experimento con un modelo de 8B, batch 32 y 8K tokens, el ciclo sincrónico tardó 300.6 s y la GPU estuvo idle un 24% del tiempo. En la versión asincrónica la GPU estuvo activa 99.4% del tiempo y el total bajó a 234.5 s. ¿Resultado? Un 22% de mejora real en tiempo de generación, sin tocar modelos ni kernels.
Idea central: desentrelazar CPU y GPU
La idea es simple enunciada y algo más fina en la práctica: preparar en CPU el batch N+1 mientras la GPU calcula el batch N. ¿Por qué no se hace así por defecto? Porque hay tres retos clave:
Recuperar control de la CPU tras lanzar trabajo en GPU.
Garantizar que los datos estén listos cuando cada operación empiece.
Construir el batch N+1 si depende de las predicciones que producirá el batch N.
Para resolverlos usamos CUDA streams, events, buffers por slots y una operación llamada carry-over. A continuación lo vas a ver paso a paso.
Streams y events: la base técnica
Un CUDA stream es una cola ordenada de operaciones GPU. Operaciones en la misma stream son secuenciales, entre streams pueden ejecutarse concurrentemente. El problema habitual es el stream por defecto en PyTorch: es sincronizante. Si usas el stream por defecto, la CPU parecerá bloquearse hasta que la GPU termine todo. Eso anula cualquier intento de paralelizar.
La receta es usar streams no por defecto: una para host-to-device (H2D), otra para cómputo (compute) y otra para device-to-host (D2H). Pero como las streams no se esperan entre sí, introducimos CUDA events.
Un event se graba en una stream con stream.record(event) y otra stream puede hacer stream.wait(event) para no empezar hasta que el event esté marcado. Importante: wait bloquea la stream en la GPU, no la CPU. Eso permite que la CPU encole trabajo y siga trabajando sin esperar.
Pipeline H2D -> compute -> D2H con eventos
El patrón por paso es:
CPU prepara inputs en host (sin stream).
Encola la copia H2D en h2d_stream y graba h2d_done.
compute_stream hace wait(h2d_done) y encola la forward.
Graba compute_done desde compute_stream.
d2h_stream hace wait(compute_done) y encola la D2H.
Finalmente, CPU sincroniza con un d2h_done_event.synchronize() para leer resultados.
Así la CPU no bloquea durante la mayor parte del ciclo; solo hay una sincronización final por batch cuando los outputs llegan al host.
Evitar corrupción de datos: doble slot y pool de memoria
Si reusas el mismo buffer en device para batch N y N+1 y la CPU empieza a sobreescribir, la GPU puede leer datos parciales. La solución práctica:
Tener dos slots de buffers (A y B) y alternar: mientras GPU procesa A, la CPU prepara B.
Esto duplica buffers, pero evita races.
Problema siguiente: si usas CUDA graphs (muy útil para latencia) cada captura queda ligada a direcciones de memoria. Con dos slots necesitarías dos graphs. Si además cada graph aloja memoria distinta, VRAM sube. La solución es usar un pool de memoria compartido: ambos graphs pueden usar el mismo pool siempre y cuando no se ejecuten en paralelo. En la práctica el uso máximo de VRAM queda cercano al mayor uso entre graphs, no su suma.
Carry-over: pasar tokens de N a N+1
Si una request aparece tanto en batch N como en N+1, el token que produce N debe ser input para N+1. Pero al preparar N+1 todavía no tienes ese token. La solución es:
Preparar N+1 con un token placeholder (por ejemplo 0).
Tras la finalización de N y antes de la forward de N+1, realizar la operación carry-over.
El carry-over usa una máscara que indica, por posición, a dónde copiar tokens producidos. En esencia:
Seleccionas tokens a llevar desde outputs de N.
Pones a cero las posiciones que no corresponden.
Truncas/padeas para el tamaño de N+1.
Sumas el tensor resultante a los input ids de N+1 (los placeholders eran cero).
Todas esas operaciones son baratas y se capturan en la CUDA graph para que no añadan overhead en tiempo crítico.
El loop asincrónico completo
Secuencia típica:
Paso 0: cold start. CPU prepara batch 0 en slot A y lo dispatcha.
Paso 1: GPU procesa batch 0 en A; CPU prepara batch 1 en B (evicción, admisión, actualizar KV cache routing, construir carry-over mask).
CPU encola H2D(B), registra y encadena events, y sigue.
GPU ejecuta pipeline: D2H de A, H2D de B, compute B cuando H2D(B) complete.
CPU sincroniza con d2h_done de A, procesa outputs, actualiza estados, arma batch 2 en A, y el ciclo continúa.
Mientras el input de N+1 esté listo en device cuando N termine, la GPU no queda idle entre batches. Normalmente la GPU sigue siendo el cuello de botella, por lo que el CPU alcanza a terminar su trabajo antes de que el GPU termine, haciendo que la superposición sea efectiva.
Resultados prácticos y números
En el experimento reportado:
Síncrono: 300.6 s total, GPU activa 76.0%, 24.0% de tiempo GPU idle.
Asíncrono: 234.5 s total, GPU activa 99.4%, mejora del 22% en tiempo.
No hubo cambios en modelos ni kernels. Fue suficiente coordinar streams, events, doble buffer y carry-over.
Implementación y recomendaciones para producción
Revisa la implementación en transformers (entrada general: continuous_batching.py, código asincrónico en ContinuousBatchingAsyncIOs).
Principales piezas a implementar:
Separar H2D, compute, D2H en streams no por defecto.
Usar events para dependencias entre streams.
Doble slots de buffers y un pool de memoria para CUDA graphs.
Capturar carry-over en las graphs.
Vigila el consumo de VRAM y el tradeoff entre menor latencia y mayor complejidad (más objetos sincronizados). En cargas de generación largas (16K+ tokens) y escenarios RL la ganancia es especialmente relevante.
Si tu pipeline usa frameworks distintos o drivers antiguos, valida que el comportamiento de streams y events sea el mismo: detalles en comportamiento del stream por defecto pueden variar.
Conclusión reflexiva
Desacoplar la preparación CPU de la ejecución GPU cambia la ecuación: ya no se trata solo de agrupar requests sino de coordinar hardware para que trabaje en paralelo. Con CUDA streams, events, buffers en slots y carry-over puedes convertir tiempo perdido en GPU en trabajo útil sin tocar modelos. Es un cambio arquitectural pequeño pero con impacto grande en costos y rendimiento.
