Resumo Técnico: Qrita - Top-k e Top-p de Alto Desempenho usando Truncamento e Seleção Baseados em Pivot

Declaração do Problema

A implementação eficiente dos algoritmos de amostragem Top-k e Top-p para Grandes Modelos de Linguagem (LLMs) com vocabulários grandes (frequentemente excedendo 100.000 tokens) permanece um desafio significativo em GPUs. Abordagens existentes enfrentam um compromisso entre eficiência computacional e correção de saída:

• Métodos baseados em ordenação: Utilizados por engines como vLLM, estes fornecem resultados exatos e determinísticos, mas incorrem em alto custo computacional e de memória devido à ordenação de vocabulário completo e acessos à memória com stride.
• Métodos estocásticos: Bibliotecas como FlashInfer e SGLang usam amostragem por rejeição para evitar ordenação, oferecendo alta vazão. No entanto, retornam apenas um único token amostrado em vez do conjunto completo de candidatos, tornando-as incompatíveis com tarefas que exigem conjuntos de candidatos determinísticos, como verificação de decodificação especulativa e aprendizado por reforço.
• Métodos de corte estatístico: Abordagens como Top-nσ e Min-p eliminam a ordenação usando limiares estatísticos (por exemplo, média e desvio padrão). Embora eficientes, alteram a distribuição de amostragem e a saída, levando a debates sobre a qualidade e coerência da geração.

Existe uma lacuna crítica para um algoritmo que entregue eficiência de memória e comportamento determinístico estrito (retornando o conjunto completo de candidatos) sem sacrificar a semântica exata de Top-k e Top-p.

Metodologia

Os autores propõem o Qrita, um algoritmo de Top-k e Top-p de alto desempenho para GPUs baseado em truncamento e seleção baseados em pivot. O Qrita estende a estratégia baseada em pivot do RTop-k (originalmente projetado para Redes Neurais de Grafos) para lidar com vocabulários grandes de LLMs, garantindo ao mesmo tempo o determinismo. O método depende de duas técnicas principais:

1. Truncamento baseado em $\sigma$ Gaussiano

Reconhecendo que os logits de LLM geralmente seguem uma distribuição quase Gaussiana onde tokens de alta probabilidade formam outliers contra um fundo "ruidoso", o Qrita emprega uma etapa de pré-filtragem.

• Mecanismo: O algoritmo aproxima a distribuição de logits como $Z \sim \mathcal{N}(\mu, \sigma^2)$. Calcula a média ($\mu$) e o desvio padrão ($\sigma$) a partir de um bloco de amostra e calcula um limiar de truncamento $t = \mu + \delta \cdot \sigma$.
• Tabelas de Busca: O coeficiente $\delta$ é determinado via tabelas de busca pré-calculadas derivadas da função quantil normal, adaptadas aos valores alvo de $k$ ou $p$.
• Resultado: Esta etapa truncar a vasta região de "ruído" do vocabulário em uma única passagem. Se o tamanho do conjunto truncado $|Z_t|$ exceder o alvo $k$ (um "acerto de truncamento"), a busca prossegue apenas neste subconjunto menor. Caso contrário, retorna ao conjunto completo de logits. Isso reduz significativamente o espaço de busca para a subsequente busca por pivot.

2. Busca Quaternária por Pivot com Tratamento de Duplicatas

Para encontrar o valor exato do pivot $\tau$ (onde exatamente $k$ elementos são $\ge \tau$) de forma eficiente e determinística:

• Busca Quaternária: Em vez de uma busca binária, o Qrita avalia três pivots (nos percentis 25º, 50º e 75º do intervalo de busca atual) em cada iteração. Isso reduz o número de iterações de busca pela metade em comparação com a busca binária.
• Tratamento de Duplicatas: Buscas por pivot padrão podem falhar ou tornar-se não determinísticas quando existem valores de logits duplicados (comuns em distribuições de baixa precisão ou cauda longa). O Qrita rastreia o valor mínimo de logit maior que o pivot ($z_{mn}$) e sua contagem de ocorrências ($N_{zmn}$).
• Condição de Terminação: A busca termina quando um pivot $\tau_n$ é encontrado tal que a contagem de elementos estritamente maiores que $\tau_n$ é menor que $k$, mas a contagem de elementos maiores ou iguais ao valor duplicado mínimo $z_{mn}$ é maior ou igual a $k$. Isso permite que o algoritmo determine precisamente quantas instâncias duplicadas do valor de fronteira devem ser retidas para satisfazer a restrição Top-k, garantindo saídas reprodutíveis e determinísticas.

