Pie: 모델 옆에서 프로그램이 도구 호출과 KV 캐시를 직접 다루는 추론 시스템

배경과 문제

LLM이 생성 도중 도구를 부르거나 여러 갈래로 추론을 펼치면, 모델의 계산 속도뿐 아니라 계산 사이의 연결 비용도 중요해진다. 논문이 비교한 방식에서는 애플리케이션 로직이 Python 클라이언트에 있고, vLLM·SGLang 서버는 추론을 맡았다. 도구를 사용할 때마다 제어가 클라이언트와 서버를 왕복하고, 문맥이 바뀌면 앞선 내용을 다시 처리하는 프리필이 필요할 수 있다. 기존 시스템이 항상 모든 KV 캐시를 버린다는 뜻은 아니다. 프리픽스 캐싱도 재계산을 줄이지만, 애플리케이션이 원하는 캐시 보존 전략을 직접 지정하기에는 제약이 있다는 문제 제기다.

Pie는 엔진이 자원과 기본 연산을 관리하고, 무엇을 언제 계산할지는 사용자 프로그램이 정하도록 역할을 나눈다. 고정된 생성 루프 대신 인퍼렛이 세분화된 API를 조합해 흐름을 만든다. 논문이 제시한 42개 API는 2025년 설계의 수치다. 현재 버전의 API 개수나 이름이 그대로라는 뜻은 아니며, 달라진 인터페이스는 뒤의 예시에서 따로 살펴본다.

기존 접근의 한계

논문은 세 가지 요구사항을 정리한다. 첫째는 애플리케이션별 KV 캐시 제어다. Tree-of-Thought나 Graph-of-Thought처럼 실행 중 추론 가지가 늘고 합쳐지는 작업에서는, 어떤 문맥을 남기고 복제하거나 버릴지 프로그램이 알아야 한다. 시스템 전체에 적용되는 LRU 축출이나 공통 앞부분을 재사용하는 프리픽스 캐싱만으로는 이 작업 구조를 충분히 반영하기 어렵다는 주장이다.

둘째는 생성 과정의 맞춤 제어다. 여러 토큰을 제안·검증하는 추측 디코딩, 문법 상태를 유지하는 제약 디코딩, 샘플링 전에 분포를 바꾸는 워터마킹을 요청별로 조합하고 싶어도 고정 루프가 제약이 될 수 있다. 셋째는 계산과 I/O의 통합이다. 도구 호출과 모델 추론을 같은 실행 흐름에 두어 외부 왕복과 불필요한 문맥 재처리를 줄이려 한다. 논문은 SGLang 등의 상위 프로그래밍 기능과 달리, KV 페이지·임베딩·샘플링 순서까지 제어하는 수준을 목표로 한다.

연구방법을 단계별로

Pie는 순전파 한 번을 세 단계로 나눈다. 입력을 임베딩으로 바꾸는 embed, 임베딩과 KV 캐시로 출력 임베딩을 계산하는 forward, 출력 임베딩에서 다음 토큰 분포를 꺼내는 sample이다. 각 단계를 API로 노출하고, 인퍼렛이 이 API를 원하는 순서로 부른다. 인퍼렛이 다루는 자원은 두 가지다. Embed는 토큰 임베딩 버퍼이고, KvPage는 PagedAttention 방식의 KV 캐시 조각으로 8~32개 토큰을 담는다. 인퍼렛은 alloc_kvpage, dealloc_kvpage로 이 자원을 명시적으로 할당하고 해제한다. 각 인퍼렛은 자기만의 가상 자원 주소 공간을 가지며, export_kvpage와 import_kvpage로 다른 인퍼렛과 KV 캐시를 공유할 수 있다.

논문의 API 42개 중 18개는 순전파와 자원 관리, 24개는 메시지 교환과 HTTP 요청 같은 런타임 작업을 맡는다. get_next_dist는 전체 어휘를 매번 옮기는 대신 상위 K개 토큰 분포를 돌려주며 기본 K는 256이다. 관련 기능은 Rust의 트레이트처럼 묶는다. 예를 들어 모델이 InputText나 InputImage를 지원하는지 인퍼렛이 실행 중 확인해 처리 방식을 바꿀 수 있다. 이처럼 모델별 기능 차이를 드러내되, 낮은 수준의 GPU 커널 구현은 추론 엔진에 맡긴다.

시스템은 세 레이어로 구성된다. 애플리케이션 레이어는 Wasmtime 기반 WebAssembly 런타임에서 인퍼렛을 실행하고 서로 격리한다. 컨트롤 레이어는 GPU가 필요 없는 API와 자원 주소 매핑을 관리하고, GPU로 갈 호출을 배치로 묶는다. 추론 레이어는 묶인 호출을 GPU에서 실행한다. 논문의 평가 구현은 PyTorch·FlashInfer를 사용하고 ZeroMQ로 레이어 사이를 연결했다. 이는 현재 README의 단일 바이너리 구현과 구분해야 한다.

