WavePP: High-Throughput Pipeline Parallel LLM Prefill under Prefix Reuse

논문 링크

WavePP — 프리픽스 재사용을 태우는 파이프라인 병렬 프리필 런타임

TL;DR

TensorRT-LLM 위에 구축한 WavePP 는 요청 admission(재사용 합의·캐시 보호·용량 예약)을 파이프라인 실행과 비동기로 겹쳐 처리 하고, wave 크기를 동적으로 조절해 파이프라인을 꽉 채운다. 동일 커널·동일 토폴로지 에서 GLM 5.2와 MiniMax M2.7의 프리필 처리량을 최대 2.91×/2.02× 까지 끌어올렸고, 28개 Kimi K3 세팅 중 동시성 8 이상의 18개 전부 에서 가장 높은 처리량을 기록했다.

핵심 아이디어

파이프라인 병렬(PP) 프리필의 진짜 병목은 GPU 연산이 아니라 “다음 요청을 얼마나 빨리 준비하느냐"다. 각 스테이지가 KV 캐시를 독립적으로 유지·축출 하는 시스템에서는 로컬 캐시 히트가 곧 글로벌 재사용 가능을 보장하지 않는다. WavePP의 핵심 주장은 하나로 요약된다: “요청 admission을 실행 경로 밖으로 빼내 GPU가 쉬는 순간을 없애면, 같은 커널로도 파이프라인 프리필 처리량을 크게 높일 수 있다.” 이 주장을 뒷받침하는 세 축이 fast reuse hint(빠른 재사용 힌트), local lease(지역 임대), capacity escrow(용량 예약) 다.

배경: 그들이 해결한 문제

에이전트 워크로드가 보편화되면서 LLM은 대화·문서·도구 출력 등 점점 더 긴 입력 을 처리한다(근거: §1). 이에 따라 prefill 단계의 비용이 급증했고, disaggregated serving에서 prefill 전용 워커를 두는 구조가 보편화되었다.

모델이 한 GPU에 안 들어갈 만큼 커지면서 등장한 병렬화 전략은 크게 세 가지다. 텐서 병렬(TP)은 레이어를 디바이스 간에 쪼개고, 파이프라인 병렬(PP)은 레이어를 연속된 스테이지 그룹으로 나누며, 컨텍스트 병렬(CP)은 입력 길이를 나눈다. PP는 여러 마이크로배치가 스테이지를 차례로 통과하며 동시에 실행되므로 처리량에 유리하지만, 가장 느린 스테이지가 전체를 지배 한다는 구조적 한계를 갖는다(근거: §3.5).

여기에 프리픽스 캐싱이 얽히면 문제가 복잡해진다. 공유 프리픽스의 캐시 상태를 재사용하면 prefill 계산을 수 배 줄일 수 있다. 예컨대 100K 컨텍스트의 90%가 캐시에 있다면 마지막 10K 토큰만 처리하면 된다(근거: §3.2). 그런데 PP에서는 각 스테이지가 자기 레이어의 로컬 캐시를 따로 갖는다. 게다가 Kimi K3 같은 하이브리드 모델은 MLA 어텐션의 밀집 KV 캐시와 KDA(선형 어텐션)의 희소 recurrent 스냅샷 을 함께 갖는데, recurrent 상태는 $S_{r'}$($r'>r$)에서 $S_r$을 되감아 복원할 수 없어 재사용 경계가 두 캐시의 교집합 에 묶인다(근거: §3.2, Fig. 1).

결국 저자가 지적하는 연구 공백은 이것이다: 기존 PP 런타임은 요청 준비(재사용 경계 합의 → 캐시 핀 → suffix 용량 확보 → 메타데이터 구축) 를 실행 경로 위에서 직렬로 수행한다. 이 준비 작업이 다음 forward pass를 지연시키면, GPU에 할 일이 남아 있어도 파이프라인이 빈다. 트리 순회는 뮤텍스를 필요로 하고, 스테이지 간 재사용 경계가 어긋나면 두 번째 순회(재보정)까지 필요해져 그 비용이 더 커진다(근거: §4.1).

