Atención Dispersa en Bloques Dispersos mediante Permutación de Tokens

Resumen

Escalar la longitud de contexto de los modelos de lenguaje grandes (LLM) ofrece beneficios significativos, pero resulta computacionalmente costoso. Este costo proviene principalmente del mecanismo de autoatención, cuya complejidad O(N^2) con respecto a la longitud de la secuencia representa un cuello de botella importante tanto para la memoria como para la latencia. Afortunadamente, la matriz de atención a menudo es dispersa, particularmente para secuencias largas, lo que sugiere una oportunidad de optimización. La atención bloque-dispersa ha surgido como una solución prometedora que divide las secuencias en bloques y omite el cálculo para un subconjunto de estos bloques. Sin embargo, la efectividad de este método depende en gran medida de los patrones de atención subyacentes, lo que puede generar una dispersión a nivel de bloques subóptima. Por ejemplo, los tokens clave importantes para las consultas dentro de un solo bloque pueden estar dispersos en numerosos otros bloques, lo que genera redundancia computacional. En este trabajo, proponemos Permuted Block-Sparse Attention (PBS-Attn), un método plug-and-play que aprovecha las propiedades de permutación de la atención para aumentar la dispersión a nivel de bloques y mejorar la eficiencia computacional del prellenado de LLM. Realizamos experimentos exhaustivos en conjuntos de datos desafiantes del mundo real de contexto largo, demostrando que PBS-Attn supera consistentemente a los métodos de atención bloque-dispersa existentes en precisión del modelo y se acerca mucho al baseline de atención completa. Impulsado por nuestros kernels personalizados de permuted-FlashAttention, PBS-Attn logra una aceleración de extremo a extremo de hasta 2.75 veces en el prellenado de contexto largo, lo que confirma su viabilidad práctica. Código disponible en 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

Atención Dispersa en Bloques Dispersos mediante Permutación de Tokens

Sparser Block-Sparse Attention via Token Permutation

Resumen

Support