Resumen Técnico: Aprovisionamiento Analítico para la Atención–FFN Desagregada en la Atención de LLM bajo Cargas de Trabajo Estocásticas

Declaración del Problema

La atención de Modelos de Lenguaje Grande (LLM) enfrenta desafíos significativos de eficiencia a medida que aumentan los tamaños de los modelos. Si bien la Desagregación de Atención–FFN (AFD) ha surgido como una arquitectura prometedora para desacoplar las capas de Atención limitadas por memoria (con estado, dominadas por la caché KV) de las capas FFN limitadas por cómputo (sin estado), determinar la proporción óptima de aprovisionamiento de recursos sigue siendo un problema abierto.

En una topología AFD (denotada como $r$A–1F, donde $r$ instancias de Atención alimentan una instancia de FFN), el rendimiento del sistema es altamente sensible a la relación $r$.

• Si $r$ es demasiado pequeña, el servidor FFN se queda sin entrada, lo que lleva a recursos de cómputo inactivos.
• Si $r$ es demasiado grande, las instancias de Atención se bloquean esperando al servidor FFN compartido, causando bloqueos a nivel de paso y burbujas en la tubería.

Los sistemas existentes suelen determinar $r$ mediante búsqueda empírica o aproximaciones deterministas ingenuas que ignoran la estocasticidad de la carga de trabajo. La dificultad central radica en la naturaleza dinámica de la carga de trabajo de Atención: a diferencia del cómputo FFN, que depende únicamente del tamaño del lote, la latencia de Atención evoluciona continuamente a medida que crecen las cachés KV y las solicitudes se reabastecen con longitudes de prompt y de decodificación aleatorias. Además, la ejecución sincronizada de múltiples trabajadores de Atención introduce un efecto de "retrasado" (straggler) gobernado por el trabajador más lento, creando una barrera de sincronización que es difícil de modelar analíticamente.

Metodología

Los autores desarrollan un marco probabilístico riguroso para derivar una regla de aprovisionamiento de forma cerrada para la relación óptima A/F ($r^*$) bajo cargas de trabajo estocásticas. La metodología aborda dos fuentes principales de aleatoriedad:

1. Evolución de la Carga de Trabajo por Ranura: Los autores modelan la carga de tokens por ranura utilizando un análisis de renovación-recompensa en tiempo discreto. Identifican una única estadística de carga de trabajo, $\theta$ (la carga de tokens estacionaria esperada por ranura), que gobierna el aprovisionamiento. Esta estadística cuenta con el efecto de "sesgo de longitud", donde las solicitudes más largas se muestrean con mayor frecuencia en pasos de decodificación aleatorios. Crucialmente, $\theta$ y su varianza $\nu^2$ pueden estimarse de forma no paramétrica a partir de trazas de solicitudes sin asumir familias de distribuciones específicas.
2. Sincronización entre Trabajadores: Para abordar el efecto de barrera donde el sistema espera al más lento de $r$ trabajadores, los autores aplican una corrección de estadística de orden gaussiana. Modelan la carga total en los $r$ trabajadores como una suma de variables aleatorias i.i.d. y aproximan la carga máxima utilizando el valor esperado del máximo de $r$ variables normales estándar ($\kappa_r$).

El análisis produce dos reglas de aprovisionamiento:

• Regla de Campo Medio ($r^*_{mf}$): Una solución de forma cerrada derivada ignorando la barrera de sincronización (asumiendo un equilibrio de carga perfecto). Esta regla se descompone en tres regímenes: cuello de botella de Atención, cuello de botella de comunicación y cuello de botella de FFN.
• Refinamiento Consciente de la Barrera ($r^*_G$): Una optimización refinada que incorpora el término de sobrecarga de sincronización ($\sqrt{B}\nu\kappa_r$) en la aproximación del tiempo de ciclo. Esto se resuelve mediante una optimización analítica unidimensional combinada con una búsqueda discreta.

Contribuciones Clave

1. Modelo de Carga de Trabajo Probabilístico: El artículo establece un marco analítico que captura la dinámica estocástica de la atención AFD con agrupación continua (continuous batching). Identifica $\theta$ como la estadística gobernante para el aprovisionamiento, válida bajo distribuciones arbitrarias de longitud de prompt y decodificación, y proporciona un estimador no paramétrico para $\theta$ y $\nu^2$ a partir de trazas.
2. Regla de Aprovisionamiento de Forma Cerrada: Los autores derivan una relación óptima de campo medio que equilibra explícitamente los cuellos de botella de Atención, comunicación y FFN. Además, proporcionan un refinamiento consciente de la barrera que cuantifica la sobrecarga de sincronización introducida por los trabajadores retrasados, mostrando que esta sobrecarga crece sub-logarítmicamente con $r$.
3. Validación Calibrada con Trazas: El equipo desarrolló un simulador AFD calibrado con trazas del mundo real y parámetros de hardware (específicamente NPUs Huawei Ascend 910C). Validaron el marco a través de diversas cargas de trabajo y configuraciones de hardware.

Resultados

La validación experimental demuestra una alta precisión en la predicción del comportamiento del sistema:

• Precisión de la Relación Óptima: La relación óptima predicha ($r^*_{mf} \approx 9.3$) coincide estrechamente con la relación óptima de simulación. La curva de rendimiento teórica sigue bien los resultados de la simulación en todas las configuraciones probadas ($r \in \{1, 2, 4, 8, 16, 24, 32\}$).
• Márgenes de Error: La relación óptima predicha coincide con la óptima de simulación dentro de un 10%.
• Impacto de los Retrasados: Se observó una brecha residual de hasta ~15% en la predicción del rendimiento entre el modelo de campo medio y la simulación a valores grandes de $r$. El análisis consciente de la barrera explica esta discrepancia, atribuyéndola a la sobrecarga de sincronización que alcanza ~11% en $r=24$.
• Sensibilidad a Parámetros: Los estudios de ablación muestran que la $r^*$ óptima aumenta con tamaños de micro-lotes ($B$) más grandes y longitudes de contexto total más largas (tanto prellenado $\mu_P$ como decodificación $\mu_D$), ya que estos factores aumentan la presión computacional en el lado de la Atención.

Significado y Afirmaciones

El artículo afirma proporcionar la primera base analítica fundamentada para el aprovisionamiento de sistemas AFD, avanzando más allá de la búsqueda empírica. Al ofrecer una cuenta compacta y calibrable de cómo la estructura de la carga de trabajo estocástica determina el aprovisionamiento, el marco permite a los diseñadores de sistemas:

• Determinar la relación óptima A/F sin una construcción extensiva del sistema o ensayo y error.
• Comprender las compensaciones entre la escalabilidad de memoria (Atención) y la escalabilidad de cómputo (FFN) bajo cargas de trabajo realistas y estocásticas.
• Cuantificar el costo de las barreras de sincronización en arquitecturas desagregadas.

Los autores señalan que, si bien la validación en sistemas reales se deja para trabajos futuros a medida que maduran las implementaciones AFD, la validación actual basada en simulación (utilizando trazas de distribuciones similares a las de producción) proporciona una confirmación empírica sólida de la teoría. El trabajo se presenta como un preprint (arXiv:2601.21351v2), indicando que es una contribución fundamental al campo emergente de la atención de LLM desagregada.