새로운 접근법: WavePP

WavePP는 이 준비 작업을 실행과 분리(decouple) 하고, 스테이지들이 캐시를 독립적으로 관리하면서도 하나의 실행 스케줄을 따르게 하는 prefill 런타임이다. 핵심 설계는 다음과 같다(근거: §4).

1) 두 스레드와 두 통신 평면. 각 랭크는 두 스레드를 가진다. WAVE 스레드는 요청 전송·재사용 합의·용량 관리를 맡고, LOOP 스레드는 executor와 캐시 블록 materialization을 맡는다. 통신도 두 평면으로 나뉜다: data plane은 wave 스케줄과 activation을, rounds plane은 요청과 admission 결과를 나른다(근거: §4.2).

2) Fast reuse hint. 순회 비용과 뮤텍스 경합을 줄이기 위해 블록을 갖지 않는 힌트 인덱스 를 둔다. 각 토큰 블록 $B_i$에 대해 체인드 프리픽스 해시를 유지한다.

밀집 어텐션 프리픽스는 gallop(2배씩 건너뛰기) + bisect(이분 탐색) 로 프로브를 $O(N)$에서 약 $O(\log N)$으로 줄이고, recurrent 스냅샷은 희소하므로 어텐션 한계 지점에서 역방향 스캔 한다. 이 힌트는 캐시를 핀하지도, 용량을 예약하지도 않으므로 대기 중 요청이 취소돼도 되돌릴 게 없다(근거: §4.3). 200K 토큰·동시성 64 스트레스 실험에서 중간 프로브 시간이 64.4 ms → 24 μs 로 줄었고, p99는 2.35 ms였다(근거: §4.3).

3) Local lease(지역 임대). 합의된 재사용 후보에 대해 실제 블록을 트리에 설치하지 않고 참조만 보호 한다. 다른 요청이 그 상태를 공유할 수는 있지만, lease가 해제될 때까지 축출할 수 없다. 스테이지 간 실제 핀 지점이 어긋나면 풍부한 쪽이 공통 지점으로 되돌린다. 이 reconciliation은 단조 감소하므로 반드시 수렴한다(근거: §4.4).

4) Capacity escrow(용량 예약). lease가 재사용을 보장해도, suffix가 들어갈 공간이 없으면 아무도 완료하지 못한다. 그래서 각 스테이지가 lease 직후 suffix 용량을 예약한다. 이 예약은 실제 블록을 할당하거나 축출 victim을 밀어내지 않고, 참조만 잡아둔 채 실제 축출·오프로드는 executor의 materialization 시점으로 미룬다(근거: §4.5). 축출 victim은 트리에 남아 있어 다른 요청이 여전히 재사용으로 점유할 수 있고, commit 시 reference count로 추가 소유자를 감지해 대체 블록을 고른다.

5) MPU(Maximum Pipeline Utilization) 스케줄링. greedy하게 꽉 채우면 웨이브 수가 줄어 유휴 스테이지가 생긴다. 2M 토큰을 budget $M$으로 채우면 2개 웨이브지만, $M/2$로 줄이면 4개가 된다(근거: Fig. 6). MPU는 남은 작업량을 보고 budget을 이산적으로 반감시킨다.

가용 웨이브 수는 $n(b)=\max\left(W/b,\ N_{>b/2}\right)$ 로 추정하고, 실행 링 $R=2P$를 공급할 수 있는 가장 큰 budget을 고른다(근거: §4.6).

작동 원리: 구체적인 예시로 살펴보기

