Resumen Técnico: InferScale

Planteamiento del Problema

Las aplicaciones de modelos de lenguaje extensos (LLM) dependen cada vez más de un contexto persistente y personalizado (por ejemplo, memorias acumuladas del usuario o historiales de conversación largos) compartido a través de múltiples solicitudes. Los sistemas de memoria actuales (como Mem0, MemGPT, Zep) abordan el costo cuadrático de pre-llenado (prefill) de contextos largos recuperando un subconjunto relevante de hechos de memoria ($k$) e inyectándolos como texto en el prompt.

Sin embargo, este enfoque de "inyección de prompt" incurre en un costo de servicio significativo: cada hecho recuperado debe ser pre-llenado nuevamente en cada solicitud. En consecuencia, el tiempo hasta el primer token (TTFT) crece de forma lineal o cuadrática con el presupuesto de recuperación ($k$), incluso cuando la memoria subyacente es idéntica en todas las solicitudes. Las técnicas existentes de reutilización de KV, como el almacenamiento en caché de prefijos (prefix caching), fallan aquí porque asumen que los tokens reutilizados ocupan posiciones fijas, mientras que los sistemas de memoria ensamblan dinámicamente diferentes subconjectos de hechos para cada solicitud, violando esta suposición.

Metodología

InferScale propone una arquitectura de recuperación e inyección nativa de GPU que reemplaza el repetido pre-llenado de prompts por la inyección directa de KV. En lugar de enviar el texto recuperado al GPU para su recomputación, InferScale precomputa las representaciones de Clave-Valor (KV) de los hechos de memoria y las inyecta directamente en el caché de atención.

El sistema se implementa como un plugin para el motor de servicio vLLM a través de su interfaz de conector KV, lo que no requiere ajuste fino (fine-tuning) del modelo ni modificaciones del motor. La arquitectura consta de dos fases:

1. Fase Offline: Codificación de Ventana de Contexto

Para abordar la pérdida de contexto entre hechos durante la codificación independiente, InferScale emplea la Codificación de Ventana de Contexto.

• Proceso: Cada hecho de memoria se codifica junto con una ventana configurable ($w$) de turnos de conversación precedentes.
• Optimización: Solo se conservan los tensores KV correspondientes al hecho objetivo; los tokens de contexto se descartan.
• Resultado: La representación KV almacenada captura el contexto de desambiguación necesario (por ejemplo, resolver pronombres) sin requerir que los tokens de contexto sean pre-llenados nuevamente en el tiempo de servicio.

2. Fase Online: RoPE Fragmentado e Inyección de KV

Para permitir la reutilización de estados KV en posiciones arbitrarias del prompt (dado que los hechos recuperados se ensamblan dinámicamente), InferScale introduce el RoPE Fragmentado (Chunked RoPE).

• Almacenamiento Pre-RoPE: Durante la codificación offline, las claves son interceptadas antes de la aplicación de las Incrustaciones de Posición Rotatorias (RoPE). Estas claves pre-rotadas y los valores brutos se almacenan en el GPU.
• Rotación Dinámica: En el tiempo de servicio, los hechos recuperados se componen en un segmento contiguo. El sistema aplica la rotación RoPE apropiada a las claves pre-almacenadas basándose en sus nuevas posiciones en el tiempo de servicio.
• Garantía Teórica: El artículo demuestra (Teorema 6) que esta re-rotación es matemáticamente exacta bajo RoPE, asegurando que las puntuaciones de atención sean idénticas a las producidas si los hechos hubieran sido pre-llenados en esas posiciones.
• Inyección: Los tensores KV compuestos se copian mediante scatter-copy directamente en el caché de KV paginado de vLLM. El modelo entonces solo realiza el pre-llenado de los tokens de consulta (query), atendiendo a la memoria inyectada como si hubiera sido computada.

Contribuciones Clave

