CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

기술 요약: CompactAttention

문제 제기
대규모 언어 모델 (LLM) 이 수십만에서 수백만 개의 토큰을 지원하는 컨텍스트 윈도우를 확장함에 따라, 단일 "원샷 (one-shot)" 프리필 (prefill) 패스를 통한 긴 시퀀스 처리는 이차적인 계산 비용과 온라인 서비스 시스템에서 디코드 요청이 멈출 수 있는 위험으로 인해 비현실적이게 되었습니다. 따라서 청크별 프리필 (chunked prefill)—축적된 키-값 (KV) 캐시를 참조하면서 고정된 크기의 청크로 입력을 순차적으로 처리하는 방식—이 vLLM 및 SGLang 과 같은 프레임워크에서 표준 서비스 전략이 되었습니다.

그러나 기존 가속화 방법들은 이러한 환경에서 상당한 한계에 직면해 있습니다:

1. 블록 희소 어텐션의 비효율성: 전통적인 블록 희소 방법 (예: XAttention, SeerAttention) 은 불규칙한 메모리 접근 및 마스크 해석의 오버헤드를 상쇄하기 위해 긴 쿼리 길이가 필요한 희소 커널에 의존합니다. 청크별 프리필에서는 쿼리 길이가 청크 크기에 의해 제한되어 ($Q \ll KV$), 이러한 커널이 병렬성과 효율성을 잃어 종종 밀집 어텐션보다 성능이 떨어집니다.
2. 패턴 탐색 오버헤드: 희소 방법은 입력에 의존하는 희소 패턴을 찾아야 합니다. 청크별 프리필에서는 이 탐색이 매 청크마다 성장하는 축적된 KV 캐시 전체에서 반복되어야 하므로, 누적 탐색 오버헤드가 치명적인 병목 현상이 됩니다.
3. 직접 KV 선택의 한계: QUOKA와 같은 방법은 부분 표본 추출된 쿼리 집합을 기반으로 KV 토큰의 하위 집합을 선택하고 밀집 어텐션을 수행함으로써 희소 커널을 피합니다. 그러나 이 접근법은 두 가지 문제에 시달립니다:
   • 커버리지 결손: 표본 추출되지 않은 쿼리에 중요한 쿼리별 KV 항목이 누락되어, 분산 정보 접근이 필요한 작업에서 정확도 저하를 초래합니다.
   • 복제 오버헤드: 토큰 수준의 선택은 어텐션 실행 전에 KV 항목을 명시적으로 수집하여 컴팩트 버퍼로 복사해야 하므로, 컨텍스트 길이와 배치 크기에 비례하는 오버헤드가 발생합니다.

방법론: CompactAttention
저자들은 CompactAttention을 제안했는데, 이는 블록 수준의 KV 선택과 어텐션 실행을 분리하는 메커니즘입니다. 2 차원 블록 희소 마스크를 희소 커널의 직접 실행 계획으로 사용하는 대신, CompactAttention 은 이를 GQA 인지 그룹별 KV 블록 테이블을 구성하기 위한 KV 선택 신호로 취급합니다.

이 방법은 두 단계로 작동합니다:

1. 선택 단계 (블록-유니온 KV 테이블 구성):

   • 시스템은 SeerAttention 또는 FlashPrefill 과 같은 경량 패턴 탐색 방법에서 생성된 2 차원 헤드별 블록 희소 마스크를 받아들입니다.
   • 현재 청크 내의 쿼리 블록 간 **Q-블록 유니온 (Q-block union)**을 적용합니다. 이는 밀집 페이지드 어텐션이 각 $(Q\text{-블록}, KV\text{-블록})$ 타일별 개별 결정이 아니라 함께 실행되는 쿼리 블록 그룹에 대한 단일 KV 블록 목록을 소비하기 때문에 필요합니다.
   • 공유 KV 그룹 (기본적으로 그룹화 쿼리 어텐션 그룹) 을 가진 쿼리 헤드 간 **그룹 내 유니온 (intra-group union)**을 적용합니다.
   • 그 결과, 원래 입력 마스크에 의해 선택된 모든 KV 블록을 보존하는 최소한의 그룹별 KV 블록 테이블 ($T_{b,g}$) 이 생성됩니다. 이 구성은 커버리지 보존을 보장합니다: 그룹 내의 어떤 쿼리 블록이나 헤드가 선택한 KV 블록은 모두 선택된 상태로 유지됩니다.

