S2O: 온라인 순열을 이용한 희소 어텐션의 조기 중단 기법

S2O: Early Stopping for Sparse Attention via Online Permutation

어텐션 연산은 시퀀스 길이에 따라 제곱으로 증가하며, 이는 장기 컨텍스트 추론에 근본적인 제약을 가합니다. 기존의 블록 단위 희소화는 지연 시간을 줄일 수 있지만, 거친 블록 구조는 고유한 희소성 한계를 발생시키며, 정교하게 설계된 방식이라 할지라도 추가적인 개선이 어렵습니다. 본 논문에서는 온라인 순열을 활용하여 희소 어텐션의 조기 중단을 수행하는 S2O 기법을 제안합니다. 메모리 시스템의 가상-물리 주소 매핑에서 영감을 받아, S2O는 FlashAttention의 실행 방식을 재검토하고 분해하여, 원래 순서대로 연속된 토큰이 아닌 불연속적인 토큰을 로드하여 추론을 수행할 수 있도록 합니다. 어텐션 히트맵의 미세한 구조에서 영감을 받아, 명시적인 순열을 온라인으로, 인덱스 기반의 이산적인 로딩 정책으로 변환합니다. S2O는 매우 가벼운 전처리 및 인덱스 재매핑 오버헤드를 통해 중요도가 높은 작은 블록 그룹에 집중합니다. 이러한 중요도 기반의 온라인 순열 로딩에 더해, S2O는 조기 중단 규칙을 도입합니다. 연산은 중요도가 높은 블록부터 낮은 블록 순으로 진행되며, 현재 블록의 점수가 특정 임계값 이하로 떨어지면 S2O는 조기에 중단하고 나머지 중요도가 낮은 블록을 건너뛰어, 효과적인 희소성을 높이고 제어된 오차 범위 내에서 연산량을 줄입니다. 결과적으로, S2O는 실질적인 희소성 한계를 크게 향상시킵니다. Llama-3.1-8B 모델에서 128K 컨텍스트 환경에서, S2O는 동일한 희소성 수준에서 단일 연산의 MSE를 3.82배 줄이고, 동일한 MSE 수준에서 사전 채우기 연산 밀도를 3.31배 줄입니다. 동시에, S2O는 전체 정확도를 유지하면서 어텐션 연산 속도를 7.51배, 전체 엔드투엔드 속도를 3.81배 향상시킵니다.

Attention scales quadratically with sequence length, fundamentally limiting long-context inference. Existing block-granularity sparsification can reduce latency, but coarse blocks impose an intrinsic sparsity ceiling, making further improvements difficult even with carefully engineered designs. We present S2O, which performs early stopping for sparse attention via online permutation. Inspired by virtual-to-physical address mapping in memory systems, S2O revisits and factorizes FlashAttention execution, enabling inference to load non-contiguous tokens rather than a contiguous span in the original order. Motivated by fine-grained structures in attention heatmaps, we transform explicit permutation into an online, index-guided, discrete loading policy; with extremely lightweight preprocessing and index-remapping overhead, it concentrates importance on a small set of high-priority blocks. Building on this importance-guided online permutation for loading, S2O further introduces an early-stopping rule: computation proceeds from high to low importance; once the current block score falls below a threshold, S2O terminates early and skips the remaining low-contribution blocks, thereby increasing effective sparsity and reducing computation under a controlled error budget. As a result, S2O substantially raises the practical sparsity ceiling. On Llama-3.1-8B under a 128K context, S2O reduces single-operator MSE by 3.82$\times$ at matched sparsity, and reduces prefill compute density by 3.31$\times$ at matched MSE; meanwhile, it preserves end-to-end accuracy and achieves 7.51$\times$ attention and 3.81$\times$ end-to-end speedups.