1. Inferencia de Recuperación e Inyección: Los autores demuestran que la memoria recuperada no necesita ser pre-llenada nuevamente. Prueban (Teorema 3) que la inyección de estados KV precomputados directamente en el caché de atención es exactamente equivalente a la inyección de prompts para transformadores de decodificador causal, siempre que la memoria se inyecte en las posiciones previstas.
2. RoPE Fragmentado: Un mecanismo novedoso para la reutilización de KV independiente de la posición. Al almacenar las claves antes de la codificación rotatoria y re-rotarlas en el momento de la inyección, el sistema permite que los hechos codificados una vez sean inyectados en posiciones arbitrarias sin aproximación (Teorema 6).
3. Codificación de Ventana de Contexto: Una técnica para recuperar la precisión perdida por la codificación independiente de hechos. Al codificar los hechos con una ventana de contexto precedente durante la fase offline, el sistema preserva las interacciones entre hechos sin incurrir en costos de recomputación en el tiempo de servicio.
4. Diseño de Sistema Nativo de GPU: InferScale integra la recuperación semántica residente en GPU (usando el índice Jasper) con el almacenamiento de KV reutilizable. Esto mantiene tanto el índice de recuperación como el almacén de KV en el GPU, eliminando la transferencia de texto recuperado vía PCIe y el pre-llenado repetido.

Resultados

Evaluado en el benchmark LoCoMo a través de tres modelos de pesos abiertos (Llama-3.1-8B, Mistral-7B, Qwen2.5-7B) frente al estado del arte Mem0:

• Latencia (TTFT): InferScale desacopla la latencia de servicio del presupuesto de recuperación. A medida que $k$ aumenta de 5 a 50, el TTFT de InferScale permanece casi constante (aumentando solo ~4% para Llama-3.1-8B), mientras que el TTFT de Mem0 aumenta un 106%. Con $k=50$, InferScale logra una aceleración de 3.6–4.8× en el TTFT comparado con Mem0.
• Rendimiento (Throughput): Bajo carga concurrente (hasta 100 usuarios), InferScale exhibe un escalado casi lineal, logrando un rendimiento 3.7–4.5× mayor que Mem0. Mem0 se satura prematuramente debido al cuello de botella computacional del pre-llenado repetido.
• Precisión: La codificación independiente ($w=0$) degrada inicialmente la precisión. Sin embargo, con la Codificación de Ventana de Contexto ($w \ge 20$), InferScale recupera casi toda la precisión perdida, alcanzando un 60.3% en Llama-3.1-8B con $k=50$ frente al 63.3% de Mem0. La brecha de precisión se limita principalmente a tareas de razonamiento de múltiples saltos que requieren una atención compleja entre hechos.
• Escalabilidad: El sistema escala a modelos más grandes (Qwen3-14B) con beneficios de latencia similares. El almacén de KV puede transferirse a la memoria DRAM del host para superar los límites de capacidad de la memoria del GPU, añadiendo solo ~3 ms al TTFT mientras mantiene una ventaja de latencia de extremo a extremo de 2.3× sobre Mem0.

Significado y Reivindicaciones

El artículo afirma que InferScale cambia fundamentalmente el primitivo de servicio para LLMs personalizados. Al pasar de la inyección de capa de tokens a la inyección de capa de atención, elimina el costo cuadrático de pre-llenado asociado con el contexto recuperado.

Los autores enfatizan que:

• Exactitud: El mecanismo de inyección de KV (con RoPE Fragmentado) no es una aproximación, sino una equivalencia exacta a la inyección de prompts para las posiciones de consulta.
• Cero Sobrecarga: El sistema no requiere ajuste fino del modelo, ni modificaciones del motor, ni cambios en los procesos de recuperación existentes.
• Desacoplamiento: Desacopla con éxito la latencia de servicio condicionada por la memoria del tamaño del contexto recuperado, permitiendo un servicio personalizado de alto rendimiento y baja latencia incluso con presupuestos de recuperación grandes.

El trabajo se distingue de los sistemas de memoria de producción (que aún pagan costos de pre-llenado) y de las técnicas existentes de reutilización de KV (que dependen de diseños de prompts fijos). Posiciona la inyección de capa de atención como un primitivo por defecto para condicionar LLMs con contextos estáticos o semi-estáticos.