DoSA가 스파스 어텐션의 토큰 선택 비용을 인증 기반으로 줄인다

무엇인가

스파스 어텐션은 소수의 키만 방문하지만, 그 소수를 찾아내는 과정은 여전히 모든 키를 모든 인덱서 헤드로 점수화해야 할 수 있다. DeepSeek Sparse Attention(DSA)은 멀티헤드 인덱서로 점수를 매긴 뒤 Top-K를 고르는데, 점수 하나에 내적과 ReLU, 부호 있는 가중 합이 들어간다. 논문은 약 100만 토큰에서 인덱서가 프리필 GPU 시간의 73.0%(GLM-5.2), 90.8%(DeepSeek-V3.2)를 차지한다고 측정한다. 논리적 행렬 FLOPs 비중은 각각 47.4%, 83.2%에 불과한데 시간 비중은 훨씬 크다는 점에서, 점수 계산뿐 아니라 선택과 메모리 트래픽까지 줄여야 한다는 문제의식이 나온다. GVR, DeepSelect, LiteTopK 같은 기존 시스템은 점수 융합과 스트리밍으로 선택 비용을 줄이지만 여전히 모든 키에 대해 전 헤드 평가가 필요하고, MISA처럼 적은 헤드로 스크리닝한 뒤 재랭킹하는 방식은 이미 제외된 승자를 되살릴 수 없다.

어떻게 동작하나

DoSA의 핵심 관찰은 Top-K 선택에 후보를 탈락시키기에 충분한 증거만 있으면 되고, 전체 점수가 필요하지는 않다는 것이다. 가시 키 중 K개 이상을 시드 집합 S로 잡으면 그 K번째 최대 점수 τ가 실제 Top-K 임계값의 하한이 된다. 어떤 키의 점수 상한 U_i가 τ보다 작으면 그 키는 Top-K에 들 수 없다. 문제는 이 상한을 점수 자체보다 훨씬 싸게, 그러면서도 타이트하게 만드는 것이다. DoSA는 일부 헤드 집합 R만 직접 계산하고 나머지는 상한으로 처리한다. 양수 가중치 헤드는 버리면 안전하지 않으므로, 각 헤드를 유지된 앵커 헤드에 음이 아닌 계수 β_h로 투영하고 잔차 d_h의 노름만큼 여유를 더한다. 이 투영은 이미 계산된 내적을 재사용하며, 단일 앵커 형태의 상한 중 잔차 노름을 가장 작게 만든다. 생략된 양수 헤드들의 잔차는 서로 같은 키를 공유한다는 점을 이용해 결합 허용량 E_P^joint로 더 조일 수 있고, 이는 추가 키 읽기 없이 기하 구조만 활용한다. 음수 가중치 헤드는 점수에서 빼지는 항이므로, 활성화 패턴에서 과반 지지(p_h ≥ 1/2)를 받는 헤드들을 하나의 집계 방향 a_A로 합쳐 빼지는 크기의 하한을 만든다. 모든 음수 헤드를 넣지 않는 이유는, 비활성 헤드의 음수 내적이 활성 헤드의 양수 내적을 상쇄해 추정치를 오히려 느슨하게 만들 수 있기 때문이다.

무엇과 다른가

최종 인증식은 U_i = Σ_{j∈R} w̃_j [q_j^T k_i]_+ + E_P ||k_i||_2 − [a_A^T k_i]_+ 이며, s_i ≤ U_i가 성립한다. 여기서 w̃_j는 투영 계수로 접힌 가중치이고, 이는 스크리닝 계획에만 쓰이며 모델 파라미터를 바꾸지 않는다. 실행은 네 단계다. 1단계에서 유지 헤드와 보조 행을 키 타일 전체에 재사용하며 상한을 계산하고, 2단계에서 U_i < τ인 키를 버리고 생존자 ID를 압축한다. 3단계에서 생존 키만 모아 원래 H개 헤드로 전체 점수를 계산하고, 4단계에서 원래 ReLU와 가중치를 적용해 s_i ≥ τ인 점수만 시드 레코드와 함께 선택에 넘긴다. 상한은 최종 점수를 대체하지 않으며, 두 필터 모두에서 경계값 동점은 살아남는다. 시드 임계값이 유효하고 살아있는 후보를 잃지 않는다는 조건에서 Top-K 선택의 정확성이 보존된다.

어떻게 쓰나

GPU 커널은 인증 기반 가지치기가 만드는 불규칙한 생존자 집합을 다루기 위해 설계됐다. 비슷한 유지 헤드 예산을 가진 쿼리를 고정 폭 타일로 묶어 Tensor Core MMA를 쓰고, 유지 쿼리 행과 접힌 가중치, 두 개의 FP8 보조 행을 한 번 준비해 키 타일마다 재사용한다. 전용 워프가 비동기 키 로드와 MMA를, 소비 워프가 TMEM 누산기에서 상한을 평가하며, 워프 ballot으로 생존 결정을 압축 마스크로 만든다. 정밀화 단계에서는 네 쿼리 그룹 안에서 두 개 이상의 상한을 통과하고 공통 가시 범위에 있는 키를 공유 스트림으로 보내 키 로드를 한 번만 수행한다. 쓰기 단계에서는 임계값을 통과한 점수만 단조 점수 코드와 정확한 키 ID로 기록해 HBM 트래픽과 선택 입력을 줄인다. 디코딩에서는 스텝당 쿼리가 적어 키 재사용이 제한적이므로, 최근 스텝에서 선택된 키를 증인으로 삼아 현재 쿼리 점수로 Top-K 임계값의 보수적 하한을 얻는 시간적 일관성 기법을 쓴다.

