SSA는 극한 희소성에서 밀집 어텐션과의 격차를 줄인다
Stratified Sparse Attention: Reducing the Sparse-to-Dense Gap at Extreme Sparsity
무엇인가
이 논문은 긴 컨텍스트 LLM 추론에서 KV 캐시를 매 디코드 단계마다 읽는 비용이 병목이 되는 문제를 다룬다. 희소 어텐션은 캐시의 일부만 접근해 비용을 줄이지만, 기존 방식은 두 갈래의 한계를 갖는다. 프록시 점수로 상위 토큰을 결정적으로 고르는 top-k·top-p 계열은 높은 점수의 토큰은 잡지만 선택 집합 밖의 잔여 어텐션 질량을 버린다. 반대로 샘플링은 커버리지를 넓히지만 분산이 커지고, top-k와 균등 샘플링을 섞는 하이브리드도 극한 희소성에서는 밀집 어텐션과 큰 격차를 보인다. SSA는 이 격차를 줄이기 위해 결정적 검색과 점수 기반 샘플링을 하나의 프레임워크로 묶는다.
어떻게 동작하나
SSA의 핵심은 토큰을 결정적으로 유지할 층과 확률적으로 탐색할 잔여 층으로 나누는 것이다. 프록시 점수 q_i를 확률로 바꾼 뒤 ∑ min(1, λ q_i)=k를 만족하는 λ를 찾고, λ q_i가 1 이상인 토큰을 고정 집합 H로 유지한다. 남은 예산 m=k-|H|는 잔여 집합 R에서 중요도 샘플링에 쓴다. 이 규칙 때문에 고정 집합 크기를 별도로 정하지 않아도 프록시 분포가 결정적 할당과 잔여 샘플링의 비율을 자동으로 정한다. 논문은 이 틀이 top-k, top-p, top-k+샘플링 하이브리드를 특수 사례로 포함한다고 주장한다.
무엇과 다른가
잔여 샘플링에서는 R 위에서 프록시 확률을 정규화해 제안 분포 p_i를 만들고, m번의 다항 추출을 한 뒤 중복을 제거해 집합 S를 얻는다. 선택된 토큰의 포함 확률은 고정 집합에서 1이고, 잔여 토큰에서는 π_i=1-(1-p_i)^m이다. 이후 선택된 토큰에 대해서만 정확한 어텐션 점수를 계산하고 exp(ℓ_i)/π_i로 보정한 가중치를 사용해 정규화 가중합을 낸다. 중복 추출은 한 번만 유지하므로 최종 선택 크기는 k 이하가 된다. 이론적으로 논문은 잔여 질량 A_R과 프록시 품질 ρ에 대해 오차가 A_R ρ/√m 규모로 줄어드는 bound를 제시하고, ε-δ 보장을 위한 표본 크기 m=O(A_R^2 ρ^2 ε^-2 log(1/δ))를 제시한다.
어떻게 쓰나
커널 구현은 긴 컨텍스트에서 메모리 대역폭 병목을 줄이는 데 초점을 둔다. 각 키를 16바이트 PQ 코드로 표현해 헤드 차원 256의 fp16 키-값 쌍 1KiB보다 훨씬 적은 바이트로 프록시 점수를 계산하고, radix top-k로 고정 집합을 고른다. 잔여 샘플링은 1024토큰 청크, 32토큰 서브타일, 최대 32토큰 스캔으로 좁혀 가는 계층적 샘플러를 쓰고, 행별 해시 테이블로 중복을 제거한 뒤 1/π_i로 재가중한다. π_i는 10^-4로 클램프해 과도한 가중치를 막지만 작은 편향을 도입한다. 최종 희소 디코드는 paged KV 캐시에서 선택된 키와 값을 모아 online softmax로 계산하고, 쿼리 헤드를 여러 스레드 블록으로 나눈 뒤 log-sum-exp로 합친다. 논문은 최적화된 GPU 구현이 적응적 예산 분할 대신 고정 heavy 비율을 사용하며, 적응형 알고리즘의 동등한 효율 구현은 향후 과제라고 밝힌다.
전제와 한계
실험은 RULER 32K-128K, LongBench 10K-16K, LOFT 32K-128K, AIME 32K 토큰 생성에서 수행됐다. 모델은 Llama-3.1-8B-Instruct, Qwen3.5-4B, Qwen3.5-9B, Ministral-3-3B-Instruct-2517, Ministral-3-8B-Instruct-2517이고, 베이스라인은 vAttention, PQCache, Oracle-top-k다. RULER-HARD 50배 희소성에서 SSA는 모든 모델과 컨텍스트 길이에서 PQ와 vAttention(PQ)을 앞섰다. Ministral-3-8B에서는 32K에서 81.8%에서 84.2%로, 128K에서 63.8%에서 69.3%로 올라 밀집 어텐션 성능의 94%와 85%를 회복했다. 32K에서 Llama-3.1-8B는 PQ 72.51에서 SSA 83.67로 올랐고 밀집은 88.76이었으며, Ministral-3-3B는 70.81에서 82.05로 올랐고 밀집은 86.69였다. 다만 Oracle-top-k는 128K의 Llama-3.1-8B와 Ministral-3-3B, 그리고 Ministral-3-8B의 모든 컨텍스트 길이에서 여전히 더 강했다. Qwen3.5-4B와 Qwen3.5-9B에서는 모든 희소 방식이 밀집과 2점 이내에 머물렀다.
AIME 2025 avg@16에서 Qwen3.5-9B 50배 희소성일 때 SSA는 59.8로 Oracle-top-k 48.3, PQ 36.7, vAttention(PQ) 35.8을 앞섰다. 100배 희소성에서는 SSA가 51.5를 유지한 반면 vAttention(PQ)는 3.5로 떨어졌다. 낮은 희소성에서는 vAttention(PQ)가 5배에서, 그리고 Qwen3.5-4B의 10배에서 더 나은 경우도 있었다. heavy 비율 h를 0.2에서 0.8로 둔 하이브리드 할당은 61.0-64.0을 기록해 순수 샘플링 40.2-42.5와 순수 top-k 49.4-53.8을 크게 앞섰다. LOFT에서는 Qwen3.5-9B와 Qwen3.5-4B 모두에서 SSA가 비-oracle 방식 중 가장 낮은 어텐션 오차를 보였고, subspan-EM은 방식 간 차이가 최대 0.04였다. LongBench의 Qwen3.5-27B에서는 50배와 20배 희소성 모두 모든 희소 방식이 밀집과 1점 이내였다.
디코드 효율은 컨텍스트와 배치가 커질수록 유리해진다. 배치 1, 128K에서는 SSA가 FlashInfer보다 느려 0.60-0.74배 속도에 그쳤다. 그러나 1M 컨텍스트에서는 50배, 20배, 10배 희소성에서 각각 1.79배, 1.39배, 1.07배 빨라졌고, 배치 16의 50배 희소성에서는 2.96배까지 올랐다. 커널 최적화는 배치 1에서 참조 구현 대비 2.95-3.80배 개선을 냈고, 2단계 샘플러는 추출 수가 늘수록 청크 전용 라우팅보다 1.21-3.16배 샘플링 비용을 줄였다. 1M 컨텍스트와 h=0.4에서 샘플링 지연은 평가된 희소성 수준에서 0.30-1.18ms였다.
개발자 관점에서 SSA는 긴 컨텍스트 디코드가 메모리 대역폭에 묶여 있고, 프록시 점수를 제공하는 근사 top-k 계열을 이미 쓰거나 새로 붙일 수 있는 경우에 특히 관련이 있다. 다만 짧은 컨텍스트나 배치 1에서는 선택·샘플링 오버헤드가 이득을 상쇄할 수 있으므로 컨텍스트 길이, 배치 크기, 희소성, heavy 비율을 함께 튜닝해야 한다. 저자들이 밝힌 한계도 분명하다. 최적화 커널은 고정 heavy 비율을 사용해 이론적 적응형 분할과 다르고, π_i 클램프는 작은 편향을 만들며, 배치 1의 128K에서는 FlashInfer보다 느리다. 또한 향후 과제로 사용자 제어 가능한 근사 품질 보장을 갖춘 검증 가능한 SSA를 제안한다.