Move the Query, Not the Cache: Characterizing Cross-Instance Latent Attention Redistribution Across GPU Fabrics

Resumen Técnico: Mover la Consulta, No el Caché

Declaración del Problema
A medida que los Modelos de Lenguaje de Gran Escala (LLMs) escalan, las cargas de trabajo de servicio dependen cada vez más de mecanismos de Atención de Cabezal Múltiple Latente (MLA) y atención dispersa (sparse attention). En estas arquitecturas, el caché de Clave-Valor (KV) se comprime en vectores latentes estrechos o se selecciona mediante indexadores dispersos, creando una "asimetría de bytes": una única fila de consulta (p. ej., $\approx 1$ KB) es órdenes de magnitud más pequeña que el bloque de caché KV al que atiende (p. ej., megabytes).

Cuando el corpus de interés (un repositorio de código canónico, documentos legales o una base de conocimiento empresarial compartida) excede la Memoria de Alto Ancho de Banda (HBM) de una sola GPU, debe particionarse entre múltiples instancias. Esto requiere atención entre instancias, donde una consulta de decodificación en una GPU debe acceder a bloques KV que residen en otra. El reflejo predominante en la industria es mover el caché (fetch): traer los bloques KV remotos al solicitante, integrarlos en el caché local y realizar la atención localmente. Sin embargo, este enfoque incurre en un costo significativo de "adaptación de posición" (re-rotar las codificaciones posicionales en caché al nuevo desplazamiento) y una gran transferencia de datos. El artículo postula que bajo MLA y selección dispersa, la aritmética puede invertirse: podría ser más barato enrutar la consulta hacia el poseedor, computar la atención parcial allí y fusionar el pequeño resultado de vuelta.

Metodología
Los autores caracterizan este compromiso en un clúster real de H100 multi-nodo conectado mediante NDR-200 InfiniBand con RDMA iniciado por el dispositivo (IBGDA) habilitado. Su metodología involucra:

1. Descomposición de Costos: Modelan la redistribución entre instancias como una función de costo aditiva: $T_{redist} = T_{probe} + T_{transfer} + T_{compute} + T_{return} + T_{merge}$.
2. Comparación de Primitivas: Comparan tres primitivas:
   • Fetch (Traer): Extraer todo el bloque KV, adaptar las posiciones (splice) y atender localmente.
   • Route (Enrutar): Enviar la consulta al poseedor, computar la atención parcial y fusionar el resultado mediante softmax online.
   • Local: Recomputar el bloque desde cero (re-prefill).
3. Caracterización de Hardware: Miden la latencia y el ancho de banda de IBGDA para cargas de trabajo con forma de atención (lotes de consultas a escala de kilobytes) y los comparan contra rutas de proxy de CPU. También analizan el impacto de la topología (mismo leaf vs. cruce de spine), el tipo de fabric (NVLink, PCIe, InfiniBand) y la concurrencia.
4. Verificación de Corrección: Verifican que la atención parcial enrutada, cuando se fusiona mediante softmax online, reproduce los resultados de la atención de instancia única dentro de los márgenes de ruido de bf16, incluso para escenarios de selección dispersa.

Contribuciones Clave

• Modelo de Costo Consciente de la Topología: Los autores derivan un modelo de costo de forma cerrada para la atención entre instancias que se ajusta a las mediciones reales de IBGDA en H100 con un $\approx 7\%$ de Error Absoluto Porcentual Medio (MAPE) para el despacho por lotes ($M_q \ge 512$). El modelo identifica que el costo está dominado por una latencia de sondeo fija ($\approx 16 \mu s$) y un ancho de banda efectivo independiente de la carga útil ($\approx 25$ GB/s para el despacho de un solo bloque), en lugar del ancho de banda máximo del enlace.
• Caracterización del Régimen de RDMA Iniciado por el Dispositivo (IBGDA): A diferencia del trabajo previo sobre el despacho de MoE (donde las rutas de proxy de CPU a veces ganan para mensajes diminutos), este artículo establece que para cargas de trabajo de atención ($\approx 2$ KB), el RDMA iniciado por el dispositivo es superior, reduciendo la latencia por cada fetch en un 40% en comparación con las rutas mediadas por el host.
• Predicado de Selección de Forma Cerrada: Proporcionan una regla de decisión para que los planificadores elijan entre route, fetch o local basándose en la forma de la solicitud ($M_q$), el tamaño del bloque ($c_t$) y las características del fabric.
  • Route gana en el decodificado (pequeño $M_q$) porque evita el costo de la adaptación de posición de $\approx 3$ ms requerido por el fetch y evita el costo de cómputo de local.
  • Fetch solo se vuelve viable cuando se amortiza el costo de la adaptación (splice) sobre muchos pasos locales subsecuentes.
