Estratégias de cache

Os vetores chave-valor (KV) são usados para calcular as pontuações de atenção. Em modelos autorregressivos, as pontuações KV são calculadas a cada vez, pois o modelo prevê um token de cada vez. Cada previsão depende dos tokens anteriores, o que significa que o modelo executa os mesmos cálculos a cada vez.

A cache de KV armazena esses cálculos para que possam ser reutilizados sem a necessidade de recálculo. O cacheamento eficiente é crucial para otimizar o desempenho do modelo, pois reduz o tempo de computação e melhora as taxas de resposta. Consulte o documento Caching para uma explicação mais detalhada sobre o funcionamento de um cache.

O Transformers oferece várias classes Cache que implementam diferentes mecanismos de cache. Algumas dessas classes Cache são otimizadas para economizar memória, enquanto outras são projetadas para maximizar a velocidade de geração. Consulte a tabela abaixo para comparar os tipos de cache e use-a para ajudá-lo a selecionar o melhor cache para o seu caso de uso.

Tipo de Cache	Suporta camadas deslizantes	Suporta offloading	Suporta torch.compile()	Uso de memória esperado
Cache Dinâmica	Sim	Sim	Não	Médio
Cache Estático	Sim	Sim	Sim	Alto
Cache Quantizado	Não	Não	Não	Baixo

Este guia apresenta as diferentes classes Cache e mostra como usá-las para geração.

Cache padrão

A DynamicCache é a classe de cache padrão para todos os modelos. Ela permite que o tamanho da cache cresça dinamicamente para armazenar um número crescente de chaves e valores à medida que a geração avança.

Observe que, para modelos que usam atenção com janela deslizante (Mistral, Gemma2,…) ou atenção em blocos (Llama4), o cache deixará de crescer quando as camadas que usam esses tipos de atenção atingirem seu tamanho máximo (o tamanho da janela deslizante ou do bloco).

Desative o cache configurando use_cache=False em 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)

As classes de cache também podem ser inicializadas primeiro e, em seguida, passadas ao parâmetro past_key_values dos modelos. Isso pode ser útil para um controle mais granular ou para usos mais avançados, como o cache de contexto.

Na maioria dos casos, é mais fácil definir a estratégia de cache no 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)

Cache de tamanho fixo

O DynamicCache padrão impede que você aproveite a maioria das otimizações just-in-time (JIT), pois o tamanho do cache não é fixo. As otimizações JIT permitem minimizar a latência em detrimento do uso de memória. Todos os tipos de cache a seguir são compatíveis com otimizações JIT, como torch.compile, para acelerar a geração.

Um cache de tamanho fixo (StaticCache) pré-aloca um tamanho máximo específico para os pares de kv. É possível gerar até o tamanho máximo do cache sem a necessidade de modificá-lo. No entanto, ter um tamanho fixo (geralmente grande) para os estados de chave/valor significa que, durante a geração, muitos tokens serão efetivamente mascarados, pois não devem participar da atenção. Assim, esse truque permite compilar facilmente a etapa de decodificação com compile, mas acarreta um desperdício de tokens no cálculo da atenção. Como tudo, trata-se de um trade-off que deve ser muito bom se você gerar com várias sequências de comprimentos mais ou menos iguais, mas pode ser subótimo se, por exemplo, você tiver 1 sequência muito grande e, em seguida, apenas sequências curtas (já que o tamanho fixo do cache seria grande, muito seria desperdiçado para as sequências curtas). Certifique-se de entender o impacto se você for usá-lo!

Quanto ao DynamicCache, observe que, para modelos que usam atenção com janela deslizante (Mistral, Gemma2,…) ou atenção em blocos (Llama4), o cache nunca será maior que o tamanho da janela deslizante/bloco nas camadas que utilizam esses tipos de atenção, mesmo que o comprimento máximo especificado seja maior.

Você pode habilitar o StaticCache configurando cache_implementation="static" em generate(). Isso também ativará a compilation automática da etapa de decodificação para as estratégias de decodificação greedy e sample.

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"

