Estrategias de caché

Los vectores clave-valor (KV) se utilizan para calcular las puntuaciones de atención. En los modelos autorregresivos, las puntuaciones KV se calculan cada vez porque el modelo predice un token a la vez. Cada predicción depende de los tokens anteriores, lo que significa que el modelo realiza los mismos cálculos cada vez.

Una cache de KV almacena estos cálculos para que puedan reutilizarse sin necesidad de recalcularlos. Una caché eficiente es crucial para optimizar el rendimiento del modelo, ya que reduce el tiempo de cómputo y mejora las tasas de respuesta. Consulte el documento Caching para obtener una explicación más detallada sobre cómo funciona una caché.

Transformers ofrece varias clases Cache que implementan distintos mecanismos de caché. Algunas de estas clases Cache están optimizadas para ahorrar memoria, mientras que otras están diseñadas para maximizar la velocidad de generación. Consulte la tabla siguiente para comparar los tipos de caché y utilícela para seleccionar la mejor opción para su caso de uso.

Tipo de caché	Admite capas deslizantes	Admite el descargado	Admite torch.compile()	Uso de memoria esperado
Cache dinámica	Sí	Sí	No	Medio
Caché estática	Sí	Sí	Sí	Alto
Cache cuantizada	No	No	No	Bajo

Esta guía le presenta las diferentes clases Cache y le muestra cómo usarlas para la generación.

Caché predeterminada

La DynamicCache es la clase de caché predeterminada para todos los modelos. Permite que el tamaño de la caché crezca dinámicamente para almacenar un número creciente de claves y valores a medida que avanza la generación.

Tenga en cuenta que, en los modelos que utilizan atención de ventana deslizante (Mistral, Gemma2,…) o atención por bloques (Llama4), la caché dejará de crecer cuando las capas que usan estos tipos de atención alcancen su tamaño máximo (el tamaño de la ventana deslizante o del bloque).

Desactive la caché configurando use_cache=False en generate().

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

model.generate(**inputs, do_sample=False, max_new_tokens=20, use_cache=False)

Las clases de caché también se pueden inicializar primero y luego pasarlas al parámetro past_key_values de los modelos. Esto puede ser útil para un control más detallado o para usos más avanzados, como la caché de contexto.

En la mayoría de los casos, es más fácil definir la estrategia de caché en el parámetro cache_implementation.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

past_key_values = DynamicCache(config=model.config)
out = model.generate(**inputs, do_sample=False, max_new_tokens=20, past_key_values=past_key_values)

Caché de tamaño fijo

El DynamicCache por defecto impide aprovechar la mayoría de las optimizaciones just-in-time (JIT) porque el tamaño de la caché no es fijo. Las optimizaciones JIT permiten minimizar la latencia a costa del uso de memoria. Todos los siguientes tipos de caché son compatibles con optimizaciones JIT como torch.compile para acelerar la generación.

Una caché de tamaño fijo (StaticCache) preasigna un tamaño máximo específico para los pares kv. Puedes generar hasta el tamaño máximo de la caché sin necesidad de modificarla. Sin embargo, tener un tamaño fijo (generalmente grande) para los estados clave/valor significa que, durante la generación, muchos tokens se enmascararán en realidad, ya que no deberían participar en la atención. Por lo tanto, este truco permite compilar fácilmente la etapa de decodificación con compile, pero conlleva un desperdicio de tokens en el cálculo de la atención. Como todo, se trata de un equilibrio que será muy ventajoso si generas con varias secuencias de longitudes más o menos similares, pero puede ser subóptimo si, por ejemplo, tienes 1 secuencia muy grande y luego solo secuencias cortas (ya que el tamaño fijo de la caché sería grande, se desperdiciaría mucho en las secuencias cortas). ¡Asegúrate de comprender el impacto si decides usarla!

En cuanto a DynamicCache, tenga en cuenta que, para los modelos que utilizan atención de ventana deslizante (Mistral, Gemma2,…) o atención por bloques (Llama4), la caché nunca será mayor que el tamaño de la ventana deslizante/bloque en las capas que usan estos tipos de atención, incluso si la longitud máxima especificada es mayor.

You can enable StaticCache by configuring cache_implementation="static" in generate(). This will also turn on automatic compilation of the decoding stage for greedy and sample decoding strategies.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("Hello, my name is", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="static")
tokenizer.batch_decode(out, skip_special_tokens=True)[0]
"Hello, my name is [Your Name], and I am a [Your Profession] with [Number of Years] of"

