sarathi-servechunked-prefillsllm-servingvllmgpu-schedulinginference-optimizationtutorial

Sarathi-Serve Guide: Chunked-Prefills & LLM Serving

Sarathi-Serve Architecture: Chunked-Prefills, Stall-Free Scheduling, and Piggybacked Decodes

The Tail-Latency Crisis in Multi-Turn LLM Serving

In high-concurrency artificial intelligence infrastructure, automated data fetching pipelines, and real-time enterprise research platforms powered by AiFetchly, serving Large Language Models (LLMs) requires balancing two opposing performance metrics: throughput (tokens per second per dollar) and latency Service Level Objectives (SLOs) (Time-to-First-Token and Inter-Token Latency).

Whether executing complex deep research queries, parsing massive web documents through high-speed web scraping, or powering interactive streaming chatbots, modern generative workloads exhibit extreme structural diversity in prompt lengths. A single serving cluster routinely processes concurrent requests where user prompts range from 50 tokens (concise conversational queries) to 32,000+ tokens (document summarization, repository analysis, and few-shot reasoning contexts).

However, modern continuous batching engines—including standard vLLM (Orca / PagedAttention) and TensorRT-LLM—encounter a fundamental systems dilemma known as the prefill-decode interference bottleneck.

The Traditional Continuous Batching Dilemma: Iteration N: Pure Decodes Batch (Low compute intensity, memory bandwidth bound) ┌──────────────┬──────────────┬──────────────┬──────────────┐ │ Decode Req 1 │ Decode Req 2 │ Decode Req 3 │ Decode Req 4 │ Duration: ~25 ms (Smooth ITL) └──────────────┴──────────────┴──────────────┴──────────────┘ │ ▼ A 4,096-token prefill arrives in queue! Iteration N+1: Monolithic Prefill Injected (Compute bound, massive GPU monopolization) ┌───────────────────────────────────────────────────────────┐ │ New Arrival: Full 4,096-Token Prompt Prefill Computation │ Duration: ~480 ms (STALL!) └───────────────────────────────────────────────────────────┘ │ ▼ Impact: All ongoing decodes are paused for nearly half a second! P99 Inter-Token Latency (ITL) jumps from 25ms to 505ms -> Catastrophic SLO violation!

1. Prefill vs. Decode: An Architectural Mismatch

Large language model inference comprises two radically distinct operational phases:

  • Prefill (Prompt Processing): Compute-bound parallel matrix multiplication. The entire user prompt is evaluated in parallel, achieving high arithmetic intensity and near-saturation of GPU Tensor Cores ($>100\text{ TFLOP/s}$). Evaluating a prompt of length $S$ scales quadratically $\mathcal{O}(S^2)$ in attention operations.
  • Decode (Autoregressive Token Generation): Memory-bandwidth-bound sequential execution. Generating one token per request requires transferring the entire model parameter set (e.g., 140 GB for an unquantized 70B model in FP16) from High-Bandwidth Memory (HBM) into on-chip cache registers. With batch size 1, Tensor Cores remain largely idle, throttled by the GPU memory bus bandwidth.

2. Head-of-Line Blocking and Severe ITL Jitter

In traditional continuous iteration-level scheduling, when a long prompt arrives, the scheduler either:

  1. Prioritizes the Prefill (Prefill Priority): The server schedules the full prompt in an upcoming iteration. This monopolizes the GPU for hundreds of milliseconds or seconds, starving all active decoding streams. Users experiencing smooth token streaming suddenly witness stuttering and frozen text generation.
  2. Prioritizes the Decodes (Decode Priority): The server delays the incoming prefill until ongoing decodes finish or yields compute in small slices, causing Time-to-First-Token (TTFT) to spike into unacceptable latencies.

This structural clash was empirically dissected by researchers from Georgia Institute of Technology and Microsoft Research in their seminal paper, Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve (USENIX OSDI 2024). Standard schedulers force cluster operators to make an unpalatable compromise: drastically over-provision expensive GPU hardware to maintain low tail latencies, or maximize throughput at the cost of severe interactive latency degradation.

The Sarathi-Serve Paradigm: Chunked-Prefills and Stall-Free Scheduling

