Prefill y decode desagregados

Para generar tokens de salida a partir de una instrucción de entrada, la inferencia de LLM se divide en dos etapas: prefill y decode. La etapa de prefill se ejecuta sobre los tokens de entrada, completando las cachés KV, antes de entrar en la etapa de decode que genera tokens uno a uno. Si bien un paso de decode típicamente se ejecuta durante decenas de milisegundos, el prefill tarda sustancialmente más. Si se ejecutan en los mismos dispositivos, mezclar el prefill con el decode degrada el rendimiento del decode. En este artículo exploramos una solución establecida consistente en el prefill y decode desagregados, ejecutándolos en dispositivos separados para maximizar tanto el rendimiento del prefill como las latencias del decode.

Rendimiento de Prefill frente a Decode

En un motor de servicio de LLM típico, el planificador de lotes selecciona las solicitudes a procesar en cada paso de ejecución de un modelo. Al ejecutarse en un solo dispositivo o nodo, tanto las solicitudes de prefill como las de decode se procesan por lotes conjuntamente. El costo de la atención, que se agrega a lo largo de la longitud de la secuencia, aumenta tanto para el prefill como para el decode, de forma proporcional a la longitud de las entradas en la caché KV (kv_len). Las solicitudes de decode suelen reenviar un solo token (qo_len=1), con un costo mínimo a través de otras capas que operan de forma independiente sobre los tokens de una secuencia. Las solicitudes de prefill reenvían miles o decenas de miles de tokens con un costo significativo a través de capas densas (gran qo_len).

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

Ampliar

Cuando se mezclan, los lotes que contienen solicitudes de prefill incurren en latencias más altas a través del paso de avance, lo que afecta negativamente el rendimiento de decode de toda la instancia. Si bien mezclar solicitudes de prefill con solicitudes de decode o emplear prefill fragmentado puede mejorar ligeramente el rendimiento de decode, es difícil mantener un rendimiento de prefill suficiente para procesar suficientes solicitudes en una instancia para maximizar el rendimiento de decode. En el caso de modelos grandes, con longitudes de salida típicas, para mantener un tamaño de lote grande para el decode, el prefill debe realizarse con la frecuencia suficiente para que degrade significativamente la latencia promedio y provoque interrupciones en la salida.

Ampliar

Ampliar

Estos problemas pueden abordarse utilizando un conjunto separado de nodos para realizar el prefill y el decode. Al asociar un nodo de prefiller con múltiples nodos de decoder, se pueden programar suficientes solicitudes para el prefill a fin de maximizar el rendimiento y mantener un número suficientemente grande de solicitudes simultáneas en los nodos de decoder para también maximizar el rendimiento de decode. Los nodos de prefiller completan las cachés KV, que luego se transfieren a los nodos de decoder. Dado que los decoders ya no tienen que detenerse para el prefill, las latencias se vuelven mucho más deterministas, ya que el impacto general del crecimiento de kv_len de las solicitudes activas es mucho menos pronunciado. El costo se paga en un aumento en el Tiempo hasta el Primer Token (TTFT), ya que la transferencia de cachés KV a través de la red puede tomar de decenas a cientos de milisegundos.

KV Messenger

En Perplexity, nuestra implementación para el prefill y decode desagregados se construye en torno a un KV messenger que interactúa con el motor de LLM para orquestar las transferencias de caché KV desde los nodos de prefiller a los nodos de decoder a través de una red. En el lado del prefiller, el messenger acepta solicitudes de los nodos de decoder, entregándolas al planificador de lotes y realizando un seguimiento de la ejecución del paso de avance para enviar las cachés KV con la menor latencia posible. En el lado del decoder, después de que se asignan las páginas que no se pueden desalojar, el messenger bloquea la solicitud de ser programada para el decode hasta que se notifique la finalización de las transferencias de la caché KV y del contexto del decoder.

Ampliar

La desagregación del prefill requiere conexiones de alto rendimiento y baja latencia, por lo que nuestra implementación está adaptada para RDMA, admitiendo tanto EFA como controladores de interfaz de red (NIC) ConnectX. El KV Messenger está construido sobre libfabric, utilizando nuestros envoltorios fabric-lib para proporcionar abstracciones de baja latencia de mayor nivel 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 NIC conectados directamente para copiar datos desde el nodo de prefiller al nodo de decoder.

