Joint Subspace Sparse Attention for Efficient Long-Context Decoding

Sparse attention reduces the KV-cache traffic that dominates long-context LLM decoding, but practical methods largely operate at coarse granularity. Token-level selection offers finer control, yet presents two challenges: low-dimensional scoring may lose query--key structure needed for accurate ranking, while gathering scattered selected tokens can offset the resulting latency benefit. We introduce Joint Subspace Sparse Attention (JSSA) to address both challenges. For selection quality, JSSA learns a shared low-dimensional query--key subspace by minimizing token-score preservation error. We upper-bound the resulting non-convex objective with a tractable marginal objective that admits a closed-form spectral solution; direct refinement of the exact objective provides no measurable improvement in the tested settings. For execution efficiency, we develop a gather-free token-level indexed-attention kernel that supports distinct per-KV-head selections under the standard grouped-query attention (GQA) KV-cache layout. Across standard GQA, hybrid sliding-window/global attention, and Multi-head Latent Attention (MLA) models, JSSA consistently remains close to full-KV accuracy while avoiding the substantial generation-length inflation and architecture-dependent failures observed with competing methods. On an NVIDIA B200 GPU, JSSA achieves up to $6.43\times$ the decoding throughput of full-KV attention at 128K context. JSSA achieves reasoning accuracy comparable to model-native DSA without model training and reaches 96--99\% of its decoding throughput even at 700B scale large models.

한 줄 요약

JSSA는 토큰 수준의 희소 어텐션을 통해 128K 컨텍스트에서 최대 6.43× 가속화를 달성한 훈련 없이 사용 가능한 방법이다.

핵심 기여도

**공유된 쿼리-키 서브스페이스 학습**: 토큰 점수 보존 오류를 최소화하여 정확한 선택을 보장.
**가져오기 없는 인덱스 어텐션 커널**: GQA 기반 KV-cache 레이아웃에서 각 KV 헤드별 선택을 지원.
**다양한 모델 아키텍처에서의 성능**: GQA, MLA, 하이브리드 어텐션 모델에서 full-KV 정확도에 근접.
**700B 규모 모델에서 96–99%의 DSA 수준 성능**: 훈련 없이도 모델 네이티브 DSA와 유사한 추론 정확도.

핵심 아이디어

JSSA는 토큰 수준의 희소 어텐션을 통해 KV-cache 트래픽을 줄이는 동시에 정확도를 유지하는 문제를 해결한다. 기존 토큰 수준 선택 방법은 쿼리-키 구조를 유지하지 못하거나, 선택된 토큰을 모으는 과정에서 지연이 발생하는 문제가 있었다. JSSA는 이 두 문제를 동시에 해결하기 위해 **공유된 쿼리-키 서브스페이스**를 학습하여 토큰 점수 보존 오류를 최소화한다. 이는 비볼록 최적화 문제를 **닫힌 형태의 스펙트럼 해법**으로 근사화함으로써 효율적으로 해결된다. 또한, **가져오기 없는 인덱스 어텐션 커널**을 통해 GQA 레이아웃에서 각 헤드별 선택을 직접 처리하여 추가적인 메모리 트래픽을 방지한다.

기술적 접근법

**공유된 쿼리-키 서브스페이스 학습**: 토큰 점수 보존 오류를 최소화하는 비볼록 목적함수를 **닫힌 형태의 스펙트럼 해법**으로 근사화.
**가져오기 없는 인덱스 어텐션 커널**: GQA KV-cache 레이아웃에서 각 KV 헤드별 선택을 지원하며, 선택된 토큰을 별도로 모으지 않음.
**모델 아키텍처**: GQA, MLA, 하이브리드 로컬/글로벌 어텐션 모델에서 테스트.
**하드웨어**: NVIDIA B200 GPU에서 평가.
**수치**: 128K 컨텍스트에서 최대 6.43× 가속화.

주요 결과

**정확도**: GQA, MLA, 하이브리드 어텐션 모델에서 full-KV 정확도에 근접.
**성능**: 128K 컨텍스트에서 full-KV 대비 최대 6.43× 높은 decoding throughput.
**대규모 모델 성능**: 700B 파라미터 모델에서 DSA 대비 96–99%의 decoding throughput 달성.
**추론 정확도**: DSA와 유사한 수준의 정확도 유지, 훈련 없이도.

의의 및 한계

JSSA는 훈련 없이도 토큰 수준의 희소 어텐션을 구현하면서도 GQA 및 MLA 모델에서 full-KV 정확도에 근접하는 성능을 보인다. 이는 대규모 모델에서의 추론 성능과 효율성을 동시에 확보할 수 있는 중요한 기술적 진전이다. 그러나, JSSA는 특정 아키텍처에 최적화된 DSA와 비교했을 때 **완전히 동일한 성능을 보장하지는 않으며**, 일부 모델 구조에서는 **아키텍처 의존적 실패**가 발생할 수 있다는 한계가 있다. 또한, **정확도 향상을 위한 직접적인 목적함수 최적화는 실험적으로 유의미한 개선을 제공하지 않았다**는 점도 주목할 만하다.

실용적 활용

JSSA는 대규모 언어 모델에서의 **장문 컨텍스트 추론**을 효율적으로 처리할 수 있는 방법으로, 특히 **실시간 추론**이나 **저비용 추론 인프라**에 적용 가능하다. GQA 및 MLA 기반 모델에서의 활용이 유망하며, **700B 이상의 대규모 모델에서도 높은 성능**을 유지하므로 클라우드 기반 대형 모델 서비스에 적합하다.