DashAttention alcança 75% de esparsidade mantendo precisão da atenção completa

Uma equipe de pesquisa de Edimburgo, Cambridge, Instituto Superior Técnico e Tsinghua publicou DashAttention, um substituto direto para a atenção hierárquica top-k usada em NSA e InfLLMv2 que mantém o estágio de seleção KV totalmente diferenciável. Todo o pipeline de dois estágios pode agora ser ajustado ponta-a-ponta. O artigo (arXiv 2605.18753, postado em 18 de maio) inclui um kernel GPU Triton e quatro checkpoints MiniCPM-4-8B ajustados finamente liberados no Hugging Face para comparação direta.

O problema central: NSA e InfLLMv2 aplicam uma etapa de seleção de bloco top-k coarse, e então executam atenção softmax padrão apenas sobre os blocos selecionados. Top-k assume que toda query precisa do mesmo número de blocos de tokens relevantes. A operação de seleção discreta corta o fluxo de gradiente entre os estágios coarse e fine, então o scorer coarse não consegue aprender com a perda de atenção downstream. Os dois estágios treinam com objetivos desalinhados.

DashAttention substitui top-k por α-entmax, uma transformação sparse diferenciável que generaliza softmax. Tokens cujos scores caem abaixo de um limiar adaptativo por query recebem exatamente zero probabilidade. O limiar é determinado pela distribuição de entrada em vez de um k fixo. Uma query com um bloco de contexto dominante atende a fewer blocos que o máximo configurado; uma query com muitos chunks relevantes atende a mais. A saída de α-entmax atua como uma ponderação prévia para a atenção softmax de segundo estágio. Como α-entmax é diferenciável, gradientes fluem do estágio softmax completamente através do seletor de blocos. Os autores chamam essa propriedade de "não-dispersiva" — a massa de atenção se concentra em blocos relevantes em vez de se espalhar por blocos irrelevantes.

DashAttention iguala a qualidade de atenção completa em 75% de esparsidade e alcança uma fronteira Pareto melhor que NSA e InfLLMv2. A lacuna aumenta em regimes de alta esparsidade. A precisão dos competidores se degrada mais rapidamente acima de intervalos de esparsidade confortáveis, enquanto a seleção adaptativa de DashAttention retém mais precisão em escala. Benchmarks rodaram em MiniCPM-4-8B base. Todas as quatro variantes — atenção completa, InfLLMv2, NSA e DashAttention — estão disponíveis no Hugging Face sob a organização fasa-org. Os autores relatam um speedup sobre FlashAttention-3 em tempo de inferência; o multiplicador exato requer a Seção 5 do artigo.

A implementação Triton é instalável via pip install -e . de github.com/fasa-org/dash-attention. A interface envolve queries, keys, values e um vetor de classificação por cabeça (head_cls) para scoring coarse. Opcionalmente retorna uma máscara active_blocks indicando quais blocos foram selecionados por query, útil para analisar esparsidade real em tempo de execução. Os flags chunk_size e estimate_diagonal controlam granularidade de bloco e correção diagonal para normalização α-entmax. GQA é suportado via enable_gqa=True.

Antes da implantação em produção, duas lacunas permanecem. O artigo está em pré-revisão por pares e reivindicações de eficiência precisam verificação além do abstract. Dados de TTFT e throughput de geração de tokens não são divulgados em materiais públicos. A etapa de seleção de bloco α-entmax adiciona uma chamada kernel extra comparada a métodos top-k mais simples; overhead em comprimentos de contexto curtos (abaixo de 16k) é não caracterizado. O vetor de classificação head_cls é uma adição arquitetônica que não existe em checkpoints Llama-family padrão. Adotar DashAttention requer ajuste fino — não pode ser trocado como uma otimização em tempo de serving sem retreinamento. Os quatro checkpoints 8B liberados fornecem um ponto de partida reprodutível, mas custos de ajuste fino em escala de produção não são divulgados.

Se você está ajustando finamente para tarefas de contexto longo acima de 32k e NSA ou InfLLMv2 está em sua lista de seleção, a seleção diferenciável e esparsidade adaptativa justificam uma comparação direta em sua suite de avaliação. Aguarde os números exatos de eficiência antes de dimensionar orçamento de GPU.