세미샤어KV 의미 유사 프롬프트를 위한 토큰 수준 LSH 기반 KV캐시 공유

초록

SemShareKV는 의미적으로 유사하지만 어휘가 다른 프롬프트 간에 KV 캐시를 재사용하기 위해 토큰 임베딩에 로컬리티-센시티브 해싱(LSH)과 로테리 포지션 임베딩(RoPE)을 결합한 프레임워크이다. 토큰 수준의 퍼지 매칭을 통해 기존 캐시를 재배열하고, 얕은 층에서는 재계산을, 깊은 층에서는 캐시 재사용을 우선시함으로써 5k 토큰 입력 시 최대 6.25배 가속과 42% 메모리 절감 효과를 달성한다. 품질 저하는 거의 없으며, 요약 데이터셋을 중심으로 실험하였다.

상세 분석

SemShareKV는 대규모 언어 모델(LLM) 추론 시 KV‑Cache가 차지하는 메모리와 연산 비용을 줄이기 위해 “의미 기반 캐시 공유”라는 새로운 패러다임을 제시한다. 기존 연구는 동일 프롬프트 내부에서 토큰 압축이나 접두사 재사용에 초점을 맞추었지만, 실제 서비스에서는 서로 다른 어휘를 사용하면서도 같은 주제·내용을 다루는 프롬프트가 빈번히 발생한다. 이러한 상황을 해결하기 위해 저자는 세 가지 핵심 관찰을 기반으로 설계 결정을 내렸다.

1. 고편차 토큰(High‑Deviation, HD) 일관성 – 실험적으로 얕은 층부터 깊은 층까지 HD 토큰 위치가 높은 상관관계를 보였으며, 이는 특정 토큰이 여러 층에서 지속적으로 큰 변화를 일으킨다는 것을 의미한다. 따라서 재계산이 필요한 토큰을 효율적으로 식별할 수 있다.

2. 깊은 층의 집중적 어텐션 – Attention Recovery(AR) 지표를 통해 깊은 층일수록 전체 어텐션 스코어 중 소수의 토큰에 집중한다는 사실을 확인했다. 이는 깊은 층에서는 전체 토큰을 모두 재계산할 필요가 없으며, 중요한 토큰만 유지하면 된다는 설계 근거가 된다.

3. 깊은 층의 중복 정보 – 다양한 토큰 보존 전략(Uniform, Exponential‑Increase, Exponential‑Decay)을 퍼플렉시티로 평가했을 때, Exponential‑Decay(깊은 층에서 보존 비율을 낮춤) 방식이 가장 낮은 퍼플렉시티를 기록했다. 이는 깊은 층이 상대적으로 중복된 정보를 많이 포함하고 있음을 시사한다.

위 관찰을 바탕으로 SemShareKV는 두 단계의 전략을 채택한다. 재계산 전략에서는 얕은 층에서 대부분의 토큰을 새로 계산하고, 깊은 층에서는 HD 토큰만 선택적으로 재계산한다. 보존 전략에서는 얕은 층에 더 많은 KV‑쌍을 유지하고, 깊은 층에서는 불필요한 KV를 삭제해 메모리를 절감한다.

핵심 기술은 토큰‑레벨 LSH 매칭과 RoPE 기반 위치 인코딩이다. 프롬프트를 입력받으면 먼저 각 토큰의 컨텍스트 임베딩(E‑Cache)을 CPU 메모리에 저장한다. 새로운 타깃 프롬프트가 들어오면 저장된 E‑Cache와 LSH 거리 기반 유사도 계산을 통해 가장 유사한 레퍼런스 프롬프트를 선택한다. 선택된 레퍼런스의 KV‑Cache를 GPU에 로드한 뒤, RoPE를 적용해 위치 정보를 임베딩에 주입한다. 이후 LSH를 이용해 타깃 토큰과 레퍼런스 토큰을 매핑하고, 매핑 결과에 따라 레퍼런스 KV‑Cache를 토큰 순서에 맞게 재배열한다. 첫 번째 트랜스포머 층에서는 전체 토큰을 완전 재계산하고, 재계산 결과와 재배열된 KV‑값을 L2 노름으로 비교해 HD 토큰을 식별한다. 이후 층에서는 식별된 HD 토큰만 재계산하고, 어텐션 점수가 낮은 토큰은 동적으로 캐시에서 제거한다.

실험은 MultiNews, XSum, CNN/DailyMail 등 다양한 요약 데이터셋을 사용했으며, 모델은 Mistral‑7B, LLaMA‑3.1‑8B, MPT‑7B 세 종류를 대상으로 했다. 5k 토큰 입력 기준으로 평균 6.25배의 추론 속도 향상과 42%의 GPU 메모리 절감 효과를 보였으며, ROUGE‑L 점수와 인간 평가에서 품질 저하는 0.2% 이하로 미미했다. Ablation 실험에서는 RoPE 없이 LSH만 적용했을 때 매핑 정확도가 15% 감소하고, Exponential‑Decay 보존 전략을 사용하지 않을 경우 메모리 절감 효과가 30% 이하로 떨어지는 것을 확인했다.

한계점으로는 레퍼런스 프롬프트 저장 비용과 LSH 매칭 비용이 프롬프트 길이가 매우 길어질 경우 병목이 될 수 있다는 점, 그리고 RoPE가 적용되지 않은 모델(ALiBi 기반)에서는 위치 정합성이 떨어져 추가적인 보정이 필요하다는 점을 언급한다. 향후 연구에서는 멀티‑GPU 환경에서 분산 캐시 관리, 동적 LSH 파라미터 튜닝, 그리고 비지도 방식으로 레퍼런스 프롬프트를 자동 클러스터링하는 방법을 탐색할 예정이다.

전반적으로 SemShareKV는 “의미 기반 캐시 재사용”이라는 새로운 효율성 차원을 열어, 장문 입력을 다루는 실제 서비스 환경에서 LLM 추론 비용을 크게 낮출 수 있음을 실증하였다.