2. 실행 단계 (제로-복제 페이지드 어텐션):

   • 토큰 수준의 선택 방법과 달리, CompactAttention 은 블록 단위로 작동합니다.
   • 표준 시퀀스 우선 레이아웃 대신 KV-헤드 우선 메모리 레이아웃 $[B, H_{kv}, L, D]$을 사용합니다. 이를 통해 각 $(배치, KV\text{-헤드}, 블록)$ 삼중항이 연속된 메모리 영역으로 처리될 수 있습니다.
   • 구성된 블록 테이블은 표준 페이지드 밀집 어텐션 커널에 메타데이터(특히 kv_indptr 및 kv_indices) 로 전달됩니다.
   • 커널은 명시적인 복사나 새로운 버퍼로의 압축 없이 원래 KV 캐시에서 선택된 KV 블록을 **인플레이스 (in-place)**로 접근합니다.

주요 기여

• 선택과 실행의 분리: 이 논문은 어떤 KV 블록을 선택할지라는 논리와 어떻게 실행할지라는 논리를 분리하는 설계를 도입하여, 경량 희소 패턴 탐색 방법의 재사용을 가능하게 하면서 희소 커널의 비효율성은 피할 수 있게 합니다.
• 블록-유니온 구성: 2 차원 희소 마스크를 실행 가능하고 GQA 인지적이며 최소적이면서도 커버리지를 보존하는 블록 테이블로 변환하기 위한 Q-블록 및 그룹 내 유니온의 독창적인 활용.
• 제로-복제 페이지드 실행: KV-헤드 우선 레이아웃과 페이지드 어텐션 메타데이터를 활용함으로써, QUOKA 와 같은 토큰 수준 선택 방법의 고유한 명시적 KV 복사 오버헤드를 제거합니다.

실험 결과
청크별 프리필 환경에서 LLaMA-3.1-8B-Instruct 및 Qwen3-30B-A3B-Instruct-2507을 평가한 결과:

• 속도 향상: 128K 컨텍스트 길이를 가진 H200 GPU 에서 CompactAttention 은 밀집 어텐션 기준 대비 최대 2.72 배의 어텐션 속도 향상 및 1.96 배의 엔드투엔드 속도 향상을 달성했습니다. 이는 $Q \ll KV$ 환경에서 종종 밀집 어텐션보다 성능이 떨어지는 블록 희소 기준 및 복사 오버헤드에 제한된 QUOKA 보다 훨씬 뛰어납니다.
• 정확도: CompactAttention 은 RULER 및 LongBench V2 벤치마크에서 밀집 어텐션과 유사한 정확도를 유지합니다. 특히 분산 정보 접근이 필요한 작업 (예: Multi-key NIAH) 에서 QUOKA 의 쿼리 부분 표본 추출로 인한 정확도 저하가 발생하는 경우, CompactAttention 은 QUOKA 보다 우수한 성능을 보입니다.
• 효율성: 제거 실험 (Ablation studies) 은 메타데이터 구성 오버헤드를 고려하더라도 제로-복제 페이지드 실행 전략이 희소 커널 실행 및 명시적 KV 압축 모두보다 더 효율적임을 확인시켜 주었습니다.

의의
이 논문은 청크별 프리필에서 희소 어텐션의 주요 병목 현상이 단순히 어떤 KV 블록을 선택할지가 아니라, 선택된 블록이 어떻게 실행되는지에 있다고 주장합니다. 블록-유니온 구성과 제로-복제 실행을 통해 희소 선택을 효율적인 밀집 페이지드 어텐션 커널과 호환되게 함으로써, CompactAttention 은 더 빠른 장문맥 LLM 서비스를 위한 실용적인 경로를 제시합니다. 이는 기존 희소 어텐션 방법들이 청크별 프리필 시스템에서 채택되지 못하게 했던 $Q \ll KV$ 환경의 특정 비효율성들을 해결합니다.