KAIST elimina a calibração do decoding em MoE e acelera blocos em até 2,69× frente ao vLLM

O decoding de Mixture-of-Experts em batches pequenos esbarra na memória, não no cálculo: a cada token, é preciso carregar os pesos dos experts roteados. Os pesquisadores do KAIST AI Geunhyeon Kim e Sung Ju Hwang mostraram que uma técnica popular de aceleração, a seleção de canais importantes dentro de cada expert, pode funcionar sem dados de calibração. O leitor deles, construído apenas sobre os pesos do modelo, preserva 98,1% da precisão do modelo completo lendo 17,8% dos pesos dos experts. Os blocos MoE com kernels em Triton rodam de 1,60 a 1,90 vez mais rápido que o vLLM otimizado em batch one.

O limiar promete um orçamento, mas não o cumpre

A ideia da seleção de canais é simples. Em cada expert, a maioria dos canais intermediários quase não influencia um token específico. Um seletor prevê quais canais importam e lê da memória apenas os pesos deles. Normalmente, escolhem-se os canais cujo score fica acima de um limiar, ajustado em uma amostra offline.

Os autores apontam duas fraquezas nesse esquema. A primeira: o limiar não fixa o volume de leitura, porque o número de canais que passam pelo corte varia de entrada para entrada. O Prox, com limiares calibrados para um orçamento de 18%, leu em média 12,9% nos benchmarks do estudo. Quem planeja latência e throughput acaba com um número diferente do esperado.

A segunda: o ganho de precisão para um dado orçamento também não é garantido.

Top-k no lugar do limiar

A alternativa dispensa dados por completo: para cada token, mantêm-se os k canais com os maiores scores. O orçamento de leitura é conhecido de antemão e coincide com o declarado. Em comparações com os mesmos scores, a versão top-k do seletor igualou ou superou a versão com limiar em 40 de 43 casos, com orçamentos de 13% e 18%. Nas tarefas de geração, venceu em todas as 23 comparações.

Há ainda um efeito qualitativo. Limiares calibrados levaram o Gemma 4 a entrar em loops de repetição, e as versões top-k evitam isso. A conclusão dos autores: a calibração do limiar não entrega nenhuma das duas coisas que justificam seu uso, ou seja, nem um orçamento de leitura conhecido, nem um ganho estável de precisão nesse orçamento.

Um leitor feito só de pesos

A segunda parte do trabalho é um leitor próprio, que não precisa de dados de calibração nem de fine-tuning. Ele funciona em dois passos:

• os canais passam por uma triagem grosseira com base em um esboço ternário (ternary sketch) dos pesos;
• os canais descartados são compensados por um mapeamento de baixo rank, também construído a partir dos pesos do modelo.

Principais números em três modelos MoE pós-treinados:

• precisão: 98,1% do modelo completo;
• volume de leitura dos pesos dos experts roteados: 17,8%;
• velocidade dos blocos MoE com kernels em Triton em batch one: de 1,60 a 1,90 vez maior que a do vLLM otimizado;
• em batches maiores, a aceleração chega a 2,69 vezes nos mesmos três modelos;
• o Prox, calibrado para 18%, leu na prática 12,9%;
• top-k contra limiar: 40 de 43 comparações com orçamentos de 13% e 18%, 23 de 23 em geração.

Para quem isso é útil

O maior ganho vai para quem roda grandes modelos MoE em hardware com memória limitada. Em uma discussão no Hugging Face, um dos autores afirmou que o problema do pruning de ativações é agudo em GPUs com memória lenta, como o DGX Spark. Nesse tipo de hardware, a largura de banda da memória, e não os FLOPS, define a velocidade de geração de tokens. Segundo ele, o próximo regime de esparsidade virá justamente da esparsidade de ativações no serviço de modelos em dispositivos com limite rígido de banda.

A vantagem prática de dispensar a calibração é que o método pode ser aplicado a qualquer modelo open-weights recém-lançado, logo após a publicação. Não é preciso buscar uma amostra representativa, e não há o risco de um limiar ajustado em um domínio quebrar a geração em outro, como aconteceu com o Gemma 4. Um orçamento de leitura fixo também simplifica as contas: os 17,8% previstos se traduzem em economia previsível de tráfego de memória e de latência.

O código foi publicado no repositório calibration-free-moe, no GitHub. Para equipes que escrevem seus próprios kernels em Triton ou otimizam inferência local, é um ponto de partida pronto. Vale testá-lo nos próprios modelos, principalmente em cenários interativos com batch one, onde o gargalo de memória é maior.