Como acelerar SLMs reutilizando o cache key-value do prefixo do prompt

Por que isso importa agora

Modelos de linguagem pequenos (SLMs, na casa de 0,5B a 3B de parâmetros) estão virando a escolha óbvia para automação estreita: classificar tickets, extrair campos, categorizar texto. Eles rodam em CPU, custam centavos e respondem rápido. O gargalo não é o modelo — é o código ao redor, que na maioria dos casos trata cada chamada como um evento isolado e recomputa o prompt inteiro toda vez. Este tutorial mostra a técnica que quebra esse desperdício: reusar o cache key-value (KV) do prefixo do prompt.

O que você ganha (e o que não ganha)

✅ O que a técnica entrega:

• Redução de ~57% no tempo total de execução no benchmark do tutorial.
• As mesmas predições do loop ingênuo — é otimização de computação pura, sem mudar o comportamento do modelo.
• Ganho que cresce com o tamanho das instruções: quanto mais longo o bloco estático, melhor o trade.

⚠️ O que ela não entrega:

• Não melhora a qualidade da saída — só reduz o custo de computação.
• Exige cuidado com o split do prompt em fronteira “token-clean” e com o recorte do cache.
• Só vale para prompts majoritariamente estáticos (automação estreita), não para conversas totalmente novas.

Requisitos

A intuição em uma frase

Em prompts de automação estreita, a instrução, a taxonomia e os exemplos somam centenas de tokens — e só uma cauda curta muda de item para item. No exemplo do tutorial, 87% do prompt é byte a byte idêntico a cada chamada. Os Transformers calculam um vetor de chave e um de valor para cada token em cada camada, e eles dependem apenas dos tokens à esquerda. Para um prefixo fixo, esses vetores são idênticos em toda chamada: calcular uma vez e guardar reduz o pre-fill de cada item aos tokens que realmente mudaram.

Passo a passo

Passo 1 — O baseline (re-codificar tudo)

Primeiro, o cenário ingênuo: um prompt few-shot realista, re-codificado por inteiro a cada ticket. A decisão usa scoring restrito (comparar o logit do primeiro token de cada label), então resta otimizar apenas o pre-fill:

pip install torch transformers accelerate

full = tokenizer(prefix_text + suffix_text(ticket),
                 add_special_tokens=False, return_tensors="pt").to(model.device)
with torch.no_grad():
    logits = model(**full).logits[0, -1, :]
    pred = LABELS[int(logits[label_first_ids].argmax())]

Resultado do baseline: 184,85 segundos para 600 tickets (~0,3 s por ticket).

Passo 2 — Populando o cache uma única vez

Rode o prefixo pelo modelo exatamente uma vez e guarde o cache resultante com DynamicCache:

prefix_cache = DynamicCache()
with torch.no_grad():
    model(input_ids=prefix_ids,
          attention_mask=torch.ones_like(prefix_ids),
          past_key_values=prefix_cache,
          use_cache=True)

Passo 3 — Classificando apenas a cauda

Depois, cada ticket alimenta só os seus próprios tokens, reaproveitando o cache. Três detalhes são obrigatórios para o resultado não sair enganosamente errado:

# A máscara cobre o prefixo EM cache + os tokens novos
attention_mask = torch.ones((1, prefix_len + suffix_len), device=model.device, dtype=torch.long)
# Os tokens novos começam na posição prefix_len, não em 0
cache_position = torch.arange(prefix_len, prefix_len + suffix_len, device=model.device)

with torch.no_grad():
    out = model(input_ids=suffix_ids,
                attention_mask=attention_mask,
                past_key_values=prefix_cache,
                cache_position=cache_position,
                use_cache=True)
    label = LABELS[int(out.logits[0, -1, label_first_ids].argmax())]

# Sem o crop, o segundo ticket atenderia aos tokens do primeiro
prefix_cache.crop(prefix_len)

Passo 4 — Resultado

Com o cache reutilizado: 80,07 segundos para os mesmos 600 tickets (~0,13 s por ticket). Redução total de ~57%, com predições idênticas em todos os tickets.

Os detalhes que fazem a diferença

• DynamicCache guarda os tensores de chave/valor por camada do prefixo; passá-lo como past_key_values diz ao modelo que essas posições já foram computadas.
• A máscara de atenção precisa cobrir o prefixo em cache e os tokens novos — a largura é prefix_len + suffix_len, mesmo passando só suffix_len ids.
• cache_position informa que os tokens novos começam no offset prefix_len, para as rotary embeddings baterem com o prompt completo.
• prefix_cache.crop(prefix_len) não é opcional: o forward append os keys/values do sufixo ao cache; sem o crop, o cache cresceria sem limite e o segundo ticket atenderia aos tokens do primeiro.
• Split token-clean: divida o prompt em fronteira limpa (fim de linha ou delimitador do chat template). Tokenizar as duas metades separadamente pode gerar sequência diferente da concatenação se o corte cair no meio de uma palavra — e aí as chaves em cache não correspondem ao que o modelo teria visto. O script de baseline tem um assert que checa isso diretamente.
• Use torch.no_grad(), não torch.inference_mode(): tensores criados em inference mode carregam uma flag que atrapalha o crop() do cache.

Tabela de resultados

Casos de uso reais

• Classificação de tickets de suporte em lotes grandes com a mesma taxonomia.
• Extração de campos de documentos com schema fixo.
• Moderação de conteúdo aplicando as mesmas regras a muitos itens.
• Roteamento de intenção em chatbots com instruções longas e estáveis.

Troubleshooting: erros comuns

• ❌ Resultados diferentes entre cache e re-codificação → Causa: split não token-clean. Solução: mova a fronteira para um delimitador; valide com o assert de equivalência.
• ❌ Cache cresce sem limite e degrada → Causa: falta de crop(prefix_len). Solução: recorte o cache após cada item.
• ❌ Embeddings rotacionais erradas → Causa: cache_position ausente ou incorreto. Solução: informe o offset prefix_len.
• ❌ Erro ao manipular o cache → Causa: tensores criados em inference_mode. Solução: use no_grad().
• ❌ Sem ganho de performance → Causa: prefixo pequeno em relação ao sufixo. Solução: o ganho escala com a proporção de conteúdo estático; instruções longas recompensam mais.

FAQ

• Isso muda o comportamento do modelo? Não. As predições devem ser idênticas — é otimização de computação pura.
• Quando não vale a pena? Quando o prefixo é curto ou muda a cada chamada; aí não há quase nada a reutilizar.
• Funciona com qualquer modelo? O padrão é geral (qualquer modelo com cache KV), mas o benchmark usou Qwen2.5-0.5B-Instruct.
• Preciso de GPU? Não — o exemplo roda em CPU, no M2 de 24 GB.
• Por que escrever o chat template à mão? Para poder dividir o prompt em uma fronteira conhecida; o apply_chat_template() esconderia esse ponto.

Onde isso leva

Um modelo pequeno deixa de ser um “compromisso” e vira a resposta óbvia para automação estreita quando o código ao redor para de tratar cada chamada como evento isolado. O prompt é quase o mesmo toda vez — e uma vez que o seu loop sabe disso, o SLM passa a competir por custo com soluções muito mais pesadas.