To dismantle this trade-off, Sarathi-Serve introduces two complementary architectural innovations: Chunked-Prefills and Stall-Free Scheduling.

The Sarathi-Serve Co-Scheduling Engine: Original Large Prompt (e.g., 2,048 Tokens) ┌───────────────────────────────────────────────────────────────────────────┐ │ Chunk 1 (512 tokens) │ Chunk 2 (512 tokens) │ Chunk 3 (512) │ Chunk 4 (512)│ └──────────────────────┴──────────────────────┴───────────────┴──────────────┘ │ │ │ │ ▼ ▼ ▼ ▼ ┌──────────────────────┬──────────────────────┬───────────────┬──────────────┐ │ Batch Iteration t │ Batch Iteration t+1 │ Iteration t+2 │ Iteration t+3│ ├──────────────────────┼──────────────────────┼───────────────┼──────────────┤ │ Chunk 1 (512 tokens) │ Chunk 2 (512 tokens) │ Chunk 3 (512) │ Chunk 4 (512)│ │ + 16 Piggyback Decodes│ + 16 Piggyback Decodes│+ 16 Decodes │+ 16 Decodes │ ├──────────────────────┴──────────────────────┴───────────────┴──────────────┤ │ Deterministic Execution Budget: Exactly ~45 ms per iteration (Zero Jitter!) │ └────────────────────────────────────────────────────────────────────────────┘

1. Chunked-Prefills: Equalizing Token Budgets

Instead of treating prompt prefill as an indivisible monolithic job, Sarathi-Serve slices incoming prompts into equal, uniform token chunks (e.g., chunks of size $C = 512$ or $C = 1,024$ tokens).

When processing Chunk $k$ of a prompt:

  1. The model computes Key-Value tensors for tokens within Chunk $k$.
  2. It queries past Key-Value states from Chunks $1 \dots k-1$ stored in the KV cache using causal attention masks.
  3. The newly generated KV states for Chunk $k$ are committed to the page cache.
  4. Only upon finishing the final Chunk does the request transition into the decode state and emit its first generated output token.

Because each chunk executes in predictable, bounded GPU time, long prompts no longer trigger massive multi-hundred-millisecond interruptions.

2. Piggybacking Decodes onto Chunked Prefills

Crucially, chunking prefill prompts enables an optimal computational symbiosis: piggybacking.

Pure decode batches are memory-bandwidth-bound (underutilizing Tensor Core arithmetic units), while pure prefill batches are compute-bound (saturating ALUs). By co-scheduling a compute-heavy prefill chunk alongside a batch of memory-bound decode requests within the same GPU forward pass:

  • The decode tokens “piggyback” on the memory transfers and compute pipelines of the prefill chunk.
  • High-Bandwidth Memory reads for model weights are amortized across both the prefill chunk and decode tokens.
  • Overall GPU compute utilization increases significantly without increasing total iteration latency.

3. Stall-Free Scheduling

In conventional batching, memory exhaustion can force the scheduler to preempt running decodes, swapping their KV caches to CPU host memory or discarding them entirely. Sarathi-Serve formulates a stall-free schedule that guarantees ongoing decode requests are never starved or preempted due to incoming prefill arrivals, enforcing strict upper bounds on P99 Inter-Token Latency.

Deep Dive: Mathematical Formulation and Attention Mechanics

To understand how Sarathi-Serve executes attention across heterogeneous chunks and decode tokens simultaneously, we examine the underlying mathematical structure of hybrid forward passes.

Attention Complexity Across Mixed Batches

Consider a mixed batch containing $B_d$ decode requests and a chunked prefill of length $C$ with an accumulated prefix context of length $L$.

For a decode request $i$ with historical sequence length $S_i$, the query tensor has sequence dimension 1: $$Q_{\text{decode}}^{(i)} \in \mathbb{R}^{1 \times d_k}$$ The Key and Value tensors encompass the entire historical cache plus the current token: $$K_{\text{decode}}^{(i)}, V_{\text{decode}}^{(i)} \in \mathbb{R}^{(S_i + 1) \times d_k}$$

