ReplaySSM repensa armazenamento de estado e acelera decodificação de modelos híbridos de IA em até 1,96x
Uma técnica chamada ReplaySSM, desenvolvida pelo Dao AI Lab, promete acelerar a decodificação de modelos híbridos de inteligência artificial ao mudar a forma como o estado interno das camadas de modelo de espaço de estados é armazenado na memória. Em vez de gravar o estado completo a cada passo de geração, o método registra as entradas recentes e reconstrói o estado somente quando necessário, reduzindo pela metade o tráfego de memória em cada etapa. Os detalhes técnicos foram apresentados no blog do Dao AI Lab, do pesquisador Tri Dao.
Modelos híbridos intercalam camadas de modelo de espaço de estados, conhecidas pela sigla SSM, com algumas camadas de atenção tradicionais. NVIDIA Nemotron 3, Qwen3.5 até Qwen3.8, Kimi-K3 e GLM-5.3-Flash foram construídos dessa forma, com Mamba-2, Gated DeltaNet (GDN) ou Kimi Delta Attention (KDA) ocupando o papel da camada SSM. A proposta dessas camadas é tornar a decodificação rápida: elas usam um estado de tamanho fixo no lugar de um cache chave-valor (KV) que cresce continuamente conforme o texto se alonga.
Apesar dessa vantagem, a decodificação com camadas SSM continua limitada pela largura de banda de memória. Em cada passo, o sistema lê o estado recorrente da memória de alta largura de banda das placas de vídeo (HBM), atualiza os valores e escreve tudo de volta. Esse estado é grande — no Mamba-2, suas dimensões dependem do número de cabeças, da dimensão de cada cabeça e da dimensão do estado em cada camada — e a escrita de volta domina o tempo de processamento quando a decodificação roda com lote de tamanho 1.
O ReplaySSM ataca exatamente esse ponto. Ele transforma a sequência de leitura e escrita do estado completo em uma operação de leitura acompanhada de uma gravação apenas periódica. Como a atualização recorrente é linear, o estado pode ser reconstruído sob demanda a partir de um ponto de controle mais um pequeno buffer circular contendo os registros dos últimos passos. A saída pode ser lida sem que a matriz de estado completa precise ser materializada nos passos intermediários.
O efeito prático é direto: o estado inicial ainda é lido a cada passo, mas a escrita do estado completo passa a ocorrer apenas a cada intervalo fixo, e não a todo momento. O tráfego de memória por passo cai de leitura mais escrita, aproximadamente o dobro do tamanho do estado, para somente leitura, cerca de uma vez o tamanho do estado, além de inserções mínimas no buffer circular. Esses registros são desprezíveis diante do estado de valores e chaves, o que representa uma redução aproximada de metade no tráfego.
Importante: o método não altera o cálculo executado pelo modelo. A mesma recorrência de estado cumulativo é mantida, mudando apenas o momento em que o estado completo é materializado e gravado. Segundo os autores, o resultado não é um novo comportamento do modelo, mas a mesma computação percorrendo um caminho de memória diferente, com tráfego menor, verificação especulativa mais barata e capacidade de atender mais requisições simultâneas.
O ganho mais expressivo aparece na decodificação especulativa, técnica em que rascunhos de tokens são gerados, verificados e depois confirmados ou descartados. Na implementação existente do vLLM, framework de inferência de código aberto, o caminho especulativo chega a ficar abaixo da decodificação padrão em tamanhos de lote típicos de serviço. Com o ReplaySSM, a escrita de estado completo deixa de ocorrer a cada passo de verificação e passa a acontecer a cada lote de tokens confirmados. Como o método guarda as entradas brutas de cada rascunho, reverter propostas rejeitadas se torna apenas um movimento de ponteiro no buffer, sem necessidade de regravar o estado.
Os números medidos cobrem duas famílias de modelos e tamanhos que vão de 4 a 550 bilhões de parâmetros. No Nemotron-3, o ReplaySSM alcançou aceleração de 1,43x a 1,84x no kernel e de 1,20x a 1,48x de ponta a ponta. No Qwen3.5, os ganhos foram de 1,43x a 1,64x no kernel e de 1,20x a 1,27x de ponta a ponta. Somando os dois cenários, a decodificação padrão do vLLM ficou até 1,48x mais rápida, chegando a 1,43x nos modelos MoE (mistura de especialistas) de grande porte.
Na decodificação especulativa, o avanço foi ainda maior: entre 1,87x e 1,96x sobre a decodificação padrão do vLLM. O método também sustentou de 3,0 a 3,3 vezes mais requisições simultâneas sob um orçamento fixo de memória. O ganho de ponta a ponta é menor que o do kernel porque o ReplaySSM atinge apenas a parte do SSM, enquanto a atenção, as multiplicações de matrizes e o restante da execução permanecem inalterados.
A implementação é um fork de pesquisa do vLLM. A troca feita pelo método — recomputar um pouco para reduzir o tráfego na memória — segue o mesmo equilíbrio entre leitura e recomputação que o FlashAttention popularizou. Os ganhos de vazão foram observados em placas H100 e B300, incluindo os modelos Nemotron-Nano-4B e Qwen3.5 nas versões de 4B e 122B. Uma proposta para incorporar a técnica ao vLLM já está em discussão aberta no repositório do projeto no GitHub.
O ReplaySSM mostra que o gargalo da decodificação híbrida está menos na computação e mais na forma de armazenar o estado. Ao registrar entradas em vez de regravar estados, mantém o resultado idêntico, acelera a decodificação padrão em até 1,48x, quase dobra o desempenho especulativo e amplia a concorrência sob o mesmo limite de memória. A explicação completa da técnica está disponível no blog do Dao AI Lab.
