토큰 수준의 선택 정확도와 실행 효율을 동시에 해결한 희소 어텐션

핵심 요약

• 긴 문맥에서 LLM 추론 시 KV-캐시 통신량을 줄이는 희소 어텐션 기법들은 주로 높은 단위의 선택만 가능하던 한계가 있었다.
• 저해상도 점수 부여로는 쿼리-키 구조 정보가 손실되어 정확한 랭킹이 어렵고, 선택된 토큰들을 모으는 과정에서 레이턴시 이득이 상쇄되는 문제가 있었다.
• Joint Subspace Sparse Attention (JSSA)은 공유 저차원 쿼리-키 부분공간을 학습해 토큰 점수 보존 오차를 최소화함으로써 선택 품질을 개선한다.
• 이 비볼록 최적화 문제를 닫힌 형태의 스펙트럼 해를 갖는 미분 가능한 상한으로 상한화하며, 대규모 모델에서 추가 정제는 측정 가능한 개선을 주지 않는다.
• Grouped-Query Attention (GQA) KV-캐시 레이아웃에서 KV-헤드별로 서로 다른 선택을 지원하는 수집-자유 토큰 레벨 인덱싱 어텐션 커널을 개발해 실행 효율을 높였다.
• 표준 GQA, 하이브리드 슬라이딩 윈도우/글로벌 어텐션, Multi-head Latent Attention (MLA) 모델 전반에서 JSSA는 완전 KV 어텐션에 가까운 정확도를 유지하면서도 경쟁 방법들에서 나타나는 생성 길이 팽창과 아키텍처 의존 실패를 피한다.
• NVIDIA B200 GPU에서 128K 토큰 문맥 시 완전 KV 어텐션 대비 6.43배 높은 디코딩 처리량을 달성한다.
• 모델 학습 없이 모델 자체 차별 선택 어텐션(DSA)만큼의 추론 정확도를 얻으면서도 700B 규모 대형 모델에서도 96~99% 수준의 디코딩 처리량에 도달한다.