For the chunked prefill request, the query tensor spans the entire current chunk length $C$: $$Q_{\text{chunk}} \in \mathbb{R}^{C \times d_k}$$ However, the Key and Value matrices span both the preceding prefix tokens $L$ and the current chunk tokens $C$: $$K_{\text{chunk}}, V_{\text{chunk}} \in \mathbb{R}^{(L + C) \times d_k}$$

The hybrid attention computation is unified through block-diagonal sparse masking: $$\text{Attention}(Q_{\text{batch}}, K_{\text{batch}}, V_{\text{batch}}) = \text{softmax}\left(\frac{Q_{\text{batch}} K_{\text{batch}}^T + M}{\sqrt{d_k}}\right) V_{\text{batch}}$$

where the attention mask $M$ enforces strict causal and isolation constraints:

$$ M_{i, j} = \begin{cases} 0 & \text{if token } j \text{ belongs to the valid prefix or causal history of token } i \ -\infty & \text{if token } j \text{ belongs to another request or future token within the same chunk} \end{cases} $$

The Unified Hybrid Attention Matrix: Prefix Tokens (L) Chunk Tokens (C) Decode Tokens (Bd) ┌───────────────────┬──────────────────┬────────────────────┐ Chunk Queries │ Dense Attention │ Causal Truncated│ -inf (Cross-req │ (C tokens) │ to Past KV Cache│ Lower-Triangular│ Isolation) │ ├───────────────────┼──────────────────┼────────────────────┤ Decode Queries│ -inf (Isolated) │ -inf (Isolated) │ Diagonal Sparse │ (Bd tokens) │ │ │ Self + Hist Cache │ └───────────────────┴──────────────────┴────────────────────┘

By utilizing high-performance variable-length attention kernels (such as FlashAttention-2 or FlashInfer), Sarathi-Serve computes this hybrid attention layout in a single fused kernel without materializing zero-padded tensors.

Determining the Optimal Chunk Size $C^*$

Selecting the prefill chunk size $C$ presents a trade-off:

  • Small Chunk Size ($C < 256$): Lower TTFT interference, but reduced GPU compute utilization and increased kernel launch overhead per prompt.
  • Large Chunk Size ($C > 2048$): Higher peak throughput, but larger iteration latency steps that re-introduce decode jitter.

Sarathi-Serve models the GPU execution time $T(C, B_d)$ as a function of chunk tokens $C$ and active decode requests $B_d$: $$T(C, B_d) = \alpha_{\text{weight_read}} \cdot \frac{W_{\text{model}}}{\text{BW}{\text{HBM}}} + \max\left( \frac{2 \cdot C \cdot W{\text{model}}}{\text{Peak FLOPS}}, \frac{2 \cdot B_d \cdot W_{\text{model}}}{\text{Peak FLOPS}} \right) + T_{\text{attn}}(C, L, B_d)$$

The optimal chunk size $C^$ saturates the GPU compute roofline precisely up to the target SLO budget $T_{\text{target}}$: $$C^ = \arg\max_C \left&#123; C ;\middle|; T(C, B_d) \le T_{\text{SLO_budget}} \right&#125;$$

Empirically, on NVIDIA A100 (80GB SXM4) and H100 architectures, $C^* = 512$ or $C^* = 1024$ tokens represents the sweet spot, yielding predictable iterations between 35ms and 55ms while saturating tensor core pipelines.

End-to-End Implementation: A Chunked-Prefill Co-Scheduler in Python

To demonstrate the core mechanisms of Sarathi-Serve, we construct a production-grade Python reference implementation using PyTorch and dataclass abstractions. This scheduler manages prompt chunking, hybrid batch formation, KV-cache prefix tracking, and piggybacked token generation.

