Agent Memory Below the Prompt: Persistent Q4 KV Cache for Multi-Agent LLM Inference on Edge Devices

Este trabajo presenta un sistema para inferencia multi-agente en dispositivos de borde que persiste la memoria de atención (KV cache) de cada agente en formato cuantizado a 4 bits en disco, permitiendo su restauración directa y eliminando la necesidad de recomputación, lo que reduce el tiempo hasta el primer token hasta en 136 veces y multiplica por cuatro la cantidad de agentes que caben en la memoria disponible.

Lo esencial

Imagina que tienes un equipo de cinco expertos (un programador, un revisor, un planificador, etc.) trabajando juntos en tu computadora portátil para resolver un problema complejo. Cada experto necesita recordar toda la conversación que han tenido hasta el momento para no perder el hilo.

El problema es que tu computadora (específicamente una MacBook con chip M4) tiene una memoria limitada, como una mesa de trabajo pequeña. Si intentas poner a los cinco expertos trabajando al mismo tiempo con sus notas completas, la mesa se llena y no caben todos.

¿Qué pasa normalmente?
Cuando un experto necesita trabajar, pero su espacio en la mesa está ocupado por otro, tienes que:

1. Sacar sus notas de la mesa y guardarlas en un archivador (el disco duro).
2. Traer las notas del siguiente experto.
3. El problema: Si el experto que sacaste vuelve a necesitar trabajar, no puedes simplemente sacar sus notas del archivador y seguir. Tienes que volver a leer y procesar toda la conversación desde el principio (como si el experto hubiera olvidado todo y tuviera que aprenderlo de nuevo). Esto tarda mucho tiempo (unos 15 segundos), haciendo que el sistema se sienta lento y torpe.

La solución de este papel: "La Memoria Persistente Cuántica"

Los autores proponen un sistema inteligente que funciona como un archivador mágico y ultra-rápido. Aquí está la analogía paso a paso:

1. El Archivador Mágico (Persistencia en Disco)

En lugar de tirar las notas al suelo cuando un experto termina, el sistema las guarda en un formato muy compacto en el disco duro (como guardar un libro gigante en un formato de bolsillo).

• La magia: Cuando el experto vuelve, el sistema no lo hace "reaprender" todo. Simplemente saca el libro del bolsillo y lo abre exactamente en la página donde se quedó.
• Resultado: Lo que antes tardaba 15 segundos en "releer", ahora tarda menos de medio segundo en "reabrir".

2. El Formato de Bolsillo (Cuantización Q4)

Para que quepan más expertos en la mesa pequeña, el sistema comprime las notas.

• Imagina que las notas originales están escritas con letras grandes y gruesas (formato FP16). Ocupan mucho espacio.
• El sistema las reescribe con una tinta muy fina y eficiente (formato Q4, o "4 bits").
• El resultado: Ahora caben 4 veces más expertos en la misma mesa. Puedes tener 12 expertos trabajando en lugar de solo 3, sin que la mesa se desborde.

3. El Baile de los Expertos (Interleaving)

El sistema es muy astuto en cómo organiza el trabajo.

• Mientras el Experto A está escribiendo su respuesta (lo cual toma unos segundos), el sistema aprovecha ese tiempo para traer las notas del Experto B desde el archivador.
• Como el tiempo de escribir es más largo que el tiempo de traer las notas, el usuario nunca nota la espera. Es como si el sistema tuviera "memoria infinita" porque nunca pierde el hilo de la conversación.

4. ¿Por qué es importante?

• Privacidad: Todo ocurre en tu computadora. Nadie envía tus conversaciones a la nube.
• Velocidad: Hace que las aplicaciones de inteligencia artificial multi-agente se sientan instantáneas, incluso con conversaciones muy largas.
• Eficiencia: Permite que dispositivos personales (como laptops) hagan cosas que antes solo podían hacer servidores gigantes en la nube.

