QKVShare: Quantized KV-Cache Handoff for Multi-Agent On-Device LLMs

기술 요약: QKVShare

문제 제기
에지 장치에 배포된 다중 에이전트 대규모 언어 모델 (LLM) 시스템은 중요한 병목 현상에 직면해 있습니다: 한 에이전트가 컨텍스트를 하위 에이전트에게 전달할 때, 수신 에이전트는 일반적으로 전체 재-프리필 (re-prefill) 을 통해 잠재 상태를 처음부터 재구성해야 합니다. 이 과정은 지연 시간이 매우 깁니다. 반면, 정밀도 전체 (FP16) 키 - 값 (KV) 캐시를 공유하는 것은 제한된 RAM 을 가진 장치에서는 메모리 용량을 초과할 정도로 비현실적입니다. 왜냐하면 캐시 크기는 컨텍스트 길이에 비례하여 선형적으로 증가하기 때문입니다. 기존 솔루션들은 지연 시간 (재-프리필) 이나 메모리 (FP16 공유) 중 하나를 낭비하는 반면, 현재 양자화 방법들은 단일 에이전트 추론에 초점을 맞추고 에이전트 경계를 넘어선 양자화 캐시의 상호작용을 다루지 않습니다.

방법론
QKVShare 는 세 가지 구성 요소를 중심으로 에이전트 간 양자화 KV 캐시 인도를 위한 프레임워크를 제안합니다:

1. 토폴로지 인식 비트 할당: 시스템은 "Don't Waste Bits" 컨트롤러를 교차 에이전트 기능을 포함하도록 확장합니다. 이는 다음 점수의 조합에 기반하여 토큰별 비트 너비 ({2, 4, 8, FP16} 중 선택) 를 할당합니다:
   • 로컬 중요도 ($I(t)$): 빈도, 품질 점수, 어텐션 분산, 엔트로피에서 유도됩니다.
   • 전송 가능성 ($T(t)$): 하위 에이전트의 수요를 추정하는 새로운 신호로, 공유 시스템 프롬프트와 같은 세그먼트 수준 사전 지식과 이전에 공유된 토큰과의 어텐션 앵커 중첩을 통해 계산됩니다.
2. CacheCard 표현: 에이전트 간 전송을 위해 양자화된 KV 그룹, 토큰 수준 비트 너비 할당, 메타데이터 (시퀀스 길이, 모델 ID, 위치 오프셋) 를 패키징하는 소형 아티팩트입니다.
3. 주입 경로: HuggingFace Transformers 와의 호환성을 위해, 현재 프로토타입은 수신 에이전트의 순전파 (forward pass) 전에 양자화된 CacheCard 에서 FP16 past_key_values 를 재구성합니다. 이는 완전히 융합된 양자화 어텐션 커널의 개발을 향후 작업으로 미루면서 지연 시간 측정을 가능하게 합니다.

주요 기여

• 교차 에이전트 토큰 스코어링: 로컬 충실도와 토폴로지의 다음 에이전트의 특정 요구 사항 모두에 기반하여 비트를 할당하는 확장된 컨트롤러.
• 양자화 인도 프로토콜: CacheCard 형식의 정의와 양자화된 상태 전송을 가능하게 하는 수신 측 주입 메커니즘.
• 측정된 프로토타입 분석: 멀티홉 추론 정확도와 인도 시간 (Time-To-First-Token, TTFT) 을 포함하는 Llama-3.1-8B-Instruct 에 대한 종단 간 평가 및 지연 시간의 단계별 분해.

실험 결과
이 평가는 12GB NVIDIA RTX 5070 Ti 노트북 GPU 에서 25 홉의 순차적 에이전트 체인과 1K8K 토큰의 합성 컨텍스트 벤치마크를 사용하여 150 개의 GSM8K 문제를 활용했습니다.

• 정확도 (멀티홉):
  • 균일 양자화: 홉 깊이가 증가함에 따라 불안정성을 보이며, 특히 4 비트 정밀도에서 두드러집니다.
  • 적응형 양자화: 더 깊은 홉과 더 높은 예산 설정에서 균일 양자화보다 명확한 개선을 보여줍니다 (예: 약 8 비트에서 적응형 토폴로지 인식 할당은 5 홉에서 82.67% 정확도를 달성한 반면, 균일 Q8 은 71.33% 였습니다).
  • 토폴로지 대 로컬: 약 4 비트에서 토폴로지 인식 컨트롤러는 로컬 전용 적응형 컨트롤러보다 미미한 순이익 (+2 승리) 을 보였으나 일관되게 우월하지는 않았습니다. 논문은 공격적인 비트 예산에서 토폴로지 인식의 이점을 확정된 승리가 아닌 열린 가설로 취급합니다.
• 지연 시간 (TTFT):
  • QKVShare 는 모든 테스트된 컨텍스트에서 전체 재-프리필에 비해 TTFT 를 크게 줄였습니다. 8K 컨텍스트의 경우 QKVShare 는 397.1ms 를 기록한 반면 재-프리필은 1029.7ms 였습니다.
  • 병목 분석: 단계별 타이밍은 현재 지연 시간이 CacheCard 생성이나 주입 자체가 아닌 주입 후 생성 (수신 순전파) 에 의해 지배됨을 보여줍니다.

의의 및 주장
이 논문은 QKVShare 를 에지 장치 다중 에이전트 추론을 위한 유망한 시스템 방향성으로 위치시키며, 양자화 KV 인도가 재-프리필의 높은 지연 시간을 피하면서도 메모리 제약을 관리한다고 주장합니다.

그러나 저자들은 현재 증거에 대해 겸손한 입장을 유지합니다:

• 프로토타입은 재-프리필을 피하는 것이 TTFT 를 실질적으로 개선함을 명확히 합니다.
• 적응형 혼합 정밀도 할당은 반복적인 인도 하에서 유망하지만, 토폴로지 인식 기능의 특정 가치가 로컬 전용 적응보다 4 비트 예산에서 특히 결론적으로 분리되지는 않았습니다.
• 현재 결과는 혼합 런타임 스택 (PyTorch/HuggingFace 대 llama.cpp) 과 주입을 위한 역양자화 단계에 대한 의존성으로 인해 방향성 있는 결과입니다.
• 논문은 토폴로지 인식 주장을 완전히 검증하기 위해 더 강력한 컨트롤러 제거 실험과 (예: 네이티브 융합 Q4->Q4 커널과 같은) "사과와 사과" 런타임 비교가 필요함을 명시적으로 지적합니다.