PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Aaryam Sharma

arXiv:2609.35263v1cs.DCcs.AI

TL;DR

Independent stage caches make common-prefix reuse and timely admission difficult to coordinate without disrupting pipeline execution. WavePP overlaps admission with execution, protects shared reusable state, reserves suffix capacity, and dynamically plans chunks; it improves throughput across most tested settings and leads measured Kimi K3 comparisons at higher concurrency.

  • Problem

    When pipeline stages retain and evict cache state independently, a local cache hit may not be globally reusable, while admission coordination can delay execution and leave the pipeline underoccupied.

    Source: §3.4 · p. 7

  • Method

    WavePP overlaps request admission with pipeline execution, agrees on reusable state, protects it, reserves suffix capacity, and dynamically plans request chunks.

    Source: §1 · p. 2, §4.2 · p. 9, §4.6 · p. 14

  • Results

    37 of 40 GLM 5.2 and MiniMax M2.7 settings improved; WavePP also had the highest measured throughput in all 18 Kimi K3 settings at concurrency eight or higher.

    Source: §1 · p. 3, §7 · p. 29

  • Takeaways & Limitations

    The results support preparing requests early and keeping enough work available to sustain pipeline execution.

    Source: §7 · p. 29

    The evaluation uses TP1×PP4 and TP1×PP8, and MPU decreases performance in some experiments.

    Source: §6 · p. 28

Abstract

from arXiv · show

Pipeline parallelism can improve prefill throughput by processing multiple request chunks concurrently across different stages of the model. However, keeping the pipeline fully utilized requires efficient scheduling and request preparation. In systems where stages retain and evict cache state independently, a local cache hit does not guarantee that the same prefix can be reused across the pipeline. Here, coordination overhead can impede request admission cadence and thus reduce overall throughput. In this paper, we present WavePP, a prefill runtime built on top of TensorRT-LLM that addresses these challenges by overlapping request admission with pipeline execution. WavePP asynchronously finds a prefix that can be reused across all stages, protects the cached state, and reserves space for the remaining input while earlier requests continue to execute. It subsequently plans the chunk sizes of each request dynamically to maximize pipeline fill. Each stage then completes the local preparation before executing the request. In the same system and pipeline topology, WavePP improves TensorRT-LLM's prefill throughput in 37 of 40 tested settings on GLM 5.2 and MiniMax M2.7. At concurrency 128 with high cache reuse, these changes increase throughput by factors of 2.91 and 2.02, respectively. Across 28 Kimi K3 settings, WavePP also has the highest measured throughput in all 18 settings at concurrency eight or higher, compared with tensor/expert-parallel and pipeline-parallel baselines from TRT-LLM, SGLang, and vLLM.

1. Introduction

Longer inputs and larger models increase the demands on inference, while pipeline parallelism depends on timely admission and balanced work to keep stages occupied. WavePP coordinates cache reuse across independently managed stages, overlaps admission with execution, and dynamically plans chunks; it improves throughput in most tested settings.

1. Introduction

  • Pipeline parallelism processes microbatches concurrently across stages, but online serving requires enough independent, balanced work to avoid idle stages.Microbatch costs vary with cache hits, prompt lengths, and chunk tails.Source: §1 · p. 2
  • With independently retained and evicted stage caches, a local hit does not ensure a prefix is reusable across the pipeline, so agreement and capacity reservation must happen early.Doing this on the main execution path can delay the next forward pass and reduce pipeline occupancy.Source: §1 · p. 2
  • WavePP overlaps admission with pipeline execution, agrees on reusable state, protects it, and reserves suffix capacity before stages finish local preparation independently.Its protocol delays eviction until reserved space is needed.Source: §1 · p. 2
  • 37 of 40 GLM 5.2 and MiniMax M2.7 settings improved with WavePP on the same TensorRT-LLM PP4 kernels; at concurrency 128, short-suffix throughput under cache pressure rose 2.91× and 2.02×.Across 18 Kimi K3 settings at concurrency eight or higher, WavePP had the highest measured throughput.Source: §1 · p. 3

2. Related Works

Prior work addresses pipeline scheduling, execution overlap, cache reuse, and model placement through varied designs. WavePP focuses on distributed admission and cache coordination within a fixed pipeline configuration.

2.1. Pipeline Parallelism

  • For inference PP, prior work chunks and schedules requests [4], overlaps preparation with execution [13], and supports pipeline execution in vLLM and SGLang; WavePP emphasizes distributed admission.WavePP agrees on reusable state, protects it, and reserves suffix capacity outside the executor loop while stages manage local caches.Source: §2.1 · p. 3

