HeteroFold, 서로 다른 모델 간 KV Cache 공유에 도전

멀티 에이전트 LLM 시스템은 계획, 검색, 실행, 검증 같은 역할을 서로 다른 모델에 배분할 수 있습니다. 모델 패밀리마다 강점이 다르기 때문에 이기종 조합은 유용하지만, 통신 방식이 텍스트에 머물면 비용이 커집니다. 송신 에이전트가 이미 처리한 긴 컨텍스트를 수신 에이전트가 다시 읽고 Prefill해야 하기 때문입니다.

KV Cache 재사용은 이 중복 계산을 줄이는 방법입니다. 모델은 컨텍스트를 처리하면서 Key와 Value 중간 상태를 저장하고, 이후 토큰 생성에서 이를 재사용합니다. 그러나 KV Cache는 모델 독립적인 데이터가 아닙니다. 토크나이저, 레이어 수, 어텐션 구조, KV 차원, 내부 표현 공간이 모델 패밀리마다 다르므로 한 모델의 캐시를 다른 모델에 그대로 넣을 수 없습니다.

논문은 이러한 문제를 해결하기 위해 HeteroFold를 제안합니다. 이 방법은 먼저 서로 다른 토크나이저에서 발생하는 토큰 경계를 맞추고, 다른 아키텍처의 레이어 대응 관계를 정렬합니다. 이어 송신 모델의 K/V 상태를 수신 모델의 표현 공간으로 매핑하고, 수신 모델의 어텐션 패턴과 출력을 유지하도록 보정합니다. 송신 모델과 수신 모델은 모두 동결되며, 두 모델의 내부 상태를 연결하는 매핑을 학습하는 방식입니다.

실험은 Llama, Qwen, Ministral 사이의 6개 전송 방향을 대상으로 진행됐습니다. 긴 컨텍스트와 짧은 컨텍스트, 멀티 에이전트 설정이 함께 평가됐습니다. 저자 보고에 따르면 HeteroFold는 4개 긴 컨텍스트 벤치마크 모두에서 가장 좋은 캐시 전송 성능을 보였고, 짧은 컨텍스트의 대부분 설정에서도 비교 방법보다 우수했습니다. 32K 컨텍스트에서 Llama 3.1 8B에서 Ministral 3 14B로 전송하는 경우, 수신 측 기본 Prefill보다 약 10.7배 빨랐으며 Dense Latent와 KV Ridge보다 1.18~1.47배 빠른 것으로 보고됐습니다. 멀티 에이전트 벤치마크에서는 텍스트 통신과 비슷한 결과도 제시됐습니다.

HeteroFold의 의미는 KV Cache를 특정 모델 내부에 갇힌 임시 상태가 아니라 모델 사이에서 교환할 수 있는 계산 자산으로 바라본다는 데 있습니다. 여러 에이전트가 긴 문서, 작업 이력, 환경 상태를 반복적으로 공유하는 상황에서는 중복 계산과 지연을 줄이고 모델 조합의 선택 폭을 넓힐 가능성이 있습니다.

다만 실제 서비스에 적용하려면 캐시 변환과 보정에 드는 비용, 캐시 전송량, 평가에 포함되지 않은 작업 분포에서의 안정성을 따져야 합니다. 제공된 자료만으로 모든 모델 규모와 운영 환경에서 동일한 이득이 발생한다고 단정할 수는 없습니다. HeteroFold는 서로 다른 모델 경계를 넘어 KV Cache를 직접 공유할 수 있다는 가능성을 보여주는 연구 방향으로 이해하는 것이 적절합니다.

출처: Hugging Face Daily Papers