4개 스테이지(PP4) 파이프라인에 200K 토큰 프롬프트가 도착했다고 하자. 90%가 캐시되어 있지만, 각 스테이지의 실제 재사용 가능 지점은 다르다.

  flowchart LR
    A["요청 도착 (200K 토큰)"] --> B["힌트 조회: 로컬 재사용 지점 탐색"]
    B --> C["올-랭크 min: 공통 후보 합의"]
    C --> D["lease: 공통 프리픽스 + 마지막 스냅샷 핀"]
    D --> E["escrow: suffix 용량 예약 (commit)"]
    E --> F["각 스테이지 로컬 materialization"]
    F --> G["rank 0이 chunk 계획 → wave 실행"]

1. 힌트 조회. 각 랭크가 자기 힌트 인덱스를 프로브해 로컬 재사용 지점을 찾는다. 스테이지 0은 176K, 1은 180K, 2는 164K, 3은 172K까지 재사용 가능하다고 하자. 후보는 올-랭크 최솟값인 164K 가 된다. 아직 캐시를 핀하지 않는다.
2. lease. 164K 프리픽스와 그 지점의 마지막 recurrent 스냅샷을 각 스테이지가 핀한다. 다른 요청이 이 블록들을 점유할 수는 있어도 축출은 못 한다.
3. escrow. 남은 suffix는 $200\text{K}-164\text{K}=36\text{K}$ 토큰이다. 각 스테이지가 36K 토큰 분량의 캐시 블록을 예약한다. 블록을 실제로 할당하거나 victim을 밀어내지는 않는다 — 이 일은 뒤로 미뤄 GPU 실행과 겹친다.
4. materialization. 각 스테이지가 자기 로컬 캐시 트리에 블록을 실제로 설치한다. 스테이지 0은 다른 스테이지가 준비를 끝내길 기다리지 않고 준비되는 즉시 실행을 시작할 수 있다.
5. wave 실행. rank 0이 suffix 36K를 chunk로 쪼갠다. $M=16{,}384$면 3개 웨이브(16K+16K+4K)지만, MPU가 $b=8{,}192$로 줄이면 더 많은 웨이브로 쪼개 4개 스테이지를 고르게 채운다.

여기서 핵심은 순서의 재배치 다. 기존 런타임은 “합의 → 핀 → 예약 → 설치"를 요청이 실행될 때 직렬로 수행해 파이프라인이 기다린다. WavePP는 “힌트(비차단) → 합의 → lease/escrow(참조만) → 실행 오버랩 중 materialization"으로 재배치해, GPU가 이전 요청을 실행하는 동안 다음 요청의 준비를 병행한다(근거: Fig. 3).

성능 검증: 주요 결과

1) 동일 커널·동일 토폴로지에서의 개선 (GLM 5.2 / MiniMax M2.7)

두 모델 모두 NVFP4 가중치, 4개 B300 GPU, PP4 토폴로지, TensorRT-LLM 1.3.0rc26 커널로 고정하고 런타임만 교체했다(근거: §5.2). 결과는 40개 세팅 중 37개 에서 WavePP가 우세했다.

(근거: Table 1, 단위 input tokens/s. cached 토큰 포함)

가장 극적인 지점은 짧은 suffix + 높은 동시성 이다. 요청이 빨리 끝날수록 런타임이 새 요청을 준비해야 하는 속도가 빨라지는데, 여기서 admission 지연이 곧바로 처리량을 깎아먹는다. $c=128$에서 WavePP는 2.91×(GLM) 와 2.02×(MiniMax) 처리량을 달성했고, 4K suffix 워크로드의 중간 prefill 완료 시간은 GLM에서 약 71%, MiniMax에서 약 55% 줄었다(근거: §5.2).

주목할 점은 역방향 사례도 존재 한다는 것이다. long-prefix reuse의 $c=8$과 4K suffix MiniMax $c=16$에서는 WavePP가 오히려 낮았다. 저자는 이를 “GPU 계산이 지배적이어서 admission 절감 효과가 상대적으로 작은 구간"이라고 설명한다(근거: §5.2).

2) Kimi K3: TP 대 PP, 그리고 라이브러리 횡단 비교