2.3. Other Parallelism Strategies

  • Layer redistribution and placement methods address pipeline configuration [25] [26], while tensor, sequence, or expert parallelism divides work within layers [5] [42] [43] [44] [45]; WavePP studies admission in a fixed configuration.These strategies can complement pipeline parallelism across layer groups.Source: §2.1 · p. 3, §2.3 · p. 4

2.2. Cache Reuse and Serving Runtimes

  • Serving systems reuse cached prefixes through block-based KV storage and radix trees [2] [3]; WavePP adopts TensorRT-LLM caching, with leases and capacity escrows protecting prefixes and reserving resources.Prior work also studies distributed cache storage [16], hybrid cache policies [14] [15], and disaggregated serving [33] [34].Source: §2.2 · p. 4

3. Preliminaries

The preliminaries describe model-state representations, prefix reuse, disaggregated prefill and decode, and pipeline parallelism. They establish why reuse boundaries and execution must be coordinated across stages.

3.1. Transformer Architectures

  • Attention layers retain KV representations, while linear attention uses a recurrent state independent of sequence length; Kimi K3’s KDA combines delta-rule updates with finer-grained gating [50] [51].Linear attention’s recurrent state captures the prefix, and later recurrent states generally cannot be rewound to recover earlier ones.Source: §3.1 · p. 5, §3.2 · p. 6

3.2. Reuse

  • Prefix caching reuses stored blocks, but recurrent caches retain sparse snapshots, so independently evolving cache types require a reusable boundary common to both.Figure 1 illustrates dense KV storage alongside four recurrent snapshots on the same prefix tree; recurrent state must match the exact prefix.Source: §3.2 · pp. 5–6, Figure 1 · p. 6

3.3. Disaggregated Serving

  • Disaggregated serving separates prefill and decode workers, transferring prefill-produced state to the decode worker; this paper focuses on dedicated prefill workers.Token throughput and prefill completion time are the primary metrics.Source: §3.3 · pp. 6–7

3.4. Parallelism Strategies

  • Pipeline parallelism partitions layers across ordered stages that own their layers’ cache state, but throughput is limited by the slowest stage and uneven online work can leave resources idle.PP therefore needs enough independent work to fill stages while balancing successive microbatches; cache consistency and admission overhead add challenges.Source: §3.4 · p. 7

4. Design of WavePP

WavePP coordinates reusable prefixes and suffix capacity across independently managed stage caches while overlapping request preparation with pipeline execution. Its scheduler dynamically sizes request chunks to keep the pipeline occupied.

4.1. A Minimum Implementation

  • A minimum admission path establishes a common reusable prefix, pins that state, reserves suffix capacity, materializes cache blocks, and then chunks and executes requests.Ranks need a common reuse boundary, while concurrent requests must protect reusable state and sufficient suffix capacity to avoid stalling.Source: §4.1 · p. 8

4.2. Decoupling Admission

  • Separate WAVE and LOOP threads let admission negotiate leases and suffix capacity while execution continues; each stage materializes locally before the wave runs.Stage 0 can start once locally ready without waiting for later stages to finish installation.Source: §4.2 · p. 9, Figure 3 · p. 10

4.3. Fast Reuse Hints

  • The sharded hint index lets queued requests agree on a reuse candidate without pinning cache or reserving suffix capacity; selected requests later validate and pin the actual prefix.The index can be stale, so an authoritative final walk remains necessary.Source: §4.3 · pp. 10–12
  • In a 200K-token, concurrency-64 stress experiment, optimized sharded lookup reduced median probe time from 64.4 ms to 24 μs, while hint discovery reached a p99 of 2.35 ms.The result is for the fast reuse-hint lookup.Source: §4.3 · p. 12

4.4. Local Leases

  • A local lease pins the agreed reusable state, while FCFS lease limits and token headroom constrain how many requests can hold cache state concurrently.Leases are reconciled to a common endpoint when ranks pin different reuse endpoints.Source: §4.4 · p. 12, Figure 5 · p. 13

4.5. Capacity Escrow

  • Capacity escrow reserves suffix space without immediately evicting victims; local materialization performs the allocation and deferred eviction or offloading.If the reuse endpoint is reduced, commit rechecks and reserves the increased suffix demand before committing.Source: §4.5 · pp. 13–14

