딥시크 희소 어텐션 DSA 원리: 라이트닝 인덱서와 상위 2048 토큰 선택 구조 해설

핵심 메시지

• 표준 어텐션은 실제로 중요한 소수의 토큰을 찾기 위해 모든 토큰을 다 확인하며, 12만 8천 토큰에서는 레이어마다 수십억 번의 점수 비교가 발생한다.
• 딥시크 희소 어텐션은 값싼 라이트닝 인덱서로 모든 과거 토큰에 점수를 매긴 뒤 상위 2048개만 남기며, 12만 8천 토큰 기준으로 1.6%에 해당한다.
• 디코딩 단계의 병목은 연산이 아니라 KV 캐시를 HBM에서 온칩 SRAM으로 옮기는 데이터 이동이다.
• MLA가 캐시 항목 하나의 크기를 줄이는 기법이라면, DSA는 불러오는 항목의 개수를 줄이는 기법이라는 점에서 서로 다른 축을 담당한다.
• 인덱서 자체는 여전히 제곱 복잡도이고 한 번 놓친 토큰은 그 위치에서 영구히 제외되는 등, 구조적 한계도 분명하게 남아 있다.

쉽게 이해하기

출발점은 어텐션 히트맵의 관찰이다. 각 행은 하나의 질의 토큰이 전체 시퀀스에 대해 갖는 분포인데, 대부분의 칸은 어둡고 값이 0에 가까우며 행마다 몇 개만 밝게 빛난다. 그렇다면 밝은 칸만 계산하면 되지 않느냐는 발상이 자연스럽게 나오지만 여기에 순환이 있다. 어느 칸이 밝은지는 다 계산해 봐야 알 수 있고, 다 계산하는 것이야말로 피하려던 일이기 때문이다.

비용의 성격은 단계에 따라 다르다. 입력 프롬프트를 처리하는 프리필은 전체 어텐션 행렬을 병렬로 만들기 때문에 연산 바운드이고 시퀀스 길이의 제곱에 비례한다. 반면 토큰을 하나씩 만들어 내는 디코딩은 메모리 대역폭 바운드다. 새 토큰마다 HBM에 있는 KV 캐시 전체를 온칩 SRAM으로 끌어와야 해서, 12만 8천 토큰에서는 매 스텝 기가바이트 단위 데이터가 오가고 연산 코어는 자주 놀게 된다.

DSA의 토대는 다중 헤드 잠재 어텐션(MLA)이다. 토큰마다 완전한 키와 값을 저장하는 대신 훨씬 작은 압축 잠재 벡터 하나만 캐시에 두고, 필요할 때 학습된 업프로젝션으로 복원한다. 보통 5~10배를 아끼며, 회전 위치 임베딩과 충돌하는 문제는 위치 성분을 별도 스트림으로 분리해 푼다. 다만 MLA는 항목 하나의 크기를 줄일 뿐 항목의 개수는 그대로여서, 문맥에 L개의 토큰이 있으면 여전히 매 스텝 L개를 읽어야 한다.

DSA는 그 앞단에 두 단계 게이트를 붙인다. 인덱서는 작은 헤드 여러 개를 쓰고 키는 모든 헤드가 공유하며, ReLU 내적의 가중합으로 각 과거 토큰의 점수를 만든다. 선택기는 그중 점수가 높은 2048개 위치만 남기고 나머지를 버리는 기계적인 역할만 한다. 그다음 실제 어텐션은 선택된 토큰에 대해서만 전정밀도로 수행되고, 소프트맥스 분모 역시 전체가 아니라 선택된 집합에 대해서만 더해진다. 선택은 질의 위치마다 독립적으로 이뤄져 토큰마다 고르는 이웃이 다르다.

학습은 처음부터가 아니라 이미 학습을 마친 V3.1 Terminus 위에 두 단계로 얹는다. 1단계는 모델 가중치를 얼린 채 인덱서만 1000스텝, 21억 토큰 동안 학습시키며 KL 발산으로 원본 어텐션 분포를 흉내 내게 한다. 2단계에서는 희소 마스크를 켜고 전체 가중치를 9437억 토큰 규모로 함께 학습하되, 인덱서와 본체의 그래디언트 경로를 완전히 분리한다. 결과적으로 12만 8천 토큰에서 희소 MLA는 전체 MLA 대비 64배 적은 연산을 쓰고, 오픈소스로 공개된 Flash MLA는 H800에서 프리필 최대 640 테라플롭스, 디코딩 410 테라플롭스를 낸다.

주요 인사이트

• 설계의 핵심은 값싼 연산의 결과 위에서 비싼 연산을 돌린다는 분업이다. 인덱서는 정확한 점수를 낼 필요가 없고 순위만 잘 매기면 되며, 선택된 집합에 대한 어텐션은 DSA가 없을 때와 똑같은 품질의 계산이다.
• 인덱서와 본체의 그래디언트를 분리한 이유가 흥미롭다. 두 경로가 이어져 있으면 본체가 언어를 잘 처리하도록 배우는 대신 자신의 어텐션 패턴을 인덱서가 맞히기 쉽게 만드는 쪽으로 학습되는 병리적 결합이 생길 수 있기 때문이다.
• 로그-로그 그래프로 보면 인덱서는 전체 MLA와 기울기가 같은 제곱 복잡도지만 상수가 약 50배 작아 아래쪽에 평행하게 놓인다. 반면 희소 MLA는 기울기가 1인 선형이어서 문맥이 길어질수록 격차가 벌어진다.
• 10만 토큰 부근에는 교차점이 있다. 그보다 짧은 구간에서는 제곱항이 충분히 억눌린 인덱서가 오히려 선형인 희소 MLA보다 싸고, 그 위로 올라가면 관계가 뒤집힌다.
• V3.2는 두 종류의 희소성을 겹쳐 쓴다. 전체 6850억 파라미터 중 토큰마다 370억만 활성화하는 전문가 혼합 구조가 하나이고, 문맥의 1.6%만 참조하는 어텐션 희소성이 다른 하나다. 둘 다 위치가 아니라 내용에 따라 학습으로 결정된다.

자주 묻는 질문

관련 AI 소식