전제와 한계

실험은 GLM-5.2(인덱서 헤드 32개)와 DeepSeek-V3.2(64개), Top-K=2048에서 수행됐다. 마이크로벤치는 단일 GPU의 GLM-5.2, 엔드투엔드는 8-way 텐서 병렬에 프리필 청크 32K(GLM-5.2)와 16K(DeepSeek-V3.2)를 쓴다. 워크로드는 InfiniteBench의 영어·중국어 도서와 코드, Project Gutenberg, SQLite 소스코드이며 컨텍스트는 128K에서 약 1M 토큰까지다. 베이스라인은 SGLang v0.5.18과 LiteTopK다. DoSA는 60개 마이크로벤치 구성 전부에서 가장 빨랐고, SGLang 대비 1.12~2.74배, LiteTopK 대비 1.17~1.69배 빨랐다. 512K 이상에서는 모든 구성에서 2배를 넘었고, 1M에서는 각각 2.47~2.74배와 1.47~1.69배였다. 이득은 쿼리 청크 크기보다 컨텍스트 길이에 더 민감하다. 영어 입력 16K 쿼리 기준으로 128K에서 1.19배였던 속도 향상이 1M에서 2.58배가 되지만, 1M에서 청크를 4K에서 32K로 8배 늘려도 2.47배에서 2.69배로만 변한다. 디코딩 인덱서는 256K, 512K, 약 1M에서 각각 1.06배, 1.11배, 1.26배 빨라졌다. 가지치기율(1 − N_keep/N_screen)은 각 도메인 21개 레이어 중 20개에서 95%를 넘었고 도메인 평균 97.2~97.5%였으며, 레이어 1에서만 89~93%로 가장 낮았다. 엔드투엔드 TTFT는 1M 토큰에서 SGLang 대비 GLM-5.2 1.86배, DeepSeek-V3.2 1.94배, LiteTopK 대비 1.26배와 1.50배를 기록했다. 인덱싱 비중이 128K의 23.6%에서 1M의 73.0%(GLM-5.2)로 커지면서 암달 법칙에 따라 전체 이득도 함께 커진다는 설명이 붙는다.

어블레이션은 LOFT의 SciFact 검색 예제(16K 쿼리, 896K 키, 21개 레이어 평균)에서 이뤄졌다. 스크리닝만 적용하면 지연이 37.4%, 스파스 쓰기만 적용하면 23.4% 줄고, 둘을 함께 쓰면 47.3% 줄어 베이스 대비 1.90배가 된다. 전체 헤드 평가를 그대로 두고 융합만 한 LiteTopK류와의 차이가 곧 인증 기반 제거가 없애는 계산량이다. 특히 음수 집계 항을 제거하면 후보 생존율이 2.65%에서 95.06%로 뛰어올라, 부호 있는 인덱서 점수에서 음수 기여를 타이트하게 묶는 것이 사실상 가지치기의 핵심임을 보여준다.

실무 관점에서 이 논문은 롱컨텍스트 서빙 파이프라인에서 인덱서가 병목인 경우에 바로 검토할 만하다. 공개된 구현은 https://github.com/llm-sys/flash-indexer 에 있다. 다만 DoSA는 원래 라우팅 규칙을 그대로 보존하는 무손실 최적화이므로, 인덱서 구조나 Top-K 규칙 자체를 바꾸는 기법과는 조합 방식이 다르다. 도입 시에는 자신의 워크로드에서 시드 임계값이 얼마나 타이트하게 형성되는지, 즉 가지치기율과 스크리닝 비용의 균형을 먼저 확인해야 한다. 컨텍스트가 짧거나 상한이 느슨하면 스크리닝 비용을 상쇄할 만큼의 절감이 나오지 않을 수 있다.

저자들이 밝힌 한계도 분명하다. DoSA는 여전히 모든 키를 스캔하며, 짧은 컨텍스트나 느슨한 상한에서는 절감이 스크리닝 비용을 넘지 못할 수 있다. 수치적 보장을 위해 원래 스코어러보다 보수적으로 평가해야 하며, 동점으로 묶인 유효한 Top-K 집합이 항상 동일한 생성 텍스트를 낳지는 않는다. 측정된 것은 계산상의 이득이고, 모델이 학습된 컨텍스트를 넘어서는 과제 정확도 향상은 아니다. 향후 과제로는 인증 여유(slack)를 학습 시점에 정규화해 건너뛰기 쉬운 인덱서를 유도하는 것, 그리고 개별 키 단위 인증을 키 블록 단위 상한으로 확장해 키 벡터를 읽지 않고도 블록을 건너뛰는 것이 제시된다.