4.6. Pipeline Scheduling

  • Maximum Pipeline Utilization adjusts the wave budget to target enough waves for pipeline fill, accounting for remaining and in-flight work while limiting unnecessary chunk splitting.The Figure 6 example illustrates that a smaller budget can create more waves, but is schematic rather than a measured speedup.Source: §4.6 · pp. 14–15, Figure 6 · p. 15

5. Evaluation

The evaluation measures prefill throughput and completion time across cold, reused-prefix, and mixed workloads. Within matched TensorRT-LLM PP4 configurations, WavePP generally improves throughput, with the largest gains under high cache pressure and short suffixes.

5.1. Metrics

  • Aggregate prompt throughput counts all input tokens, including cached tokens, while completion time includes admission and queueing.Throughput is completed-request input tokens divided by run wall time; reported completion-time statistics include p50 and p95 for Kimi K3 comparisons.Source: §5.1 · p. 16

5.2. Pipeline Parallelism in TensorRT-LLM

  • The matched-kernel PP4 comparison uses one configuration per system across workloads and reports medians from three independent server boots.It covers cold, 90%-reuse, mixed, and 4,096-token-suffix workloads, with the latter following cold traffic and its cache-eviction history.Source: §5.2 · pp. 16–17, §5 · p. 16, Table 1 · p. 19
  • Cold and mixed traffic favors WavePP in every setting, with throughput gains of 1.9–9.1% for GLM and 7.5–10.7% for MiniMax.With 90% reuse of shorter inputs, gains at c=16 and c=32 are 11.5% and 17.8% for GLM, and 9.1% and 10.4% for MiniMax.Source: §5.2 · p. 17, Figure 7 · p. 17
  • For 4K suffixes, TensorRT-LLM PP4 throughput falls from c=32 to c=128 while WavePP throughput rises, as more requests compete for cache space.Overlapping reuse preparation and suffix-capacity reservation with earlier execution helps keep new requests ready as others finish.Source: §5.2 · p. 18, Figure 8 · p. 18
  • WavePP4 beats TensorRT-LLM PP4 in 37 of 40 settings; at c=128 after cache pressure, it reaches 2.91× and 2.02× their throughput on GLM and MiniMax.These factors apply to the 4,096-token suffix workload after earlier cold traffic filled the cache.Source: §5.2 · pp. 17–18, Figure 8 · p. 18, Table 1 · p. 19

5.3. Kimi K3 Experimental Setup

The Kimi K3 experiments compare WavePP with tensor/expert-parallel and pipeline-parallel serving implementations on a shared eight-GPU, two-node setup. The workload suite varies input length, reuse, traffic mix, and concurrency.

  • Kimi K3 is a 2.8-trillion-parameter hybrid KDA/MLA MoE model, evaluated on eight GB300 GPUs across two four-GPU nodes.WavePP and pipeline baselines use TP1×PP8; tensor baselines use TP8/EP8.Source: §5.3 · p. 19
  • The comparison includes TRT-LLM TP8/EP8 and PP8, plus SGLang and vLLM TP8/EP8 and PP8 implementations.WavePP and the TRT-LLM TP reference share Kimi K3 kernels; the SGLang configurations use v0.5.18.Source: §5.3 · p. 19
  • Workloads cover 65.5K, 131K, and 262K inputs, seeded prefix reuse, and mixed traffic with 75% warm 131K and 25% cold 262K requests.All systems receive the same request sequence, with the stated vLLM PP8 context-length adjustment; all requests completed successfully.Source: §5.3 · p. 20, Table 3 · p. 20

5.4. Kimi K3 Overall Performance

Across the Kimi K3 workload grid, WavePP leads most often overall and in every setting at concurrency eight or higher. Its geometric-mean throughput also exceeds both external PP8 baselines across the full grid.

  • At low concurrency, WavePP can trail some baselines when there is less independent work to fill the pipeline.The passage notes similar gaps for other PP methods.Source: §5.4 · p. 20
  • Across the full grid, WavePP improves geometric-mean throughput by 9.2% over SGLang PP8 and 15.7% over vLLM PP8.SGLang PP8 itself improves geometric-mean throughput by 26.8% over SGLang TP8/EP8.Source: §5.4 · p. 20
  • WavePP has the highest measured throughput in 21 of 28 Kimi K3 settings, including all 18 settings at c≥8.The figure identifies the highest-throughput system in each cell and reports WavePP’s difference from the strongest TP8/EP8 and external PP8 results.Source: §5.4 · p. 20, Figure 9 · p. 21