En resumen:
Este papel presenta un sistema que permite a tu computadora portátil manejar múltiples agentes de IA trabajando juntos sin volverse loca por falta de memoria. Lo hace guardando el "estado mental" de cada agente de forma comprimida en el disco duro y recuperándolo al instante cuando es necesario, eliminando la necesidad de volver a empezar desde cero cada vez. Es como tener una memoria de elefante que cabe en un bolsillo y que nunca olvida nada, incluso si apagas la computadora.

Resumen técnico

Resumen Técnico: Memoria Persistente de Agentes para Inferencia LLM en Dispositivos de Borde

1. El Problema: Gestión de Memoria en Sistemas Multi-Agente en el Borde

Los sistemas de LLM multi-agente en dispositivos de borde (como un MacBook Pro M4 Pro) enfrentan un cuello de botella crítico: la gestión de la memoria RAM.

• Limitación de Hardware: Los dispositivos de borde tienen RAM soldada y fija (ej. 24 GB). Tras reservar memoria para el modelo y el sistema operativo, el presupuesto para la caché KV (Key-Value) es limitado (aprox. 10.2 GB).
• Incapacidad de Acomodar Múltiples Agentes: Con un formato de precisión FP16, un dispositivo con 10.2 GB de presupuesto de caché solo puede mantener simultáneamente a 3 agentes con un contexto de 8K tokens. Un flujo de trabajo de 10 agentes obliga a expulsar y recargar constantemente las cachés.
• Costo de la Recarga (Re-prefill): Sin persistencia, cada vez que un agente es expulsado de la memoria y necesita ser reactivado, el sistema debe realizar un prefill completo (procesar todo el historial de conversación a través del modelo).
  • En un M4 Pro, esto toma 15.7 segundos por agente para un contexto de 4K.
  • En un escenario de 10 agentes, esto genera tiempos muertos inaceptables (ej. 78.5 segundos de latencia acumulada).
• Fallo de Soluciones Actuales: Las soluciones existentes como vLLM o SGLang, diseñadas para GPUs de centros de datos, dependen de la memoria volátil. En el borde, la falta de persistencia hace que las reinicializaciones tras reinicios del servidor o falta de memoria sean prohibitivas.

2. Metodología y Diseño del Sistema

El sistema propuesto, Agent Memory, resuelve este problema mediante la persistencia de la caché KV en disco utilizando un formato cuantizado de 4 bits (Q4), eliminando la necesidad de recalcular el prefijo.

Componentes Clave:

1. Pool de Bloques Persistente (Persistent Block Pool):
   • Divide la caché KV en bloques fijos de 256 tokens.
   • Almacena la caché de cada agente de forma aislada en disco en formato safetensors (cuantizado a Q4).
   • Permite que la caché sobreviva a reinicios del servidor y cambios de modelo, actuando como una "memoria virtual" para el estado de atención.
2. Pipeline de Cuantización Q4:
   • Utiliza un formato de 4 bits (datos uint32 empaquetados + escalas y sesgos bfloat16) con un tamaño de grupo de 64.
   • Eficiencia: Reduce el uso de memoria en un 72% comparado con FP16 (ratio 0.281). Esto permite alojar 4 veces más agentes en la misma memoria.
   • La inferencia se realiza directamente sobre tensores cuantizados Q4 mediante kernels de atención fusionados en MLX.
3. Inferencia Cuantizada por Lotes (BatchQuantizedKVCache):
   • Implementa un programador que alterna entre agentes durante el prefijo y entrelaza los pasos de decodificación.
   • Permite la inferencia concurrente de múltiples agentes con sus respectivas cachés cuantizadas, ocultando la latencia de carga de disco detrás de la generación de tokens de otros agentes.
4. Inyección de Contexto Trans-Fase (Cross-Phase Context Injection):
   • En flujos de trabajo multi-fase (ej. negociación, debate), el sistema acumula el estado de la caché KV entre fases sin recalcular el prefijo.
   • Utiliza coincidencia de texto a nivel de caracteres para detectar extensiones de contexto (EXTEND match), permitiendo reutilizar la caché cargada desde el disco.

3. Contribuciones Principales

