IndexCache: 계층 간 인덱스 재사용을 통한 희소 어텐션 가속화

왜 중요한가

긴 문맥을 처리할 때 병목이 되는 희소 어텐션의 인덱서 연산을 획기적으로 줄이는 방법을 제시한다. 계층 간의 유사성을 활용해 성능 저하 없이 추론 속도를 높이고 비용을 절감할 수 있어 실무적 가치가 크다.

핵심 기여

IndexCache 아키텍처 설계

레이어를 Full 레이어와 Shared 레이어로 구분하여, 일부 레이어에서만 인덱스를 계산하고 나머지는 이를 공유하는 구조를 통해 연산 중복을 제거했다.

Greedy 레이어 선택 알고리즘

가중치 업데이트 없이도 언어 모델 손실값(LM Loss)을 최소화하는 최적의 인덱서 유지 레이어 조합을 찾아내는 데이터 기반 검색 방식을 도입했다.

Multi-layer Distillation 기법

인덱서가 자신이 서빙하는 여러 레이어의 평균적인 어텐션 분포를 동시에 학습하도록 하여, 공유 구조에서도 높은 정확도를 유지하게 했다.

대규모 모델 실증 및 가속

30B 및 744B(GLM-5) 모델에서 인덱서 연산의 75%를 제거하면서도 성능 저하 없이 최대 1.82배의 프리필 속도 향상을 달성했다.

핵심 아이디어 이해하기

Transformer의 Self-Attention은 모든 토큰 간의 관계를 계산하므로 시퀀스 길이의 제곱에 비례하는 연산량이 필요하다. 이를 해결하기 위해 DeepSeek Sparse Attention(DSA) 같은 기법은 Indexer라는 가벼운 모듈을 사용해 중요한 상위 k개 토큰만 골라내어 연산량을 줄이지만, 이 Indexer 자체도 모든 레이어에서 독립적으로 실행되며 제곱 단위의 연산량을 소모하는 한계가 있다.

연구진은 인접한 레이어들이 선택하는 중요한 토큰의 집합이 70~100% 일치한다는 계층 간 중복성에 주목했다. IndexCache는 모든 레이어에서 인덱서를 돌리는 대신, 일부 레이어(Full)에서만 인덱스를 계산하고 나머지 레이어(Shared)는 이를 그대로 재사용함으로써 중복된 연산을 제거한다.

단순히 일정 간격으로 공유하는 것이 아니라, 언어 모델의 손실값(LM Loss)을 최소화하는 최적의 레이어 조합을 찾아내거나, 인덱서가 여러 레이어의 요구사항을 동시에 만족하도록 학습시키는 방식을 통해 정확도 손실을 거의 없앴다. 결과적으로 긴 문맥 처리 시 가장 큰 병목이었던 인덱서 연산을 75%까지 줄여 전체 추론 속도를 대폭 개선했다.

방법론

IndexCache는 전체 N개의 레이어를 Full(F)과 Shared(S) 역할로 구분한다. F 레이어는 자체 인덱서를 실행해 새로운 top-k 인덱스를 생성하고 이를 캐시에 저장하며, S 레이어는 가장 가까운 이전 F 레이어의 캐시된 인덱스를 그대로 가져와 어텐션 연산을 수행한다.

최적의 레이어 패턴을 찾기 위해 Greedy Search 알고리즘을 사용한다. 모든 레이어를 F로 시작하여, 하나씩 S로 바꿨을 때 캘리브레이션 데이터셋에서의 LM Loss 증가가 가장 적은 레이어를 선택해 나가는 방식이다. [패턴 후보 입력 → 순방향 전파 연산 → LM Loss 출력 → 최소 손실 패턴 확정] 과정을 통해 가중치 업데이트 없이도 효율적인 구조를 결정한다.

학습 기반 방식에서는 Multi-layer Distillation Loss를 도입한다. 특정 F 레이어의 인덱서가 자신이 서빙하는 모든 레이어의 평균적인 어텐션 분포를 예측하도록 학습시킨다. [각 레이어의 어텐션 분포 입력 → 평균값 계산 → KL Divergence 연산 → 인덱서 가중치 갱신] 과정을 통해 인덱서가 여러 레이어에 범용적인 토큰을 선택할 수 있게 한다.

주요 결과

30B DSA 모델 실험 결과, 인덱서의 75%(1/4 유지)를 제거했을 때 200K 문맥 길이에서 프리필(Prefill) 속도는 1.82배, 디코딩(Decode) 속도는 1.48배 향상되었다. 이는 인덱서 연산이 긴 문맥에서 전체 지연 시간의 상당 부분을 차지하기 때문이다.

성능 유지 측면에서, Greedy Search로 찾은 패턴은 단순 균등 간격(Uniform) 방식보다 월등히 우수했다. 1/4 유지 패턴에서도 원래 DSA 모델과 거의 동일한 벤치마크 점수를 기록했으며, 특히 AIME 2025와 GPQA-Diamond 같은 추론 태스크에서는 오히려 성능이 소폭 향상되는 정규화 효과도 관찰되었다.

