Prompt Caching

La caché de instrucciones, también llamada caché de prefijos, almacena los resultados intermedios que calcula un modelo de lenguaje grande (LLM) para el inicio de una instrucción y los reutiliza en solicitudes posteriores que empiezan exactamente con el mismo texto. Muchas solicitudes comparten un prefijo largo e invariable —como una instrucción de sistema, definiciones de herramientas, ejemplos, un documento extenso o la conversación hasta el momento— al que sigue una parte nueva y breve. Cuando se encuentra una coincidencia en la caché, el servidor procesa solo los nuevos tokens, lo que reduce el tiempo hasta el primer token de salida y, en las API alojadas, rebaja el precio de la entrada almacenada en caché a una fracción de la tarifa habitual.

Cómo funciona la caché de instrucciones#

Un LLM de transformador responde en dos fases. Durante el prellenado, cada capa calcula un vector de clave y otro de valor para cada token de la instrucción; durante la decodificación, cada token nuevo presta atención a todas esas claves y valores mediante el mecanismo de atención. En una misma solicitud, la caché KV conserva estos tensores para evitar que cada token generado vuelva a calcular la instrucción. La caché de instrucciones los conserva entre solicitudes. Como la atención de estos modelos es causal, las claves y los valores de un token dependen solo de los tokens anteriores; por tanto, dos instrucciones con el mismo prefijo tienen claves y valores idénticos para ese prefijo, y la copia almacenada puede sustituir al nuevo cálculo.

El siguiente ejemplo muestra la reutilización con GPT-2 de Hugging Face Transformers. Primero realiza el prellenado de una instrucción de sistema de 409 tokens y, después, responde a dos preguntas procesando solo sus propios tokens sobre el prefijo almacenado en caché. Por último, compara el resultado con un recálculo completo:

import copy

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")

# A long system prompt shared by every request, followed by different user questions
system = "You answer questions about a camera store. " + "Returns are accepted within 30 days with a receipt. " * 40
questions = [" Question: Can I return a lens?", " Question: Do you ship abroad?"]
prefix_ids = tokenizer(system, return_tensors="pt").input_ids

with torch.no_grad():
    # Prefill the shared prefix once and keep its key-value cache
    prefix_cache = model(prefix_ids).past_key_values
    for question in questions:
        ids = tokenizer(question, return_tensors="pt").input_ids
        # Cache hit: only the question's tokens run through the model
        cached = model(ids, past_key_values=copy.deepcopy(prefix_cache)).logits[0, -1]
        # No cache: the whole prompt is computed from scratch
        full = model(torch.cat([prefix_ids, ids], dim=1)).logits[0, -1]
        same = torch.allclose(cached, full, atol=1e-3)
        print(f"Computed {ids.shape[1]} of {prefix_ids.shape[1] + ids.shape[1]} tokens, same next-token scores: {same}")

Imprime Computed 8 of 417 tokens, same next-token scores: True y Computed 7 of 416 tokens, same next-token scores: True: cada solicitud procesó con el modelo menos del 2 % de su instrucción y obtuvo las mismas puntuaciones para el siguiente token que el cálculo completo.

Los motores de servicio de modelos de código abierto aplican este mecanismo automáticamente. vLLM mantiene la caché KV en bloques de tamaño fijo, siguiendo la disposición de memoria de PagedAttention, reutiliza los bloques del prefijo coincidente más largo para cada solicitud nueva y expulsa los bloques menos utilizados cuando escasea la memoria de GPU. RadixAttention de SGLang mantiene los prefijos almacenados en caché en un árbol radix, de modo que las solicitudes con cualquier parte inicial en común puedan reutilizarlos. El artículo de SGLang informa de un rendimiento hasta 6,4 veces superior al de otros motores de inferencia de última generación. El artículo Prompt Cache, de In Gim y sus colegas y presentado en MLSys 2024, amplía la reutilización más allá de los prefijos exactos a módulos de instrucciones declarados, como plantillas y documentos, y comunica reducciones del tiempo hasta el primer token que van de 8 veces en GPU a 60 veces en CPU.

Caché de instrucciones en las API de LLM#

Las API alojadas ofrecen el mismo mecanismo de pago. La caché de instrucciones de OpenAI está activada de forma predeterminada en los modelos compatibles, mientras que la caché de instrucciones de Claude, de Anthropic, se activa con el campo cache_control, ya sea una vez por solicitud o en un máximo de cuatro bloques de contenido. Anthropic cobra una décima parte del precio base de entrada o menos por la lectura de la caché, y 1,25 veces el precio base por escribir en la caché predeterminada de cinco minutos; cada coincidencia reinicia ese plazo de cinco minutos. Ambos proveedores solo aceptan prefijos idénticos: Anthropic compara las herramientas, después la instrucción de sistema y, por último, los mensajes, así que modificar una definición de herramienta invalida todo lo almacenado después. La guía de OpenAI recomienda colocar primero las instrucciones estables y el material de referencia, y dejar para el final los detalles específicos de cada usuario.

Las cargas de trabajo con prefijos largos compartidos son las que más se benefician:

• Agentes de IA: Cada paso vuelve a enviar las definiciones de las herramientas y todo el historial, que incluye la instrucción del paso anterior, la última acción y su resultado.
• Chatbots: Cada turno de una conversación amplía el prompt del turno anterior, por lo que el historial se almacena en caché hasta el mensaje nuevo.
• Preguntas y respuestas sobre documentos: Muchas preguntas sobre un contrato, informe o repositorio extenso reutilizan el documento procesado, y los sistemas de recuperación lo reutilizan cada vez que vuelven a recuperar los mismos pasajes.
• Prompts con pocos ejemplos: Se calculan una vez decenas de ejemplos fijos y se reutilizan para cada entrada.

Los modelos de visión de vocabulario abierto reutilizan el cálculo del prompt de forma similar. Ultralytics YOLOE codifica una vez los prompts de clase, como «carretilla elevadora», con un codificador de texto MobileCLIP cuando se llama a set_classes(), y almacena las incrustaciones resultantes en el modelo, de modo que las predicciones posteriores omiten por completo el codificador de texto; save_prompt_embeddings() las escribe en un archivo NPZ que una sesión posterior carga con load_prompt_embeddings(), en lugar de volver a ejecutar el codificador.

Caché de prompts y términos relacionados#

• Caché KV: Reutiliza claves y valores entre los tokens de una generación. La caché de prompts los reutiliza entre solicitudes independientes.
• Caché semántica: Almacena respuestas completas y devuelve una cuando una consulta nueva significa lo mismo que otra anterior, sin ejecutar el modelo. La caché de prompts sigue ejecutando el modelo con la parte nueva del prompt y genera una respuesta nueva.
• Compresión de prompts: Acorta el propio prompt, lo que resulta adecuado para texto que cambia con cada solicitud. Ambas técnicas se combinan bien: la parte fija del prompt se almacena en caché y la parte variable se comprime.
• Ingeniería de prompts: Determina qué dice el prompt. El almacenamiento en caché solo cambia cuánto vuelve a calcular el servidor, y un acierto deja intactas las puntuaciones de los siguientes tokens, salvo por el redondeo en coma flotante.