Prefill y Decodificación Desagregados

Para generar tokens de salida a partir de un aviso de entrada, la inferencia de LLM se divide en dos etapas: pre-rellenar y decodificar. El pre-rellenado se ejecuta en los tokens de entrada, poblando los cachés KV, antes de entrar a la etapa de decodificación que genera tokens uno por uno.

Mientras que un solo paso de decodificación típicamente dura decenas de milisegundos, el pre-rellenado toma sustancialmente más tiempo. Si se ejecuta en los mismos dispositivos, mezclar pre-rellenado con decodificación degrada el rendimiento de decodificación. En este artículo exploramos una solución establecida en forma de pre-rellenado y decodificación desagregados, ejecutándolos en dispositivos separados para maximizar tanto el rendimiento de pre-rellenado como las latencias de decodificación.

Rendimiento de Pre-Rellenado vs Decodificación

En un motor de servicio LLM típico, el programador de lotes selecciona solicitudes para procesar en cada paso de ejecución de un modelo. Cuando se ejecuta en un solo dispositivo o nodo, tanto las solicitudes de pre-rellenado como las de decodificación se agrupan juntas. El costo de la atención, que se agrega a lo largo de la longitud de la secuencia, crece tanto para el pre-rellenado como para la decodificación, proporcionalmente a la longitud de las entradas en el caché KV (kv_len). Las solicitudes de decodificación generalmente avanzan un solo token (qo_len=1), a un costo mínimo a través de otras capas que operan independientemente sobre los tokens de una secuencia. Las solicitudes de pre-rellenado avanzan miles o decenas de miles de tokens a un costo significativo a través de capas densas (gran qo_len).

La latencia de un paso de avance está más influenciada por el número de tokens independientes que pasan por capas densas (qo_len) que por el número de tokens recuperados del caché KV durante la atención (kv_len). La atención puede paralelizarse tanto por el número de solicitudes como por el kv_len proporcional a las longitudes de las secuencias, logrando una buena utilización. El pre-rellenado está limitado por el cómputo: debido a que qo_len es alto, los núcleos GEMM pueden asignar bloques suficientes a lo largo de la dimensión M para utilizar completamente las capacidades de cómputo de las GPU modernas. La decodificación está limitada por la memoria: debido a los tamaños de lote típicamente bajos, el número de entradas a lo largo de M suele ser pequeño, suficiente para solo un bloque. Si bien los núcleos Split-K GEMM pueden mejorar la utilización de SM para tamaños de lote de tokens bajos, los cachés y las unidades de multiplicación de matrices típicamente permanecen infrautilizadas.

Cuando se mezclan juntos, los lotes que contienen solicitudes de pre-rellenado incurren en latencias más altas durante el paso de avance, afectando negativamente el rendimiento de decodificación de toda la instancia. Aunque mezclar solicitudes de pre-rellenado con solicitudes de decodificación o emplear pre-rellenado fragmentado puede mejorar ligeramente el rendimiento de decodificación, es difícil mantener un rendimiento de pre-rellenado suficiente para procesar suficientes solicitudes en una instancia para maximizar el rendimiento de decodificación. En el caso de modelos grandes, con longitudes de salida típicas, para mantener un tamaño de lote grande para la decodificación, el pre-rellenado debe realizarse con la frecuencia suficiente como para degradar significativamente la latencia promedio y causar tartamudeo en la salida.

Estos problemas se pueden abordar utilizando un conjunto separado de nodos para realizar el pre-rellenado y la decodificación. Al asociar un nodo de pre-rellenado con múltiples nodos de decodificación, se pueden programar suficientes solicitudes para el pre-rellenado para maximizar el rendimiento y mantener un número suficientemente grande de solicitudes concurrentes en los nodos de decodificación para también maximizar el rendimiento de decodificación. Los nodos de pre-rellenado pueblan los cachés KV, que luego se transfieren a los nodos de decodificación. Dado que los decodificadores ya no tienen que detenerse para el pre-rellenado, las latencias se vuelven mucho más determinísticas, ya que el impacto general del creciente kv_len de solicitudes activas es mucho menos pronunciado. El costo se paga en un aumento en el Tiempo para Primer Token (TTFT), ya que la transferencia de cachés KV a través de la red puede tomar decenas a cientos de milisegundos.

KV Messenger

En Perplexity, nuestra implementación para el pre-rellenado y la decodificación desagregados se basa en un mensajero KV que interactúa con el motor LLM para orquestar las transferencias de cachés KV desde los nodos de pre-rellenado a los nodos de decodificación a través de una red. En el lado del pre-rellenado, el mensajero acepta solicitudes de los nodos de decodificación, enviándolas al programador de lotes y rastreando la ejecución del paso de avance para despachar los cachés KV con la menor latencia posible. En el lado del decodificador, después de que se asignan páginas no expulsables, el mensajero bloquea que la solicitud sea programada para decodificación hasta que se notifique la finalización de las transferencias de cachés KV y contexto del decodificador.

