Resumo Técnico: FlashTrie

Declaração do Problema

A recuperação generativa reformula a recuperação de documentos como uma tarefa de sequência-para-sequência, mapeando consultas diretamente para identificadores de documentos (docIDs). Embora esse paradigma ofereça uma modelagem ponta a ponta compacta, o atendimento online é limitado por orçamentos de latência rigorosos. O procedimento de decodificação, particularmente para decodificadores não autorregressivos (NAR), frequentemente gera identificadores inválidos que requerem filtragem via busca restrita.

As abordagens de estado da arte atuais utilizam a busca de feixe (beam search) restrita por trie para garantir que os identificadores gerados correspondam a uma biblioteca predefinida. No entanto, as implementações existentes dependem de representações de trie baseadas em ponteiros (ex: MARISA padrão) executadas em CPUs. Esses designs sofrem de:

1. Acesso irregular à memória: Levando a uma baixa utilização de hardware e alta latência.
2. Paralelismo limitado: À medida que as larguras de feixe crescem, a travessia baseada em CPU e a validação de candidatos tornam-se o gargalo dominante de atendimento.
3. Ineficiência em GPUs: As tries sucintas tradicionais (como a MARISA) são otimizadas para chaves de nível de caractere e mapeiam mal para arquiteturas de GPU, causando divergência de warp e acesso de memória não coalescente.
4. Sobrecarga de Orquestração: A aplicação de restrições é frequentemente implementada fora do caminho central de decodificação da GPU, exigindo sincronização repetida entre host e dispositivo por passo de decodificação.

O artigo aborda três questões fundamentais: Podem as tries sucintas ser adaptadas para vocabulários de nível de token sem sacrificar a eficiência de espaço? A busca de feixe restrita pode ser executada inteiramente na GPU? Sistemas como este podem atender às restrições de latência de produção sem degradar a qualidade da recuperação?

Metodologia: FlashTrie

O FlashTrie é um framework de decodificação restrita nativo de GPU que otimiza conjuntamente a representação de restrições e o cálculo de decodificação. Ele substitui a orquestração baseada em CPU por um kernel CUDA cooperativo totalmente no dispositivo.

1. Estrutura de Dados: Trie Sucinta Consciente de Inteiros

O FlashTrie redesenha o layout da trie MARISA para suportar grandes vocabulários de tokens inteiros (mais de 1 milhão de tokens) mantendo uma pequena pegada de memória adequada para a Memória de Alta Largura de Banda (HBM) da GPU.

• LOUDS Estreito: Os autores estendem a representação LOUDS (Level-Order Unary Degree Sequence) para tokens inteiros de 32 bits.
• Compressão de Bits: Para reduzir o uso de memória, o layout divide os rótulos dos nós e os índices de sufixo em uma parte baixa de 2 bytes (armazenada no array base) e uma parte alta empacotada em bits (armazenada em um array secundário para nós que a requeiram).
• Benefícios: Este "estreitamento" reduz pela metade o tamanho do array dominante, encolhe o índice de restrição (encaixando 800 milhões de palavras-chave em ~3 GB) e melha a localidade de memória ao reduzir o número de bits buscados por acesso.

2. Algoritmo: Busca de Feixe Cooperativa no Dispositivo

O FlashTrie executa toda a busca de feixe restrita dentro de um único kernel CUDA persistente, eliminando a orquestração host-per-step.

• Paralelismo de Dois Níveis:
  • Paralelismo de Feixe (Beam-Parallel): Feixes ativos através de uma consulta são mapeados para diferentes blocos de threads (CTAs).
  • Paralelismo de Proposta (Proposal-Parallel): Dentro de cada feixe, as threads realizam validação paralela. Cada thread testa uma das top-$K$ propostas do modelo de linguagem contra a trie usando busca binária paralela sobre rótulos de filhos ordenados.
• Fluxo de Trabalho do Kernel:
  1. Expansão: Os feixes ativos são expandidos usando as top-$K$ propostas.
  2. Validação: As threads realizam busca binária nos rótulos dos filhos da trie para validar prefixos.
  3. Poda (Pruning): Um kernel de seleção top-$B$ paralelo substitui o tradicional heap de CPU. Ele utiliza um buffer de candidatos de append-only e ordenação paralela para podar os feixes até a largura alvo, evitando contenção de lock e perseguição de ponteiros (pointer chasing).
