멀티 LoRA 에이전트를 위한 효율적인 KV 캐시 공유 LRAgent

초록

멀티‑LoRA 기반 다중 LLM 에이전트는 동일한 백본을 공유하지만 각 에이전트가 별도의 KV 캐시를 유지해 메모리·연산 비용이 크게 늘어난다. 저자는 에이전트 간 캐시 차이가 LoRA 어댑터 출력에 의해 주도된다는 사실을 발견하고, 캐시를 ‘베이스 캐시’와 ‘LR 캐시’(저‑랭크 어댑터 캐시)로 분리한다. BaseShared는 베이스 캐시만 공유하고 LR 캐시는 개별 저장하며, BaseLRShared는 다운‑프로젝션을 공유해 LR 캐시까지 공유한다. 또한 Flash‑LoRA‑Attention 커널을 설계해 저‑랭크 캐시를 풀 차원으로 복원하는 비용을 최소화한다. 실험 결과 메모리 사용량과 지연시간이 크게 감소하면서도 정확도는 비공유 기준에 근접한다.

상세 분석

본 논문은 멀티‑LoRA 구조를 활용한 다중 LLM 에이전트 시스템에서 발생하는 KV 캐시 중복 문제를 근본적으로 해결하고자 한다. 전통적인 멀티‑LoRA 설정에서는 각 에이전트가 동일한 프리트레인 백본을 공유하지만, 역할별 LoRA 어댑터가 다르기 때문에 추론 시마다 별도의 KV 캐시(키·밸류)를 생성한다. 이는 특히 툴‑증강(long‑context) 트래젝터리를 다루는 에이전트에서 메모리 사용량을 급격히 증가시키고, 동일한 컨텍스트를 여러 번 재계산하게 만들어 지연시간을 악화시킨다.

저자는 먼저 동일한 입력 시퀀스에 대해 에이전트 간 베이스 캐시(Y_base = X·W₀)의 코사인 유사도가 0.95 이상으로 매우 높고, 어댑터에 의해 추가되는 ΔY = (X·A)·B는 크기가 작지만 서로 거의 직교한다는 실험적 관찰을 제시한다(표 1, 그림 1). 이는 전체 캐시 차이가 어댑터 출력에 의해 주도된다는 의미이며, 베이스 캐시만 공유하고 어댑터 부분만 별도로 관리하면 정확도 손실을 최소화하면서 메모리 절감이 가능함을 시사한다.

이에 기반해 LRAgent은 KV 캐시를 두 부분으로 분해한다. 첫 번째는 프리트레인 가중치 W₀만을 사용해 계산되는 ‘베이스 캐시’이며, 이는 모든 에이전트가 동일하게 재사용한다. 두 번째는 LoRA 어댑터의 다운‑프로젝션 결과 X·A(차원 l×r)를 저장하는 ‘LR 캐시’이다. LR 캐시는 저‑랭크 형태이므로 메모리 비용이 크게 감소한다. 런타임 시에는 LR 캐시와 각 에이전트의 업‑프로젝션 B를 곱해 ΔY를 복원한다.

두 가지 공유 전략이 제안된다. BaseShared는 베이스 캐시만 공유하고, 각 에이전트는 자체 LR 캐시를 유지한다. 이는 메모리 사용량을 (1 + 1/N)배에서 1 + 1/N배로 감소시킨다(여기서 N은 에이전트 수). BaseLRShared는 최근 연구에서 제안된 ‘Shared‑A’ 구조를 활용해 모든 에이전트가 동일한 다운‑프로젝션 A를 사용하도록 설계한다. 이 경우 LR 캐시 자체도 공유 가능해져 메모리와 연산 모두 추가 절감된다. 특히, 이미 처리된 컨텍스트에 대해 다른 에이전트가 재사용할 때, 베이스 캐시와 LR 캐시 모두를 재활용함으로써 KV 계산을 완전히 건너뛸 수 있다.

LR 캐시를 풀 차원으로 복원하는 과정에서 발생할 수 있는 연산 오버헤드를 최소화하기 위해, 저자는 Flash‑LoRA‑Attention 커널을 설계했다. 이 커널은 기존 FlashAttention의 블록‑스파스 매트릭스 연산 흐름을 재배열해, LR 캐시를 직접 전치·곱셈 없이 바로 어텐션 연산에 삽입한다. 결과적으로 메모리 이동 비용과 불필요한 매트릭스 확장을 회피하면서도 GPU 메모리 대역폭을 최적화한다.

실험에서는 LLaMA‑3.1‑8B‑Instruct와 MiniSTR‑8B‑Instruct를 기반으로 3‑role(플래닝, 액션, 리플렉션) 에이전트를 구성하고, HotpotQA, GSM‑8K 등 장기 컨텍스트가 요구되는 벤치마크에서 평가했다. BaseShared와 BaseLRShared 모두 비공유 설정 대비 메모리 사용량을 각각 38%·55% 절감했으며, 토큰당 처리량(throughput)은 0.92×~0.97× 수준으로 거의 동일했다. 정확도 측면에서는 전체 캐시와 비교해 0.3% 이하의 차이만 보였으며, 특히 Shared‑A 기반 BaseLRShared는 어댑터 간 차이를 최소화해 정확도 손실을 거의 없앴다.

결론적으로 LRAgent은 멀티‑LoRA 에이전트 환경에서 KV 캐시 중복을 구조적으로 제거함으로써, 메모리·연산 효율성을 크게 향상시키면서도 성능 저하를 최소화한다. 코드와 모델은 공개되어 향후 다양한 멀티‑에이전트 시스템에 적용 가능할 것으로 기대된다.