Mantenha os tensores de geração na CPU

Backends de compilador como Neuron e TPU traçam seu modelo em um grafo de computação fixo. O loop de geração mantém tensores (sequência de saída, attention_mask, position_ids) que crescem em um token a cada etapa. O compilador re-traca o grafo sempre que esses tensores mudam de forma no acelerador, o que torna a geração mais lenta.

generate() move apenas os tensores que forward consome para o dispositivo do modelo, logo antes de cada chamada de forward. A saída é movida de volta para corresponder ao dispositivo da entrada. Passe suas entradas na CPU para manter o controle dos tensores em crescimento do loop fora do acelerador. O gráfico compilado permanece estável e, como a saída segue o dispositivo da entrada, a saída gerada permanece na 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

Cache offloading

O cache KV pode ocupar uma parte significativa da memória e se tornar um gargalo para a geração de contexto longo. Caches eficientes em memória focam em trocar velocidade por uma redução no uso de memória. Isso é especialmente importante para grandes modelos de linguagem (LLMs) e quando seu hardware tem restrições de memória.

O offloading do cache economiza memória de GPU ao mover o cache KV das camadas do modelo, exceto uma, para a CPU. Apenas o cache da camada atual é mantido na GPU durante a iteração forward do modelo pelas camadas. Ele pré-carrega de forma assíncrona o cache da próxima camada e envia de volta para a CPU o cache da camada atual após o cálculo de atenção.

Você pode considerar o offloading se tiver uma GPU pequena e estiver recebendo erros de falta de memória (OOM).

Você pode notar uma pequena degradação no throughput de geração em comparação com um cache totalmente no dispositivo, dependendo do seu modelo e das opções de geração (tamanho do contexto, número de tokens gerados, número de feixes, etc.). Isso ocorre porque mover os estados de chave/valor de um lado para o outro requer algum trabalho.

O offloading está disponível tanto para DynamicCache quanto para StaticCache. Você pode ativá-lo configurando cache_implementation="offloaded" para a versão dinâmica, ou cache_implementation="offloaded_static" para a versão estática, em GenerationConfig ou generate(). Além disso, você também pode instanciar seu próprio DynamicCache ou StaticCache com a opção offloading=True e passar esse cache em generate ou no forward do seu modelo (por exemplo, past_key_values=DynamicCache(config=model.config, offloading=True) para um cache dinâmico).

Observe que as duas classes Cache mencionadas acima possuem uma opção adicional ao instanciá-las diretamente, offload_only_non_sliding. Esse argumento adicional decide se as camadas que usam atenção de janela deslizante/chunk (se houver) também serão descarregadas. Como essas camadas geralmente são curtas de qualquer forma, pode ser melhor evitar o descarregamento delas, pois o descarregamento pode incorrer em uma penalidade de velocidade. Por padrão, essa opção é False para DynamicCache e 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.

O exemplo abaixo mostra como você pode recorrer a um cache descarregado caso a memória se esgote:

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)

Cache quantizado

O QuantizedCache reduz os requisitos de memória ao quantizar os valores de KV para uma precisão menor. O QuantizedCache atualmente suporta três backends de quantização:

hqq suporta os tipos de dados int1, int2, int3, int4 e int8.
quanto suporta os tipos de dados int2 e int4. Este é o backend de quantização padrão.
fp8 armazena os valores de KV em ponto flutuante de 8 bits (float8_e4m3fn). Ele não requer nenhuma dependência extra e quantiza por tensor, em vez de por grupo, portanto rejeita nbits, axis_key, axis_value, q_group_size e residual_length.

A quantização do cache pode prejudicar a latência se o comprimento do contexto for curto e houver memória de GPU suficiente disponível para a geração sem ativar a quantização do cache. Tente encontrar um equilíbrio entre eficiência de memória e latência.

Ative o QuantizedCache configurando cache_implementation="quantized" em GenerationConfig. O backend de quantização, juntamente com quaisquer parâmetros adicionais de quantização, é especificado no dicionário cache_config passado para generate(). Você deve usar os valores padrão para esses parâmetros adicionais, a menos que esteja enfrentando erros de memória insuficiente. Nesse caso, considere diminuir o comprimento residual.