744B 규모의 GLM-5 모델에 적용했을 때도 유사한 경향이 확인되었다. 1/2 인덱서 유지 시 성능 저하 없이 전체 추론 속도가 약 1.2배 향상되었으며, 이는 대규모 프로덕션 모델에서도 IndexCache의 확장성과 실용성이 높음을 입증한다.

기술 상세

IndexCache는 DSA(DeepSeek Sparse Attention)의 핵심 구성 요소인 Lightning Indexer의 중복성을 해결한다. DSA는 O(L²)의 인덱서와 O(Lk)의 코어 어텐션으로 구성되는데, 문맥이 길어질수록 인덱서의 비중이 커지는 문제를 계층 간 인덱스 공유로 해결한다.

수학적으로 Multi-layer Distillation Loss는 각 레이어의 KL Divergence 합으로 정의되지만, 이는 각 레이어 어텐션 분포의 평균(Centroid)을 타겟으로 하는 단일 증류 손실과 그래디언트가 동일함을 증명했다. 이를 통해 구현 효율성을 높이면서도 여러 레이어를 아우르는 합의된(consensus) top-k를 학습할 수 있다.

레이어 선택의 민감도는 네트워크의 초기 및 전이 영역에서 높게 나타났다. Greedy Search는 이러한 임계 레이어(critical layers)를 자동으로 식별하여 인덱서를 유지함으로써, 단순한 균등 간격 공유가 유발하는 누적 오차와 성능 저하를 방지한다.

구현 측면에서 Shared 레이어는 추가적인 GPU 메모리를 거의 사용하지 않는다. F 레이어에서 계산된 인덱스 텐서를 임시 버퍼에 저장하고 S 레이어에서 이를 참조하기만 하면 되므로, 기존 DSA 할당량 외의 오버헤드가 거의 없다.

한계점

1/8 이하의 극단적인 인덱서 유지 비율에서는 Greedy Search를 사용하더라도 긴 문맥 성능의 유의미한 하락이 관찰되었다. 또한, 본 연구는 DSA 아키텍처를 중심으로 검증되었으므로 다른 형태의 희소 어텐션이나 동적 토큰 선택 기법에 대한 일반화 가능성은 추가 연구가 필요하다.

실무 활용

긴 문맥을 처리하는 에이전트 워크플로우나 대규모 문서 분석 시스템에서 추론 비용을 낮추고 속도를 높이는 데 즉시 적용 가능하다. 특히 DeepSeek-V3나 GLM-5와 같이 희소 어텐션을 사용하는 최신 LLM 아키텍처에 최적화되어 있다.

긴 문맥(200K+)을 사용하는 법률/의료 문서 분석 서비스의 추론 가속
멀티스텝 에이전트의 실시간 응답 속도 개선 및 서버 비용 절감
제한된 GPU 메모리 환경에서 긴 시퀀스를 처리해야 하는 온프레미스 LLM 배포
DeepSeek 기반 모델의 추론 엔진(vLLM, SGLang 등) 최적화 플러그인

코드 공개 여부: 공개

코드 저장소 보기

키워드

Sparse Attention(희소 어텐션)DeepSeek Sparse Attention(DeepSeek 희소 어텐션)Cross-layer Redundancy(계층 간 중복성)Inference Optimization(추론 최적화)Long-context LLM(긴 문맥 대형 언어 모델)

용어 해설

희소 어텐션(Sparse Attention): — 모든 토큰 간의 관계를 계산하는 대신 중요한 일부 토큰만 선택하여 연산하는 기법이다. 연산 복잡도를 시퀀스 길이의 제곱에서 선형 수준으로 낮추어 긴 문맥 처리를 가능하게 하는 핵심 기술이다.
KV 캐시(KV Cache): — 이전 추론 단계에서 계산된 Key와 Value 벡터를 메모리에 저장해 두었다가 재사용하는 기술이다. 매 단계마다 전체 시퀀스를 다시 계산할 필요가 없게 하여 LLM의 생성 속도를 비약적으로 높인다.
지식 증류(Knowledge Distillation): — 복잡한 모델(교사)의 출력 분포를 더 단순한 모델(학생)이 모방하도록 학습시키는 기법이다. 본 논문에서는 여러 레이어의 어텐션 정보를 하나의 인덱서가 학습하도록 하는 데 활용된다.
프리필(Prefill): — LLM 추론의 첫 단계로, 입력된 프롬프트 전체를 한꺼번에 처리하여 초기 상태를 생성하는 과정이다. 문맥이 길어질수록 이 단계에서 발생하는 연산 병목 현상이 심화된다.

IndexCache: 계층 간 인덱스 재사용을 통한 희소 어텐션 가속화

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

LLM 추론 엔지니어링과 AI 반도체 최적화 기법 심화 강의

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 기사

LLM 추론 엔지니어링과 AI 반도체 최적화 기법 심화 강의