Mantener los tensores de generación en la CPU

Los backends de compilador como Neuron y TPU trazan su modelo en un grafo de cómputo fijo. El bucle de generación mantiene tensores (secuencia de salida, attention_mask, position_ids) que crecen un token en cada paso. El compilador vuelve a trazar el grafo cada vez que estos tensores cambian de forma en el acelerador, lo que ralentiza la generación.

generate() solo mueve a la unidad del modelo los tensores que forward consume, justo antes de cada llamada a forward. La salida se mueve de vuelta para que coincida con la unidad de la entrada. Pasa tus entradas en CPU para mantener la contabilidad de los tensores que crecen en el bucle fuera del acelerador. El grafo compilado se mantiene estable y, como la salida sigue la unidad de la entrada, la salida generada permanece en CPU.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", device_map="auto")

# leave the inputs on CPU instead of calling .to(model.device).
inputs = tokenizer("The French Bread Law states", return_tensors="pt")

# generate runs forward on the model device but returns the output on the input_ids device.
output = model.generate(**inputs, do_sample=False, max_new_tokens=20)
print(output.device)
cpu

Descarga de caché

La caché KV puede ocupar una parte significativa de la memoria y convertirse en un cuello de botella para la generación de contexto largo. Las cachés eficientes en memoria se centran en intercambiar velocidad por un menor uso de memoria. Esto es especialmente importante para los modelos de lenguaje grandes (LLM) y si tu hardware tiene restricciones de memoria.

El descargado de la caché ahorra memoria de GPU al mover la caché KV de las capas del modelo, excepto una, a la CPU. Solo se mantiene en la GPU la caché de la capa actual durante la iteración forward del modelo sobre las capas. Se preobtiene de forma asíncrona la caché de la siguiente capa y se envía de vuelta a la CPU la caché de la capa actual después del cálculo de atención.

Es posible que quieras considerar el offloading si tienes una GPU pequeña y estás recibiendo errores de memoria insuficiente (OOM).

Es posible que note una pequeña degradación en el rendimiento de generación en comparación con una caché completamente en el dispositivo, dependiendo de su modelo y de las opciones de generación (tamaño del contexto, número de tokens generados, número de haces, etc.). Esto se debe a que mover los estados de clave/valor de un lado a otro requiere cierto trabajo.

El offloading está disponible tanto para DynamicCache como para StaticCache. Puedes activarlo configurando cache_implementation="offloaded" para la versión dinámica, o cache_implementation="offloaded_static" para la versión estática, ya sea en GenerationConfig o en generate(). Además, también puedes instanciar tu propio DynamicCache o StaticCache con la opción offloading=True y pasar esta caché en generate o en el forward de tu modelo (por ejemplo, past_key_values=DynamicCache(config=model.config, offloading=True) para una caché dinámica).

Tenga en cuenta que las 2 clases Cache mencionadas anteriormente tienen una opción adicional al instanciarlas directamente, offload_only_non_sliding. Este argumento adicional decide si las capas que usan atención de ventana deslizante/por bloques (si las hay) también se descargarán. Dado que estas capas suelen ser cortas de todos modos, puede ser mejor evitar descargarlas, ya que la descarga puede suponer una penalización en la velocidad. De forma predeterminada, esta opción es False para DynamicCache y True para StaticCache.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, dtype=torch.float16, device_map="auto")
inputs = tokenizer("Fun fact: The shortest", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=23, cache_implementation="offloaded")
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
Fun fact: The shortest war in history was between Britain and Zanzibar on August 27, 1896.

El siguiente ejemplo muestra cómo puedes recurrir a una caché descargada si te quedas sin memoria:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from accelerate import Accelerator

def resilient_generate(model, *args, **kwargs):
    oom = False
    device = Accelerator().device
    torch_device_module = getattr(torch, device, torch.cuda)
    try:
        return model.generate(*args, **kwargs)
    except torch.OutOfMemoryError as e:
        print(e)
        print("retrying with cache_implementation='offloaded'")
        oom = True
    if oom:
        torch_device_module.empty_cache()
        kwargs["cache_implementation"] = "offloaded"
        return model.generate(*args, **kwargs)

ckpt = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(ckpt, dtype=torch.float16, device_map="auto")
prompt = ["okay "*1000 + "Fun fact: The most"]
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
beams = { "num_beams": 40, "num_return_sequences": 20, "max_new_tokens": 23, "early_stopping": True, }
out = resilient_generate(model, **inputs, **beams)
responses = tokenizer.batch_decode(out[:,-28:], skip_special_tokens=True)

