Stratified Sparse Attention acelera a decodificação em 2,96 vezes com contexto de 1M

Uma equipe de pesquisadores propôs o Stratified Sparse Attention (SSA), um método que acessa apenas uma pequena fração do KV-cache e perde bem menos qualidade do que os esquemas esparsos anteriores. No RULER-HARD com o Ministral-3-8B, com atenção 50 vezes mais esparsa, o SSA eleva a precisão de 81,8% para 84,2% em contexto de 32K e de 63,8% para 69,3% em 128K. Isso corresponde a 94% e 85% do resultado da atenção densa (dense) completa. Com contexto de um milhão de tokens, os kernels CUDA otimizados aceleram a decodificação em até 2,96 vezes em relação ao FlashInfer.

Onde a atenção esparsa perde qualidade

Em contextos longos, o gargalo da inferência deixa de ser tanto o cálculo e passa a ser a leitura do KV-cache, que cresce a cada passo da decodificação. A atenção esparsa reduz esse custo ao consultar só uma parte do cache. Aqui há uma bifurcação.

O primeiro caminho é a seleção determinística. Top-k ou top-p escolhem os tokens com as maiores pontuações. Isso captura com confiabilidade o essencial, mas todo o restante da atenção, que individualmente é pequeno e somado pode ser relevante, simplesmente é descartado. O segundo caminho é a amostragem. Ela cobre melhor a cauda da distribuição, mas paga o preço da variância: o resultado oscila mais de uma execução para outra.

Os métodos híbridos combinam top-k com amostragem uniforme. Os autores mostram que, com esparsidade agressiva, a distância para a atenção densa continua grande. A amostragem uniforme é cega: gasta o orçamento em tokens que quase não influenciam a resposta.

Como o SSA divide o cache em estratos

O SSA monta um esquema unificado com duas partes. Os tokens com pontuações altas são lidos de forma determinística. O restante do cache não é descartado nem escolhido ao acaso: seus tokens são amostrados com probabilidades proporcionais às estimativas proxy de atenção. Para que a contribuição dos tokens escolhidos não distorça o resultado, o SSA aplica uma correção pela probabilidade de inclusão de cada um na amostra. Assim, recupera-se a massa de atenção que fica fora do conjunto selecionado.

A ideia se aproxima da amostragem estratificada clássica: a região importante é coberta por inteiro, e o restante é estimado a partir de uma subamostra com pesos. Os autores ressaltam que o SSA é uma generalização. Top-k, top-p e o híbrido de top-k com amostragem são casos particulares dele, de modo que o método pode ser visto como um arcabouço geral para toda essa família.

Números e aceleração

A avaliação usou quatro benchmarks (RULER, LongBench, LOFT e AIME) e cinco modelos de 3B a 27B parâmetros. Com a mesma esparsidade, o SSA supera de forma consistente os métodos de referência. Principais resultados:

• RULER-HARD, Ministral-3-8B, 32K: 81,8% → 84,2% contra o baseline prático mais forte, 94% da atenção densa;
• RULER-HARD, Ministral-3-8B, 128K: 63,8% → 69,3%, 85% da atenção densa;
• modo: atenção com esparsidade de 50x;
• velocidade: até 2,96x na decodificação em relação ao FlashInfer com contexto de 1M;
• código aberto: github.com/skylight-org/SSA-Paper-Code.

Há ainda uma segunda parte importante no resultado. O método de maior precisão costuma exigir operações mais complexas, e a aceleração pode desaparecer. Aqui, os autores mostram, com kernels CUDA otimizados, que o ganho de qualidade não é consumido pelo overhead. O aumento de precisão em 128K é maior que em 32K (5,5 p.p. contra 2,4 p.p.). Ou seja, o método ajuda mais onde o contexto é mais longo e a esparsidade pesa mais.

O que muda para quem opera esses sistemas

Para as equipes que servem contexto longo, o KV-cache é um item direto de custo em memória de GPU e largura de banda. Se uma redução de 50 vezes nos acessos ao cache preserva de 85% a 94% da qualidade da atenção densa em tarefas difíceis, os engenheiros ganham mais uma alavanca entre custo e precisão. Isso é especialmente interessante para cenários agênticos e para trabalho com grandes bases de código, em que o contexto cresce rápido.

As limitações também são claras. Parte dos números se refere a um único modelo, o Ministral-3-8B, e em 128K ainda há uma defasagem visível em relação ao modo denso: 85% é um bom resultado, mas não é paridade. Para uso em produção, será preciso testar o método nas tarefas da própria equipe e verificar a compatibilidade dos kernels com o stack em uso. O código aberto, por outro lado, reduz a barreira: é possível reproduzir os experimentos e comparar o SSA com os esquemas top-k atuais sem esperar por uma implementação nos frameworks mais populares.