배치를 묶는 단위는 커맨드 큐다. 큐는 API 호출 사이의 의존 관계와 우선순위를 알려 준다. 수평 배치는 서로 다른 큐의 같은 종류 호출을, 수직 배치는 같은 큐에서 연속된 같은 종류 호출을 묶는다. 단, 같은 KvPage에 동시에 쓰는 것과 같은 충돌이 없어야 한다. 작업 보존 스케줄링은 GPU가 바쁜 동안 호출을 모으고, GPU가 준비되면 쌓인 호출을 보내 계산 자원이 놀지 않게 한다.

결과와 숫자의 의미

논문 평가는 GCP g2-standard-32의 NVIDIA L4 GPU 24GB 한 장에서 Llama 3의 1B·3B·8B 모델을 사용했다. 가중치·KV 캐시·활성값은 BF16이며, vLLM v0.6.0·SGLang v0.4.4와 비교했다. 세 시스템 모두 FlashInfer 백엔드를 사용해 커널 차이를 줄였고, 지연 시간은 캠퍼스 네트워크의 원격 Python 클라이언트에서 측정했다. 다음 결과는 이 과거 버전과 단일 GPU 조건에 해당하며, 2026년 최신 엔진이나 대형 분산 서비스의 우열을 확정하지 않는다.

Figure 6은 1B 모델로 ReACT의 웹 API 호출 8회, CodeACT의 코드 실행 8회, Swarm의 에이전트 간 I/O 32회를 수행한다. Pie의 지연은 각각 4.27·3.18·6.14초, 처리량은 29.94·40.18·5.21 agents/s였다. 이 단위는 생성 토큰 수가 아니라 초당 처리한 에이전트 수다. 저자들은 ReACT에서 지연 최대 15% 감소와 처리량 최대 30% 증가를 보고한다. 그래프는 워크로드마다 가장 긴 지연과 가장 높은 처리량을 각각 기준으로 정규화했으므로 두 막대의 비율을 같은 분모로 읽으면 안 된다. 외부 상호작용이 2회 미만이면 차이가 거의 없었고, 평가한 범위에서는 I/O가 많을수록 격차가 커졌다.

Figure 7의 vLLM 대비 처리량 3.5배는 별도 함수 호출 워크로드의 결과다. 자주 쓰는 API 문서의 KV 캐시를 남기고, 결과를 기다리지 않아도 되는 호출은 생성 중 호출 형식이 확인되는 즉시 병렬로 보내며, 한 번만 사용하는 API 명세의 캐시는 이후 문맥에서 제외한다. 세 최적화를 모두 적용한 Pie를 Python으로 구현한 vLLM 워크플로와 비교했고, 3.5배는 에이전트 128개 지점의 값이다. API 사용 빈도와 반환값의 필요성을 아는 애플리케이션이 직접 최적화한 결과로, Pie를 켜기만 하면 얻는 향상이 아니다.

프로그래밍 가능성에는 비용이 따른다. 논문은 일반 텍스트 생성처럼 Pie의 이점이 없는 작업에서 출력 토큰당 생성 시간(TPOT)을 vLLM과 비교했다. Llama 3 8B에서는 vLLM 64.06ms 대비 Pie 65.59ms로 1.53ms(2.39%) 느렸고, 3B에서는 30.30ms 대비 32.01ms로 1.71ms(5.64%), 1B에서는 16.83ms 대비 18.75ms로 1.92ms(11.41%) 느렸다. 증가율의 분모는 vLLM의 TPOT다. 추가 시간은 1.5~1.9ms로 비슷하므로, 모델이 클수록 비율이 작아진다. 8B 모델·32개 동시 인퍼렛에서 원인을 분해한 Table 3에서는, 가장 큰 몫은 GPU에서 샘플링을 순전파와 겹쳐 실행하지 못하는 비용(토큰당 1.320ms)이었다. Wasm 처리 자체의 0.001ms가 전체 런타임 비용을 뜻하는 것은 아니다.

구체적인 예시

논문의 자기회귀 루프 예시는 인퍼렛의 역할을 보여 준다. 먼저 프롬프트를 토큰화하고 임베딩·KV 페이지를 할당한 뒤 프리필을 수행한다. 다음에는 토큰 분포를 받아 가장 높은 값을 고르고, 선택한 토큰을 다시 임베딩해 다음 생성 단계로 넘기는 과정을 반복한다. 끝나면 자원을 해제한다. 이 저수준 제어를 매번 손으로 쓸 필요는 없다. 지원 라이브러리를 이용하면 Context를 만들고 프롬프트를 채운 뒤 generate_until로 생성하는 세 단계로 줄일 수 있다.