Para o backend hqq, recomendamos definir os parâmetros axis_key e axis_value como 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 o backend quanto, recomendamos definir os parâmetros axis_key e axis_value como 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

O backend fp8 não aceita nenhum parâmetro adicional e gera um erro se for fornecido um dos parâmetros por 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

Cache de encoder-decoder

EncoderDecoderCache foi projetado para modelos encoder-decoder. Ele gerencia os caches de autoatenção e de atenção cruzada para garantir o armazenamento e a recuperação dos pares kv anteriores. É possível definir individualmente um tipo de cache diferente para o encoder e para o decoder.

Esse tipo de cache não requer nenhuma configuração. É um wrapper simples em torno de 2 Caches, conforme descrito acima, que serão usados de forma independente diretamente pelo modelo.

Caches específicos de modelo

Alguns modelos possuem uma maneira única de armazenar pares kv passados ou estados que não é compatível com nenhuma outra classe de cache.

Modelos Mamba, como o Mamba, exigem um cache específico, pois o modelo não possui um mecanismo de atenção nem estados kv. Portanto, eles não são compatíveis com as classes Cache acima.

Geração iterativa

Um cache também pode funcionar em configurações de geração iterativa, nas quais há interação contínua com um modelo (chatbots). Assim como na geração regular, a geração iterativa com um cache permite que um modelo processe conversas em andamento de forma eficiente, sem precisar recalcular todo o contexto a cada etapa.

Para geração iterativa com cache, comece inicializando uma classe de cache vazia e, em seguida, você pode inserir seus novos prompts. Mantenha o histórico do diálogo com um chat template.

O exemplo a seguir demonstra o Llama-2-7b-chat-hf. Se você estiver usando um modelo de chat diferente, o apply_chat_template() pode processar as mensagens de forma diferente. Ele pode remover tokens importantes, dependendo de como o template Jinja é escrito. Para modelos de chat multimodais, consulte o guia iterative chatting with cache para saber como processar imagens ou áudio.

Por exemplo, alguns modelos usam tokens especiais <think> ... </think> durante o raciocínio. Esses tokens podem se perder durante a re-codificação, causando problemas de indexação. Você pode precisar remover ou ajustar manualmente tokens extras das conclusões para manter a estabilidade.

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 um cache

Use crop() para reverter tokens de um cache, como quando uma continuação candidata é rejeitada durante a geração iterativa. Passe o número de tokens a remover como um inteiro negativo.

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

Camadas de atenção com janela deslizante e de atenção linear descartam os estados anteriores à medida que avançam, portanto, por padrão, não há para onde voltar. Chame activate_past_recording() antes da passagem de frente para que elas retenham esses estados. Sem isso, crop levanta um RuntimeError assim que uma camada de janela deslizante preenche sua janela, e imediatamente para camadas de atenção linear.

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

crop(0) não limpa o cache. Ele não remove tokens da sequência em cache e é uma operação inerte para camadas de atenção completa. Para camadas de janela deslizante e de atenção linear, ele pode descartar informações em cache que a próxima chamada do modelo não usará. Isso reduz o uso de memória sem reverter 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
                            └───┘

Agora, ambas as camadas relatam um comprimento de sequência de 3, e a geração retoma a partir da posição 3.

Passar 0 não remove nenhum token da sequência. É uma operação sem efeito para camadas de atenção completa, mas para camadas de janela deslizante e de atenção linear, ela ainda descarta os estados registrados que a próxima passagem para frente não lerá, o que libera memória sem desfazer nada.

Preencher um cache (cache de prefixo)

Em algumas situações, você pode querer preencher um Cache com pares de kv para um determinado prompt de prefixo e reutilizá-lo para gerar diferentes sequências.

O exemplo abaixo inicializa um StaticCache e, em seguida, faz o cache de um prompt inicial. Agora você pode gerar várias sequências a partir do prompt pré-preenchido.

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