2.8T 파라미터의 하이브리드 KDA/MLA MoE 모델인 Kimi K3에서 8개 GB300 GPU(4-GPU 노드 2개)로, WavePP(TP1×PP8)를 TRT-LLM/SGLang/vLLM의 TP8/EP8 및 PP8과 비교했다(근거: §5.3). 28개 세팅 중 21개, 그리고 동시성 8 이상의 18개 전부 에서 WavePP가 최고 처리량을 기록했다(근거: §5.4).

• Cold long-context: TRT-LLM TP8/EP8 대비 기하평균 처리량이 65.5K/131K/262K에서 각각 +58.8% / +49.1% / +37.8%. $c=32$에서 44,643 / 35,878 / 28,031 input tok/s(근거: §5.5).
• Seeded reuse(90%): 131K $c=32$에서 289,096 tok/s, 262K $c=32$에서 218,743 tok/s. 동시성 8 이상에서 처리량과 p50 완료 시간 모두 일관되게 우위(근거: §5.6).
• Mixed(75% warm-131K / 25% cold-262K): TP 대비 기하평균 +46.1~+56.8%, PP 대비 +5.0%(SGLang)/+14.6%(vLLM). $c=32$에서 61,190 tok/s(근거: §5.7).

흥미로운 관찰은 PP가 prefill에서 TP보다 유리하다 는 점이다. SGLang PP8이 SGLang TP8/EP8보다 기하평균 처리량이 +26.8% 높았고, WavePP는 여기서 한 단계 더 올라갔다(근거: §5.4).

3) 캐시 유지율과 Ablation

프리픽스 유지 실험에서 WavePP의 good-hit rate(기대 프리픽스의 80% 이상 재사용)는 88~100% 로, TP8/EP8의 73~83%를 모든 도착률에서 앞섰다(근거: §5.8). 이는 admission을 실행과 분리한 덕에 캐시가 새 요청에 밀려나기 전에 재사용이 확보되기 때문이다.

Ablation은 메커니즘별 기여를 분리한다(근거: §5.9, Fig. 14).

• Cold 131K $c=16$: R1→R2에서 8,826 → 29,692 tok/s (+236.4%). 이후 변화는 미미. MPU는 p50을 69.69→43.03초로 줄였지만 처리량은 +1.4%에 그침.
• Reuse 131K $c=8$: lease/escrow(R3)가 96,796 → 149,021 (+54.0%), MPU(R4)가 → 211,214 (+41.7%).

즉, cold에서는 admission·transport가, reuse에서는 캐시 조정과 스케줄링이 각각 결정적 이다. 다만 MPU가 항상 이기는 건 아니다: reuse 262K $c=32$에서는 MPU가 처리량을 -9.7% 낮췄다(그래도 TP8/EP8보다 +24.5% 높음). planner-driven packing이 특정 워크로드에서는 정적 packing보다 못할 수 있음을 시인한 셈이다.

우리의 관점: 강점, 한계, 그리고 이 연구가 중요한 이유

강점. 이 논문의 가장 큰 공헌은 문제의 재정의 다. 프리필 처리량 최적화를 “커널 연산 가속"이 아니라 “admission cadence” 문제로 보고, hint → lease → escrow → materialization이라는 개념적 프레임워크를 제시한 점은 실무 서빙 시스템 설계에 직접 이식 가능하다. 또 동일 커널·동일 토폴로지 로 런타임만 바꿔 비교해, 커널 성능과 스케줄링 성능을 깔끔하게 분리한 실험 설계가 돋보인다. SGLang PP8의 재현성 체크(day-0 리포트 대비 +5.7%)까지 넣은 것도 신뢰도를 높인다.

한계. 첫째, MPU는 chunk의 실제 실행 시간을 추정하지 않고 토큰 수에만 의존하는 이산적 budget 조절이라, 똑같은 크기의 chunk도 어텐션 프리픽스 길이에 따라 비용이 다르다는 점을 반영하지 못한다. 실제로 일부 세팅에서 역회귀 를 일으켰다. 둘째, TP1×PP4·PP8만 검증해 TP2×PP4나 PP16, 불균형 레이어 배치에서의 동작은 미지수다. 셋째, FCFS 스케줄링과 단일 replica만 다뤄 우선순위·마감시한·캐시 인지 라우팅 같은 실무 요구를 남겨뒀다. 마지막으로, 이는 TensorRT-LLM에 국한된 구현이어서 다른 런타임(vLLM/SGLang)으로의 이식성은 입증되지 않았다(근거: §6).