Detalhes de Implementação

• Framework: Implementado usando Triton, uma linguagem de programação para GPU.
• Autotuning: Emprega uma estratégia de "Autotune de Tamanho de Bloco em Níveis do Triton". Como o tamanho do vetor varia drasticamente entre o conjunto completo de logits e o conjunto truncado, o compilador pode selecionar tamanhos de bloco ótimos (por exemplo, blocos maiores para conjuntos completos, menores para conjuntos truncados) para maximizar a utilização de hardware.
• Pipeline: O algoritmo processa lotes em paralelo, aplicando truncamento Gaussiano, seguido pela busca quaternária e, finalmente, geração de máscara de tratamento de duplicatas para produzir a saída final.

Contribuições Principais

1. Amostragem Determinística de Alto Desempenho: O Qrita fornece o conjunto completo de candidatos Top-k/Top-p (diferente de FlashInfer/SGLang) mantendo a semântica exata de algoritmos baseados em ordenação, tornando-o adequado para decodificação especulativa e RL.
2. Redução do Espaço de Busca: Ao aproveitar as propriedades estatísticas dos logits, o truncamento Gaussiano $\sigma$ reduz drasticamente o espaço de busca, evitando a ordenação de vocabulário completo.
3. Busca por Pivot Robusta: A introdução da busca quaternária combinada com tratamento explícito de duplicatas garante comportamento determinístico e reduz contagens de iteração, abordando as ineficiências de métodos padrão baseados em pivot em vocabulários grandes.
4. Integração: O Qrita é implementado como o amostrador Top-k e Top-p padrão para o caminho de execução em GPU do vLLM.

Resultados

Os autores avaliaram o Qrita contra baselines incluindo ordenação ingênua do PyTorch, amostrador padrão do vLLM, FlashInfer e SGLang em GPUs NVIDIA H100, RTX4090 e AMD MI355X usando vários modelos (Llama, Qwen, GPT-OSS, Gemma).

• Vazão: O Qrita alcança até 1,4$\times$ de melhoria na vazão de serviço ponta a ponta sobre a baseline do vLLM no H100, com algumas configurações mostrando até 2$\times$ de melhoria na RTX4090. Supera significativamente as baselines baseadas em ordenação enquanto corresponde à saída de algoritmos baseados em ordenação.
• Uso de Memória: O Qrita reduz o uso de memória em aproximadamente 50% comparado a métodos baseados em ordenação, principalmente evitando a alocação de grandes buffers intermediários necessários para ordenação.
• Estudos de Ablação:
  • Truncamento Gaussiano: Proporciona o maior aumento de velocidade, reduzindo a latência em ~74% comparado à busca no conjunto completo de logits.
  • Busca Quaternária: Embora ligeiramente mais lenta que a busca binária em casos de "acerto de truncamento" devido à redução do espaço de busca, proporciona uma redução significativa de latência (~2,3ms) em casos de fallback onde o vocabulário completo deve ser buscado.
  • Tratamento de Duplicatas: Adiciona sobrecarga mínima (~3,5% de aumento de latência), mas é crítico para correção e determinismo.
  • Autotuning: Reduz a latência de execução em quase 30% comparado ao uso de um único tamanho de bloco fixo.

Significado e Alegações

O artigo afirma que o Qrita preenche a lacuna entre eficiência de memória e comportamento determinístico estrito para decodificação de vocabulário grande. Afirma que, ao evitar a ordenação de vocabulário completo e aproveitar a seleção baseada em pivot com pré-filtragem estatística, o Qrita oferece uma solução de alto desempenho para decodificação de LLM compatível com tarefas avançadas que exigem conjuntos completos de candidatos (por exemplo, decodificação especulativa). Os autores posicionam o Qrita não como um substituto para as semânticas estatísticas de Top-k/Top-p, mas como uma implementação eficiente que as preserva exatamente, diferentemente de métodos de corte estatístico heurísticos ou amostradores estocásticos de token único. A integração no vLLM como amostrador padrão destaca sua utilidade prática em ambientes de produção.