AKV, 에이전트의 기록을 바꾸지 않고 KV 캐시를 관리하다

검색, 도구 호출, 관찰, 행동을 반복하는 LLM 에이전트의 실행 궤적은 빠르게 길어집니다. 궤적이 모델의 최대 컨텍스트 길이를 넘으면 어떤 기록을 남길지 결정해야 합니다. 하지만 문제는 입력 토큰 수에만 있지 않습니다. 지금까지 계산한 궤적에 대응하는 KV 캐시를 저장하고, 메모리 계층 사이에서 읽고 쓰며, 다음 요청에서 복원하는 비용도 함께 커집니다. 장시간 동작하는 에이전트에서는 KV 캐시의 저장과 I/O가 주요 서빙 병목이 될 수 있습니다.

기록을 줄이면 재계산이 발생한다

기존의 컨텍스트 엔지니어링은 오래된 기록을 잘라내거나 요약문으로 바꿔 입력을 줄이는 방식이 일반적입니다. 그러나 Transformer의 후속 상태는 앞부분의 계산에 의존합니다. 프리픽스를 수정하면 남겨 둔 텍스트가 동일하더라도 그 이후에 만들어진 캐시가 무효화될 수 있습니다. 결국 보존한 문맥을 다시 프리필해야 합니다.

재계산을 줄이려고 기록을 많이 남기면 메모리와 I/O 비용이 커집니다. 반대로 너무 많이 줄이면 작업 정확도가 낮아질 수 있습니다. 논문은 이 문제를 텍스트 기록 자체를 수정하는 대신, 현재 어텐션에 참여하는 캐시를 관리하는 방식으로 풀고자 합니다.

drop과 repos의 역할

‘캐시 엔지니어링’에서는 원래의 에이전트 궤적을 그대로 유지하면서 후속 계산에 사용할 캐시만 조절합니다. 핵심 연산은 다음과 같습니다.

• drop: 선택한 캐시를 이후 어텐션에서 제외합니다. 기록에서 해당 텍스트를 삭제할 필요는 없습니다.
• repos: 남겨 둔 키의 위치 인코딩을 조정해 캐시 관리 이후에도 활용할 수 있도록 합니다.
• 정책 기반 관리: 컨텍스트 길이가 임계치에 도달한 뒤 한 번에 요약하는 대신, 도구 응답 같은 작업 경계에서 캐시를 점진적으로 정리합니다.

이는 전체 텍스트 궤적을 기록으로 보존하면서도, 모델이 현재 참조하는 작업 집합은 작게 유지하는 접근입니다. 에이전트의 장기 기록과 즉시 사용할 메모리를 분리할 수 있다는 점이 핵심입니다.

단순한 프리픽스 매칭의 한계

직접 캐시를 조작하면 새로운 서빙 문제가 생깁니다. 같은 토큰 시퀀스라도 과거에 어떤 drop이나 repos 연산을 거쳤는지에 따라 최종 캐시 상태가 달라질 수 있습니다. 텍스트 프리픽스만 비교하는 시스템은 호환되지 않는 캐시를 같은 것으로 판단해 잘못 재사용할 가능성이 있습니다.

연구팀은 이를 해결하기 위해 AKV 서빙 엔진을 함께 설계했습니다. 외부에는 기존 서빙 API와 호환되는 무상태 인터페이스를 제공하지만, 내부에서는 상태를 가진 캐시 재사용, 축출, 복원을 관리합니다. 호출 방식은 크게 바꾸지 않으면서도 캐시 연산 이력을 고려해 안전한 재사용을 수행하려는 구조입니다.

실험 결과와 시사점

실험에서 도구 응답을 기준으로 캐시를 순차적으로 버리는 rolling-drop 정책은 임계치에 도달하면 요약하는 방식과 비교해 평균 벤치마크 점수를 1.7점 높이고 평균 비용을 27.7% 낮췄습니다. AKV를 적용했을 때 출력 처리량은 최대 1.2배, P95 요청 지연시간은 최대 1.5배 개선됐습니다. 또한 재귀적 자기 개선으로 캐시 정책을 탐색하면 전문가가 설계한 정책보다 비용을 추가로 7.1% 줄일 수 있다는 결과도 제시됐습니다.

이 연구가 주는 큰 시사점은 에이전트의 메모리가 전체 텍스트 기록과 같을 필요가 없다는 것입니다. KV 캐시를 실행 중 관리되는 자원으로 보고, 도구 사용 단계와 작업 진행 상황에 맞춰 활성 범위를 조절할 수 있습니다. 다만 캐시를 과도하게 제거하면 나중에 필요한 정보가 사라질 수 있고, repos 같은 위치 조정도 모델과 작업에 따라 다르게 작동할 수 있습니다. 따라서 AKV는 모든 상황에 적용되는 단일 해법이라기보다, 장기 실행 에이전트를 위한 새로운 서빙 설계 방향으로 보는 것이 적절합니다.

출처: Hugging Face Daily Papers