기술 요약: CoSA - 프록시-커널 공동 설계된 희소 어텐션을 통한 롱 컨텍스트 추론 가속화

1. 문제 정의

셀프 어텐션 메커니즘의 이차(quadratic) 계산 비용은 현대 대규모 언어 모델(LLM)의 롱 컨텍스트 추론을 매우 비싸게 만듭니다. 희소 어텐션(Sparse Attention)은 중요한 쿼리-키(Query-Key) 상호작용만을 선택적으로 계산함으로써 해결책을 제시하지만, 기존 방식들은 특히 엄격한 희소성 예산(sparsity budget) 하에서 효율성과 충실도(fidelity) 사이의 근본적인 트레이드오프에 직면해 있습니다.

현재의 방법들은 일반적으로 다음 두 가지 범주로 나뉩니다:

1. 프록시 기반 접근법 (Proxy-Driven Approaches): 커널이 실행되기 전에 이진 희소 마스크(binary sparse mask)를 예측하기 위해 경량 프록시를 사용합니다. 중간 정도의 예산에서는 효과적이지만, 예산이 타이트해지면 프록시는 필연적으로 중요한 블록을 놓치게 되어 심각한 정확도 저하를 초래합니다.
2. 인-커널 스키핑 접근법 (In-Kernel Skipping Approaches): 커널 내에서 정확한 Softmax 통계량을 사용하여 무시할 수 있는 블록을 건너뜁니다. 이는 정확하지만 두 가지 한계가 있습니다:
   • 밀집 프록시 요구 (Dense Proxy Requirement): 스킵 결정을 내리기 위해 모든 블록에 대해 전체 쿼리-키(QK) 상호작용을 계산해야 하며, 이는 지배적인 이차적 QK 비용을 제거하지 못합니다.
   • 보수적인 스키핑 (Conservative Skipping): "버킷 효과(bucket effect)"(단 하나의 아웃라이어 행이 스킵을 방해하는 현상)와 글로벌 최댓값 대신 러닝 맥시멈(running maximum, 글로벌 최댓값의 하한선)을 사용하는 특성 때문에, 이러한 방법들은 과도하게 보수적인 스킵 규칙을 적용하여 잠재적인 속도 향상을 활용하지 못합니다.

식별된 핵심 병목 현상은 프록시와 백엔드 커널 사이의 단절입니다. 프록시는 실행 순서를 고려하지 않는 정적 이진 마스크를 생성하는 반면, 커널은 최적의 희소성을 저해하는 고정된 순서로 실행됩니다.

2. 방법론: CoSA

저자들은 프록시-커널 공동 설계(proxy-kernel co-design) 원칙에 기반한 2단계, 훈련 불필요(training-free) 블록 희소 어텐션 메커니즘인 CoSA를 제안합니다. CoSA는 **커널 인식 프록시(Kernel-Aware Proxy, KAP)**와 **순서 기반 스키핑 커널(Ordered-Skipping Kernel, OSK)**을 결합하여 블록 선택과 실행 순서를 공동으로 최적화합니다.

2.1 핵심 설계 원칙

• 순서 불변성 (Order Invariance): 온라인 소프트맥스(Online-Softmax, OSM) 누적이 키 블록을 방문하는 순서와 독립적이라는 성질을 활용하여, CoSA는 최종 출력을 변경하지 않고도 블록 처리 순서를 임의로 재정렬할 수 있게 합니다.
• 계산-순서 마스크 (Computation-Order Mask): 이진 마스크 대신, CoSA는 KV 페이지를 방문해야 하는 구체적인 순서를 규정하는 마스크를 생성합니다.

2.2 1단계: 커널 인식 프록시 (KAP)

KAP는 적절한 예산 하에서 블록을 선택하고 방문 순서를 결정합니다.

• 블록 랭킹 (Block Ranking): 쿼리와 키를 다운샘플링하고, 근사 점수를 계산하며, MaxPool(키 대상) 및 MaxNorm을 적용합니다. 이 과정은 블록 내의 "행최댓값(rowmax)"을 보존하고 점수를 상대적 랭킹 기여도로 변환합니다.
• HRM 플래그 생성: "행최댓값 보유(Have Rowmax, HRM)" 블록(특정 쿼리 행의 최댓값을 포함하는 블록)을 식별합니다.
• 플래그 그룹별 정렬 (Flag-Grouped Sorting): KAP는 HRM 블록을 먼저 방문하고, 그 다음 추정 점수에 따라 정렬된 다른 블록들을 방문하도록 하는 계산-순서 마스크($M_{KAP}$)를 구성합니다. 이 순서는 커널이 탐색 중에 글로벌 최댓값에 조기에 도달하도록 보장합니다.

2.3 2단계: 순서 기반 스키핑 커널 (OSK)

OSK는 $M_{KAP}$ 마스크를 소비하여 세 가지 뚜렷한 패턴으로 희소 어텐션을 실행합니다:

