Pie: A Programmable Serving System for Emerging LLM Applications

framework 2510.24051
programmable-servingkv-cache-handlersospllm-applicationswebassemblyinferlet

Pie: A Programmable Serving System for Emerging LLM Applications — L2 #

§1 TL;DR #

Pie decomposes the monolithic LLM generation loop into fine-grained handlers (embed / forward / sample), delegating end-to-end control to user-provided Wasm programs called inferlets — achieving only 2.4% overhead on 8B models while delivering 1.3–3.4× throughput on agentic workflows via application-specific KV cache, decoding, and I/O integration.

§2 痛点 · 方法 · 结果 #

Q1 痛点 #

Existing LLM serving systems (vLLM, SGLang, TGI) enforce a monolithic prefill–decode loop with three fundamental limitations:

  1. Implicit KV cache management (R1): System-wide policies (LRU, prefix caching) prevent application-specific control needed by tree/graph-of-thought, modular caching, or attention sinks.
  2. Inflexible decoding process (R2): The tightly-coupled predict–sample loop resists per-request customization for speculative decoding, MCTS, grammar-constrained generation, or watermarking.
  3. Poor workflow integration (R3): Agentic workflows requiring external tool calls force costly round-trips — discarding KV cache between interactions and re-prefilling context.
  4. These limitations are architectural, not implementation bugs: the monolithic design couples application logic to the execution engine.

    Q2 方法 #

    Pie introduces two architectural shifts:

    1. Decomposed handlers: The generation pipeline is dismantled into independent, fine-grained service handlers (embedding, forward pass, KV cache ops, sampling). Each handler is an opaque function exposed through an API.
    2. Inferlets as first-class programs: User-provided programs (inferlets) orchestrate handlers via a 42-API surface, replacing prompts as the basic unit of service. Inferlets explicitly manage KV cache allocation/deallocation, define custom generation sequences, and integrate arbitrary I/O — all within a single-threaded, event-driven Wasm runtime.
    3. Three-layer architecture: Application layer (Wasm runtime + ILM) → Control layer (resource manager + adaptive batch scheduler) → Inference layer (GPU kernel handlers via FlashInfer).

      核心技术壁垒: The key insight is that the generation loop can be decomposed into composable, per-API-call handlers without sacrificing batch efficiency — the adaptive batch scheduler (vertical + horizontal batching with work-conserving dispatch) bridges the gap between fine-grained programmability and GPU-efficient execution. This is non-obvious because prior systems assumed decomposition would destroy batching opportunities.

      Q3 结果 #

      • Standard tasks: 3–12% latency overhead vs vLLM (2.39% TPOT on 8B, 5.64% on 3B, 11.41% on 1B).
      • Agentic workflows: 1.1–2.4× lower latency, 1.3–3.4× higher throughput on ReACT/CodeACT/Swarm.
      • Stacked optimizations: 3.5× throughput via application-specific KV export + concurrent API calls + KV masking (Fig 7).
      • Adaptive batching: 17× throughput over naive eager dispatch.
      • Expressiveness: 19 distinct LLM applications implemented as inferlets (38–255 LoC each).

      §3 架构 / 方法图 #

      Figure 2: Pie architecture — handlers exposed as independent APIs, inferlets orchestrate them

      Paper's Figure 2, verbatim (caption: "Our proposed system, Pie, dismantles the sequential generation process into independent handlers, and delegates control to user-provided programs called inferlets.").

      The architecture decomposes the traditional monolithic loop (Fig 1) into three layers. Inferlets issue API calls through command queues, which the control layer batches (vertical + horizontal) before dispatching to GPU handlers. The key difference from prior systems: programs, not prompts, are the unit of service — enabling hundreds of concurrent inferlets with distinct optimization strategies.

      Figure 3: Inferlet service workflow across three layers

      Paper's Figure 3, verbatim (caption: "Inferlet service workflow. The application layer executes inferlets that make API calls to the control layer whose batch scheduler adaptively batches these calls and forwards them to the inference layer.").

      The request lifecycle: (1) user submits Wasm binary → ILM creates inferlet, (2) inferlet issues API calls to control layer, (3) batch scheduler groups compatible calls, (4) inference layer executes batched GPU ops, (5) results flow back via event dispatcher.

      Figure 4: Batch scheduling — vertical and horizontal batching

      Paper's Figure 4, verbatim (caption: "Batch scheduling example. Horizontal batching groups calls across different command queues, while vertical batching groups consecutive calls of the same type within the same queue if they do not conflict.").

      The scheduler uses a work-conserving policy: when the GPU becomes idle, the inference layer immediately notifies the control layer via IPC to trigger batch formation, maximizing GPU occupancy.

      System scope #

      • Stage coverage: Both prefill and decode; the API decomposes them into embed → forward → sample handlers.
      • Serving vs training: Serving only (continuous batching via adaptive scheduler).
      • Parallelism: Not explicitly covered — delegated to inference layer. Single-GPU evaluation only; multi-GPU discussed as future work.
      • Deployment mode: Single node, centralized control layer.

      Scheduler #

      • Queueing discipline: Work-conserving adaptive batching — GPU busy → queue; GPU idle → immediate dispatch.
      • Batch selection: When multiple API types eligible, selects the batch whose oldest pending call waited longest. Priority-aware (higher-priority queues placed earlier, truncate from tail).
      • Vertical batching: Consecutive same-type commands within one queue.
      • Horizontal batching: Same-type commands across different queues.

      KV / Memory manager #

      • Allocation unit: PagedAttention-style KvPage (8–32 tokens per page).
      • Lifecycle: Explicit allocation/deallocation by inferlets — no system-wide eviction policy.
      • Sharing: Import/export APIs for cross-inferlet KV cache sharing.
      • Virtualization: Each inferlet has its own virtual resource address space (opaque pointers).
      • Contention policy: FCFS; terminates most recently created inferlets until resources freed.

      §4 作者证明 #

      无形式化作者证明 — 仅实证

      Pie does not present a formal throughput/latency model. The closest analytical content is the programming model abstraction (§4, three-stage view: embed → forward → sample) and the opportunity-cost breakdown (Table 3).

      Opportunity-cost decomposition (Table 3): The dominant overhead is the lack of pipelined sampling (1.32 ms/token) — all other costs (control-layer scheduling: 0.05 ms, boundary crossing: 0.007 ms, Wasm processing: 0.001 ms) are negligible. This decomposition is the strongest analytical evidence that the architecture's indirection does not create fundamental bottlenecks.

      A formal model would have clarified:

      • At what batch size does horizontal batching saturate GPU utilization?
      • What is the theoretical maximum number of concurrent inferlets before control-layer becomes the bottleneck?
      • How does the work-conserving scheduling interact with varying API call latencies?

      6 minimum checks:

      1. ✅ Overhead decomposition validated (Table 3): 1.53 ms total overhead on 8B model, with each component measured.
      2. ✅ TPOT scales with model size (Table 4): overhead percentage decreases (11.41% → 5.64% → 2.39% for 1B → 3B → 8B).
      3. ✅ Adaptive batching validated (Table 5): 17× over Eager, demonstrating work-conserving policy effectiveness.
      4. ✅ Wasm cold start measured (Fig 9): 35–81 ms, negligible vs per-token latency.
      5. ✅ Control layer scales to 896 concurrent inferlets with <30 μs per call.
      6. ✅ Stacked optimization gains multiplicative (Fig 7): each optimization contributes independently.
      7. §5 实验与数据 #

        Setup: GCP G2 instance, NVIDIA L4 (24 GB), Llama 3 (1B/3B/8B), BF16. Baselines: vLLM v0.6.0, SGLang v0.4.4, LMQL v0.7.3, StreamingLLM. All use FlashInfer backend.

        Figure 6: Agentic workflow latency and throughput

        Paper's Figure 6, verbatim (caption: "Latency and throughput of LLM agents hosted by different serving systems. Numbers are normalized to the longest latency or the greatest throughput in each case.").

        Pie outperforms baselines on all three agentic patterns. The advantage is proportional to the ratio of I/O interactions to total tokens — no difference at <2 external interactions, gap widens linearly. On smaller models (1B, 3B), round-trip elimination dominates; on larger models (8B+), KV cache retention across interactions avoids costly re-prefills.

        Figure 7: Stacked application-specific optimizations

        Paper's Figure 7, verbatim (caption: "Performance gains via applying workload-specific optimizations to the simple agentic workflow. Stacked optimizations further improve the performance.").

        Three optimizations stack multiplicatively: (1) retain frequently-used API doc KV cache via export_kvpage, (2) concurrent API calls upon detecting callable signature, (3) drop one-use API spec KV cache via mask_kvpage. Combined: 3.5× throughput over baseline Python workflow on vLLM. This demonstrates Pie's core thesis — generic heuristics leave enormous performance on the table for heterogeneous workloads.

        Figure 8: Generation strategy comparison

        Paper's Figure 8, verbatim (caption: "Latency and throughput of example LLM inference techniques hosted by different serving systems.").

        On deliberate prompting (ToT, RoT, GoT, SkoT): up to 28% latency reduction, 34% throughput improvement. Advantage from program-controlled KV cache reuse (more precise than implicit management). On attention-level techniques (attention sink, windowed, hierarchical): 1.5× lower latency, 30× higher throughput vs StreamingLLM (partially influenced by kernel library differences).

        Where Pie loses: Beam search latency slightly worse than vLLM (Fig 8). 11.41% overhead on 1B models (Table 4) is significant for latency-sensitive small-model deployments.

        Workload characterization #

        Workload regimePieBaselineWhy
        Standard text completionComparable (~2–12% overhead)Baseline performanceDecomposition overhead without programmability benefit
        Agentic with many I/O interactions1.3–3.4× throughputPoor (round-trips, re-prefills)KV cache retention + integrated I/O eliminates round-trips
        Tree/graph reasoning with branching28–34% betterAdequate but suboptimalExplicit KV cache fork/reuse beats implicit policies
        Small models (1B) with simple tasksSlightly worse (11.4% overhead)BaselineFixed overhead dominates at low per-token latency

        §6 论证链 #

        StepClaimEvidenceDepends on
        1Monolithic loop is fundamentally inflexible for R1/R2/R3§2.2: beam search nearly removed from vLLM; round-trip KV discard; MCTS/grammar require invasive modifications
        2Decomposing into handlers + inferlets satisfies R1–R3Table 2: 19 distinct applications implemented in 22–255 LoC, covering all three requirementsStep 1
        3Adaptive batch scheduling preserves GPU efficiency despite decompositionTable 5: 17× over Eager, 8–40% over single-dimension batching; Table 3: dominant overhead is pipelining (1.32 ms), not architectureStep 2
        4Wasm provides adequate isolation with negligible overheadTable 3: 0.001 ms Wasm overhead; Fig 9: 35–81 ms cold start; 896 concurrent inferlets supportedStep 2
        5End-to-end performance competitive on standard tasks, superior on emerging tasksTables 3–4: 2.4–11.4% overhead; Figs 6–8: 1.1–3.4× improvements on agentic/reasoning workloadsSteps 3, 4

        §7 实现 cross-reference #

        Open source: https://github.com/pie-project/pie

        • 13,650 SLOC total: 11,640 lines Rust (core system + support library) + Python GPU handlers.
        • Wasm runtime: wasmtime with WASI for system interfaces.
        • GPU kernels: PyTorch + FlashInfer, communication via ZeroMQ (IPC).
        • Support library: Rust procedural macros + async runtime + Context abstraction for automatic KvPage management.
        • Native C++/CUDA implementation available (10–30% lower latency) but only supports Forward + InputText traits.

        关键实现细节:

        1. Pooled allocation for Wasm instances: wasmtime's pooled allocation pre-allocates virtual memory for up to 1,000 inferlet instances, enabling fast warm starts (10–50 ms with cached JIT binaries). Without this, cold start JIT compilation would add 35–81 ms per inferlet launch.
        2. Python single-threaded deserialization is the bottleneck: The inference layer's main overhead comes from Python-side API call deserialization (10–300 μs per call depending on concurrency), not from architectural indirection. A production deployment would benefit from Rust/C++ inference handlers.
        3. API & usability #

          • User-facing API: Custom 42-call API organized by Rust-like traits (Allocate, Forward, InputText, OutputText, Tokenize). Not OpenAI-compatible.
          • Languages: Any language compiling to Wasm (C++, Rust, Python via Wasm compilation).
          • Config surface: Minimal system-level config (GPU memory allocation at startup). Per-application logic entirely in inferlet code.
          • Migration cost: High — requires rewriting application logic as inferlets. Not a drop-in replacement for vLLM/SGLang.

          Deployment context #

          • Serving stage: Both prefill and decode.
          • Concurrency regime: Tested up to 896 concurrent inferlets — mid to high concurrency.
          • Hardware affinity: Evaluated on L4 (24 GB); multi-GPU support acknowledged as future work. Control layer currently centralized.
          • Ecosystem integration: Standalone system — not a plugin to vLLM/SGLang. Would need complete migration.
          • Migration path: Replace entire serving stack; rewrite application logic as Wasm inferlets.