Sparse‑K가 llama.cpp에서 수백만 개의 Attention 연산을 줄이는 방법
Source: Dev.to
어텐션이란 무엇일까?
Sparse‑K에 들어가기 전에, 어텐션 메커니즘이 실제로 무엇을 하는지, 그리고 왜 계산 비용이 큰지 이해하는 것이 중요합니다.
언어 모델이 다음과 같은 문장을 받았을 때:
“The cat sat on the gray carpet,”
단순히 단어들을 읽는 것이 아니라 의미와 관계, 그리고 각 단어가 다른 단어와 어떻게 연결되는지를 파악해야 합니다. 이를 위해 각 단어는 문장 내 모든 다른 단어와 얼마나 강하게 연관되는지를 평가합니다. 이렇게 하면 모든 토큰이 다른 모든 토큰에 대해 점수를 계산하는 조밀한 연결망이 만들어집니다.
수식:
(N) 토큰 → (N \times N) 어텐션 점수 계산.
이와 같은 2차 성장(quadratic growth)이 어텐션을 특히 긴 시퀀스에서 비용이 많이 들게 만드는 원인입니다.
왜 이것이 성능 문제인가?
계산량이 2차적으로 증가하기 때문입니다 — (O(N^{2})).
2048 토큰 길이의 시퀀스는 단일 어텐션 레이어에서 400만 개가 넘는 점수 계산을 필요로 합니다.
이 매트릭스를 자세히 살펴보면, 대부분의 연결은 최종 예측에 의미 있게 기여하지 않으며, 많은 연결이 약하거나 무관합니다. 시각적으로 보면, 모든 토큰이 모든 다른 토큰에 어텐션을 두는 조밀한 웹 대신, 실제로는 몇 개의 의미 있는 연결만 필요합니다.
Sparse‑K는 실제로 무엇을 하는가?
Sparse‑K는 간단하지만 강력한 아이디어를 도입합니다: 토큰 간 모든 쌍 관계를 평가하는 대신, 각 토큰에 대해 상위 K(top‑K) 가장 의미 있는 연결만 유지하고 나머지는 무시합니다.
- 어텐션 메커니즘의 구조는 그대로 유지됩니다; 처리되는 데이터 양만 줄어듭니다.
- 이로 인해 계산량이 크게 감소하고 메모리 압력이 낮아지며 전체 추론 파이프라인이 빨라집니다.
어텐션 점수에서 Sparse 마스크로
표준 어텐션에서는 유사도 점수((Q \times K))를 계산한 뒤 전체 매트릭스가 생성됩니다. Sparse‑K는 추가 단계를 삽입합니다:
- Top‑K 선택 – 각 토큰에 대해 가장 높은 K개의 점수만 남깁니다.
- 마스킹 – 나머지 값은 마스크를 적용해 “제거”합니다.
- 음의 무한대 – 마스크된 값은 매우 큰 음수((-\infty))로 설정되어 Softmax에 영향을 주지 않게 합니다.
그 결과 Softmax와 이후 모든 연산은 실제로 중요한 연결에 대해서만 수행됩니다.
이미지는 각 토큰과 모든 다른 토큰 사이의 어텐션 점수 계산을 보여줍니다. 각 토큰에 대해 상위 K개의 점수가 강조됩니다.
이 그림은 Top‑K 선택 후 최종 Sparse‑K 어텐션 마스크를 나타냅니다.
왜 마스크가 적절한 해결책인가?
마스크는 이미 어텐션 메커니즘의 기본 요소입니다. 모든 트랜스포머 모델은 다음과 같은 여러 종류의 마스크에 의존합니다:
- 인과 마스크(Causal masks) – 모델이 미래 토큰을 참조하지 못하도록 방지합니다.
- 시퀀스 구분 마스크(Sequence separation masks) – 서로 다른 입력 구간을 구분합니다.
Sparse‑K는 새로운 메커니즘을 도입하는 것이 아니라 기존 마스킹 시스템에 자연스럽게 통합됩니다:
- Top‑K 선택을 나타내는 추가 마스크를 만든다.
- 이를 기존 기본 마스크와 결합한다.
- 이미 존재하는 어텐션 루틴 내에서 사용한다.
이 접근 방식은:
- 현재 파이프라인에 매끄럽게 맞춰져 호환성을 깨뜨리지 않는다.
- 계산 그래프 구조를 그대로 유지한다.
- 최적화된 어텐션 구현(예: Flash Attention)을 재사용할 수 있게 한다. 이러한 구현은 이미 마스크를 인터페이스의 일부로 받아들이기 때문이다.
따라서 마스크를 사용함으로써 Sparse‑K는 기존 최적화를 그대로 활용하면서 효율성, 호환성, 유지보수성을 보장할 수 있습니다.
llama.cpp에의 통합
Sparse‑K는 어텐션 블록에 직접 삽입되며, 계산 흐름의 자연스러운 지점에 배치됩니다:
- 유사도 점수 (Q \times K)를 계산한다.
- Sparse‑K 마스크(Top‑K 선택)를 만든다.
- 남은 어텐션 단계 전에 결합된 마스크를 적용한다.
모델의 다른 부분—임베딩 레이어, MLP 연산, 멀티‑헤드 구성—은 전혀 변경되지 않습니다. 이렇게 함으로써 Sparse‑K는 모델에서 가장 계산 비용이 큰 부분을 목표로 하면서도 나머지 아키텍처는 그대로 유지합니다.