CompactAttention: 블록-유니온 KV 선택을 통한 청크 프리필 가속화

초록

청크 기반 프리필(chunked prefill)은 장기 컨텍스트 대규모 언어 모델을 위한 널리 채택된 서빙 전략이 되었지만, 이 체계에서 효율적인 어텐션 계산은 여전히 어려운 과제로 남아 있다. 기존의 희소 어텐션(sparse attention) 방법은 주로 원샷 프리필(one-shot prefill)을 위해 설계되었으며, 청크 프리필에서는 효율적으로 적용되지 않는다: 블록 희소 커널(block-sparse kernel)은 쿼리 길이가 청크 크기로 제한될 때 효율성이 떨어지고, 미세 패턴 탐색은 모든 청크에서 누적된 KV 캐시에 대해 반복될 때 비용이 많이 든다. 최근 청크 프리필을 직접 대상으로 하는 QUOKA는 희소 커널 오버헤드를 피하지만, 쿼리 서브샘플링(query-subsampled) 기반의 토큰 수준 KV 선택에 의존하여 쿼리별 KV 항목을 놓칠 수 있고 명시적인 KV 복사 오버헤드를 초래한다. 이러한 한계를 해결하기 위해, 우리는 블록 유니온 KV 선택(Block-Union KV Selection)에 기반한 청크 프리필 어텐션 메커니즘인 CompactAttention을 제안한다. CompactAttention은 2D 블록 희소 마스크를 직접적인 희소 커널 실행 계획이 아닌 KV 선택 신호로 취급하고, 이를 Q-블록 유니온(Q-block union)과 그룹 내 유니온(intra-group union)을 통해 GQA 인식(GQA-aware) 그룹별 KV 블록 테이블로 변환한다. 이 구성은 페이지 실행 제약 하에서 입력 마스크에 의해 선택된 모든 KV 블록을 보존하는 최소 블록 테이블을 생성하며, 선택된 KV 블록이 명시적 KV 압축 없이 제자리에서 접근될 수 있게 한다. LLaMA-3.1-8B-Instruct에서 CompactAttention은 RULER 벤치마크에서 정밀 어텐션에 가까운 정확도를 유지하면서, 청크 프리필 환경에서 128K 컨텍스트 길이 기준 최대 2.72배의 어텐션 속도 향상을 제공한다.

English

Chunked prefill has become a widely adopted serving strategy for long-context large language models, but efficient attention computation in this regime remains challenging. Existing sparse attention methods are primarily designed for one-shot prefill and do not translate efficiently to chunked prefill: block-sparse kernels lose efficiency when the query length is limited by the chunk size, while fine-grained pattern search becomes costly when repeated over the accumulated KV cache at every chunk. QUOKA, a recent method that directly targets chunked prefill, avoids sparse-kernel overhead but relies on query-subsampled, token-level KV selection, which can miss query-specific KV entries and introduce explicit KV-copy overhead. To address these limitations, we propose CompactAttention, a chunked-prefill attention mechanism based on Block-Union KV Selection. CompactAttention treats 2D block-sparse masks as KV-selection signals rather than direct sparse-kernel execution plans, and converts them into GQA-aware per-group KV block tables through Q-block union and intra-group union. This construction produces the minimal block tables that preserve all KV blocks selected by the input masks under paged execution constraints, enabling selected KV blocks to be accessed in place without explicit KV compaction. On LLaMA-3.1-8B-Instruct, CompactAttention maintains accuracy close to dense attention on the RULER benchmark while delivering up to 2.72times attention speedup at 128K context length under chunked prefill.