CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

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.

📋 한 줄 요약

**[LLM 서빙 시스템]** Chunked prefill 환경에서 Block-Union KV Selection으로 어텐션을 가속하는 CompactAttention을 제안하여 LLaMA-3.1-8B에서 128K 컨텍스트 기준 최대 2.72배 속도 향상을 달성했다.

🎯 핵심 기여도

청크 단위 prefill에 특화된 새로운 sparse attention 메커니즘 CompactAttention 제안
2D 블록 sparse 마스크를 직접 실행 계획이 아닌 KV 선택 신호로 재해석하는 관점 제시
Q-block union 및 intra-group union을 통해 GQA(Grouped-Query Attention) 인식 per-group KV block table 구성 방법론 정립
명시적 KV 복사 없이 paged execution 제약 하에서도 효율적으로 동작하는 실용적 구현 제공

💡 핵심 아이디어

기존 sparse attention 기법들은 일회성 prefill에 맞춰 설계되어 chunked prefill에서 비효율적이다. 블록 sparse 커널은 청크 크기로 query 길이가 제한되면 효율이 떨어지고, fine-grained 패턴 탐색은 매 청크마다 누적 KV cache에 반복되어 비용이 크다. CompactAttention은 sparse 마스크를 실행 계획이 아닌 "어느 KV 블록을 선택할지"의 신호로 보고 최소 블록 테이블만 구성한다.

🔬 기술적 접근법

Block-Union KV Selection: 입력 마스크가 선택한 모든 KV 블록을 보존하는 최소 블록 테이블 생성
Q-block union: query 블록 단위로 선택 KV를 합집합 처리
Intra-group union: GQA의 그룹 내부에서 KV 선택을 통합
Paged execution: 선택된 KV 블록을 명시적 압축 없이 in-place로 접근

📊 주요 결과

LLaMA-3.1-8B-Instruct에서 RULER 벤치마크 dense attention 대비 근접 정확도 유지
128K 컨텍스트에서 chunked prefill 기준 최대 2.72배 attention 속도 향상
QUOKA 등 기존 방식의 KV-copy 오버헤드 제거

💭 의의 및 한계

긴 컨텍스트 LLM 서빙의 핵심 병목인 chunked prefill 어텐션 효율을 실용적으로 개선한 점에서 의미가 있다. 다만 LLaMA-3.1-8B-Instruct 단일 모델 평가에 머물러 다양한 GQA 구성·모델 크기에 대한 일반화 검증은 추가 연구가 필요하다.

🚀 실용적 활용

긴 문서 분석, 대화형 에이전트, RAG 등 long-context LLM 추론 서비스의 prefill 단계 비용 절감에 직접 적용 가능하며, GPU 메모리·지연시간이 중요한 프로덕션 서빙 스택에 통합하기 적합하다.