JSSA, 공동 부분공간과 Gather-free 커널로 장문맥 디코딩 가속

장문맥 LLM 추론에서 비용을 키우는 요인은 어텐션의 연산량만이 아닙니다. 매 디코딩 단계마다 긴 KV-cache를 읽어야 하므로 메모리 트래픽도 핵심 병목이 됩니다. 컨텍스트가 길어질수록 full-KV 어텐션이 이동시키는 데이터가 늘어나기 때문에, 필요한 토큰만 선택하는 희소 어텐션이 중요한 대안으로 떠오르고 있습니다. 하지만 토큰 단위 선택은 두 가지 문제를 안고 있습니다. 저차원 점수 계산이 쿼리와 키의 구조를 충분히 보존하지 못할 수 있고, 선택된 토큰이 메모리 곳곳에 흩어지면 gather와 인덱싱 비용이 희소화의 이점을 상쇄할 수 있습니다.

‘Joint Subspace Sparse Attention for Efficient Long-Context Decoding’은 이 두 문제를 선택 알고리즘과 실행 커널을 함께 설계하는 방식으로 다룹니다.

핵심 아이디어

• 공유 쿼리-키 부분공간 학습. JSSA는 쿼리와 키를 공통의 저차원 공간에서 표현하고, 토큰 점수 보존 오차를 최소화하도록 부분공간을 학습합니다. 단순히 차원을 줄이는 것이 아니라, 어떤 토큰이 중요한지 순위를 매기는 데 필요한 관계를 유지하는 것이 목표입니다.
• 계산 가능한 목적함수. 정확한 점수 보존 목적은 비볼록 문제입니다. 연구진은 이를 상계하는 다루기 쉬운 주변 목적을 만들고, 폐형식 스펙트럴 해를 얻었습니다. 논문은 테스트한 환경에서 정확한 목적을 추가로 직접 정제해도 측정 가능한 개선이 없었다고 설명합니다.
• Gather-free 인덱스 어텐션. 토큰 선택 결과가 KV-cache에 분산되어 있으면 일반적인 gather가 추가 메모리 비용을 만듭니다. JSSA는 토큰 단위 indexed-attention 커널을 설계해 이 비용을 줄이고, 표준 Grouped-Query Attention(GQA) KV-cache 구조에서 KV head마다 서로 다른 선택 결과를 지원합니다.

보고된 실험 결과

평가는 표준 GQA 모델, 슬라이딩 윈도와 전역 어텐션을 결합한 하이브리드 구조, Multi-head Latent Attention(MLA) 모델에서 진행됐습니다. 제공된 요약에 따르면 JSSA는 full-KV 어텐션에 가까운 정확도를 유지하면서, 일부 경쟁 방법에서 나타난 생성 길이의 큰 증가와 아키텍처별 실패를 피했습니다.

NVIDIA B200에서 컨텍스트 길이 128K를 사용했을 때는 full-KV 어텐션보다 최대 6.43배 높은 디코딩 처리량을 달성했다고 보고했습니다. 또한 추가 모델 학습 없이 모델 자체 DSA와 비슷한 추론 정확도를 보였으며, 700B 규모 모델에서도 DSA 디코딩 처리량의 96~99%에 도달했다고 합니다. 다만 이 수치는 논문이 제시한 실험 조건의 결과이며, 모든 모델과 하드웨어에 동일하게 적용된다는 의미는 아닙니다.

의미와 한계

JSSA의 차별점은 선택 정확도와 메모리 접근을 하나의 시스템 문제로 본다는 데 있습니다. 점수 모델만 개선하면 불규칙한 읽기 문제가 남고, 커널만 최적화하면 잘못된 토큰 선택으로 정확도가 떨어질 수 있습니다. JSSA는 두 층을 함께 다뤄 장문맥 디코딩에서 큰 비중을 차지하는 KV-cache 트래픽을 직접 줄이려 합니다.

실제 효과는 어텐션 구조, 선택 토큰 수, GPU 커널 구현, 하드웨어에 따라 달라질 수 있습니다. 제공된 자료에는 선택 비율, 지연 시간의 세부 구성, 다양한 가속기에서의 종합 비교가 포함되어 있지 않습니다. 따라서 서비스에 적용하려면 대상 모델과 실제 컨텍스트 분포를 기준으로 별도 검증이 필요합니다.

출처: Hugging Face Daily Papers