KV Cache

La caché KV (caché de clave-valor) es una técnica de optimización fundamental que se utiliza principalmente en los modelos de lenguaje grandes (LLM) y otras arquitecturas basadas en Transformer para acelerar la latencia de inferencia y reducir los costes computacionales. En esencia, la caché KV almacena las matrices Key y Value que genera el mecanismo de atención para los tokens anteriores de una secuencia. Al guardar estos cálculos intermedios, el modelo evita volver a calcular los estados de atención de todo el historial de la conversación cada vez que genera un token nuevo. Este proceso transforma el flujo de trabajo de generación de texto, que pasa de una operación de complejidad cuadrática a una lineal y hace viables las interacciones en tiempo real con chatbots y agentes de IA.

Mecanismo y ventajas#

En un modelo Transformer estándar, para generar la siguiente palabra hay que tener en cuenta todas las palabras anteriores para comprender el contexto. Sin almacenamiento en caché, el modelo tendría que volver a calcular las relaciones matemáticas de toda la secuencia en cada paso. La caché KV resuelve este problema actuando como un banco de memoria.

• Mejora de la velocidad: Al recuperar de la memoria claves y valores precalculados, el sistema acelera considerablemente el motor de inferencia. Esto es esencial para las aplicaciones que requieren una latencia baja, como la inferencia en tiempo real en bots de atención al cliente.
• Eficiencia de recursos: Aunque aumenta el uso de memoria (VRAM), reduce significativamente el cálculo (FLOPs) necesario por token. Este equilibrio suele gestionarse mediante técnicas como la cuantización de modelos o la paginación, de forma parecida a cómo los sistemas operativos gestionan la RAM.
• Contexto ampliado: La gestión eficiente de la caché KV permite a los modelos manejar una ventana de contexto más amplia, lo que les permite procesar documentos largos o mantener conversaciones coherentes durante periodos prolongados.

Aplicaciones en el mundo real#

La caché KV es un componente fundamental para implementar la IA generativa moderna, pero sus principios también se aplican a la visión artificial (CV).

1. Chatbots generativos: Servicios como ChatGPT o Claude dependen en gran medida del almacenamiento en caché KV. Cuando un usuario hace una pregunta de seguimiento, el modelo no vuelve a leer todo el historial del chat desde el principio. En su lugar, añade la nueva entrada a los estados almacenados en caché del turno anterior, lo que permite responder casi al instante.

2. Comprensión de vídeo: En las tareas de comprensión de vídeo, los modelos procesan los fotogramas secuencialmente. Al igual que ocurre con los tokens de texto, las características visuales de los fotogramas anteriores pueden almacenarse en caché para ayudar al modelo a seguir objetos o reconocer acciones sin volver a procesar todo el historial del vídeo. Esto es especialmente relevante para el reconocimiento de acciones, donde el contexto temporal es crucial.

Gestión eficiente de la memoria#

A medida que los modelos crecen, el tamaño de la caché KV puede convertirse en un cuello de botella y consumir gigabytes de memoria de GPU. Los avances recientes se centran en optimizar este almacenamiento.

• PagedAttention: Inspirado en la memoria virtual de los sistemas operativos, PagedAttention (introducido por vLLM) permite almacenar la caché KV en bloques de memoria no contiguos. Esto reduce la fragmentación y permite usar tamaños de lote mayores durante el servicio de modelos.
• Cuantización de la caché KV: Para ahorrar espacio, los desarrolladores suelen aplicar precisión mixta o cuantización int8 específicamente a los valores almacenados en caché. Esto reduce la huella de memoria y permite que los dispositivos de IA en el borde con RAM limitada ejecuten modelos capaces.
• Almacenamiento en caché de prompts: Una técnica relacionada que calcula una vez los estados KV de un prompt de sistema estático (por ejemplo, «Eres un asistente de programación útil») y los reutiliza en muchas sesiones de usuario distintas. Es una función clave para optimizar a gran escala los flujos de trabajo de ingeniería de prompts.

Diferenciación de conceptos relacionados#

Es útil distinguir la caché KV de otros términos relacionados con el almacenamiento en caché y la optimización:

• Caché KV frente a almacenamiento en caché de prompts: La caché KV suele referirse a la memoria dinámica, token a token, que se utiliza durante un único flujo de generación. El almacenamiento en caché de prompts se refiere específicamente a guardar el estado procesado de una instrucción de entrada fija para reutilizarlo en varias llamadas de inferencia independientes.
• Caché KV frente a incrustaciones: Las incrustaciones son representaciones vectoriales de datos de entrada (texto o imágenes) que capturan su significado semántico. La caché KV almacena las activaciones (claves y valores) derivadas de estas incrustaciones dentro de las capas de atención, específicamente para generar secuencias.
• Caché KV frente a pesos del modelo: Los pesos del modelo son los parámetros estáticos aprendidos de la red neuronal. La caché KV está formada por datos dinámicos y temporales que se generan durante la pasada hacia delante de una secuencia de entrada concreta.

Ejemplo: contexto en modelos de visión#

Aunque el almacenamiento en caché KV es más conocido en el procesamiento del lenguaje natural, el concepto de mantener el estado también se aplica a los modelos de visión avanzados. En el ejemplo siguiente, simulamos la idea de transferir el estado (contexto) en una situación de seguimiento de vídeo con Ultralytics YOLO26. En este caso, el rastreador mantiene la identidad de los objetos entre fotogramas, algo conceptualmente similar a cómo una caché mantiene el contexto entre tokens.

from ultralytics import YOLO

# Load the Ultralytics YOLO26 model
model = YOLO("yolo26n.pt")

# Track objects in a video, maintaining identity state across frames
# The 'track' mode effectively caches object features to link detections
results = model.track(source="https://ultralytics.com/images/bus.jpg", show=False)

# Print the ID of the tracked objects
if results[0].boxes.id is not None:
    print(f"Tracked IDs: {results[0].boxes.id.numpy()}")

Los desarrolladores que quieran gestionar conjuntos de datos e implementar modelos optimizados pueden usar la plataforma Ultralytics, que simplifica el flujo de trabajo desde la anotación de datos hasta la implementación de modelos eficiente. Para profundizar en los mecanismos de atención, bibliotecas como PyTorch proporcionan los componentes básicos en los que se implementan estos mecanismos de almacenamiento en caché.