DoSA, 근사 없이 희소 어텐션의 Top-K 선택 가속

들어가며

희소 어텐션은 긴 컨텍스트에서 관련성이 높은 일부 키만 읽어 어텐션 비용을 낮추는 접근입니다. 하지만 실제 추론 시스템에서는 필요한 키를 찾는 과정 자체가 병목이 될 수 있습니다. 인덱서는 최종적으로 선택될 후보가 소수에 불과하더라도, 모든 후보에 대해 여러 인덱서 헤드의 점수를 계산한 뒤 Top-K를 결정할 수 있습니다. 후보 풀이 커질수록 이 탐색 비용은 희소 접근이 제공하는 이점을 일부 상쇄합니다.

DoSA(Dots Sparse Attention)는 이 문제를 무손실 후보 선별로 다룹니다. 모든 후보를 완전히 계산하는 대신, 가장 유리한 경우를 가정해도 현재 Top-K 경계를 넘을 수 없는 후보를 먼저 증명하고 제거합니다. 최종 선택 결과는 바꾸지 않은 채, 결과에 영향을 줄 수 있는 후보에만 전체 계산을 적용하는 방식입니다.

핵심 작동 방식

• 정확한 점수보다 충분한 배제 증거를 사용: 후보의 가능한 최대 점수가 현재 선택 임계값보다 낮다고 확인되면 정확한 점수를 계산할 필요가 없습니다. Top-K에서 중요한 것은 모든 후보를 완전히 평가하는 것이 아니라, 탈락을 확실히 입증하는 것입니다.
• 시드 집합으로 임계값의 하한을 확보: DoSA는 먼저 시드 후보 집합에서 현재 승자 임계값의 하한을 얻습니다. 이후 후보가 이 기준을 넘을 가능성이 있는지 검사합니다.
• 양의 기여와 음의 기여를 구분해 제한: 음이 아닌 헤드 투영과 결합 잔차 경계를 이용해 양의 기여가 커질 수 있는 정도를 상한으로 묶습니다. 음의 헤드가 만드는 차감량은 활성값 정보를 반영한 집계 경계로 제한합니다. 이 두 종류의 경계를 결합해 후보의 자격 유지 여부를 판정합니다.
• 정확한 계산 뒤에도 비용을 절감: 선별을 통과한 후보만 전체 점수를 계산한 뒤, 정확한 점수 필터링을 적용해 기록과 최종 선택에 필요한 작업을 추가로 줄입니다.
• 정확한 Top-K 보장: 논문은 이 인증 절차와 원래 계산과 동일한 Top-K를 얻기 위한 조건을 정식화합니다. 따라서 근사 검색으로 정확도를 희생하는 방법이 아닙니다.

결과와 영향

GPU 구현은 영어, 중국어, 코드 워크로드에서 평가되었습니다. 보고된 결과에 따르면 전체 Indexer는 SGLang 대비 최대 2.74배, LiteTopK 대비 최대 1.69배 빨라졌습니다. 엔드투엔드 평가에서는 GLM-5.2의 첫 토큰 도달 시간이 최대 1.94배 개선됐고, DeepSeek-V3.2에서는 최대 1.50배 개선됐습니다. 성능 향상은 점수 계산 감소에만 그치지 않고, 후보 기록과 이후 Top-K 선택 작업의 축소에서도 발생합니다.

DoSA의 핵심 의미는 희소성을 최종 읽기 단계보다 앞선 후보 탐색 단계에까지 적용했다는 데 있습니다. 시스템이 결국 소수의 토큰만 사용한다면, 그 토큰을 찾는 과정에서도 결과에 영향을 주지 않는 계산을 건너뛸 수 있어야 한다는 관점입니다. DoSA는 이를 경험적 휴리스틱이 아니라 점수 경계에 기반한 검증 가능한 절차로 구현합니다.

효과는 모델과 작업 조건에 따라 달라질 수 있습니다. 논문은 표현 공간의 기하학적 특성과 선택 후보 사이의 점수 마진이 성능을 좌우한다고 분석합니다. 경계가 촘촘하고 선택 차이가 작으면 전체 계산이 필요한 후보가 많아질 수 있으며, 반대로 상한이 타이트하고 후보 간 차이가 크면 조기 제거가 더 잘 작동합니다. 제시된 배속은 평가된 구현, 모델, 워크로드에서의 최대값입니다. 그럼에도 정확한 Top-K 의미를 유지하면서 “불필요함을 증명한 뒤 계산한다”는 설계는 긴 컨텍스트 추론을 위한 실용적인 최적화 방향을 보여줍니다.

출처: Hugging Face Daily Papers