• Sincronização: Uma barreira em nível de grid sincroniza os passos, permitindo que o kernel execute por $T$ passos de decodificação em um único lançamento.

Principais Contribuições

1. Índice de Restrição Residente na GPU: Um novo layout de trie com compressão de bits que mantém bibliotecas de restrições de escala bilionária inteiramente na HBM da GPU, reduzindo paradas de memória (memory stalls).
2. Kernel CUDA Cooperativo: Um modelo de execução de kernel único que realiza expansão de feixe, validação e poda no dispositivo, removendo o overhead de lançamentos repetidos de kernels e orquestração do host.
3. Primitivas Conscientes de GPU: A substituição da travessia irregular baseada em ponteiros por busca binária paralela e seleção top-$K$ livre de heap, melhorando significativamente a utilização de warp e reduzindo a divergência.
4. Escalabilidade: O sistema escala a decodificação restrita para 800 milhões de palavras-chave em uma única GPU, um regime anteriormente viável apenas em configurações offline.

Resultados Experimentais

O sistema foi avaliado em um modelo NAR de produção com um vocabulário de 2,2 milhões de tokens e uma trie de restrição de 800 milhões de palavras-chave, rodando em uma GPU NVIDIA A100 80GB.

Métricas de Desempenho

• Latência: O FlashTrie reduz a latência de busca na trie para menos de 3 ms (p99) mesmo com uma largura de feixe ($BW$) de 1000.
  • Com $BW=1000$, o FlashTrie atinge uma latência média de 1,91 ms, comparado a 46,30 ms para a linha de base otimizada de CPU (MARISA-Opt).
  • Isso representa uma aceleração de 24× sobre a linha de base de CPU multithread altamente otimizada.
• Throughput: Sob atendimento em lote (batching) ($b=8$), o FlashTrie atinge o pico de 1.621 consultas/segundo para $BW=1000$, entregando um throughput 32×–71× maior que a linha de base de CPU.
• Construção: O FlashTrie constrói o índice 3,4× mais rápido e usa 22% menos espaço em disco que o MARISA-Opt.

Qualidade de Recuperação

• Paridade com CPU: O FlashTrie iguala a qualidade de recuperação da linha de base de CPU (MARISA-Opt) dentro de 0,001 de precisão em todas as larguras de feixe.
• Ganhos de Qualidade via Largura de Feixe: Como o FlashTrie mantém baixa latência, ele permite o uso de larguras de feixe muito maiores (até 1000) que são inviáveis na CPU. Aumentar a largura do feixe de 100 para 1000 melhora a Precisão@100 de 0,53 para 0,78.

Teste A/B Online

Em um experimento online de larga escala em um motor de busca patrocinada comercial:

• Configuração: O grupo de tratamento usou FlashTrie com $BW=600$, enquanto o grupo de controle usou MARISA-Opt com $BW=200$ (o máximo viável para CPU dentro dos limites de latência).
• Resultados: O FlashTrie entregou um aumento de receita estatisticamente significativo de +0,71% e um aumento de 0,17–0,20% nas taxas de clique (CTR) sem degradar a qualidade do anúncio (as taxas de defeito permaneceram não significativas).
• Latência: O tratamento manteve orçamentos de latência rigorosos, reduzindo a latência p95 em 32% em comparação com o limite do controle.

Significância e Alegações

O artigo argumenta que a decodificação restrita não precisa permanecer um gargalo de CPU. Ao redesenhar as estruturas de dados para execução paralela e mover toda a lógica de aplicação de restrições para a GPU, o FlashTrie permite:

1. Escalabilidade em Tempo Real: Permite que sistemas de recuperação generativa escalem o tamanho do feixe em até 5× em aplicações críticas de latência (como busca patrocinada) sem violar os orçamentos de atendimento.
2. Resolução do Trade-off Qualidade-Desempenho: Desacopla a qualidade da recuperação das restrições de latência, permitindo que os sistemas aproveitem larguras de feixe maiores para melhor seleção de candidatos enquanto mantêm baixa latência.
3. Aplicabilidade Geral: A abordagem é aplicável a qualquer tarefa de geração restrita com grandes vocabulários estáticos, como vinculação de entidades ou recuperação de catálogo de produtos.

Os autores concluem que o FlashTrie demonstra um padrão mais amplo onde estruturas de dados tradicionalmente implementadas em CPUs podem se tornar altamente eficazes em GPUs quando seus padrões de acesso são redesenhados para execução em lote e paralela.