5.5. Cold Long-Context Prefill

On cold Kimi K3 inputs, WavePP’s throughput advantage over TRT-LLM TP8/EP8 narrows as input length grows. At concurrency eight or higher, its p50 completion time is lower than the tensor baseline and both PP baselines in every setting.

  • Across 65.5K, 131K, and 262K cold inputs, WavePP’s geometric-mean gains are 2.2%, 4.2%, and 8.0% over vLLM PP8, and 9.4%, 4.5%, and 1.9% over SGLang PP8.At c=1, WavePP is below the best baseline at every input length and remains below vLLM PP8 for cold 65.5K at c=4.Source: §5.5 · p. 21
  • At c≥8, WavePP p50 completion time is 31.6–49.0% below TRT-LLM TP8/EP8 and below both PP baselines in every setting.Its p95 is also below both PP baselines except cold 65.5K at c=8, where it is 1.6% higher than vLLM PP8.Source: Figure 10 · p. 22, Table 5 · p. 23, §5.5 · p. 23
  • WavePP improves cold-input geometric-mean throughput over TRT-LLM TP8/EP8 by 58.8%, 49.1%, and 37.8% for 65.5K, 131K, and 262K inputs, respectively.At c=32, it sustains 44,643, 35,878, and 28,031 input tok/s for those lengths.Source: §5.5 · p. 21, Table 4 · p. 22

5.6. Seeded Prefix Reuse

With 90% prefix reuse, WavePP’s throughput advantage depends on concurrency and remaining input length. It outperforms the listed PP baselines at concurrency eight or higher for both input lengths, while completion-time results are less consistent at p95.

5.6. Seeded Prefix Reuse

  • From c=8 onward, WavePP has the highest throughput at both 131K and 262K input lengths; at c=1, all pipeline-parallel baselines trail tensor-parallel baselines, while WavePP leads PP methods.For 90%-cached requests at c=1, WavePP is 63.3% and 51.4% below vLLM TP8/EP8 for 131K and 262K inputs, respectively. At c=4 for 262K inputs, it is 15.3% above TRT-LLM TP8/EP8 and 25.8% above SGLang PP8.Source: §5.6 · p. 23
  • At 131K reuse and c≥8, WavePP improves throughput by 36.8–55.4% over TRT-LLM TP8/EP8.Across all five concurrencies, its geometric-mean throughput is 4.1% below TRT-LLM TP8/EP8 but 11.6% and 24.9% above SGLang PP8 and vLLM PP8.Source: §5.6 · p. 23, Table 6 · p. 24
  • For 262K reuse, WavePP improves geometric-mean throughput by 11.1% over TRT-LLM TP8/EP8, 22.2% over SGLang PP8, and 45.3% over vLLM PP8.At c=32, it serves 218,743 input tok/s versus 184,753 input tok/s for SGLang PP8.Source: §5.6 · pp. 23–24, Table 6 · p. 24
  • At c≥8, WavePP reduces p50 completion time by 19.1–42.2% versus TRT-LLM TP8/EP8, while its p95 is 2.4–3.2% higher than the lower-PP baseline in three settings.Its p95 is lower than both PP baselines in the other three settings.Source: §5.6 · p. 24, Table 7 · p. 25

5.7. Mixed Traffic and Completion Latency

On mixed warm and cold traffic, WavePP leads the compared systems in throughput at all three concurrencies, with mixed latency results. In the prefix-retention sweep, it also records higher good-hit rates at every tested arrival rate.

5.7. Mixed Traffic and Completion Latency

  • WavePP has the highest throughput at all three mixed-traffic concurrencies, with geometric-mean gains of 46.1% over TRT-LLM TP8/EP8, 56.8% over SGLang TP8/EP8, and 50.0% over vLLM TP8/EP8.Against PP baselines, its geometric-mean throughput is 5.0% above SGLang and 14.6% above vLLM.Source: §5.7 · p. 25, Table 8 · p. 26
  • For mixed traffic, WavePP reduces p50 completion time by 9.5–33.1% and p95 by 22.5–43.4% versus the lowest corresponding TP baseline values.Its p95 is 14.3% and 11.8% higher than the lowest PP baseline at c=16 and c=32, respectively.Source: §5.7 · p. 25, Table 9 · p. 26