""" sarathi_serve_scheduler.py Production Reference Implementation: Chunked-Prefill and Decode Co-Scheduler Architectural concept based on Sarathi-Serve (USENIX OSDI 2024). """ from dataclasses import dataclass, field from enum import Enum from typing import List, Optional, Tuple, Dict import torch import torch.nn as nn import time class RequestPhase(Enum): PREFILL = "PREFILL" DECODE = "DECODE" COMPLETED = "COMPLETED" @dataclass class ServingRequest: request_id: str prompt_tokens: List[int] max_new_tokens: int created_at: float = field(default_factory=time.time) # State tracking phase: RequestPhase = RequestPhase.PREFILL processed_prefill_tokens: int = 0 generated_tokens: List[int] = field(default_factory=list) inter_token_latencies: List[float] = field(default_factory=list) last_token_timestamp: float = 0.0 @property def total_prompt_length(self) -> int: return len(self.prompt_tokens) @property def remaining_prefill_tokens(self) -> int: return self.total_prompt_length - self.processed_prefill_tokens def is_prefill_complete(self) -> bool: return self.processed_prefill_tokens >= self.total_prompt_length def is_finished(self) -> bool: return len(self.generated_tokens) >= self.max_new_tokens @dataclass class ScheduledIterationBatch: # Chunked prefill slice chunked_req: Optional[ServingRequest] chunk_tokens: List[int] chunk_prefix_length: int # Piggybacked decode batch decode_reqs: List[ServingRequest] decode_tokens: List[int] class SarathiCoScheduler: """ Sarathi-Serve Co-Scheduler: Slices large incoming prompt requests into uniform token chunks and co-schedules them with concurrent decode requests. """ def __init__(self, chunk_size: int = 512, max_batch_tokens: int = 2048): self.chunk_size = chunk_size self.max_batch_tokens = max_batch_tokens # Scheduling Queues self.waiting_prefill_queue: List[ServingRequest] = [] self.active_decode_queue: List[ServingRequest] = [] self.completed_requests: List[ServingRequest] = [] def add_request(self, req: ServingRequest) -> None: self.waiting_prefill_queue.append(req) def schedule_next_iteration(self) -> Optional[ScheduledIterationBatch]: """ Forms the next iteration batch by combining at most ONE chunked prefill with all pending decodes within memory budget constraints. """ if not self.waiting_prefill_queue and not self.active_decode_queue: return None scheduled_chunk_req: Optional[ServingRequest] = None scheduled_chunk_tokens: List[int] = [] chunk_prefix_len: int = 0 # 1. Select and slice next prefill request if tokens budget allows if self.waiting_prefill_queue: curr_req = self.waiting_prefill_queue[0] remaining = curr_req.remaining_prefill_tokens tokens_to_take = min(self.chunk_size, remaining) start_idx = curr_req.processed_prefill_tokens end_idx = start_idx + tokens_to_take scheduled_chunk_tokens = curr_req.prompt_tokens[start_idx:end_idx] chunk_prefix_len = curr_req.processed_prefill_tokens scheduled_chunk_req = curr_req # 2. Select decode requests that fit into remaining token budget scheduled_decodes: List[ServingRequest] = [] scheduled_decode_tokens: List[int] = [] available_tokens = self.max_batch_tokens - len(scheduled_chunk_tokens) for req in self.active_decode_queue: if len(scheduled_decodes) < available_tokens: scheduled_decodes.append(req) # Decode input token is the most recently generated or last prompt token last_tok = req.generated_tokens[-1] if req.generated_tokens else req.prompt_tokens[-1] scheduled_decode_tokens.append(last_tok) else: break return ScheduledIterationBatch( chunked_req=scheduled_chunk_req, chunk_tokens=scheduled_chunk_tokens, chunk_prefix_length=chunk_prefix_len, decode_reqs=scheduled_decodes, decode_tokens=scheduled_decode_tokens ) def step_iteration(self, batch: ScheduledIterationBatch) -> None: """ Updates request states following physical forward step execution. """ now = time.time() # Update chunked prefill request if batch.chunked_req is not None: batch.chunked_req.processed_prefill_tokens += len(batch.chunk_tokens) if batch.chunked_req.is_prefill_complete(): # Pop from prefill queue and promote to decode queue promoted_req = self.waiting_prefill_queue.pop(0) promoted_req.phase = RequestPhase.DECODE promoted_req.last_token_timestamp = now # Emit initial token from prefill completion promoted_req.generated_tokens.append(42) # Simulated emitted token self.active_decode_queue.append(promoted_req) # Update decode requests finished_decode_indices = [] for i, req in enumerate(batch.decode_reqs): # Record Inter-Token Latency (ITL) itl = (now - req.last_token_timestamp) * 1000.0 # ms req.inter_token_latencies.append(itl) req.last_token_timestamp = now # Emit newly decoded token req.generated_tokens.append(100 + len(req.generated_tokens)) if req.is_finished(): req.phase = RequestPhase.COMPLETED self.completed_requests.append(req) finished_decode_indices.append(req) # Remove completed requests from active decode queue for req in finished_decode_indices: if req in self.active_decode_queue: self.active_decode_queue.remove(req) # Execution Simulation if __name__ == "__main__": scheduler = SarathiCoScheduler(chunk_size=512, max_batch_tokens=2048) # 1. Add 4 interactive short decode requests (simulated already in decode phase) for i in range(4): req = ServingRequest( request_id=f"decode_short_{i}", prompt_tokens=[1] * 64, max_new_tokens=8, phase=RequestPhase.DECODE, processed_prefill_tokens=64, generated_tokens=[10], last_token_timestamp=time.time() ) scheduler.active_decode_queue.append(req) # 2. Add a massive 2,048-token document extraction prefill request large_req = ServingRequest( request_id="doc_extract_large_0", prompt_tokens=[99] * 2048, max_new_tokens=10 ) scheduler.add_request(large_req) print(f"=== Initialized Sarathi-Serve Scheduler ===") print(f"Active Decodes: {len(scheduler.active_decode_queue)} | Queued Prefills: {len(scheduler.waiting_prefill_queue)}") # Execute scheduler loop iteration = 1 while scheduler.waiting_prefill_queue or scheduler.active_decode_queue: batch = scheduler.schedule_next_iteration() if batch is None: break prefill_info = f"Prefill Chunk ({len(batch.chunk_tokens)} tokens)" if batch.chunked_req else "None" print(f"Iteration {iteration:02d} -> {prefill_info} co-scheduled with {len(batch.decode_reqs)} Decodes.") # Simulate ~40ms forward execution time.sleep(0.04) scheduler.step_iteration(batch) iteration += 1 print("\n=== All Requests Completed ===") for req in scheduler.completed_requests: avg_itl = sum(req.inter_token_latencies) / max(len(req.inter_token_latencies), 1) print(f"Req: {req.request_id} | Gen Tokens: {len(req.generated_tokens)} | Avg ITL: {avg_itl:.2f} ms")

