Resumen Técnico: Las Solicitudes de una Pluma Deben Agruparse

Declaración del Problema

La inferencia de Modelos de Lenguaje Grande (LLM) está fundamentalmente limitada por la memoria durante la fase de decodificación auto-regresiva. Cada token generado requiere atender la caché de Clave-Valor (KV) de todos los tokens anteriores, lo que hace necesaria una exploración completa de la caché KV en la memoria de la GPU en cada paso. Aunque trabajos anteriores se han centrado en maximizar el tamaño del lote para mejorar la utilización de la GPU, este artículo identifica un compromiso crítico y pasado por alto: la homogeneidad de los prefijos.

En cargas de trabajo del mundo real (por ejemplo, indicaciones con pocos ejemplos, indicaciones del sistema, RAG), las solicitudes a menudo comparten prefijos significativos. Los motores de inferencia existentes (como vLLM y SGLang) utilizan árboles radices o estructuras de hash para detectar estos prefijos compartidos y reutilizar las cachés KV, principalmente para reducir la huella de memoria. Sin embargo, los planificadores actuales priorizan la maximización del tamaño del lote o la minimización de las expulsiones de memoria sin considerar el impacto de la heterogeneidad de los prefijos en los patrones de acceso a la memoria.

Los autores demuestran que los lotes heterogéneos (que contienen solicitudes con diferentes prefijos) interrumpen la localidad espacial y temporal en los accesos a la DRAM de la GPU, lo que lleva a una mala utilización del ancho de banda y a una reducción del rendimiento. Por el contrario, los lotes homogéneos en prefijos (donde todas las solicitudes comparten un prefijo común) permiten una recuperación anticipada eficiente por parte del hardware y la reutilización de caché. El artículo argumenta que los planificadores existentes no reconocen que un lote moderadamente pequeño y homogéneo puede superar a uno más grande y heterogéneo. Además, los mecanismos utilizados para detectar el compartir prefijos (por ejemplo, recorridos de árboles radices) incurren en una sobrecarga de CPU prohibitiva, a menudo comparable o superior al tiempo de ejecución de la GPU.

Metodología

El artículo introduce Feather, un planificador consciente de los prefijos diseñado para optimizar el compromiso entre el tamaño del lote y la homogeneidad de los prefijos. Feather integra dos componentes principales:

1. Árbol de Hash Fragmentado (CHT)

Para abordar la alta sobrecarga de CPU de los recorridos tradicionales de árboles radices, Feather emplea una estructura de datos ligera llamada Árbol de Hash Fragmentado.

• Mecanismo: En lugar de comparaciones token a token, las solicitudes se representan como vectores de hashes acumulativos sobre fragmentos de tokens de tamaño fijo.
• Eficiencia: El CHT mantiene un "conjunto de trabajo" de hashes para el lote activo y un montículo mínimo de solicitudes en espera ordenadas por su "conteo faltante" (el número de fragmentos aún no cubiertos por el lote activo).
• Complejidad: Esto permite al planificador identificar la mejor solicitud candidata para agregar a un lote en un tiempo de $O(\log W)$ (donde $W$ es el número de solicitudes en espera), una mejora significativa sobre la complejidad de $O(W \times T)$ o $O(E)$ de los recorridos de árboles radices en profundidad.
• Operaciones: El CHT admite una inserción eficiente, la búsqueda de la mejor solicitud, la adición de una solicitud al lote activo y la eliminación de solicitudes completadas, utilizando actualizaciones perezosas del montículo y mantenimiento incremental de la punta para minimizar la sobrecarga.

2. Política de Agrupación por Aprendizaje por Refuerzo (RL)

Feather utiliza una política basada en RL para determinar el punto óptimo de parada para la formación de lotes.

• El Compromiso: Agregar más solicitudes aumenta la utilización de la GPU (tamaño de lote más grande) pero puede diluir la homogeneidad de los prefijos (reduciendo la localidad de memoria).
• Estado y Acción: El planificador observa el tamaño actual del lote, la pérdida en la profundidad del prefijo compartido ($\Delta$) y el número de solicitudes en espera que comparten la nueva profundidad de prefijo ($w$). Decide si Agregar otra solicitud o Detenerse y despachar el lote.
• Políticas: Los autores evalúan tres políticas: una heurística estática, un Bandido Contextual (UCB) y Q-Learning. El objetivo es maximizar la recompensa, definida como el rendimiento de decodificación.
• Adaptabilidad: La política de RL aprende a formar lotes que sean "justos": lo suficientemente grandes para utilizar la GPU pero lo suficientemente homogéneos para maximizar la localidad de memoria. Se adapta dinámicamente a la intensidad de la carga de trabajo y a los patrones de compartir prefijos.