1. 마스크 스킵 (Mask-Skipped): KAP에 의해 선택되지 않은 블록은 **KV 페이지 리매핑(remapping)**을 통해 물리적으로 건너뜁니다. 이를 통해 표준 희소 커널에서 발생하는 제어 흐로(control-flow) 및 동기화 비용을 제거합니다.
2. 인-커널 스킵 (In-Kernel Skipped): 1단계에서 유지된 블록에 대해, OSK는 정확한 온라인-소프트맥스 통계량을 사용하여 2단계 스킵을 수행합니다. HRM 블록을 먼저 방문했기 때문에(KAP의 순서 덕분), 러닝 맥시멈이 빠르게 글로벌 최댓값에 도달합니다. 이를 통해 스킵 조건(로컬 맥시멈과 러닝 맥시멈 간의 차이 기반)을 더 공격적이고 안전하게 충족할 수 있으며, "버킷 효과"를 피할 수 있습니다.
3. 계산 (Computed): 두 필터를 모두 통과한 블록은 완전히 계산됩니다.

3. 주요 기여

• 단절의 규명: 본 논문은 프록시 설계와 커널 실행 사이의 격리가 높은 정확도를 유지하면서 높은 희소성을 달로 달성하는 데 있어 근본적인 장벽임을 밝혀냈습니다.
• KAP 설계: 이진 마스크가 아닌 계산-순서 마스크를 출력하는 새로운 프록시를 제안합니다. 이는 적절한 예산 하에서 블록을 선택하고, 행최댓값 값을 포함하는 블록을 우선시하는 실행 순서를 규정합니다.
• OSK 설계: KAP 마스크를 바탕으로 페이지 리매핑을 통해 물리적으로 페이지를 건너뛰고, 재정렬된 탐색을 통해 개선된 러닝 맥시멈 통계량을 활용하여 공격적인 인-커널 스키핑을 수행하는 최적화된 커널을 제안합니다.
• 실증적 검증: 다양한 LLM 백본(Qwen3-8B, Llama-3.1-8B-Instruct)과 벤치마크에 대해 기존 최신 기술(MInference, FlexPrefill, XAttention)보다 낮은 예산에서 더 높은 정확도를 달성함을 입증했습니다.

4. 실험 결과

실험은 4K에서 128K 컨텍스트 길이에 걸쳐 RULER(합성 태스크)와 LongBench-v2(실제 추론 태스크)에서 수행되었습니다.

• 정확도: CoSA는 기존 희소 방법들을 일관되게 능가합니다. Qwen3-8B를 사용한 LongBench-v2에서 CoSA는 가장 낮은 예산(15%)에서 작동하면서도 가장 높은 평균 정확도(36.51%)를 달성했습니다. 또한 "w/o CoT"(직접 모드)와 "w/ CoT"(추론 모드) 모두에서 견고한 성능을 유지했습니다.
• 효율성:
  • 어텐션 속도 향상: CoSA는 128K 컨텍스트 길이에서 어텐션 계산 시 4.93배의 속도 향상을 달성합니다.
  • 종단 간 지연 시간 (End-to-End Latency): 128K 컨텍스트에서 TTFT(Time-to-First-Token)를 밀집 어텐션 대비 미미한 성능 저하만으로 2.53배 단축합니다.
• 절제 연구 (Ablation Studies): 논문은 KAP(순서 기반 선택)와 OSK(페이지 리매핑 + 인-커널 스킵)의 조합이 시너지 효과를 낸다는 것을 보여줍니다. KAP만 사용하는 것은 정확도를 높이고, 인-커널 스키핑을 추가하면 희소성이 증가하며, 페이지 리매핑을 추가하면 예산을 더욱 타이트하게 가져가면서도 정확도를 회복하고 개선할 수 있습니다.

5. 의의 및 주장

본 논문은 CoSA가 프록시와 커널을 분리된 구성 요소로 취급하는 것에서 벗어나 공동 설계된 시스템으로 전환하는 패러다임 시프트를 나타낸다고 주장합니다. 계산-순서 마스크로 이 둘을 연결함으로써, CoSA는 현재의 희소 어텐션 방법들이 겪고 있는 효율성-충실도 트레이드오프를 해결합니다.

저자들의 주장은 다음과 같습니다:

1. 공격적인 희소성 달성 가능: HRM 블록을 전면에 배치함으로써 커널이 더 많은 블록을 안전하게 스킵할 수 있게 하여, 이전보다 더 낮은 희소성 예산에서도 높은 성능을 낼 수 있습니다.
2. 물리적 실행의 중요성: 좋은 마스크를 갖는 것만으로는 충분하지 않습니다. 이론적인 속도 향상을 실현하고 불필요한 동기화 비용을 피하기 위해서는 커널이 메모리 접근(페이지 리매핑을 통해)을 물리적으로 재정렬해야 합니다.
3. 확장성: CoSA는 특히 어텐션의 이차 비용이 지배적인 롱 컨텍스트(128K+)에서 효과적이며, RAG(검색 증강 생성) 및 자율 에이전트 애플리케이션을 위한 실질적인 지연 시간 감소를 제공합니다.

논문은 CoSA가 현재 프리필링(prefilling) 단계에 집중하고 있지만, 향후 연구를 통해 서로 다른 계산 패턴을 가진 디코딩(decoding) 단계로 이러한 공동 설계 원칙을 확장할 것임을 밝히며 마무리합니다.