Desagregar el pre-rellenado requiere conexiones de alto rendimiento y baja latencia, por lo tanto, nuestra implementación está adaptada para RDMA, compatible con ambos EFA y Controladores de Interfaz de Red (NICs) ConnectX. El mensajero KV se construye sobre libfabric, utilizando nuestros envoltorios fabric-lib para proporcionar abstracciones de baja latencia de nivel superior sobre las primitivas de Acceso Directo a Memoria Remota (RDMA), implementando transferencias eficientes de páginas y metadatos, junto con señalización de baja latencia. En segundo plano, fabric-lib coordina una GPU y sus NICs directamente conectadas para copiar datos del nodo de pre-rellenado al nodo de decodificación.

Una vez recibido, el nodo de pre-rellenado asigna un conjunto correspondiente de páginas KV fuente y programa la solicitud para pre-rellenado usando su motor local. Para minimizar la latencia, las transferencias no esperan el paso de avance: en cambio, las copias de páginas KV se inician tan pronto como el modelo termina de añadir entradas de caché KV al caché KV para capas individuales. Dado que las solicitudes de pre-rellenado pueden fragmentarse, el programador de lotes notifica al mensajero KV de los fragmentos actualmente programados antes de la ejecución. Para soportar gráficos CUDA mientras se puede rastrear las capas, el mensajero mantiene un hilo dedicado que sondea un contador incrementado después de la proyección de salida de la atención. El contador se mantiene solo en el nodo inicial en un entorno fragmentado: aunque las entradas del caché KV son válidas después de añadir y antes de la atención, la proyección de salida se reduce entre rangos, sincronizándolos implícitamente. Una vez que se observa un cambio en el contador, se notifica al mensajero y llama a fabric-lib para iniciar la transferencia de una capa.

Después de que la transferencia del último fragmento esté completa, cualquier metadato adicional también se copia: la decodificación especulativa o MTP requiere que se muevan logits y estados ocultos al decodificador. Estas copias también se realizan a través de RDMA, hacia y desde búferes pre-asignados.

Tras la finalización de todas las transferencias pendientes del último fragmento, el nodo de pre-rellenado libera las páginas KV y completa la solicitud. El nodo decodificador no se notifica explícitamente: en cambio, utiliza contadores inmediatos para llevar un seguimiento del número de operaciones completadas. El número de operaciones RDMA en el lado del pre-rellenado es proporcional al número de páginas transferidas. Tras la finalización del número conocido de copias de páginas y contexto, fabric-lib llama al mensajero KV para indicar que una solicitud está lista para decodificación. El mensajero libera cualquier contexto y entrega la solicitud al motor LLM.

Transferencias de Cachés KV Fragmentados

Si el pre-rellenado y el decodificador dependen del paralelismo tensorial (TP) y fragmentan o replican las cachés KV de manera idéntica, un único motor de transferencia coordina múltiples dispositivos para enviar y recibir las páginas de todas las réplicas. Para poder utilizar un único mensajero y motor de transferencia a pesar de que el ejecutor del modelo esté replicado en múltiples dispositivos y procesos, se utiliza cuMem y cuMemImportFromShareableHandle para asignar la memoria del dispositivo que respalda las cachés KV y mapearla en el proceso principal. El motor de transferencia inspecciona la topología del nodo para encontrar los NICs y las CPUs en el nodo NUMA más cercano para utilizarlos en las transferencias de cada una de las porciones de caché KV.

Si la fuente y el destino se fragmentan de manera idéntica, las transferencias son triviales ya que hay una asignación uno a uno de los dispositivos y páginas de la fuente y el destino. En esta situación, el fragmentado ayuda implícitamente a las latencias de transferencia: al usar más GPUs, se pueden emplear más NICs asociadas, alcanzando una utilización de ancho de banda casi completa. Sin embargo, si hay una discrepancia, el motor de transferencia debe dividir o reconstruir páginas dependiendo de la relación entre las porciones de fuente y destino.

Si el pre-rellenador divide el caché KV entre más dispositivos, las páginas completas se reconstruyen en el decodificador enviando las mitades correspondientes desde los dispositivos del pre-rellenado. Si el decodificador tiene más fragmentaciones, recibe páginas de múltiples fuentes. El decodificador necesita conocer el esquema de fragmentación del pre-rellenado para poder calcular el número de escrituras RDMA que se espera recibir. Si se involucra replicación, el pre-rellenador agrupa los dispositivos en conjuntos de réplicas que replican el caché KV completo dentro de ellos mismos. Los conjuntos de réplicas de destino se asignan aleatoriamente a uno de los conjuntos de origen para utilizar todos los dispositivos disponibles para iniciar las escrituras RDMA.

