Atenção Esparsa em Blocos Dispersos via Permutação de Tokens

Resumo

A escalabilidade do comprimento de contexto em grandes modelos de linguagem (LLMs) oferece benefícios significativos, mas é computacionalmente dispendiosa. Este custo deriva principalmente do mecanismo de autoatenção, cuja complexidade O(N²) em relação ao comprimento da sequência representa um grande garganto para memória e latência. Felizmente, a matriz de atenção é frequentemente esparsa, particularmente para sequências longas, sugerindo uma oportunidade de otimização. A atenção esparsa por blocos surgiu como uma solução promissora que particiona sequências em blocos e ignora a computação para um subconjunto desses blocos. No entanto, a eficácia deste método é altamente dependente dos padrões de atenção subjacentes, o que pode levar a uma esparsidade em nível de bloco subótima. Por exemplo, tokens-chave importantes para consultas dentro de um único bloco podem estar dispersos por inúmeros outros blocos, levando a redundância computacional. Neste trabalho, propomos a Permuted Block-Sparse Attention (PBS-Attn), um método plug-and-play que aproveita as propriedades de permutação da atenção para aumentar a esparsidade em nível de bloco e melhorar a eficiência computacional do pré-preenchimento de LLMs. Realizamos experimentos abrangentes em conjuntos de dados desafiadores do mundo real com contexto longo, demonstrando que a PBS-Attn supera consistentemente os métodos existentes de atenção esparsa por blocos em precisão do modelo e se aproxima muito da linha de base de atenção completa. Alimentado pelos nossos kernels personalizados de permuted-FlashAttention, a PBS-Attn alcança uma aceleração de ponta a ponta de até 2,75 vezes no pré-preenchimento de contexto longo, confirmando sua viabilidade prática. Código disponível em https://github.com/xinghaow99/pbs-attn.

English

Scaling the context length of large language models (LLMs) offers significant benefits but is computationally expensive. This expense stems primarily from the self-attention mechanism, whose O(N^2) complexity with respect to sequence length presents a major bottleneck for both memory and latency. Fortunately, the attention matrix is often sparse, particularly for long sequences, suggesting an opportunity for optimization. Block-sparse attention has emerged as a promising solution that partitions sequences into blocks and skips computation for a subset of these blocks. However, the effectiveness of this method is highly dependent on the underlying attention patterns, which can lead to sub-optimal block-level sparsity. For instance, important key tokens for queries within a single block may be scattered across numerous other blocks, leading to computational redundancy. In this work, we propose Permuted Block-Sparse Attention (PBS-Attn), a plug-and-play method that leverages the permutation properties of attention to increase block-level sparsity and enhance the computational efficiency of LLM prefilling. We conduct comprehensive experiments on challenging real-world long-context datasets, demonstrating that PBS-Attn consistently outperforms existing block-sparse attention methods in model accuracy and closely matches the full attention baseline. Powered by our custom permuted-FlashAttention kernels, PBS-Attn achieves an end-to-end speedup of up to 2.75times in long-context prefilling, confirming its practical viability. Code available at https://github.com/xinghaow99/pbs-attn

Atenção Esparsa em Blocos Dispersos via Permutação de Tokens

Sparser Block-Sparse Attention via Token Permutation

Resumo

Support