La transferencia KV basada en push de vLLM reduce la latencia en inferencia desagregada

Amazon Web Services ha introducido un mecanismo de transferencia de clave-valor (KV) basado en push para el motor de inferencia open-source vLLM. Esta actualización se dirige a arquitecturas de prefill-decode desagregadas, con el objetivo de reducir significativamente el tiempo hasta el primer token (TTFT) y mejorar el rendimiento durante la generación de tokens. La funcionalidad fue validada en instancias Amazon EC2 P5en utilizando el modelo Qwen3-32B.

Qué ocurrió

La inferencia de modelos de lenguaje grandes consta de dos fases distintas: prefill, que procesa el prompt de entrada para construir la caché KV, y decode, que genera los tokens de salida utilizando dicha caché. En el servicio desagregado, estas fases se ejecutan en instancias de motor separadas para permitir el escalado y la optimización independientes. Sin embargo, esta separación requiere transferir la caché KV desde la instancia de prefill a la instancia de decode antes de que pueda comenzar la generación, lo que hace que la eficiencia de esta transferencia sea crítica para la latencia general.

Anteriormente, vLLM utilizaba un método basado en pull mediante el conector NVIDIA Inference Xfer Library (NIXL). En este modo predeterminado, la instancia de decode espera a que termine el prefill, recibe metadatos de un router, asigna memoria y luego extrae los datos. Estos pasos ocurren secuencialmente después de completar el prefill, añadiendo una demora significativa al camino crítico. El nuevo conector basado en push cambia este flujo superponiendo las tareas de coordinación con el propio cálculo del prefill.

Los ingenieros de AWS desarrollaron y probaron esta funcionalidad en clústeres conectados con AWS Elastic Fabric Adapter (EFA), que soporta Acceso Directo a Memoria Remota (RDMA). Aunque inicialmente se construyó para instancias AWS Trainium, la contribución es agnóstica respecto al acelerador y funciona también con GPUs. Las pruebas de referencia muestran que bajo alta carga, el modo push puede reducir el TTFT hasta tres veces en comparación con el modo pull, mejorando además el tiempo por token de salida (TPOT) aproximadamente un 30 por ciento.

Cómo funciona

La arquitectura basada en push elimina los cuellos de botella de serialización permitiendo que la instancia de decode registre sus bloques de memoria destino con antelación. En lugar de esperar a que termine el prefill para prepararse para recibir los datos, el lado de decode envía una notificación de registro a la instancia de prefill tan pronto como asigna los bloques locales. Esto significa que el destino es conocido antes de que los datos fuente estén listos.

En el lado de prefill, un hilo de fondo dedicado gestiona la preparación de la transferencia. Una vez finalizado el cálculo del prefill y llega el registro desde el decodificador, el bucle principal del motor entrega los bloques KV a este hilo de fondo. El hilo establece la conexión NIXL, mapea los ID de bloque y emite una operación NIXL WRITE para copiar los datos directamente en la memoria registrada del decodificador. Esto mantiene la carga pesada de preparación de descriptores y envío de transferencias fuera de los bucles principales críticos en latencia de ambas instancias.

El sistema utiliza un protocolo de encuentro (rendezvous) para manejar las diferencias de temporización. Si el registro llega antes de que termine el prefill, el productor lo retiene. Si el prefill termina primero, retiene los bloques hasta que llega el registro. El evento que ocurra en segundo lugar dispara la transferencia inmediata. El decodificador cuenta las notificaciones de finalización de cada rango productor y solo procede a la generación de tokens una vez que han llegado todos los datos esperados.

Detalles clave

• Ganancia de rendimiento: Reducción de hasta 3x en TTFT y mejora del 30% en TPOT para Qwen3-32B en despliegues 1P1D a altas tasas de solicitud.
• Hardware probado: Instancias Amazon EC2 P5en con EFA, utilizando GPUDirect RDMA para la transferencia directa de memoria GPU a GPU.
• Mecanismo: Utiliza operaciones NIXL WRITE iniciadas por la instancia de prefill, en lugar de NIXL READ iniciadas por la instancia de decode.
• Multihilo: Un hilo de fondo en el worker de prefill gestiona la configuración de la transferencia, evitando estancamientos en el bucle principal de ejecución del modelo.
• Compensación (Trade-off): El modo push requiere seleccionar la instancia de decode más temprano en el proceso, lo que puede reducir la flexibilidad para el enrutamiento consciente de la carga en comparación con el modo pull.
• Configuración: Los usuarios habilitan esto configurando tanto los servidores de prefill como los de decode para usar NixlPushConnector en lugar del predeterminado NixlConnector.

Por qué importa

Para los ingenieros que construyen servicios LLM de alto rendimiento, el servicio desagregado es una estrategia común para gestionar el uso de recursos eficientemente. Sin embargo, la sobrecarga de mover grandes cachés KV entre nodos puede anular los beneficios si no se maneja cuidadosamente. El modelo tradicional de pull introduce un periodo de espera síncrono donde la instancia de decode está inactiva, esperando metadatos y luego extrayendo datos. Esta demora se vuelve pronunciada a medida que crecen las longitudes de contexto y aumentan las tasas de solicitud, impactando directamente en la latencia percibida por el usuario.

Al trasladar el trabajo de coordinación a un hilo de fondo y superponerlo con el cálculo del prefill, el modelo push asegura que la instancia de decode esté lista para recibir datos en el momento en que estén disponibles. Esto reduce el tiempo de inactividad en el camino crítico, conduciendo a tiempos de respuesta más rápidos para los usuarios. También evita que el worker de decode se atasque con tareas de configuración de transferencia, permitiéndole mantener velocidades consistentes de generación de tokens incluso bajo alta concurrencia.

Este desarrollo destaca la importancia de las optimizaciones de infraestructura de bajo nivel en el servicio de IA. A medida que los modelos crecen y los contextos se alargan, la eficiencia del movimiento de datos entre nodos de cómputo se vuelve tan importante como la velocidad de las propias multiplicaciones de matrices. Los ingenieros que desplieguen vLLM en AWS o infraestructuras similares habilitadas para RDMA deberían considerar esta opción para maximizar el rendimiento.

Qué puedes hacer

• Evalúa tu despliegue actual de vLLM para ver si el servicio de prefill-decode desagregado es apropiado para tu carga de trabajo.
• Prueba el NixlPushConnector en un entorno de staging con patrones de tráfico realistas para medir las mejoras en TTFT.
• Asegúrate de que tu infraestructura soporte RDMA, como AWS EFA, para aprovechar las funciones de acceso directo a memoria.
• Monitoriza la compensación entre la selección temprana de la instancia de decode y las necesidades de balanceo de carga dinámico en tu configuración de router.
• Deshabilita el caché de prefijos durante las pruebas de referencia para aislar el rendimiento de la transferencia, ya que el caché puede ocultar los costes de transferencia.
• Revisa la documentación de vLLM para las banderas de configuración específicas requeridas para cambiar de modo pull a modo push.