기술 요약: PTStore (Prefix Tensor Store)

문제 정의

대규모 언어 모델(LLM) 추론 워크로드는 고성능 컴퓨팅(HPC) 데이터 센터에서 에너지 소비와 자원 요구량 측면에서 훈련을 능가하며 지배적인 부하가 되었습니다. LLM 추론은 두 단계로 구성됩니다: 프리필(prefill)(입력 프롬프트를 병렬로 처리)과 디코드(decode)(토큰을 순차적으로 생성). 어텐션 메커니즘의 중복 계산을 피하기 위해 시스템은 Key-Value (KV) 캐시를 사용하여 중간 결과들을 저장합니다.

vLLM과 같은 최첨단 런타임은 단일 GPU 또는 노드 내에서의 KV 캐싱은 최적화하지만, 규모가 커질 때 다음과 같은 중대한 한계에 직면합니다:

1. 교차 노드 재사용의 부재: 기존 시스템은 분산된 컴퓨팅 노드 간에 메모리를 집계하는 데 실패하는 경우가 많습니다. 한 노드의 요청이 다른 노드의 요청과 접두사(prefix)를 공유하더라도, 두 번째 노드는 캐시된 텐서를 재사용하는 대신 접두사를 다시 계산합니다.
2. 메타데이터 및 레이턴시 병목 현상: 분산 캐싱을 시도하는 접근 방식(예: LMCache, EvoStore)은 원격 메모리 액세스로 인한 높은 I/O 오버헤드나 복잡한 메타데이터 동기화(예: 단일 노드를 넘어 확장되는 Radix-Attention)로 인해 어려움을 겪는 경우가 많습니다.
3. 메모리 제약: 개별 GPU 메모리는 거대한 컨텍스트 윈도우를 수용하기에 불충분하며, 호스트 메모리나 SSD로 오프로딩하는 것은 캐싱의 이점을 상쇄할 정도의 레이턴시를 유발합니다.

핵심 과제는 과도한 I/O나 메타데이터 오버헤드 없이, 여러 컴퓨팅 노드에 걸쳐 분산된 수많은 GPU 간에 확장 가능하고 저지연인 KV 캐시 접두사 재사용을 가능하게 하는 것입니다.

방법론: PTStore 아키텍처

PTStore(Prefix Tensor Store)는 인기 있는 KV 캐시 접두사를 분산하고 복제함으로써 이러한 한계를 해결하기 위해 설계된 분산형 복제 텐서 스토어입니다. 이 시스템은 각 컴퓨팅 노드가 로컬 호스트 메모리와 SSD를 집계하여 로컬 및 원격 GPU 클라이언트에 서비스를 제공하는 서버를 실행하는 클라이언트-서버 모델을 채택합니다.

주요 설계 원칙

1. 증분형 텐서 저장 (Trie 구조):

   • PTStore는 전체 KV 블록을 저장하는 대신, 새로운 객체와 이전에 저장된 객체들의 가장 긴 공통 접두사(LCP) 사이의 **증분 차이(incremental differences, 텐서)**를 저장합니다.
   • 이를 통해 접사를 텐서 수준의 입도(granularity)로 구현하여, 접두사가 시간이 지남에 따라 서로 다른 방향으로 중복 없이 성장할 수 있게 합니다(trie와 유사함).
   • 통합된 메타데이터: 비용이 많이 드는 분산 trie 탐색을 피하기 위해, PTStore는 평면적인(flat) 메타데이터 구조를 사용합니다. 각 객체의 메타데이터에는 고유한 텐서 ID 목록이 포함됩니다. 로드 작업은 이 ID들을 순회하며 로컬 복제 캐시에 존재하는지 확인하고, 누락된 경우 소유자(owner) 서버로부터 원격으로 가져옵니다.

2. 복제를 활용한 분산 계층형 캐싱:

   • 소유 캐시(Owned Cache): 특정 서버가 책임지는 증분 텐서를 저장합니다.
   • 복제 캐시(Replication Cache): 액세스 로컬리티를 개선하기 위해 서버에 "핫(hot)"한(인기 있는) 접두사 복사본을 로컬에 저장합니다.
   • 트레이드오프 관리: 시스템은 소유 캐시와 복제 캐시 사이의 구성 가능한 임계값을 관리합니다. 시스템은 검색 속도와 저장 용량 사이의 균형을 맞추기 위해, 느린 저장소로 플러싱해야 하는 소유 텐서를 제거하는 것보다 재요청이 가능한 복제 텐서를 우선적으로 폐기하도록 합니다.

3. 액세스 패턴 인식 기반 제거(Eviction):

   • PTStore는 LRU(Least Recently Used) 대신 접두사 구조에서 초기 텐서들이 더 자주 액세스된다는 점을 고려하여, GDSF에서 파생된 빈도 기반 제거 정책을 사용합니다.
   • 이는 크기 대 빈도의 트레이드오프를 고려하여, 작은 빈도의 텐서가 가져오기 비용이 큰 큰 텐서를 밀어내지 않도록 보장합니다.