Al recibir, el nodo de prefiller asigna un conjunto correspondiente de páginas KV de origen y programa la solicitud para el prefill utilizando su motor local. Para minimizar la latencia, las transferencias no esperan al paso de avance: en su lugar, las copias de páginas KV se inician tan pronto como el modelo termina de agregar entradas de caché KV a la caché KV para capas individuales. Dado que las solicitudes de prefill pueden fragmentarse, el planificador de lotes notifica al KV messenger sobre los fragmentos programados actualmente antes de la ejecución. Para admitir gráficos CUDA mientras se puede realizar un seguimiento de las capas, el messenger 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 principal en un entorno fragmentado: aunque las entradas de caché KV son válidas después de la adición y antes de la atención, la proyección de salida se reduce a través de rangos, sincronizándolos implícitamente. Una vez que se observa un cambio en el contador, se notifica al messenger y este llama a fabric-lib para iniciar la transferencia de una capa.

Ampliar

Después de que se completa la transferencia del último fragmento, también se copian los metadatos adicionales: el decodificado especulativo o MTP requieren que los logits y los estados ocultos se muevan al decoder. Estas copias también se realizan a través de RDMA, hacia y desde búferes preasignados.

Tras la finalización de todas las transferencias pendientes del último fragmento, el nodo de prefiller desasigna las páginas KV y completa la solicitud. El nodo de decoder no recibe una notificación explícita: en su lugar, utiliza contadores inmediatos para realizar un seguimiento del número de operaciones completadas. El número de operaciones RDMA en el lado del prefiller es proporcional al número de páginas transferidas. Tras la finalización del número conocido de copias de páginas y de contexto, fabric-lib llama al KV messenger para indicar que una solicitud está lista para el decodificado. El messenger desasigna cualquier contexto y entrega la solicitud al motor de LLM.

Transferencias de caché KV fragmentadas

Si el prefiller y el decoder dependen del Paralelismo de Tensores (TP) y fragmentan o replican las cachés KV de forma 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 messenger y motor de transferencia a pesar de que el ejecutor del modelo se replica a través de múltiples dispositivos y procesos, se utilizan cuMem y cuMemImportFromShareableHandle para asignar la memoria del dispositivo que respalda las cachés KV y para asignarla al proceso principal. El motor de transferencia inspecciona la topología del nodo para encontrar los NIC y las CPU en el nodo NUMA más cercano para usar en las transferencias de cada una de las porciones de la caché KV.

Si el origen y el destino se fragmentan de forma idéntica, las transferencias son triviales ya que existe una asignación uno a uno desde los dispositivos y páginas del origen y el destino. En esta situación, la fragmentación ayuda implícitamente a las latencias de transferencia: al usar más GPU, se pueden emplear más NIC asociados, acercándose a la utilización completa del ancho de banda. Sin embargo, si hay una discrepancia, el motor de transferencia debe dividir o reconstruir las páginas dependiendo de la proporción entre las porciones de origen y destino.

Ampliar

Si el prefiller divide la caché KV entre más dispositivos, las páginas completas se reconstruyen en el decoder enviando las mitades correspondientes desde los dispositivos del prefiller. Si el decoder tiene más fragmentos, recibe páginas de múltiples fuentes. El decoder necesita conocer el esquema de fragmentación del prefiller para poder calcular el número de escrituras RDMA que se espera que reciba. Si hay replicación, el prefiller agrupa los dispositivos en conjuntos de réplicas que replican la caché KV completa dentro de sí 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.

Ampliar

Las transferencias fragmentadas requieren un ligero ajuste en las cachés KV. De forma predeterminada, FlashInfer depende del diseño NHD, que ordena los tokens dentro de una página dentro de los cabezales. Dado que las cachés probablemente se fragmenten a lo largo del número de cabezales de atención, esto crea una discontinuidad dentro del cabezal. Las transferencias RDMA no admiten implícitamente escrituras escalonadas, lo que requiere una operación por cabezal 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 del cabezal antes del número de tokens. Esto garantiza la continuidad, permitiendo copiar una página con una sola escritura.

Decodificado especulativo