현재 main README의 Python 예제는 인터페이스가 달라졌다. WorkingSet이 KV 페이지 등을 관리하고, ForwardPass의 epilogue에 전달한 토큰 선택 함수는 한 번 추적된 뒤 GPU에서 실행되며 호스트는 결과 토큰을 받는다. 논문에서 호스트로 분포를 가져와 샘플링하던 예시와 구별된다. 다만 이 변화가 비용을 얼마나 줄였는지 보여 주는 현재 버전의 동일 조건 TPOT 표는 확보한 자료에 없다. FAQ의 raw forward는 Rust 전용이라는 설명도 README와 달라, 사용하는 릴리스와 예제를 함께 확인해야 한다.

한계와 오해하기 쉬운 점

논문은 사용자 코드와 네트워크 I/O를 엔진 안에 넣으면 공격 표면이 넓어진다고 밝힌다. 메모리 격리가 있어도 자원 고갈, 부채널 유출, 토큰 분포 악용 위험까지 해결됐다는 뜻은 아니며 포괄적 보안 검증은 논문 범위 밖이다. 또한 당시에는 단일 추론 백엔드의 중앙 스케줄러를 사용했고, KV 페이지 부족 시 가장 최근에 시작한 인퍼렛을 종료했다. 이런 당시 구현의 한계를 현재 버전의 확정 동작으로 옮겨 쓰면 안 된다.

현재 FAQ는 프로덕션 준비가 되지 않은 연구 프로토타입이며, 릴리스 사이에 API가 예고 없이 바뀔 수 있다고 안내한다. 한 호스트의 다중 GPU는 지원하지만 다중 노드 추론은 지원하지 않는다는 설명이다. 확인된 최신 안정 릴리스는 2026년 9월 22일의 0.5.2다. 로드맵에는 여전히 0.5.0의 분산 백엔드 계획이 남아 있으므로, 계획이나 main 코드 존재만으로 해당 기능이 안정 릴리스에 제공됐다고 판단할 수 없다.

0.4.0 블로그의 추측 실행(speculative execution)은 GPU가 현재 순전파를 계산하는 동안 다음 호출을 준비해 런타임 비용을 겹치는 방식이다. 여러 토큰 초안을 검증하는 추측 디코딩과 다르다. 블로그의 90%는 프로그래밍 가능한 서빙의 오버헤드를 줄였다는 주장으로, 전체 응답 시간 90% 감소가 아니다. 약 320µs의 런타임 작업과 5~8ms의 순전파도 그 글의 설명값이다. FAQ의 약 1ms/token, 논문 TPOT 비용과 조건이 같다고 볼 수 없다. 논문 Table 5에도 표와 본문의 개선율이 맞지 않는 부분이 있어 여기서는 그 배수를 성과로 사용하지 않았다.

실제로 어떻게 활용할까

활용을 검토할 만한 경우는 도구를 자주 부르면서 같은 문맥을 재사용하는 에이전트, 추론 가지 사이에 KV 캐시를 공유하는 작업, 새 샘플러를 엔진 수정 없이 실험하려는 연구다. 반대로 독립적인 짧은 요청을 반복 처리한다면 추가 비용만 생길 수 있다. 먼저 현재 서비스에서 도구 왕복과 재프리필이 얼마나 큰 비중인지 측정하고, 제거할 수 있는 비용이 인퍼렛 관리 비용보다 큰지 확인하는 것이 실용적이다.

README의 기본 사용 흐름은 설정과 모델을 준비하고, Python·JavaScript·Rust 인퍼렛을 작성한 뒤 pie run으로 실행하는 것이다. 상시 엔진은 pie serve로 띄우고 클라이언트로 프로그램을 제출한다. OpenAI·Anthropic·Google 형식의 호환 API도 있지만, 기존 클라이언트를 그 주소로 바꾸기만 해서는 애플리케이션별 캐시 제어가 자동으로 생기지 않는다. 이득을 얻으려면 인퍼렛에 어떤 문맥을 보존하고 언제 도구를 호출할지 표현해야 한다.

검증할 때는 지금 사용하는 엔진 버전, 같은 GPU와 모델, 동시 요청 수, 도구 호출 횟수와 생성 길이를 맞춰 비교해야 한다. 토큰당 시간뿐 아니라 전체 작업 지연과 완료 작업 수를 함께 봐야 한다. 여기서는 논문과 공개 문서를 대조했으며 코드를 실행해 성능을 재현하지 않았다. 도입 전에는 선택한 릴리스의 예제와 모델 지원 범위를 확인하고, 같은 업무에서 실제 이득이 유지되는지 평가해야 한다.