Como colocar um milhão de tokens em uma GPU comum usando Kimi Linear

Você já se perguntou por que trabalhar com textos longos em LLMs é tão caro? Tudo se resume ao KV-cache. Quando você alimenta um documento massivo em um modelo, sua "memória" (o cache) incha até tamanhos obscenos, devorando toda a sua memória de vídeo. A galera da MoonshotAI decidiu que era hora de fazer algo a respeito, e lançou o Kimi Linear. Isso não é apenas mais um modelo "melhorado"—é uma tentativa de repensar a arquitetura de atenção para que possamos trabalhar com um contexto de um milhão de tokens sem precisar comprar uma fazenda de servidores.

Qual é o problema com a atenção padrão

O mecanismo padrão de Full Attention (o dos transformers clássicos) é uma fera faminta. Sua complexidade cresce quadraticamente com o tamanho do texto. Quer um contexto duas vezes maior? Prepare-se para gastar quatro vezes mais recursos. Soluções populares como Flash Attention ou MLA (Multi-head Latent Attention) do DeepSeek ajudam, mas não resolvem o problema radicalmente quando se trata de sequências realmente longas.

O Kimi Linear adota uma abordagem diferente. Os desenvolvedores usaram um método híbrido combinando os pontos fortes dos transformers e estruturas semelhantes a RNNs.

Como funciona o Kimi Delta Attention

O núcleo do projeto é o mecanismo Kimi Delta Attention (KDA). Sem mergulhar em matemática pesada, é uma evolução do conceito de Gated DeltaNet. O truque principal aqui é o "esquecimento" inteligente.

Em um RNN padrão, a memória é limitada por um tamanho de estado fixo. O KDA emprega um mecanismo de gating que determina quais informações passadas devem ser retidas nesse estado comprimido e quais podem ser descartadas. Isso permite que o modelo mantenha alta precisão mesmo em distâncias vastas, onde modelos lineares convencionais começam a "derivar" e perder o fio da narrativa.

O que isso entrega na prática

Os desenvolvedores introduziram uma arquitetura onde KDA e MLA (Global Attention) são misturados em uma proporção de 3:1. Essa combinação alcançou vários resultados impressionantes:

1. Economia de memória. Os requisitos de KV-cache caíram 75%. Isso é crítico ao fazer deploy do modelo no seu próprio hardware.
2. Velocidade de geração. Em um contexto de 1 milhão de tokens, a taxa de transferência de tokens aumenta até 6x em comparação com arquiteturas padrão.
3. Contexto honesto. Testes de benchmark no RULER mostram que o modelo genuinamente "vê" e usa informações ao longo de todos os 128k (e até 1M) tokens, não apenas fingindo.

Os gráficos abaixo demonstram como o Kimi Linear (linha azul) se destaca em velocidade à medida que o tamanho do contexto cresce:

Testando na prática

A MoonshotAI não segurou e lançou os pesos do modelo no Hugging Face. Há uma versão base e uma variante Instruct com 48 bilhões de parâmetros. Graças à arquitetura de Mixture-of-Experts (MoE), apenas 3 bilhões de parâmetros são ativados durante a computação, tornando o modelo surpreendentemente leve para sua classe.

Para começar, você vai precisar do PyTorch mais recente e da biblioteca fla-core. O código parece bastante padrão para quem já trabalhou com transformers:

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

# Обычный чат-шаблон
messages = [
    {"role": "system", "content": "You are a helpful assistant provided by Moonshot-AI."},
    {"role": "user", "content": "Расскажи, в чем преимущество линейного внимания перед обычным?"}
]

input_ids = tokenizer.apply_chat_template(
    messages, 
    add_generation_prompt=True, 
    return_tensors="pt"
).to(model.device)

generated_ids = model.generate(inputs=input_ids, max_new_tokens=500)
response = tokenizer.batch_decode(generated_ids)[0]
print(response)

Se você precisa fazer deploy em produção, o modelo é compatível com vLLM. Você pode subir uma API compatível com OpenAI com um único comando no terminal, especificando um max-model-len massivo.

Vale a pena baixar

O projeto parece promissor para quem está construindo sistemas RAG ou analisando logs e documentos longos.

Quem definitivamente deveria dar uma olhada mais de perto:

• Quem atingiu o teto de memória das GPUs ao trabalhar com contextos longos.
• Desenvolvedores que se importam com a velocidade de resposta (TPOT) em tempo real.
• Pesquisadores buscando alternativas aos transformers padrão.

Por outro lado, a arquitetura é relativamente nova, e suporte em ferramentas de terceiros (como quantização ou otimizadores específicos) pode não chegar imediatamente. Mas ter kernels KDA prontos na biblioteca FLA é encorajador.

O Kimi Linear é um bom exemplo de que otimização de algoritmos ainda pode render ganhos maiores do que simplesmente empilhar mais teraflops. Se você precisa "digerir" uma biblioteca inteira ou uma base de código massiva em uma única passagem—esta é provavelmente uma das ferramentas mais interessantes disponíveis agora.