Contribuciones Clave

1. Identificación de un Compromiso Crítico: El artículo establece que la homogeneidad de los prefijos es un determinante más fuerte del rendimiento de decodificación que el tamaño del lote por sí solo. Demuestra que incluso una sola solicitud divergente en un lote puede interrumpir la localidad y reducir el rendimiento en casi 2 veces.
2. Cuantificación de la Sobrecarga de CPU: Los autores revelan que los mecanismos de detección de prefijos más avanzados (como los recorridos de árboles radices en SGLang) incurren en sobrecargas de CPU que pueden alcanzar el 50–90% de la latencia total, negando efectivamente los beneficios de la reutilización de prefijos.
3. Planificador Feather: El diseño e implementación de Feather, que integra la estructura de datos CHT con una política de agrupación basada en RL. Feather se integra tanto en vLLM como en SGLang.
4. Agnosticismo de Hardware: A diferencia de los kernels de atención conscientes de los prefijos que requieren un ajuste específico de hardware, Feather opera en la capa de planificación, lo que lo hace compatible con los kernels de atención existentes y las configuraciones de hardware.

Resultados

La evaluación se realizó en GPUs NVIDIA RTX 6000 Ada y A100 utilizando modelos como Llama 3 8B, Qwen y LongChat 13B en diversas cargas de trabajo (L-Eval, LongBench).

• Ganancias de Rendimiento: Feather logra un rendimiento de extremo a extremo 2–10 veces mayor en comparación con los planificadores existentes (FCFS de vLLM, LPM/DFS-W de SGLang) en cargas de trabajo con prefijos compartidos que van desde 1K hasta 10K tokens.
• Sensibilidad a los Prefijos: En experimentos con 800 solicitudes divididas en 8 grupos de prefijos, procesarlos en 8 lotes homogéneos de 100 solicitudes produjo un rendimiento casi 2 veces mayor que procesarlos en lotes más grandes y heterogéneos.
• Reducción de Sobrecarga: El CHT de Feather reduce la sobrecarga de detección de prefijos hasta en 1000 veces en comparación con los enfoques basados en DFS para secuencias largas. La sobrecarga de planificación se mantiene por debajo del 1% del tiempo de ejecución de la GPU, mientras que los enfoques de referencia pueden exceder el tiempo de ejecución de la GPU.
• Robustez: Feather rinde a la par con el estado actual (FCFS) cuando las cargas de trabajo carecen de compartir prefijos, asegurando que no haya degradación del rendimiento en escenarios no compartidos.
• Comparación con Kernels: Feather supera o iguala a kernels de atención especializados conscientes de los prefijos (como PAT) a pesar de no modificar el kernel en sí, demostrando que las optimizaciones a nivel de planificación pueden ser tan efectivas como las de nivel de kernel.

Significancia

El artículo afirma que maximizar el tamaño del lote ya no es suficiente para un servicio eficiente de LLM. El cuello de botella principal en las fases de decodificación modernas es el ancho de banda de memoria, que está fuertemente influenciado por la localidad de acceso. Al cambiar el enfoque de "maximizar el tamaño del lote" a "optimizar la homogeneidad de los prefijos", Feather desbloquea ganancias significativas de rendimiento sin requerir modificaciones de kernel específicas de hardware.

El trabajo destaca que el costo de encontrar prefijos compartidos (sobrecarga de CPU) ha sido un cuello de botella oculto en los sistemas actuales. El CHT ligero de Feather y su política de RL adaptativa demuestran que es posible equilibrar dinámicamente la utilización de cómputo y la localidad de memoria, lo que conduce a mejoras sustanciales tanto en el rendimiento como en la latencia para cargas de trabajo de inferencia de LLM en el mundo real. Los autores planean hacer de código abierto Feather contribuyéndolo a los repositorios de vLLM y SGLang.