Key Architectural Traits of the Implementation

  1. Granular Chunk Slicing: A 2,048-token prompt is cleanly partitioned into four 512-token increments across consecutive scheduling iterations (doc_extract_large_0 takes 4 scheduled turns to complete prefill).
  2. Zero Decode Starvation: At every step, the active decode requests continue executing concurrently without being paused or evicted.
  3. Bounded Iteration Time: By fixing the maximum token footprint ($C + B_d \le \text{Budget}$), individual step latencies remain tightly bounded, protecting downstream SLOs.

For enterprise teams integrating modern inference engines into custom microservices or building data extraction workflows with custom API integrations, this scheduling model forms the backbone of reliable real-time AI.

Production Engineering Best Practices

Deploying chunked-prefill architectures in high-throughput enterprise environments requires addressing several low-level systems trade-offs.

+-----------------------------------------------------------------------------------+ | SARATHI-SERVE PRODUCTION DEPLOYMENT CHECKLIST | +-----------------------------------------------------------------------------------+ | 1. Attention Kernel Selection: | | - Deploy FlashInfer or FlashAttention-2 with native variable-length support | | - Avoid padding ragged batches to max_len (eliminates memory overhead) | | | | 2. Paged KV-Cache Page Sizing: | | - Use 16-token or 32-token virtual blocks to prevent fragmentation | | - Pre-allocate contiguous chunk blocks during prefill stages | | | | 3. Distributed Tensor Parallelism (TP): | | - Synchronize chunk boundaries across all TP ranks | | - Execute All-Reduce for intermediate layer projections across chunk tokens | | | | 4. Prefix Caching Synergy: | | - Reuse prompt prefixes across shared system instructions | | - Chunk only the dynamic suffix of the incoming request | +-----------------------------------------------------------------------------------+

