[논문리뷰] CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

링크: 논문 PDF로 바로 열기

메타데이터

저자: Jiwon Song, Dongwon Jo, Beomseok Kang, Jae-Joon Kim

---

1. Key Terms & Definitions (핵심 용어 및 정의)

• Chunked Prefill: 긴 컨텍스트 입력을 고정된 크기의 Chunk 단위로 나누어 처리하며, 이전 Chunk들의 KV 캐시를 누적하여 활용하는 효율적인 추론 방식입니다.
• Block-Union KV Selection: 개별 쿼리 블록 및 헤드에 대한 2D 블록 스파스 마스크를 GQA(Grouped-Query Attention) 기반의 KV 블록 테이블로 통합하는 CompactAttention의 핵심 선택 메커니즘입니다.
• Zero-Copy Paged Attention: 선택된 KV 블록들을 별도의 메모리 버퍼로 복사하지 않고, 기존의 KV 캐시 레이아웃 위에서 메타데이터(페이지 테이블) 기반으로 직접 액세스하여 실행하는 최적화 기법입니다.
• KV-Head-Major Layout: 기존의 시퀀스 중심 메모리 레이아웃을 개선하여, 각 (batch, KV head, block) 단위를 독립적인 페이지로 주소 지정할 수 있도록 구성한 데이터 배치 형식입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 기존 Chunked Prefill 환경에서 Block-Sparse Attention 및 Query-Subsampled KV Selection 방식이 가진 성능 한계를 극복하기 위해 CompactAttention을 제안합니다. 기존의 Block-Sparse 기법들은 Q≪KV 상태인 Chunked Prefill 체제에서 커널 실행 효율이 저하되며, 반복적인 패턴 검색 오버헤드가 발생한다는 문제점이 있습니다 [Figure 1]. 반면, QUOKA와 같은 Query-Subsampled 방식은 특정 쿼리 종속적 KV 항목을 누락하여 정확도가 떨어지며, 토큰 단위의 KV 선택으로 인해 매번 명시적인 복사(copy) 오버헤드가 발생합니다. 결과적으로, 긴 컨텍스트 처리를 위한 효율적인 어텐션 실행 프레임워크가 절실히 요구되는 상황입니다 [Figure 2].

3. Method & Key Results (제안 방법론 및 핵심 결과)

본 논문은 KV 블록 선택 단계와 어텐션 실행 단계를 분리하여 최적화하는 CompactAttention을 제안합니다. 먼저, Q-block union 및 intra-group union을 통해 생성된 2D 스파스 마스크를 GQA를 고려한 최소 KV Block Table로 변환합니다 [Figure 3]. 이 방식은 선택된 모든 KV 블록을 보존하면서도, Zero-Copy Paged Attention을 통해 메모리 복사 없이 기존 KV 캐시에서 데이터를 즉시 가져옵니다 [Figure 4]. 실험 결과, LLaMA-3.1-8B-Instruct 모델에서 RULER 벤치마크 기준 덴스 어텐션과 유사한 정확도를 유지하면서도, 128K 컨텍스트 길이에서 최대 **2.72×**의 어텐션 속도 향상을 달성하였습니다 [Table 1]. 또한, H200 GPU 환경에서 End-to-End Latency를 크게 단축하며 기존의 FlashPrefill 대비 월등한 효율성을 입증하였습니다 [Figure 5].

4. Conclusion & Impact (결론 및 시사점)

본 연구는 Block-Union KV Selection을 통해 Chunked Prefill 환경에서 어텐션 연산의 병목을 성공적으로 해결하였습니다. CompactAttention은 기존의 스파스 커널 의존성을 제거하고 고효율 덴스 페이지드 어텐션 커널과의 호환성을 확보함으로써 성능 향상을 이끌어냈습니다. 이는 긴 컨텍스트를 다루는 대규모 언어 모델의 실시간 서비스 분야에서 매우 중요한 기술적 토대가 될 것입니다. 향후 다양한 모델 아키텍처 및 하드웨어 환경에서도 일관된 고성능을 제공할 것으로 기대됩니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.