• Análisis de Dispersión y Selección: Para la atención dispersa (p. ej., DeepSeek-V3.2, GLM-5.1), el artículo demuestra que el fetch se convierte en un "gather" disperso a través de múltiples poseedores, incurriendo en una penalización que crece con el número de poseedores. Route mantiene un costo plano porque la consulta se envía una vez, y el poseedor atiende a los bloques dispersos in situ sin re-rotación.
• Codos de Capacidad del Poseedor (Holder Capacity Elbows): Identifican cuellos de botella específicos para el lado del poseedor:
  • Codo de Copia (Copy Elbow): Un pool de $K=8$ streams de CUDA es óptimo para preparar datos para solicitudes concurrentes.
  • Codo de Cómputo (Compute Elbow): Un poseedor puede servir a $\approx 8$ solicitantes enrutados concurrentes con un aumento de latencia insignificante antes de que ocurra la saturación de cómputo.

Resultos

• Dominancia de Latencia: En tamaños de lote típicos de decodificado ($M_q = 256$), un viaje de ida y vuelta enrutado cuesta 31–48 $\mu s$ a través de varios fabrics (NVLink, PCIe, IBGDA). Esto es 60–100$\times$ más rápido que el costo de adaptación de $\approx 3$ ms de traer (fetch) un bloque.
• Eficiencia de Bytes en el Cable: El enrutamiento mueve $\ge 76\%$ menos bytes en el cable que el fetch durante el decodificado.
• Independencia del Fabric: En tamaños de decodificado, el despacho enrutado está limitado por el despacho (limitado por la capacidad de un solo bloque de hilos para emitir transferencias, $\approx 18\text{--}25$ GB/s) en lugar de estar limitado por el ancho de banda. Consecuentemente, un enlace NVLink 4.0 de 900 GB/s se comporta de manera similar a un enlace InfiniBand de 25 GB/s entre nodos para consultas de un solo bloque; el ancho de banda máximo solo importa para tamaños de lote grandes.
• Resiliencia a la Congestión: La primitiva de enrutamiento es robusta a la congestión de red. La latencia permanece plana hasta que el enlace está totalmente saturado ($K=3$ flujos concurrentes), e incluso entonces, sigue siendo un orden de magnitud más rápida que el fetch.
• Advertencia de Sobrecarga del Host: Aunque el enrutamiento gana en bytes de cable y latencia de transporte, los autores señalan que en su prototipo actual, la sobrecarga del host (cliente Python, lógica de despacho) domina el tiempo de llegada del primer token (time-to-first-token). En tamaños de lote pequeños, un simple fetch de "bytes-back" puede ganar actualmente en tiempo total debido a la menor sobrecarga del host, aunque se identifican tres reducciones de implementación específicas (respuesta colapsada, amortización de cómputo, loteo) para cerrar esta brecha.

Significancia y Reivindicaciones
El artículo afirma que el reflejo de "mover el caché" es subóptimo para el régimen emergente de atención comprimida y dispersa. Al invertir la aritmética, enrutar la consulta se convierte en la primitiva dominante para la atención entre instancias, particularmente para cargas de trabajo agénticas donde muchos sub-agentes consultan un corpus compartido e inmutable.

Los autores enfatizan que su contribución no es un nuevo algoritmo de enrutamiento (que existe en trabajos previos como DistAttention), sino una caracterización y un modelo de costo que valida el caso económico para el enrutamiento bajo condiciones específicas de hardware y arquitectura (MLA, IBGDA, selección dispersa). Argumentan que esta asimetría de bytes es la propiedad clave que motiva el uso de RDMA iniciado por el dispositivo para la atención. El artículo concluye que, si bien el enrutamiento es teóricamente superior, lograr sus beneficios de latencia de extremo a extremo en un stack de servicio de producción requiere abordar las sobrecargas del lado del host, una tarea que queda para futuros esfuerzos de ingeniería. El modelo de costo y el predicado proporcionados sirven como artefactos reutilizables para que los diseñadores de sistemas tomen decisiones de programación conscientes de la topología sin necesidad de volver a perfilar cada nueva configuración.