1. Persistencia de Caché KV en Disco: Primera implementación de un pool de bloques aislado por agente que sobrevive a reinicios, almacenado en formato Q4 en safetensors.
2. Inferencia Concurrente Cuantizada: Un mecanismo de BatchQuantizedKVCache que permite inferencia simultánea sobre múltiples cachés cuantizadas en dispositivos de borde, algo no disponible en bibliotecas estándar de MLX.
3. Abstracción de Arquitectura: El sistema soporta modelos con arquitecturas distintas (GQA densa, GQA híbrida con ventanas deslizantes y MLA - Multi-Latent Attention) mediante una especificación de modelo agnóstica (ModelCacheSpec).
4. Evaluación Empírica Exhaustiva: Validación en tres arquitecturas distintas (Gemma 3 12B, DeepSeek-Coder-V2-Lite 16B, Llama 3.1 8B) con métricas de latencia, rendimiento y calidad.

4. Resultados y Métricas

Rendimiento y Latencia (TTFT - Time To First Token):

• Reducción Masiva de Latencia: La restauración de caché desde disco reduce el TTFT en comparación con el prefill frío (re-cálculo completo).
  • Gemma 3 (32K contexto): De 172 segundos a 1.3 segundos (aceleración de 136x).
  • DeepSeek (32K contexto): De 47.3 segundos a 624 ms (aceleración de 76x).
  • Llama 3.1 (16K contexto): De 47.6 segundos a 526 ms (aceleración de 91x).
• Comparación con vLLM-MLX:
  • vLLM-MLX (caché FP16 en memoria) falla bajo presión de memoria multi-agente (no puede mantener 5 agentes a 8K+).
  • El sistema propuesto (Q4 en disco) mantiene la latencia baja incluso tras reinicios, algo que la caché FP16 volátil no puede hacer.
  • La latencia de recarga en caliente (disco) es comparable a la latencia de caché en memoria de FP16 (~500 ms), pero con persistencia.

Capacidad de Memoria:

• En un M4 Pro (10.2 GB de presupuesto de caché):
  • FP16: A 8K contexto, caben 3 agentes.
  • Q4: A 8K contexto, caben 12 agentes.
  • A 16K contexto, FP16 no puede alojar ni un solo flujo de trabajo multi-agente, mientras que Q4 permite 6 agentes.

Calidad del Modelo (Perplejidad):

• El impacto de la cuantización Q4 en la calidad es mínimo y medible:
  • Gemma 3: -0.7% (dentro del ruido de medición).
  • Llama 3.1: +2.8%.
  • DeepSeek: +3.0%.
• Estos resultados son consistentes con la literatura previa sobre cuantización de KV cache y son aceptables para la mayoría de aplicaciones de agentes.

Escenarios de Uso:

• Flujo de Trabajo de 5 Fases: La persistencia trans-fase redujo el tiempo total de pared en un 23% (Gemma) al evitar el recálculo del prefijo en fases posteriores.
• Enrutamiento Multi-Agente: En un escenario de 10 expertos, las consultas a agentes con caché "caliente" fueron 24x más rápidas que el prefill frío.

5. Significado e Impacto

Este trabajo es fundamental para la viabilidad de agentes LLM autónomos en dispositivos de borde:

• Soberanía de Datos: Permite inferencia local completa sin enviar historiales a la nube, cumpliendo con regulaciones como GDPR y HIPAA.
• Experiencia de Usuario: Transforma la latencia de "minutos" (recálculo) a "milisegundos" (recarga de disco), haciendo que los agentes multi-turno sean interactivos y responsivos.
• Eficiencia de Recursos: Demuestra que la cuantización Q4 combinada con persistencia en disco es una estrategia superior a la simple expansión de memoria en hardware de borde, permitiendo flujos de trabajo complejos que de otro modo serían imposibles.
• Arquitectura de Software: Establece un nuevo patrón de diseño donde el estado de atención (KV cache) se trata como un recurso de memoria virtual paginable entre RAM y SSD, similar a la gestión de memoria en sistemas operativos tradicionales.

El sistema es de código abierto y compatible con la API de OpenAI, lo que facilita su integración en frameworks existentes como AutoGen o CrewAI sin modificaciones en la lógica de la aplicación.