El decodificado especulativo requiere ligeros ajustes al prefill-decode desagregado. En nuestra implementación, no se permite que los nodos de prefiller muestreen 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 prefillers y decoders. En los mecanismos MTP y decodificado especulativo, realizar el prefill del modelo borrador hasta el último token implica muestrear tokens del modelo objetivo.

Ampliar

Para evitar estos problemas, el prefill no incluye el último token de la secuencia de entrada. En su lugar, se transfieren los estados ocultos o logits del prefill anterior al último token y se trata como un token de decode en el siguiente paso en el decoder. Si bien esto aumenta ligeramente las latencias, ya que se debe realizar un paso de decode completo después del prefill para emitir el primer token, la complejidad de la implementación se reduce enormemente.

Implementaciones desagregadas

Hemos implementado o experimentado con múltiples configuraciones desagregadas con diferentes modelos, para admitir el tráfico de producción o cargas de trabajo de evaluación internas. Según el tamaño y el mecanismo de atención de los modelos, elegimos esquemas de fragmentación adecuados para los nodos de prefiller y decoder a fin de utilizar mejor las GPU.

DeepSeek-R1

Con DeepSeek, consideramos implementaciones de Paralelismo de Tensores (TP) y Paralelismo de Datos (DP). Como se discutió en publicaciones anteriores del blog, las implementaciones de TP brindan una mejor latencia a costa de un menor rendimiento, lo que requiere más GPU para atender un tráfico intenso. Las implementaciones de 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 entre nodos.

DeepSeek depende de la Atención Latente de Múltiples Cabezales, comprimiendo las cachés KV. Dado que todos los cabezales KV se comprimen en un único vector latente, el TP no puede fragmentar las cachés KV, ya que en su lugar debe replicar los vectores latentes en todos los rangos. La fragmentación ocurre después de la descompresión, ya que cada rango puede extraer diferentes cabezales de la misma representación latente. En consecuencia, todos los fragmentos de caché KV son idénticos en ambos fragmentos de prefiller y decoder.

Con una configuración de TP intranodo, tanto los prefillers como los decoders se fragmentan de forma idéntica. Las transferencias se envían desde todos los rangos para utilizar plenamente todos los NIC disponibles. Sin embargo, con una implementación de DP, donde el tamaño del rango de TP es menor o cada rango de DP se asigna a una sola GPU, cualquier dispositivo de prefiller que contenga una copia replicada de la caché KV puede enviarla. Para equilibrar las solicitudes entre todos los NIC disponibles, seleccionamos aleatoriamente una GPU y un NIC para enviar la caché KV desde el prefiller al decoder.

Con prefill-decode mixto, nuestra implementación de R1 luchaba por superar constantemente los 50 TPS debido a las frecuentes interrupciones de prefill del orden de cientos de milisegundos. Por el contrario, al separar el prefill, incurrimos en una penalización de aproximadamente 100 ms en el TTFT para cada solicitud, pero un solo nodo de prefiller podía mantener tamaños de lote consistentes en 3 nodos de decoder, entregando un rendimiento superior a 90 TPS mientras manejaba una carga de aproximadamente 1 QPS por nodo de decoder. Con implementaciones de paralelismo de datos, los TPS fueron ligeramente más bajos, alrededor de 50; sin embargo, las instancias podían manejar una carga de 1 QPS por rango, con 8 rangos en un solo nodo.

Qwen3-Coder

Este modelo de 480B utiliza Atención de Consulta Agrupada (GQA), por lo que la atención se puede fragmentar fácilmente y puede beneficiarse del paralelismo de tensores sin sacrificar memoria para las cachés KV. En consecuencia, pudimos fragmentar el modelo en 8 GPU tanto para el prefill como para el decode, emparejando alrededor de 3 nodos de decoder con un solo nodo de prefiller. Dado que la atención está fragmentada, confiamos en el diseño de caché KV HND para fragmentar las cachés KV del prefiller y del decoder, emparejando los rangos del prefiller con los rangos del decoder y utilizando plenamente todos los NIC para transferir porciones en paralelo.

Citación

@online{perplexity2025DisaggregatedPrefillAndDecode,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Kevin Hu and Erin McKnight and Phi Hoang},
  title = {{Disaggregated prefill and decode}},
  date = {2025-08-01},
  url = {https://perplexity.coolkeypoint.com/hub/blog/disaggregated-prefill-and-decode}
}