Caché cuantizada

QuantizedCache reduce los requisitos de memoria al cuantizar los valores KV a una precisión inferior. QuantizedCache admite actualmente tres backends de cuantización:

hqq admite los tipos de datos int1, int2, int3, int4 e int8.
quanto admite los tipos de datos int2 e int4. Este es el backend de cuantización predeterminado.
fp8 almacena los valores de KV en punto flotante de 8 bits (float8_e4m3fn). No requiere ninguna dependencia adicional y cuantiza por tensor en lugar de por grupo, por lo que rechaza nbits, axis_key, axis_value, q_group_size y residual_length.

La cuantización de la caché puede perjudicar la latencia si la longitud del contexto es corta y hay suficiente memoria de GPU disponible para la generación sin activar la cuantización de la caché. Intente encontrar un equilibrio entre la eficiencia de memoria y la latencia.

Active QuantizedCache configurando cache_implementation="quantized" en GenerationConfig. El backend de cuantización, junto con cualquier parámetro adicional de cuantización, se especifica en el diccionario cache_config que se pasa a generate(). Deberías usar los valores predeterminados para estos parámetros adicionales a menos que estés agotando la memoria. En ese caso, considera reducir la longitud residual.

Para el backend hqq, recomendamos establecer los parámetros axis_key y axis_value en 1.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=20,
    cache_implementation="quantized",
    cache_config={
        "nbits": 4,
        "backend": "hqq",
        "axis_key": 1,
        "axis_value": 1
    }
)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. It's a great way to express myself and rel

Para el backend quanto, recomendamos establecer los parámetros axis_key y axis_value en 0.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=20,
    cache_implementation="quantized",
    cache_config={
        "nbits": 4,
        "backend": "quanto",
        "axis_key": 0,
        "axis_value": 0
    }
)
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it's loud and energetic. It's a great way to express myself and rel

El backend fp8 no acepta ningún parámetro adicional y genera un error si se le proporciona uno de los parámetros de grupo.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B-Instruct-2507")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-4B-Instruct-2507", dtype=torch.float16, device_map="auto")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

out = model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="quantized", cache_config={"backend": "fp8"})
print(tokenizer.batch_decode(out, skip_special_tokens=True)[0])
I like rock music because it is powerful and emotional. I also like classical music because it is elegant and sophisticated. I enjoy

Caché de codificador-decodificador

EncoderDecoderCache está diseñado para modelos encoder-decoder. Gestiona tanto las cachés de autoatención como las de atención cruzada para garantizar el almacenamiento y la recuperación de los pares kv anteriores. Es posible establecer un tipo de caché diferente para el encoder y el decoder de forma individual.

Este tipo de caché no requiere ninguna configuración. Es un envoltorio simple alrededor de 2 Cache como se describió anteriormente, que el modelo utilizará de forma independiente y directa.

Cachés específicos de modelo

Algunos modelos tienen una forma única de almacenar pares kv pasados o estados que no es compatible con ninguna otra clase de caché.

Los modelos Mamba, como Mamba, requieren una caché específica porque el modelo no tiene un mecanismo de atención ni estados kv. Por lo tanto, no son compatibles con las clases Cache anteriores.

Generación iterativa

Una caché también puede funcionar en configuraciones de generación iterativa donde hay interacción de ida y vuelta con un modelo (chatbots). Al igual que en la generación regular, la generación iterativa con una caché permite que un modelo maneje conversaciones en curso de manera eficiente sin recalcular todo el contexto en cada paso.

Para la generación iterativa con una caché, comience inicializando una clase de caché vacía y luego puede introducir sus nuevos prompts. Mantenga un registro del historial de la conversación con una plantilla de chat.

El siguiente ejemplo demuestra Llama-2-7b-chat-hf. Si estás usando un modelo de chat diferente, apply_chat_template() puede procesar los mensajes de manera distinta. Es posible que elimine tokens importantes dependiendo de cómo esté escrito la plantilla Jinja. Para los modelos de chat multimodales, consulta la guía iterative chatting with cache para saber cómo procesar imágenes o audio.

Por ejemplo, algunos modelos usan tokens especiales <think> ... </think> durante el razonamiento. Estos podrían perderse durante la re-codificación, lo que causaría problemas de indexación. Es posible que necesites eliminar o ajustar manualmente tokens adicionales de las completaciones para mantener la estabilidad.