Las transferencias fragmentadas requieren un ligero ajuste a las cachés KV. Por defecto, FlashInfer se basa en el diseño NHD, que ordena los tokens dentro de una página dentro de las cabezas. Dado que las cachés son más probablemente fragmentadas a lo largo del número de cabezas de atención, esto crea discontinuidad dentro de la cabeza. Las transferencias RDMA no admiten implícitamente escrituras con escalonamiento, requiriendo una operación por cada cabeza para realizar la transferencia. En cambio, para reducir el número de interacciones con libfabric, organizamos las cachés KV utilizando el diseño HND que coloca la dimensión de la cabeza antes del número de tokens. Esto asegura continuidad, permitiendo que una página sea copiada con una sola escritura.

Decodificación Especulativa

La decodificación especulativa requiere ligeros ajustes al pre-rellenado-decodificación desagregados. En nuestra implementación, los nodos de pre-rellenado no están permitidos a muestrear tokens. Dado que los modelos Sonar de Perplexity admiten salida estructurada, no queremos incurrir en la complejidad de sincronizar las implementaciones del procesador de esquemas entre pre-rellenadores y decodificadores. En los mecanismos de MTP y decodificación especulativa, el pre-rellenado del modelo preliminar hasta el último token implica muestrear tokens del modelo objetivo.

Para resolver estos problemas, el pre-rellenado no incluye el último token de la secuencia de entrada. En cambio, se transfieren estados ocultos o logits desde el pre-rellenado previo al último token y se trata como un token de decodificación en el siguiente paso en el decodificador. Aunque esto incrementa ligeramente las latencias, ya que debe realizarse un paso completo de decodificación después del pre-rellenado para emitir el primer token, la complejidad de la implementación se reduce considerablemente.

Despliegues Desagregados

Hemos desplegado o experimentado con múltiples configuraciones desagregadas con diferentes modelos, para soportar tráfico de producción o cargas de trabajo de evaluación interna. Basándonos en el tamaño y el mecanismo de atención de los modelos, elegimos esquemas de fragmentación adecuados para nodos de pre-rellenado y decodificación para utilizar mejor las GPUs.

DeepSeek-R1

Con DeepSeek, consideramos despliegues tanto de Tensor-Parallel (TP) como Data-Parallel (DP). Como se discutió en publicaciones anteriores del blog, los despliegues TP proporcionan mejor latencia al costo de menor rendimiento, requiriendo más GPUs para servir tráfico pesado. Los despliegues DP escalan mucho mejor con la carga, sin embargo, su rendimiento máximo es menor debido al costo de la comunicación entre dispositivos o nodos.

DeepSeek se basa en la Atención Latente Multi-Cabeza, comprimiendo los cachés KV. Dado que todas las cabezas KV se comprimen en un solo vector latente, TP no puede fragmentar los cachés KV, ya que debe replicar los vectores latentes en todas las posiciones. La fragmentación ocurre después de la descompresión, ya que cada posición puede extraer diferentes cabezas de la misma representación latente. En consecuencia, todas las fragmentaciones de cachés KV son idénticas tanto en fragmentaciones de pre-rellenado como decodificador.

Con una configuración TP intra-nodo, tanto los pre-rellenadores como los decodificadores se fragmentan de manera idéntica. Las transferencias se despachan desde todas las posiciones para utilizar completamente todos los NICs disponibles. Sin embargo, con un despliegue DP, donde el tamaño de rango TP es menor o cada rango DP se asigna a una sola GPU, cualquier dispositivo de pre-rellenado que tenga una copia replicada del caché KV puede despacharlo. Para equilibrar las solicitudes entre todos los NICs disponibles, seleccionamos aleatoriamente una GPU y un NIC para enviar el caché KV del pre-rellenado al decodificador.

Con el pre-rellenado-decodificación mixto, nuestro despliegue R1 estaba luchando para exceder consistentemente 50 TPS debido a interrupciones frecuentes de pre-rellenado en el orden de cientos de milisegundos. En contraste, al separar el pre-rellenado, incurrió una penalización de aproximadamente 100 ms en el TTFT para cada solicitud, pero un solo nodo de pre-rellenado podría mantener tamaños de lote consistentes en 3 nodos decodificadores, con un rendimiento superior a 90 TPS mientras manejaba una carga de aproximadamente 1 QPS por nodo decodificador. Con los despliegues DP, TPS fue ligeramente menor, alrededor de 50, sin embargo, las instancias podrían manejar una carga de 1 QPS por rango, con 8 rangos a un solo nodo.

Qwen3-Coder

Este modelo de 480B utiliza Atención de Consulta Agrupada (GQA), por lo que la atención puede fragmentarse fácilmente y puede beneficiarse del paralelismo tensorial sin sacrificar la memoria para los cachés KV. En consecuencia, pudimos fragmentar el modelo en 8 GPUs tanto para pre-rellenado como para decodificación, emparejando alrededor de 3 nodos decodificadores con un solo nodo de pre-rellenado. Dado que la atención está fragmentada, confiamos en el diseño de caché KV HND para fragmentar los cachés KV de pre-rellenador y decodificador, emparejando rangos de pre-rellenador con rangos de decodificador y utilizando todos los NICs para transferir porciones en paralelo.