5.8. Prefix Retention

  • WavePP’s good-hit rate is 88–100% across five offered-load settings, compared with 73–83% for TRT-LLM TP8/EP8, and is higher at every tested rate.The tested arrival rates range from 0.50 to 1.00 requests/s; a good hit reuses at least 80% of the expected prefix.Source: Figure 13 · p. 26, §5.8 · p. 26

5.9. Ablation Study

The ablation separates admission improvements from scheduling changes across cold and reused requests. Admission and cache-capacity mechanisms deliver large gains for reused requests, while MPU’s effect varies by workload.

5.9. Ablation Study

  • The cumulative configurations progress from a naive PP8 re-implementation through wave admission and transport, leases and capacity escrows, and MPU.MPU combines planner-driven packing, adaptive wave sizing, and a progress deadline; successive PP comparisons keep configuration settings fixed.Source: §5.9 · pp. 26–27, Table 11 · p. 27
  • On reuse 131K at c=8, leases and capacity escrows increase throughput by 54.0%, and MPU adds a further 41.7%.Throughput rises from 96,796 to 149,021 and then 211,214 input tok/s; MPU also lowers p50 from 6.51 to 4.20 s.Source: Figure 14 · p. 27, §5.9 · p. 28
  • On cold 131K at c=16, wave admission and transport raise throughput from 8,826 to 29,692 input tok/s, a 236.4% increase.MPU lowers p50 from 69.69 to 43.03 s, while throughput rises by only 1.4% and mean completion time changes from 63.18 to 62.33 s.Source: Figure 14 · p. 27, §5.9 · pp. 27–28
  • On reuse 262K at c=32, MPU reduces throughput by 9.7% to 158,238 input tok/s, though it remains 24.5% above TP8/EP8.The paper notes static packing performs better for this specific workload.Source: Figure 14 · p. 27, §5.9 · p. 28

6. Limitations and Further Directions

The discussion identifies limits in MPU’s chunk-cost awareness, the evaluated pipeline topologies, and the current scheduling and deployment scope. It outlines corresponding directions for further evaluation and policy development.

6. Limitations and Further Directions

  • MPU adjusts token budgets in discrete steps but does not estimate chunk execution time, although equal-sized chunks can have different costs.The paper suggests choosing chunk sizes using estimated execution times and observed stage progress, while keeping planning inexpensive.Source: §6 · p. 28
  • The evaluation covers TP1×PP4 and TP1×PP8; other stage counts, tensor-parallel collectives, inter-node bandwidth, and uneven layer placement remain to be examined.The paper notes that keeping requests ready cannot by itself remove persistent stage-execution imbalance from uneven layer placement.Source: §6 · p. 28
  • WavePP currently packs requests in FCFS order and reserves capacity before scheduling chunks.Alternative policies considering request priorities, deadlines, and dynamic cache state are left for further work.Source: §6 · p. 28
  • The paper leaves multi-replica cache-aware routing and evaluation under realistic traffic patterns for future work.Such evaluation could test the benefits of larger cache capacity and whether the approach improves tail latency.Source: §6 · p. 28

7. Conclusion

WavePP prepares requests while earlier requests run, coordinating cache reuse and capacity across independently managed stages. It improves measured prefill throughput across most tested GLM 5.2 and MiniMax M2.7 settings and leads the Kimi K3 comparisons at concurrency eight or higher.

7. Conclusion

  • Independent stage caches make request preparation distributed, and doing it on the execution path can delay the next wave even when GPUs have work available.Pipeline throughput therefore depends on finding reusable state and securing cache capacity while earlier requests run.Source: §7 · pp. 28–29
  • WavePP moves reuse discovery and admission ahead of execution, using hints, leases, and capacity escrows before each stage prepares its request for scheduled chunks.The mechanisms let cache management proceed independently across stages while the pipeline follows a common schedule.Source: §7 · p. 29
  • WavePP improves TensorRT-LLM prefill throughput in 37 of 40 GLM 5.2 and MiniMax M2.7 settings, and has the highest measured throughput in all 18 Kimi K3 settings at concurrency eight or higher.The GLM 5.2 and MiniMax M2.7 comparison uses the same kernels and PP4 topology.Source: §7 · p. 29
  • At concurrency 128, short-suffix requests under cache pressure gain throughput factors of 2.91 and 2.02.The ablations show admission and transport improve cold prefill, while coordinated cache admission and scheduling provide additional gains under reuse.Source: §7 · p. 29
Loading 2609.35263v1…