왜 중요한가. 에이전트 워크로드가 길어지고, 하이브리드(어텐션+recurrent) 모델이 늘어나면서 “여러 캐시 타입의 재사용 경계를 어떻게 일관되게 합의할 것인가"는 앞으로 더 중요해질 문제다. WavePP는 이 분산 admission 문제에 대해 실행과 준비의 분리 라는 명확한 원칙과 측정 가능한 이득(최대 2.91×)을 함께 제시했다는 점에서, 서빙 시스템 연구자와 엔지니어 모두에게 참조점이 될 만한 작업이다.

다음 단계는?: 앞으로의 길

저자가 직접 제안한 방향(근거: §6)에 더해, 합리적 후속 작업을 정리하면 다음과 같다.

1. 실행 시간 추정 기반의 동적 chunking. MPU의 이산적 budget 대신 chunk별 실행 시간과 스테이지 진행 상황을 관찰해 크기를 정하면, 역회귀를 줄이면서도 파이프라인 균형을 개선할 수 있다.
2. 다양한 토폴로지·불균형 배치. TP2×PP4, PP16, 그리고 스테이지 간 레이어 불균형에서 admission 오버헤드와 오버랩이 어떻게 변하는지 검증할 필요가 있다.
3. 우선순위·마감시한·캐시 상태를 고려한 스케줄링. FCFS를 넘어 tail latency와 SLO를 최적화하는 정책이 자연스러운 확장이다.
4. 멀티 replica + 캐시 인지 라우팅. 여러 prefill 워커에 걸쳐 캐시를 고려한 라우팅과 SLO 곡선을 결합하면 실제 배포에서의 tail latency 개선을 검증할 수 있다.
5. 다른 런타임으로의 이식. hint/lease/escrow 프레임워크를 vLLM·SGLang 등 다른 백엔드에 적용해 일반화 가능성을 확인하는 것도 가치 있는 후속 연구다.

논문 원문의 표

arXiv e-print 의 LaTeX 원본에서 기계적으로 옮긴 표입니다. 숫자는 논문의 값이며 모델을 거치지 않았습니다.

표 1. PP4 comparison on matched kernels. Aggregate prompt throughput (tokens/s, including cached tokens) is the median of three boots per system and model. Each system uses one configuration throughout. Throughput change $\Delta$ is $(\mathrm{WavePP}/\text{TRT-LLM}-1)\times100\%$. Green shading marks the higher throughput. Short and long reuse use 90% cached prefixes. The 4K-suffix cells follow cold traffic and retain its cache-pool eviction history.

표 2. Configurations of the Kimi K3 comparison baselines.

표 3. Kimi K3 workloads and offered concurrency. Input lengths are abbreviated in decimal thousands; the vLLM PP8 262K exception is described in the text.

표 4. Cold aggregate prompt throughput (input tok/s). Positive deltas indicate higher throughput for WavePP.

표 5. Queue-inclusive completion time for cold prefill, in seconds, with WavePP deltas against the lowest-latency TP8/EP8 and PP8 results in each cell. Negative deltas indicate lower completion time for WavePP.

표 6. Seeded-reuse aggregate prompt throughput (input tok/s). Positive deltas indicate higher throughput for WavePP.

표 7. Queue-inclusive completion time for seeded reuse, in seconds, with WavePP deltas against the lowest-latency TP8/EP8 and PP8 results in each cell. Negative deltas indicate lower completion time for WavePP.

표 8. Mixed-workload aggregate prompt throughput (input tok/s). Positive deltas indicate higher throughput for WavePP.