JSSA가 토큰 단위 희소 어텐션의 선택 품질과 실행 비용을 함께 줄인다
Joint Subspace Sparse Attention for Efficient Long-Context Decoding
무엇인가
긴 문맥 자기회귀 디코딩은 매 스텝마다 커지는 KV 캐시를 반복해서 읽는 데 병목이 걸린다. 풀 어텐션은 현재 쿼리를 전체 문맥에 대해 점수화하고 모든 값을 읽어서 메모리 트래픽이 시퀀스 길이에 선형으로 증가한다. 어텐션이 메모리 대역폭에 묶이는 이 영역에서는 KV 접근을 줄이는 것이 가장 직접적인 가속 방법이다. 희소 어텐션은 각 디코딩 스텝에서 문맥의 일부만 선택해 KV 접근을 줄이지만, 기존 실용 기법은 대부분 블록이나 페이지 단위의 거친 입자로 동작한다. 토큰 단위 선택은 더 세밀한 패턴을 표현할 수 있지만 두 가지 문제가 있다. 첫째, 저차원 점수화가 토큰 순위를 정하는 데 필요한 쿼리-키 구조를 잃을 수 있다. 둘째, 흩어진 선택 토큰을 모으는 비용이 지연 이득을 상쇄할 수 있다. JSSA는 이 두 문제를 함께 다룬다.
어떻게 동작하나
JSSA는 학습이 필요 없는 토큰 단위 희소 어텐션 방법으로, GQA와 MLA 모델을 대상으로 한다. 각 레이어와 KV 헤드마다 원본 KV 캐시와 함께 저차원 투영 키 캐시를 유지한다. 선택을 위해 JSSA는 쿼리와 키에 같은 정규직교 기저 M을 적용하는 공유 저차원 부분공간을 학습한다. GQA에서는 KV 헤드에 연결된 G개 쿼리 헤드의 평균 쿼리 q_bar를 사용해 투영 점수 s_hat = q_bar^T P k를 계산하고, 점수 오차는 s - s_hat = q_bar^T (I-P) k가 된다. 여기서 P = M^T M은 랭크 r 투영기다. JSSA는 토큰 순위를 결정하는 이중선형 점수의 기대 제곱 오차를 최소화하는 P를 고른다. 정확한 선택 목적식 L_sel(P) = tr((I-P) G_Qbar (I-P) G_K)는 비볼록이라 일반적인 닫힌 해가 없지만, 저자들은 L_sel(P) <= 1/4 L_marg(P)^2, L_marg(P) = tr((I-P)(G_Qbar + G_K))라는 상한을 유도한다. 이 주변 목적식은 M* = TopEig_r(G_Qbar + G_K)^T라는 닫힌 형태의 스펙트럼 해를 갖는다. 저자들은 M*를 그대로 사용하며, 정확한 목적식을 직접 개선해도 실험 설정에서는 측정 가능한 선택 이득이 없었다고 밝힌다. 따라서 캘리브레이션은 레이어와 KV 헤드마다 2차 모멘트를 누적하고 한 번의 고유분해를 수행하는 것으로 줄어든다.
무엇과 다른가
선택 후 실행 단계에서 JSSA는 각 KV 헤드에 대해 상위 B개 위치를 고른 뒤, 원본 전체 차원 KV 캐시의 흩어진 위치에서 직접 어텐션을 계산한다. 토큰 하나의 키나 값 슬라이스는 헤드 차원에 걸쳐 이미 연속이므로 불규칙 인덱스 자체가 장벽은 아니고, 인덱스 읽기는 버스트 크기 트랜잭션으로 남는다. 진짜 문제는 GQA 캐시 레이아웃에서 각 KV 헤드가 별도 파티션을 차지하고 서로 다른 토큰 집합을 선택할 수 있다는 점이다. 선택을 하나의 레이어 공유 집합으로 합치면 긴 문맥에서 정확도가 떨어진다. 그래서 저자들은 표준 GQA KV 캐시 레이아웃에서 KV 헤드별로 다른 선택을 지원하는, 수집 없는 토큰 단위 인덱스 어텐션 커널을 Triton으로 개발했다. 이 커널은 레이어당 한 번 실행되고 내부에서 헤드별 인덱스 집합을 해결하며, 각 쿼리 헤드를 연결된 KV 헤드가 선택한 인덱스에 매핑해 해당 키와 값을 제자리에서 읽는다. 이는 gather-then-attend에 필요한 스테이징 버퍼와 KV 헤드별 커널 실행을 모두 피한다. MLA에서는 FlashMLA의 희소 어텐션 커널을 사용한다. 점수화용 투영 캐시는 K_proj = K M^T로 유지하며, GQA에서는 FP8 투영 키 캐시가 원본 BF16 KV 캐시의 최대 6.25%를 차지하도록 r을 정한다(LLaMA, Qwen, Ministral은 r=32, Gemma4는 r=64). MLA에서는 64개 잠재 차원과 64개 RoPE 차원을 합친 128차원 점수 표현을 사용해 DSA 인덱서와 같은 용량을 맞춘다. FP8을 전반에 사용하며 BF16 대비 정확도 영향은 무시할 수 있다고 검증했다.
어떻게 쓰나
실험은 표준 GQA, 하이브리드 로컬/글로벌 GQA, MLA 아키텍처를 아우른다. LongBench는 KV 예산 256 토큰에서 LLaMA-3.1-8B-Instruct, Qwen3-14B, Ministral-3-8B-Instruct-2512로 평가했고, JSSA는 세 모델 모두에서 최고 또는 근접한 희소 결과를 내며 풀 KV와 0.67점 이내 차이를 보였다. 경쟁 베이스라인 중 어떤 것도 아키텍처 전반에서 일관되게 가장 강하지는 않았다. RULER는 32K에서 128K 문맥까지 고정 1.5625% KV 예산으로 평가했는데, JSSA가 전체적으로 풀 KV에 가장 가깝고 128K에서 가장 높은 희소 정확도를 냈다. 경쟁 기법은 긴 문맥에서 더 급격히 성능이 떨어졌다(예: LLaMA-3.1-8B-Instruct의 Loki). Ministral-3-8B-Instruct-2512에서는 JSSA가 128K에서 풀 KV와 일치하고 32K와 64K에서 2점 이내였다. PG-19 퍼플렉시티는 32K에서 128K까지 같은 고정 밀도 예산으로 측정했고, JSSA는 세 모델과 모든 문맥 길이에서 풀 KV에 가깝게 유지되며 최대 격차가 1.11 퍼플렉시티 포인트였다. 반면 Loki는 LLaMA에서 급격히 나빠지고 Ministral에서는 잘하는 등 일관성이 떨어졌고, FASA는 LLaMA 128K에서 퍼플렉시티가 튀었다.
전제와 한계
추론 평가는 KV 예산 1,024 토큰에서 표준 GQA, 하이브리드 로컬/글로벌 GQA, MLA 모델을 대상으로 했다. 각 항목은 pass@1 정확도와 풀 KV 대비 생성 길이를 포함하는데, 희소 어텐션이 나빠지면 최종 답이 맞아도 추론 궤적이 불안정해질 수 있기 때문이다. JSSA는 아키텍처 전반에서 가장 좋은 정확도-길이 절충을 제공했고, 대부분 벤치마크에서 풀 KV와 1.3점 이내 차이를 유지하면서 Qwen3-8B와 Gemma4-26B-A4B-It에서 생성 길이 팽창을 2~17%로 제한했다. TriAttention, QUEST, SeerAttention-R은 훨씬 긴 생성을 보였다. FASA와의 격차는 Gemma4에서 가장 컸는데 AIME 2025에서 12.92점 차이였고, GLM-4.7-Flash에서는 JSSA가 AIME 2025, MATH-500, GPQA-Diamond 세 벤치마크 모두에서 풀 KV와 같거나 능가했다. 처리량은 단일 NVIDIA B200 GPU에서 Qwen3-8B, Gemma4-26B-A4B-It, GLM-4.7-Flash로 측정했고 모든 희소 기법이 KV 예산 B=2,048 토큰을 사용했다. 32K와 128K 문맥에서 배치 크기를 스윕한 결과, JSSA는 풀 KV보다 일관되게 좋았고 128K에서 이점이 더 커졌다. 128K에서 JSSA는 Qwen3-8B, Gemma4-26B-A4B-It, GLM-4.7-Flash에서 각각 최대 1.51배, 1.88배, 6.43배 처리량을 냈다. GLM-4.7-Flash에서는 FASA와 같은 처리량을 내면서 추론 정확도는 훨씬 높았다.
700B 파라미터 규모에서는 모델 네이티브 DeepSeek Sparse Attention(DSA)을 갖춘 DeepSeek-V3와 GLM-5.2를 평가했다. DeepSeek에서는 풀 어텐션 모델인 DeepSeek-V3.1-Terminus에 JSSA를 적용해 continued pre-training으로 네이티브 DSA를 도입한 DeepSeek-V3.2-Exp와 비교했다. GLM-5.2에서는 JSSA가 일부 레이어에서만 선택을 호출하고 인접 레이어에 인덱스를 재사용하는 모델 고유 스케줄을 따랐다. AIME 2025와 GPQA-Diamond pass@1에서 JSSA, 풀 KV, 네이티브 DSA는 좁은 정확도 범위에 머물렀다. JSSA는 DeepSeek와 GLM-5.2 모두 AIME 2025에서 네이티브 DSA와 일치했고, GLM-5.2의 GPQA-Diamond에서는 두 방법이 동일했다. 저자들은 수치 차이가 작아 JSSA가 네이티브 DSA보다 일관되게 우세하다기보다 정확도가 비슷하다고 해석한다. 처리량에서 JSSA는 DeepSeek의 측정된 문맥 길이와 배치 크기 전반에서 네이티브 DSA의 96.1~99.2%에 도달했고 continued training 없이 이를 달성했다. 128K 문맥, 배치 크기 8에서는 풀 KV 대비 최대 3.49배 처리량을 냈다. 절제 실험에서 투영 랭크를 4에서 32까지 바꾸면 JSSA가 모든 랭크에서 가장 좋았고, 가장 빡빡한 제약에서 이득이 컸다(r=4에서 FASA와 Loki보다 LongBench 점수 +5.40, +24.48). 캘리브레이션 코퍼스(LongAlign, NarrativeQA, LooGLE)에는 거의 둔감했지만, 캘리브레이션 길이는 중요했다. 32K에서 캘리브레이션한 기저는 128K 평가에서 퍼플렉시티가 16.05~18.31로 급격히 나빠졌고, 128K에서 캘리브레이션한 기저는 11.31~11.48이었다. 따라서 캘리브레이션은 최대 배포 문맥 길이에 가까운 시퀀스를 포함해야 한다. 수집 없는 실행은 128K LLaMA-3.1-8B-Instruct에서 배치 4와 16에서 풀 KV 대비 어텐션 단계 속도 향상을 약 두 배로 늘렸다(배치 16에서 25.35배에서 65.26배). Gemma4-26B-A4B-It에서는 gather-then-attend 대비 종단 간 처리량을 1.07~2.22배 개선했고, 이득은 배치 크기가 커질수록 증가했다.
개발자 입장에서 JSSA는 이미 배포된 GQA나 MLA 기반 긴 문맥 LLM에 학습 없이 붙일 수 있는 토큰 단위 희소 디코딩 경로다. 다만 적용 전에 몇 가지를 확인해야 한다. 첫째, 캘리브레이션 데이터를 최대 배포 문맥 길이에 가깝게 준비해야 한다. 32K 캘리브레이션으로 128K를 평가하면 성능이 크게 떨어진다는 절제 결과가 있다. 둘째, 아키텍처별 적용 범위가 다르다. Gemma4처럼 하이브리드 로컬/글로벌 구조에서는 글로벌 어텐션 레이어에만 JSSA를 적용하고 로컬 레이어는 네이티브 슬라이딩 윈도우를 유지한다. GLM-5.2에서는 모델 고유의 레이어 선택 스케줄을 따른다. 셋째, GQA에서 KV 헤드별로 다른 선택을 지원하는 수집 없는 인덱스 커널이 핵심이므로, 자체 서빙 스택에 통합할 때 이 커널과 KV 캐시 레이아웃 호환성을 확인해야 한다. 넷째, 점수화는 여전히 전체 N개 후보를 훑는다. 어텐션은 선택된 B개만 접근하지만 선택 비용은 O(N)으로 남아 있어, 문맥이 극단적으로 길어지면 이 부분이 병목이 될 수 있다. 마지막으로 FP8 투영 키 캐시가 원본 BF16 KV 캐시의 최대 6.25%를 차지하므로 메모리 예산에 이 보조 캐시를 포함해야 한다.
저자들이 밝힌 한계와 전제는 다음과 같다. 선택 목적식의 상한은 다루기 쉬운 대리 목적식을 제공하지만, 그 상한이 타이트하다거나 전역 최적성을 보장하지는 않는다. M*에서 정확한 목적식을 직접 개선해도 테스트한 설정에서는 측정 가능한 이득이 없었지만, 이것이 다른 설정에서도 항상 그렇다는 보장은 없다. 캘리브레이션 길이가 성능에 민감하다는 점도 실무 전제다. 또한 GLM-5.2에는 풀 어텐션 체크포인트가 없어서 네이티브 DSA와의 비교가 제한적이고, Gemma4에서는 글로벌 어텐션 레이어에만 적용했다. 정확도는 풀 KV에 가깝지만 동일하지는 않다. LongBench에서 0.67점, PG-19에서 최대 1.11 퍼플렉시티, 추론에서 대부분 1.3점 이내 차이가 났다. 700B 규모에서는 네이티브 DSA보다 우세하다기보다 정확도가 비슷하고 처리량이 96~99% 수준이라고 저자들 스스로 해석한다. 토큰 점수화가 전체 후보를 훑는 O(N) 비용을 갖는다는 점도 긴 문맥에서 남는 구조적 부담이다.