Resumo Técnico: SpecLA

Declaração do Problema

A decodificação autorregressiva continua sendo um gargalo primário no serviço de grandes modelos de linguagem (LLMs), particularmente para cargas de trabalho limitadas por memória, onde a latência é dominada pelo acesso ao cache, movimentação de estado e overhead de lançamento de kernels. Embora a decodificação especulativa tenha acelerado com sucesso modelos Transformer ao verificar múltiplos tokens de rascunho em uma única passagem de alvo, os sistemas existentes dependem da manipulação de caches KV crescentes.

Modelos de atenção linear (ex: Gated DeltaNet ou GDN) substituem o cache KV por um estado recorrente denso de tamanho fixo. Embora isso reduza o crescimento assintótico de memória, o caminho de decodificação permanece limitado pela memória porque cada token requer a leitura, atualização e escrita de um vetor de estado denso (ex: 2 MiB por camada para o GDN-1.3B). Adaptações ingênuas de decodificação especulativa para esses modelos falham por duas razões:

1. Replay de kernel de decodificação: Verificar tokens sequencialmente preserva o custo de ida e volta (round-trip) de estado por token (leitura/escrita no HBM) que a especulação visa amortizar.
2. Reuso de kernel de prefill: Tratar tokens de rascunho como uma sequência de prefill curta incorre em altos custos de configuração (conversão de layout, transformações scan/UT) que não podem ser amortizados sobre janelas especulativas curtas, e falha em lidar com as dependências causais não regulares de rascunhos estruturados em árvore.

Consequentemente, há uma necessidade de um runtime de decodificação especulativa projetado especificamente para alvos de atenção linear baseados em estado, que aborde três desafios acoplados: verificação eficiente residente em SRAM, aceitação de estado denso sem snapshots completos e rascunho alinhado ao alvo.

Metodologia

O artigo apresenta o SpecLA, um runtime de decodificação especulativa que substitui a manipulação de sufixo de cache KV por três mecanismos conscientes do alvo:

1. Kernels de Verificação de Topologia Conscientes

Em vez de forçar todas as propostas através de caminhos padrão de decodificação ou prefill, o SpecLA adapta o cronograma de execução à topologia do rascunho:

• Verificação Serial Residente no Estado: Para candidatos do tipo cadeia, a ordem de execução é reorganizada de token-major para layer-major. Isso permite que tokens adjacentes na mesma camada executem consecutivamente, mantendo os tiles do estado recorrente residentes no SRAM (via tiling de dimensão-V) e evitando round-trips no HBM entre tokens.
• Verificação Paralela com Máscara de Árvore: Para rascunhos de ramificação, um fator completo insere uma máscara de árvore na fatoração UT do Gated DeltaNet (GDN). Isso permite a computação paralela de logits candidatos enquanto garante que as atualizações se propaguem apenas ao longo dos camos ancestrais, evitando vazamento entre irmãos.
• Verificação Híbrida Decomposta em Cadeia: Para equilibrar o baixo overhead dos kernels seriais com o paralelismo dos kernels de árvore, a árvore de rascunho é decomposta em cadeias que respeitam a dependência (usando decomposição leve-pesada). As cadeias são verificadas em paralelo usando o kernel serial residente no estado assim que seus estados de fronteira estão prontos.

2. Gerenciamento de Estado de Fator Aceito

Como candidatos rejeitados não podem ser removidos pela truncagem de um sufixo de cache, o SpecLA desacopla o progresso especulativo do estado durável:

• Buffering de Fatores: Em vez de armazenar snapshots de estado completos (intensivos em memória) ou realizar o replay de tokens (redundantes computacionalmente), o runtime armazena em buffer os fatores compactos por camada (chaves projetadas, valores, decaimento e coeficientes da regra delta) produzidos durante a verificação.
• Atualização de Estado Adiada: Em vez de lançar um kernel independente para consolidar os estados aceitos, o SpecLA funde a atualização do estado com a próxima iteração de verificação. Os fatores aceitos pendentes são aplicados aos tiles de estado que já foram carregados no SRAM no início do próximo kernel, eliminando uma passagem extra de acesso ao estado.

3. Integração de Rascunho Alinhada ao Alvo

O processo de rascunho é adaptado para a dinâmica recorrente do alvo:

• Poda Guiada por Confiança: Caminhos de baixa probabilidade são podados com base na log-probabilidade cumulativa do caminho antes da verificação para evitar desperdício de trabalho dependente de estado.
• Rascunhador Alinhado ao Alvo: Um cabeçalho de rascunho estilo EAGLE é treinado em características extraídas diretamente dos portadores de estado recorrente do alvo de atenção linear, garantindo que as pontuações de proposta reflitam a dinâmica de execução específica do alvo, em vez de estados ocultos genéricos de Transformers.

Principais Contribuições

1. Runtime SpecLA: Um sistema de decodificação especulativa especializado para modelos de atenção linear baseados em estado que supera o descompasso de abstração entre a especulação baseada em cache KV e modelos de estado recorrente.
2. Inovações de Kernel: O design de kernels conscientes da topologia (serial residente no estado, paralelo com máscara de árvore e híbrido decomposto em cadeia) que otimizam a residência de estado e o paralelismo para cargas de trabalho especulativas curtas e irregulares.
3. Estratégia de Gerenciamento de Estado: Uma abordagem inovadora usando buffering de fatores e atualizações de estado adiadas para recuperar estados aceitos de forma eficiente sem snapshots de estado completo ou replay de tokens.
4. Alinhamento de Rascunho: A integração de um rascunhador alinhado ao alvo e poda guiada por confiança para garantir que o trabalho de candidatura submetido seja útil para o alvo recorrente específico.

Resultados

Avaliado em uma GPU NVIDIA H100 usando um modelo alvo GDN-1.3B público:

• Aceleração End-to-End: O SpecLA alcança até 1.70× de aceleração sobre a decodificação autorregressiva no dataset GSM8K, 1.42× em um conjunto de prompts mistos e 1.06× no HumanEval.
• Microbenchmarks:
  • A verificação híbrida decomposta em cadeia melhora em relação ao replay raiz-folha de 1.80× a 7.11×, dependendo da forma da árvore.
  • O reuso de buffer de fator reduz a latência de recuperação de estado aceito em 2.74× a 4.28× comparado ao replay de tokens.
  • A atualização de estado adiada reduz a latência de fronteira de commit em 1.15× a 1.44×.
• Sensibilidade: A aceleração é altamente dependente da taxa de aceitação; o sistema torna-se lucrativo quando a probabilidade de aceitação excede ~0.7, com potenciais acelerações atingindo 3.71× (alvo 1.3B) ou 4.61× (projetado para escala 9B) sob aceitação oracle.

Significância

O artigo argumenta que, embora os modelos de atenção linear eliminem o crescimento do cache KV, eles não eliminam a natureza limitada pela memória da decodificação autorregressiva. O SpecLA demonstra que a decodificação especulativa pode ser efetivamente adaptada a esses modelos, mas apenas mudando fundamentalmente o runtime de manipulação de cache KV para o gerenciamento de estado consciente da topologia. O trabalho estabelece que a especulação eficiente para modelos baseados em estado exige o co-design dos kernels de verificação, do ciclo de vida do estado e da estratégia de rascunho para corresponder à dinâmica recorrente do alvo, em vez de simplesmente conectar kernels de atenção linear em frameworks especulativos existentes baseados em Transformer.