4. RDMA 인식 통합:

   • 흩어짐을 최소 최소화하기 위해, LCP에 추가되는 증분들은 소유자 서버 상의 단일 연속 영역으로 통합됩니다.
   • 로드 작업은 단일 RPC를 통해 흩어진 세그먼트들을 병렬로 가져오기 위해 **벌크 RDMA(bulk RDMA)**를 사용하여, 데이터를 전송하기 전에 연속적인 영역으로 복사하는 오버헤드를 피합니다.

주요 기여

1. 설계 원칙: 증분 텐서 저장, 통합된 메타데이터, 접두사 복제를 통합하는 분산형 저장소에 대한 일련의 고수준 원칙을 제시합니다.
2. PTStore 프로토타입: C++ 저수준 API와 vLLM과 같은 LLM 런타임과 원활하게 통합할 수 있는 Python 인터페이스를 특징으로 하는 연구용 프로토타입을 구현합니다.
3. 성능 검증: 최첨단 베이스라인들과 비교하여 I/O 오버헤드 및 엔드 투 엔드 실행 시간을 크게 감소시킨 광범ening한 실험 결과를 보여줍니다.

실험 결과

저자들은 ALCF Polaris HPC 테스트베드(560개 노드, A100 GPU)에서 두 가지 추출형 QA 워크로드인 WikiQA(긴 컨텍스트)와 SQUAD(높은 질문 볼륨)를 사용하여 PTStore를 평가했습니다. 사용된 LLM은 Mistral-7B-instruct-V2입니다.

베이스라인

• vLLM Vanilla: 교차 요청 접두사 공유 기능이 없는 표준 vLLM.
• vLLM Prefix: 로컬 접두사 공유(노드 내)를 지원하는 vLLM.
• EvoStore: 증분 저장 및 RDMA를 사용하지만 로컬 접두사 복제가 없는 분산 텐서 스토어.
• PTStore: 분산 인지 및 로컬 복제를 갖춘 제안된 시스템.

주요 결과

• 약한 확장성 (8–32 GPUs): PTStore는 EvoStore 및 vLLM Prefix보다 성능이 크게 향상되었습니다. EvoStore는 원격 접두사를 가져올 때 높은 RDMA I/O 오버헤드로 고통받은 반면, PTStore의 로컬 복제는 이를 완화하여 첫 토큰 생성 시간(TTFT)에서 "분리된 우위(detached advantage)"를 보여주었습니다.
• 시퀀스 길이 확장성 (1k–8k tokens):
  • 짧은 시퀀스(1k)의 경우, vLLM의 로컬 캐싱이 경쟁력이 있었습니다.
  • 시퀀스 길이가 길어질수록 PTStore의 이점이 커졌습니다. 8k 토큰에서 PTStore는 vLLM의 접두사 캐싱보다 거의 2배 빨랐으며, EvoStore보다 20% 더 빨랐습니다.
  • 재계산 또는 원격 I/O 비용이 로컬 전용 캐싱의 이점을 상회하기 때문에, 컨텍스트가 길어질수록 성능 격차가 벌어졌습니다.
• 효율성 이득: 긴 패시지 Q&A 데이터셋에서 PTStore는 노드 간 메모리를 집계하지 않고 KV 캐시를 재생성해야 하는 베이스라인들보다 5~6배 더 효율적으로 추론을 실행했습니다.

의의 및 주장

본 논문은 PTStore가 현재의 분산 노드 간 효율적인 KV 캐시 접두사 재사용 불가능성이라는 중요한 격차를 해결한다고 주장합니다. 증분 저장을 통한 중복 최소화, 빠른 쿼리를 위한 통합된 메타데이터, 그리고 로컬리티 최적화를 위한 복제 전략을 결합함으로써 PTStore는 다음을 가능하게 합니다:

• 클러스터 전체의 메모리를 집계함으로써 유효 KV 캐시 크기를 수 차례 확장.
• 재계산 비용이 큰 긴 컨텍스트 워크로드에서 TTFT의 대폭적인 감소.
• 이전의 분산 접근 방식들을 괴롭히던 통신 병목 현상 및 메타데이터 동기화 문제를 피하는 확장성.

저자들은 PTStore를 확장 가능한 AI 추론을 향한 기초적인 단계로 규정하며, 향후 연구는 동적 메모리 밸런싱, 머신 러닝 기반 제거 정책, 그리고 실제 대화 및 코드 완성 트레이스에 대한 LMCache 및 Mooncake와의 광범위한 벤치마킹에 집중할 것이라고 밝혔습니다.