1. Fused Ragged Batch Attention Kernels

Never construct hybrid batches by padding individual request sequences to the maximum sequence length. In mixed batches where decode requests have sequence lengths of $1$ and the prefill chunk has length $512$, padding creates massive memory inflation and wasted FLOPs. Production systems rely on ragged batch representations where sequence metadata is managed via a cumulative token offsets array (cu_seqlens).

2. Paged Memory Block Management

When a request’s prompt is chunked, its KV-cache pages are allocated progressively. Ensure the memory manager locks allocated physical blocks so they cannot be evicted before subsequent chunks are evaluated. Review our infrastructure cost breakdown to evaluate how memory efficiency translates directly to reduced GPU cluster expenditure.

3. Distributed Tensor Parallelism Coordination

In multi-GPU deployments utilizing Tensor Parallelism (TP = 2, 4, or 8), all GPU ranks must execute identical token slices simultaneously. The scheduler must run on the orchestrator rank (Rank 0) and broadcast the batch layout (chunk_id, cu_seqlens, decode_batch_size) to all peer worker ranks to avoid desynchronization barriers.

Empirical Benchmark Evaluation

The USENIX OSDI 2024 paper evaluated Sarathi-Serve across representative open-source models (LLaMA-13B, LLaMA-70B, Falcon-40B) and heterogeneous enterprise workloads (ShareGPT conversational traces, CodeX completion, and ArXiv summarization).

Performance Comparison: Sarathi-Serve vs. Standard Continuous Batching (vLLM / Orca) Workload: LLaMA-70B on 4x NVIDIA A100-80GB GPUs | Strict SLO: P99 ITL < 50ms Normalized Serving Capacity (Requests/sec under SLO) vLLM (Orca Baseline) ████████ (1.0x Baseline) Splitwise / Static ████████████ (1.5x) Sarathi-Serve (OSDI) ████████████████████████████████████████ (5.6x Capacity Boost!) P99 Inter-Token Latency (ITL) Spike During 4k Prefill Arrival vLLM (Orca Baseline) ████████████████████████████████████████ (495 ms - Severe Jitter!) Sarathi-Serve (OSDI) ████ (42 ms - Rock Solid Predictability!)

1. Throughput Under Strict SLOs

Under loose latency budgets, traditional schedulers can batch aggressively. However, as enterprise SLOs tighten (e.g., demanding that 99% of generated tokens stream out in under $50\text{ ms}$), traditional schedulers collapse: they must throttle incoming prefill requests or restrict batch sizes. Sarathi-Serve achieves 2.6x to 5.6x higher serving capacity because its predictable iteration runtime eliminates tail-latency degradation.

2. Elimination of Preemption Swapping

In standard serving environments, unexpected prefill spikes consume remaining GPU memory, triggering costly request preemptions (swapping active KV caches to CPU RAM). Because Sarathi-Serve smooths memory consumption into predictable chunk increments, request preemptions are reduced to virtually zero, saving massive PCIe interconnect bandwidth. Detailed configurations are available in our developer documentation.

Conclusion: The Foundation of Modern Inference Engines

The architectural principles introduced by Sarathi-Serve represent a permanent transition in LLM serving systems. Following its publication at USENIX OSDI 2024, the concept of chunked-prefills was rapidly incorporated into mainstream serving frameworks, including vLLM (--enable-chunked-prefill), TensorRT-LLM, and SGLang.

By replacing monolithic prompt execution with uniform chunked scheduling and co-scheduling compute-bound prefills with memory-bound decodes, inference engineers can:

  • Eradicate Inter-Token Latency spikes and ensure silky-smooth text streaming.
  • Amortize high-bandwidth memory reads, pushing GPU arithmetic utilization close to the physical hardware roofline.
  • Multiply cluster serving capacity, slashing cloud infrastructure costs while meeting mission-critical enterprise SLOs.

As generative AI workloads scale into million-token document understanding and real-time reasoning agents, fine-grained co-scheduling architectures like Sarathi-Serve will continue to define the frontier of high-performance deep learning systems. Explore how AiFetchly utilizes cutting-edge inference and data pipelines to power next-generation enterprise AI.