import torch
from transformers import AutoTokenizer,AutoModelForCausalLM, DynamicCache, StaticCache

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(model_id)

user_prompts = ["Hello, what's your name?", "Btw, yesterday I was on a rock concert."]

past_key_values = DynamicCache(config=model.config)

messages = []
for prompt in user_prompts:
    messages.append({"role": "user", "content": prompt})
    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
    input_length = inputs["input_ids"].shape[1]
    outputs = model.generate(**inputs, do_sample=False, max_new_tokens=256, past_key_values=past_key_values)
    completion = tokenizer.decode(outputs[0, input_length: ], skip_special_tokens=True)
    messages.append({"role": "assistant", "content": completion})

Recortar una caché

Use crop() para revertir tokens desde una caché, por ejemplo, cuando una continuación candidata se rechaza durante la generación iterativa. Pase el número de tokens a eliminar como un entero negativo.

# Remove the last three tokens from the cache.
past_key_values.crop(-3)

Las capas de atención lineal y de ventana deslizante descartan los estados anteriores a medida que avanzan, por lo que, de forma predeterminada, no hay nada a lo que retroceder. Llama a activate_past_recording() antes de la pasada hacia adelante para que conserven dichos estados. Sin ello, crop lanza un RuntimeError en cuanto una capa de ventana deslizante llena su ventana, y de inmediato en las capas de atención lineal.

# Enable this before the forward pass if you need to roll back tokens.
past_key_values.activate_past_recording()

crop(0) no limpia la caché. No elimina tokens de la secuencia en caché y es una operación sin efecto para las capas de atención completa. Para las capas de atención con ventana deslizante y atención lineal, puede descartar información en caché que la siguiente llamada al modelo no utilizará. Esto reduce el uso de memoria sin revertir tokens.

                    cached positions
                    ┌───┬───┬───┬───┬───┬───┐
before crop(-3)     │ 0 │ 1 │ 2 │ 3 │ 4 │ 5 │
                    └───┴───┴───┴───┴───┴───┘
                                ╰─────┬─────╯
                                 rolled back

after crop(-3)
                    ┌───┬───┬───┐
full attention      │ 0 │ 1 │ 2 │    keeps the whole prefix
                    └───┴───┴───┘
                            ┌───┐
sliding_window=2            │ 2 │    keeps sliding_window - 1
                            └───┘

Ambas capas informan ahora de una longitud de secuencia de 3, y la generación se reanuda desde la posición 3.

Pasar 0 no elimina ningún token de la secuencia. Es una operación sin efecto para las capas de atención completa, pero para las capas de atención con ventana deslizante y de atención lineal, aún descarta los estados registrados que la siguiente pasada hacia adelante no leerá, lo que libera memoria sin deshacer nada.

Prefill de una caché (caché de prefijo)

En algunas situaciones, es posible que desee llenar una Cache con pares kv para un prefijo de prompt determinado y reutilizarla para generar diferentes secuencias.

El siguiente ejemplo inicializa una StaticCache y luego almacena en caché un prompt inicial. Ahora puede generar varias secuencias a partir del prompt prellenado.

import copy
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, StaticCache

model_id = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_id, dtype=torch.bfloat16, device_map={"": 0})
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Init StaticCache with big enough max-length (1024 tokens for the below example)
# You can also init a DynamicCache, if that suits you better
prompt_cache = StaticCache(config=model.config, max_cache_len=1024)

INITIAL_PROMPT = "You are a helpful assistant. "
inputs_initial_prompt = tokenizer(INITIAL_PROMPT, return_tensors="pt").to(model.device.type)
# This is the common prompt cached, we need to run forward without grad to be able to copy
with torch.no_grad():
     prompt_cache = model(**inputs_initial_prompt, past_key_values = prompt_cache).past_key_values

prompts = ["Help me to write a blogpost about travelling.", "What is the capital of France?"]
responses = []
for prompt in prompts:
    new_inputs = tokenizer(INITIAL_PROMPT + prompt, return_tensors="pt").to(model.device.type)
    past_key_values = copy.deepcopy(prompt_cache)
    outputs = model.generate(**new_inputs, past_key_values=past_key_values,max_new_tokens=20)
    response = tokenizer.batch_decode(outputs)[0]
    responses.append(response)

print(responses)

Update on GitHub