토큰 순열을 통한 희소 블록 희소 어텐션

초록

대규모 언어 모델(LLM)의 컨텍스트 길이 확장은 상당한 이점을 제공하지만 계산 비용이 많이 듭니다. 이러한 비용은 주로 자기 주의(self-attention) 메커니즘에서 기인하며, 시퀀스 길이에 대한 O(N^2) 복잡도는 메모리와 지연 시간 측면에서 주요 병목 현상으로 작용합니다. 다행히도 주의 행렬(attention matrix)은 특히 긴 시퀀스에서 흔히 희소(sparse)한 특성을 보여 최적화 가능성을 시사합니다. 블록 희소 주의(Block-sparse attention)는 시퀀스를 블록으로 분할하고 이 블록들의 일부에 대한 계산을 생략하는 유망한 해결책으로 부상했습니다. 그러나 이 방법의 효과는 기본 주의 패턴에 크게 의존하며, 이로 인해 차선의 블록 수준 희소성이 발생할 수 있습니다. 예를 들어, 단일 블록 내의 질의(query)에 대한 중요한 키(key) 토큰들이 수많은 다른 블록에 흩어져 있을 수 있어 계산 중복이 발생합니다. 본 연구에서는 주의의 순열(permutation) 특성을 활용하여 블록 수준 희소성을 높이고 LLM 프리필링(prefilling)의 계산 효율성을 향상시키는 플러그 앤 플레이(plug-and-play) 방식인 PBS-Attn(Permuted Block-Sparse Attention)을 제안합니다. 우리는 도전적인 실제 장문 컨텍스트 데이터셋에 대한 포괄적인 실험을 수행하여 PBS-Attn이 모델 정확도 측면에서 기존 블록 희소 주의 방법들을 지속적으로 능가하며 전체 주의(full attention) 기준선(baseline)과 매우 유사한 성능을 보임을 입증했습니다. 저희가 개발한 사용자 정의 순열 플래시 어텐션(permuted-FlashAttention) 커널을 기반으로 PBS-Attn은 장문 컨텍스트 프리필링에서 최대 2.75배의 종단 간(end-to-end) 속도 향상을 달성하여 실용적인 타당성을 확인했습니다. 코드는 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

토큰 순열을 통한 희소 블록 희소 어텐션

Sparser Block-Sparse Attention via Token Permutation

초록

Support