7.8 Disaggregated Prefill/Decode & Chunked Prefill¶

Autoregressive generation has two fundamentally different compute phases: prefill, which processes the entire prompt in parallel, and decode, which generates one token at a time. These phases live inside the same GPU in almost every deployed system — but they have radically different resource profiles. Mixing them together on the same hardware leads to a class of performance problems that have quietly been the biggest source of latency waste in production LLM serving since the first continuous-batching schedulers were deployed.

This chapter tears apart the prefill/decode interference problem, explains why separating these phases onto different hardware pools can dramatically improve both latency and throughput, and covers the practical engineering — KV-cache transfer, chunked prefill scheduling, and the systems DistServe and Splitwise that have formalized these ideas. Readers already comfortable with continuous batching (Continuous Batching & Request Scheduling) and PagedAttention (PagedAttention & KV-Cache Memory Management) will find this chapter the natural next step toward cutting-edge serving system design.

The Two Phases Are Nothing Alike¶

Before we can appreciate why disaggregation helps, we need a precise picture of what each phase actually does to the hardware.

Prefill: Compute-Bound Matrix Multiplications¶

During prefill, the server processes a sequence of length \(T_p\) (the prompt length) in a single forward pass. The dominant cost is the projection matrices in each attention layer and each MLP layer, which see batched matrix multiplications of shape \([T_p, d_\text{model}]\) against \([d_\text{model}, d_\text{ff}]\). With a long prompt, \(T_p\) is large, the matrices are fat, and modern GPUs can achieve high arithmetic intensity — often north of 100–200 FLOPs/byte — pushing utilization close to the peak TFLOP ceiling.

The arithmetic intensity for a single linear layer with input \([B \cdot T_p, d]\) and weight \([d, d']\) is roughly:

Dividing through, this rearranges to \(I_\text{prefill} = \dfrac{2}{\text{bytes\_per\_element} \cdot \left(1/d' + 1/(B T_p)\right)}\): intensity rises with the token count \(B \cdot T_p\) and saturates at \(2 d' / \text{bytes\_per\_element}\). In BF16 with \(d' = 4096\), a batch of \(B T_p = 512\) tokens gives \(\approx 455\) FLOP/byte — comfortably above an H100’s ridge point of \(\approx 295\) FLOP/byte (989 TFLOP/s dense BF16 divided by 3.35 TB/s HBM), so the GEMM is firmly compute-bound. Attention contributes a further \(O(T_p^2)\) term on top, which only pushes prefill deeper into the compute-bound regime as prompts grow. Keep that ridge point in mind: it reappears in the chunk-size analysis below, where a chunk that is too small drops the same GEMMs back below it.

Decode: Memory-Bandwidth-Bound Vector-Matrix Multiplications¶

During decode, we generate one token at a time. The batch dimension is the number of concurrent sequences, \(B_d\). Each layer executes a matrix-vector multiply: shape \([B_d, d_\text{model}]\) times \([d_\text{model}, d']\). Unless \(B_d\) is in the hundreds, this is entirely memory-bandwidth-bound: we stream gigabytes of weights from HBM every step just to do a handful of FLOPs per byte.

For a 70B parameter model in BF16 (140 GB of weights), each decode step reads roughly 140 GB of weights from HBM regardless of batch size — plus the KV cache of every sequence in the batch, which unlike the weights does grow with batch size and context length, and eventually dominates at long context. On an A100 SXM with HBM bandwidth of roughly 2 TB/s, that’s on the order of 70 ms of pure bandwidth time per step — leaving essentially no room for compute to hide. An H100 (3.35 TB/s bandwidth) gets this to around 40 ms in theory, and a 2026-frontier Blackwell B200 (8 TB/s HBM3e) to roughly 17 ms — but the decode phase stays bandwidth-bound on every generation of hardware.

The key insight: prefill wants to be compute-bound and loves big batches; decode wants high bandwidth and is bottlenecked by memory, not FLOPs. These constraints point toward different hardware configurations: fewer, faster GPUs (or GPUs with high TFLOP/s) for prefill, and more, bandwidth-rich GPUs for decode.

The Interference Problem¶

Now imagine both phases run on the same GPU pool under a continuous batching scheduler (see Continuous Batching & Request Scheduling). Every iteration of the inference engine processes a mixed batch: some sequences are in prefill, others are in decode. This mixing has two serious consequences:

1. Prefill preempts decode. A long incoming prompt (say 8,192 tokens) takes tens of milliseconds to prefill. During that time, every decode-only request that was already generating is paused. The time-to-next-token for existing users spikes unexpectedly, violating service-level objectives (SLOs) on P99 latency.

2. Decode throughput degrades under mixed batches. When a decode step includes prefill tokens from new requests, the attention kernel must handle variable-length sequences with very different KV-cache patterns. GPU kernel efficiency drops; memory fragmentation increases; the effective batch size for decode shrinks.

This is the prefill-decode interference problem. It is not a theoretical concern — production teams have measured it causing P99 time-to-first-token (TTFT) and P99 inter-token latency (ITL) to exceed SLOs by 2–5x under moderate load.

Disaggregated Prefill/Decode: The Architecture¶

The core idea of disaggregation is simple to state: run prefill on a dedicated pool of instances (prefill workers), run decode on a separate pool (decode workers), and transfer the KV cache between them.

Lifecycle of a Request¶

1. A request arrives at the router with a prompt of \(T_p\) tokens.
2. The router schedules it on a prefill worker. The prefill worker runs a single forward pass over the entire prompt, producing the first output token and — critically — the full KV cache for layers \(1 \ldots L\).
3. The KV cache (a tensor of shape \([L, 2, T_p, n_\text{kv\_heads}, d_\text{head}]\)) is transmitted to a decode worker via a high-speed interconnect (NVLink, RDMA over InfiniBand, or PCIe depending on cluster topology).
4. The decode worker receives the KV cache, places it into its paged KV-cache allocator, and begins generating tokens one step at a time — without interfering with any prefill computation.
5. Generated tokens stream back to the client.

Why This Helps¶

Decode workers never see a long prefill stall. Every step they execute is a pure decode step over a batch of active sequences. Time-to-next-token becomes predictable and low. Meanwhile, prefill workers handle new requests at maximum compute efficiency — no decode sequences competing for memory bandwidth.

KV Cache Transfer: The Engineering Challenge¶

Transferring KV caches is not free. For a single request with prompt length \(T_p = 4096\), a 70B model with 80 layers, 8 KV heads (GQA), and \(d_\text{head} = 128\), storing in BF16:

Transferring that at NVLink4 speeds (~900 GB/s between two H100s on the same NVSwitch fabric) takes about 1.5 ms. Over InfiniBand HDR100 (100 Gb/s ≈ 12.5 GB/s), it takes over 100 ms — longer than the prefill itself for many prompts. The transfer mechanism fundamentally shapes the architecture:

For real deployments, the recommendation is: keep prefill and decode workers on the same NVSwitch fabric (same node or adjacent nodes connected via NVSwitch) to keep transfer latency under a few milliseconds. If that is not possible, pipeline the transfer with decode (start decoding even as later layers’ KV caches arrive) to overlap transfer and computation.

Layer-Wise Pipelining¶

A practical optimization is to transfer KV caches layer by layer as they are computed during prefill, not as a single bulk transfer at the end. The decode worker still cannot start its first step until layer \(L\)’s cache has arrived, so the win is not that decode starts earlier in the layer order — it is that the transfers of layers \(1 \ldots L-1\) are hidden underneath the prefill compute of layers \(2 \ldots L\). With bulk transfer the exposed cost is the whole \(L\)-layer cache; with layer-wise streaming it is only the last layer’s slice, roughly \(1/L\) of it, provided the per-layer transfer time is shorter than the per-layer prefill time. Concretely, for the 70B/4K example above (\(1.34\) GB over 80 layers \(\approx 17\) MB per layer), NVLink4 moves one layer’s slice in roughly 19 µs, while that layer’s prefill compute — about \(2 \times 875\text{M} \times 4096 \approx 7.2\) TFLOPs, spread over an 8-way tensor-parallel group — is on the order of a millisecond. Transfer is two orders of magnitude cheaper, so it disappears entirely. On IB HDR100 the same slice takes about 1.3 ms, i.e. comparable to the per-layer compute: the link saturates, pipelining buys you little, and the transfer becomes a first-order cost. That ratio — per-layer transfer time versus per-layer prefill time — not the pipelining trick itself, is what decides whether cross-node disaggregation is viable.

# Pseudocode: layer-wise KV transfer from prefill worker to decode worker
# (simplified, assumes a hypothetical RPC / RDMA abstraction)

import torch
from typing import Tuple

class PrefillWorker:
    def __init__(self, model, kv_sender):
        self.model = model
        self.kv_sender = kv_sender  # e.g. a RDMA/NVLink send handle

    def prefill_and_stream_kv(
        self,
        input_ids: torch.Tensor,   # [1, T_p]
        request_id: str,
    ) -> torch.Tensor:
        """
        Run prefill layer-by-layer, streaming each layer's KV
        to the paired decode worker as we go.
        Returns the first output token (greedy) so TTFT is fast.
        """
        x = self.model.embed(input_ids)          # [1, T_p, d_model]
        first_token = None

        for layer_idx, layer in enumerate(self.model.layers):
            # Standard attention + MLP forward
            x, kv_cache = layer.forward_with_kv(x)
            # kv_cache shape: [2, T_p, n_kv_heads, d_head]

            # Fire-and-forget async send — does NOT block prefill forward pass
            self.kv_sender.send_async(
                request_id=request_id,
                layer_idx=layer_idx,
                kv=kv_cache,
            )

        # Compute logits only for last position (first output token)
        logits = self.model.lm_head(x[:, -1, :])   # [1, vocab]
        first_token = logits.argmax(dim=-1)          # greedy; real systems sample
        return first_token


class DecodeWorker:
    def __init__(self, model, kv_receiver, paged_kv_manager):
        self.model = model
        self.kv_receiver = kv_receiver
        self.kv_mgr = paged_kv_manager

    def receive_kv_and_decode(
        self,
        request_id: str,
        first_token: torch.Tensor,
        max_new_tokens: int,
    ):
        """
        Wait for all layers' KV caches to arrive, then decode.
        In a real system this overlaps with prefill's later layers.
        """
        # Block until all L layers have been received
        kv_caches = self.kv_receiver.collect(request_id)
        # kv_caches: list of [2, T_p, n_kv_heads, d_head] tensors, one per layer

        # Allocate paged KV slots and copy into the page table
        slot = self.kv_mgr.allocate(request_id, kv_caches)

        generated = [first_token.item()]
        cur_token = first_token
        for _ in range(max_new_tokens - 1):
            # Pure decode step: append current token's KV to each layer's cache
            logits = self.model.decode_step(cur_token, slot)
            cur_token = logits.argmax(dim=-1)
            generated.append(cur_token.item())
            if cur_token.item() == self.model.eos_id:
                break

        return generated

Chunked Prefill: Serving Both Phases on One Pool¶

Disaggregation requires separate hardware pools and a network transfer path — a significant operational complexity. Chunked prefill is the middle-ground technique that keeps prefill and decode on the same GPU but breaks long prefills into small chunks, interleaving them with decode steps so no single iteration monopolizes the GPU.

The Basic Idea¶

Instead of processing a 16K-token prompt in one monolithic forward pass (which would stall decode for many tens of milliseconds), we split the prompt into chunks of at most \(C\) tokens each, say \(C = 512\). Each inference iteration processes:

• One chunk of the current prompt (partial prefill), contributing \(C\) tokens to the KV cache.
• All decode tokens from in-flight sequences (one new token each).

The decode sequences are never fully preempted. Their ITL (inter-token latency) increases only by the overhead of the prefill chunk, not by the full prompt length. This converts the bursty, unpredictable P99 latency problem into a smooth, bounded one.

Choosing the Chunk Size \(C\)¶

The chunk size \(C\) is a critical knob:

• Too large: each iteration still stalls decode sequences for too long. If \(C = 4096\), the P99 ITL spikes are only 4x better than without chunking.
• Too small: the prefill is broken into so many chunks that the total time to complete prefill (TTFT) grows. Each chunk incurs per-iteration overhead (kernel launches, scheduling, KV-cache bookkeeping). Additionally, attention kernels are less efficient on shorter sequences — you leave FLOP/s on the table.

“Too small” has a sharp quantitative meaning, and it is the roofline calculation from the opening section. Every prefill iteration re-reads the model’s weight matrices from HBM, so splitting a prompt into \(n = T_p / C\) chunks streams the full weight set \(n\) times instead of once. That is only free if each iteration is still compute-bound — i.e. if the per-iteration token count (the chunk’s \(C\) tokens plus the decode tokens piggybacked into the same batch) keeps the GEMMs above the hardware’s ridge point of a few hundred FLOP/byte. Using the earlier formula, that means a few hundred tokens per iteration on an H100 in BF16. Go below that and prefill itself becomes bandwidth-bound, and you pay the penalty on every chunk. This is precisely why Sarathi-Serve schedules against a token budget per iteration rather than a fixed chunk count: it packs all pending decode tokens into the batch first, then tops up with prefill tokens until the budget is reached, so every batch sits at or above the saturation point. In vLLM that budget is exactly max_num_batched_tokens.

A second subtlety: chunks are not equally expensive. Chunk \(k\) runs the same GEMMs as every other chunk but attends over \(kC\) already-cached tokens, so its attention cost grows linearly in \(k\). Summed over all chunks the attention work is \(\sum_{k} C \cdot (k{+}1) C \approx T_p^2/2\) — identical to a monolithic prefill, so chunking adds no asymptotic work — but the last chunk of a 32K prompt is far more expensive than the first. A fixed \(C\) therefore produces ITL spikes that grow through a long prefill. Production schedulers respond either by shrinking \(C\) as a prompt progresses or by budgeting on estimated chunk time rather than token count.

In practice, production systems tune \(C\) per deployment based on their latency SLOs. Values in the range \(C \in [256, 2048]\) are common. The scheduler can also make \(C\) dynamic: use large chunks when the decode batch is empty (no one is waiting) and small chunks when many decode sequences are active.

Attention on Partial KV Caches¶

A subtle implementation detail: when processing chunk \(k\) of a prompt, the attention layer must attend over the KV cache built from chunks \(0 \ldots k-1\). This is not the same as standard decode (which attends over a complete past KV cache). It is also not the same as full prefill (which attends over the full prompt). The attention mask must reflect that:

• Tokens in chunk \(k\) can attend to all previous prompt tokens (chunks \(0 \ldots k-1\)) that have been cached.
• Tokens in chunk \(k\) can attend to earlier tokens within the same chunk (causal masking within the chunk).
• Tokens in chunk \(k\) cannot attend to later prompt tokens (not yet processed).

This requires a custom attention mask and careful block layout in the KV cache. Systems like vLLM (see vLLM: Architecture, PagedAttention & Internals) implement chunked prefill by extending their paged KV cache manager to handle mid-prompt KV writes and mid-sequence mask construction.

import torch
import torch.nn.functional as F

def chunked_prefill_attention(
    q_chunk: torch.Tensor,    # [C, n_heads, d_head] — queries for current chunk
    k_full: torch.Tensor,     # [T_past + C, n_kv_heads, d_head] — all keys so far
    v_full: torch.Tensor,     # [T_past + C, n_kv_heads, d_head]
    T_past: int,              # tokens already in KV cache (from previous chunks)
    C: int,                   # chunk size
    scale: float,
) -> torch.Tensor:
    """
    Compute attention for a chunk of prefill tokens.

    The mask allows each query at position T_past+i to attend to
    positions 0 .. T_past+i (standard causal), but NOT T_past+i+1 ..
    T_past+C-1 (future tokens in the same chunk).
    """
    n_heads, d_head = q_chunk.shape[1], q_chunk.shape[2]

    # Expand GQA: if n_kv_heads < n_heads, repeat KV heads
    # (omitted for brevity — same as decode)

    # Build causal mask for the chunk against the full context
    # Shape: [C, T_past + C]
    T_total = T_past + C
    causal_mask = torch.ones(C, T_total, dtype=torch.bool)
    for i in range(C):
        # query at position T_past + i can see positions 0 .. T_past + i
        causal_mask[i, T_past + i + 1:] = False

    # Attention scores: [n_heads, C, T_total]
    q = q_chunk.transpose(0, 1)     # [n_heads, C, d_head]
    k = k_full.transpose(0, 1)      # [n_kv_heads, T_total, d_head]
    # (assume n_heads == n_kv_heads for clarity)
    scores = torch.bmm(q, k.transpose(1, 2)) * scale   # [n_heads, C, T_total]

    # Apply causal mask (broadcast over head dim)
    scores = scores.masked_fill(
        ~causal_mask.unsqueeze(0),  # [1, C, T_total]
        float('-inf'),
    )

    attn = F.softmax(scores, dim=-1)    # [n_heads, C, T_total]
    v = v_full.transpose(0, 1)          # [n_kv_heads, T_total, d_head]
    out = torch.bmm(attn, v)            # [n_heads, C, d_head]
    return out.transpose(0, 1)          # [C, n_heads, d_head]

DistServe: Formalizing Disaggregation¶

DistServe (Zhong et al., 2024) is the landmark paper that formally analyzed and implemented disaggregated prefill/decode serving. Its key contributions:

1. Quantified the interference problem with measurements showing that mixed batches cause P99 TTFT to grow proportionally with the longest prefill in the batch, and P99 ITL to grow as the decode batch is interrupted by prefill work.

2. Proposed resource allocation as an optimization problem: given a fleet of GPUs, how many should be assigned to the prefill pool vs. the decode pool? The answer depends on the workload’s prompt-to-output ratio. Long-prompt workloads (RAG pipelines, document summarization) need more prefill capacity; chatbot workloads with short prompts and long outputs need more decode capacity.

3. Implemented KV transfer via RDMA with layer-wise pipelining to overlap transfer with computation, achieving near-zero transfer overhead on NVLink-connected nodes.

4. Demonstrated SLO attainment: under tight latency SLOs, disaggregated serving sustains substantially more requests per second than mixed serving while keeping P99 TTFT and P99 ITL within bounds — the paper reports up to 7.4× on its most favourable workload; 2–4× is the more typical range practitioners report on mixed production traffic.

The optimization problem DistServe solves is (informally):

where \(r_P\) and \(r_D\) are the number of replicas (GPU groups) allocated to prefill and decode, and \(N\) is the total GPU budget.

Splitwise: Heterogeneous Hardware for Each Phase¶

Splitwise (Patel et al., ISCA 2024, Microsoft Research; arXiv preprint 2023) takes disaggregation one step further: it argues that because prefill is compute-bound and decode is memory-bandwidth-bound, you should use different GPU models for the two pools. Specifically:

• Prefill workers: use high-FLOP/s, moderately-bandwidth GPUs. In an H100/A100 world, this often means fewer GPUs with aggressive compute configurations.
• Decode workers: use high-bandwidth-memory GPUs — or even CPUs with large memory for small batches (CPU offloading). High-HBM parts shine here, from the H100’s 3.35 TB/s HBM3 up to the Blackwell B200’s ~8 TB/s HBM3e.

Splitwise also introduced the term “prompt phase” for prefill and “token phase” for decode, now widely adopted in the systems literature. Their key empirical finding: on commercial cloud deployments, the decode phase uses far fewer FLOPs per token than prefill but consumes a comparable fraction of total serving cost due to the time it spends waiting for memory bandwidth. Disaggregation with heterogeneous hardware can reduce per-token cost by routing each phase to its best-fit hardware.

This connects directly to inference economics (see Inference Economics: Latency, Throughput & Cost) — the compute-to-cost frontier is different for prefill and decode, and ignoring this difference means overpaying.

Scheduling Under Disaggregation¶

A disaggregated system requires a more sophisticated scheduler than continuous batching. The scheduler must:

1. Route new requests to prefill workers with available capacity.
2. Match prefill worker output to decode workers with available KV-cache pages.
3. Handle KV transfer back-pressure: if decode workers are full, the prefill worker must pause or queue its output — this is analogous to a producer/consumer problem.
4. Rebalance pools dynamically as the arrival rate shifts between short-prompt (decode-heavy) and long-prompt (prefill-heavy) workloads.

Priority and Preemption Policies¶

In a pure disaggregated system, decode workers can still be preempted if they run out of KV-cache memory for long sequences. The preemption options remain the same as in standard serving: swap KV cache to CPU memory or recompute from scratch (at the cost of a prefill re-run). These policies are discussed in depth in PagedAttention & KV-Cache Memory Management.

Under chunked prefill (without full disaggregation), the scheduler gains a finer-grained lever: it can dynamically adjust \(C\) to give more or less priority to prefill vs. decode. A simple policy:

def adaptive_chunk_size(
    decode_queue_depth: int,       # number of active decode sequences
    prefill_queue_depth: int,      # number of prompts waiting to start prefill
    base_chunk_size: int = 512,
    max_chunk_size: int = 4096,
    min_chunk_size: int = 128,
    decode_pressure_threshold: int = 16,
) -> int:
    """
    Return the chunk size for the next iteration.

    When many decode sequences are active (high decode_queue_depth),
    use small chunks to minimise decode ITL inflation.
    When the prefill queue is long and decode is idle, use large chunks
    to drain prompts quickly and minimise TTFT.
    """
    if decode_queue_depth == 0:
        # No decode sequences waiting; use full large chunks for fast TTFT
        return max_chunk_size

    if decode_queue_depth >= decode_pressure_threshold:
        # Lots of decode sequences in-flight; be gentle on ITL
        return min_chunk_size

    # Interpolate linearly between min and base
    ratio = decode_queue_depth / decode_pressure_threshold
    chunk = int(base_chunk_size - ratio * (base_chunk_size - min_chunk_size))
    return max(min_chunk_size, min(chunk, base_chunk_size))

Implementation in vLLM and SGLang¶

Both vLLM and SGLang have shipped production-quality chunked prefill implementations.

vLLM’s Chunked Prefill¶

vLLM introduced chunked prefill (often called “prefill chunking” in its documentation) as a scheduler-level feature. In the V0 engine it was an opt-in flag; in the V1 engine — the default since vLLM 0.8 — the scheduler is unified: there is no separate “prefill batch” and “decode batch” at all, only a per-iteration token budget that mixed prefill chunks and decode tokens draw from. Chunked prefill is therefore on by default, and the knob you actually tune is the budget:

# vllm serve configuration (YAML or CLI flags)
max_num_batched_tokens: 2048   # total tokens (prefill chunks + decode) per iteration
max_num_seqs: 256              # max concurrent sequences
# enable_chunked_prefill: true  # V0-era flag; implicit in the V1 scheduler

Internally, the V1 Scheduler (vllm/v1/core/sched/scheduler.py) tracks a num_computed_tokens counter per request and, each round, hands out tokens from the budget: one per running sequence that needs a decode step, and the remainder as prefill chunks for requests that still have prompt left. A request needs no special “prefill vs decode” state — it is simply a sequence with num_computed_tokens < num_prompt_tokens or not. This is the same design as the ChunkedPrefillScheduler sketched below, and it is why the V1 rewrite made chunked prefill essentially free to support.

max_num_batched_tokens is the single most important latency/throughput dial in a vLLM deployment: raise it for throughput (bigger, more efficient GEMMs, faster TTFT), lower it for tighter P99 ITL. vLLM’s own defaults differ by mode — a large budget when optimizing throughput, a smaller one when optimizing latency — so set it explicitly rather than inheriting whatever the release picked.

SGLang’s RadixAttention and Chunking¶

SGLang (see SGLang: RadixAttention & Structured Programs) combines chunked prefill with its RadixAttention prefix cache. A chunked prefill step can reuse prefix cache hits from earlier chunks, meaning that if two requests share a common prefix, their prefill work is shared at the chunk boundary — a multiplicative efficiency win.

Actually Running a Disaggregated Deployment¶

The pseudocode earlier in this chapter shows the mechanism; in 2026 you do not implement it yourself. Three open-source layers exist, and it is worth knowing which does what:

• The transfer engine. NIXL (NVIDIA Inference Xfer Library) is the emerging common abstraction over RDMA/InfiniBand, NVLink, and local copies; Mooncake’s Transfer Engine plays the same role and additionally backs a distributed KV store over DRAM/SSD; LMCache adds a tiered KV cache (GPU → CPU → disk) usable both for cross-instance transfer and for prefix reuse.
• The engine-side connector. vLLM exposes these through its KVConnector interface (NixlConnector, LMCacheConnectorV1, MooncakeConnector, …), configured with --kv-transfer-config. SGLang exposes them via --disaggregation-mode plus a transfer backend.
• The orchestrator. NVIDIA Dynamo supplies the piece neither engine provides: independently scalable prefill and decode pools, a KV-aware router that sends a request to the worker most likely to already hold its prefix, and pool rebalancing.

A minimal two-process vLLM deployment on one node — prefill on GPU 0, decode on GPU 1:

# Terminal 1 — prefill worker (KV producer)
CUDA_VISIBLE_DEVICES=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --port 8100 \
  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer"}'

# Terminal 2 — decode worker (KV consumer)
CUDA_VISIBLE_DEVICES=1 vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --port 8200 \
  --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer"}'

# Terminal 3 — proxy: send the prompt to :8100 with max_tokens=1 to build the KV
# cache, then forward the same request to :8200, which pulls the KV blocks over
# NIXL and streams the remaining tokens. vLLM ships worked proxy examples under
# examples/online_serving/disaggregated_serving/ in the repo.

The SGLang equivalent launches two servers with explicit roles and a small load balancer that pairs them:

# Prefill server
python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct \
  --disaggregation-mode prefill --port 30000

# Decode server
python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct \
  --disaggregation-mode decode  --port 30001

# A router/load-balancer process then pairs prefill and decode instances and
# presents a single OpenAI-compatible endpoint to clients.

A Minimal End-to-End Chunked Prefill Scheduler¶

import dataclasses
from collections import deque
from typing import List, Optional
import torch

@dataclasses.dataclass
class SequenceState:
    seq_id: int
    prompt_ids: List[int]                # full prompt token ids
    num_computed: int = 0                # how many prompt tokens have been processed
    kv_cache: Optional[torch.Tensor] = None  # accumulated KV cache
    output_ids: List[int] = dataclasses.field(default_factory=list)
    finished: bool = False

class ChunkedPrefillScheduler:
    """
    Minimal scheduler demonstrating chunked prefill logic.
    Does NOT implement actual model calls — shows scheduling decisions only.
    """

    def __init__(
        self,
        chunk_size: int = 512,
        max_decode_seqs: int = 64,
        max_batched_tokens: int = 4096,
    ):
        self.chunk_size = chunk_size
        self.max_decode_seqs = max_decode_seqs
        self.max_batched_tokens = max_batched_tokens

        self.waiting: deque[SequenceState] = deque()      # not yet started
        self.prefilling: deque[SequenceState] = deque()   # partially prefilled
        self.decoding: List[SequenceState] = []           # in decode phase

    def add_request(self, seq_id: int, prompt_ids: List[int]):
        self.waiting.append(SequenceState(seq_id=seq_id, prompt_ids=prompt_ids))

    def schedule(self) -> dict:
        """
        Produce a batch descriptor for the next forward pass.
        Returns a dict describing which sequences to process and how.
        """
        budget = self.max_batched_tokens
        batch = {"decode": [], "prefill_chunks": []}

        # --- Step 1: schedule decode sequences (highest priority) ---
        for seq in self.decoding:
            if not seq.finished:
                # Each decode sequence consumes 1 token of budget
                if budget >= 1:
                    batch["decode"].append(seq.seq_id)
                    budget -= 1

        decode_budget_used = self.max_batched_tokens - budget

        # --- Step 2: schedule prefill chunks with remaining budget ---
        # First, continue any partially prefilled sequences
        next_prefilling = deque()
        for seq in self.prefilling:
            remaining = len(seq.prompt_ids) - seq.num_computed
            chunk = min(remaining, self.chunk_size, budget)
            if chunk <= 0:
                next_prefilling.append(seq)
                continue

            batch["prefill_chunks"].append({
                "seq_id": seq.seq_id,
                "token_ids": seq.prompt_ids[seq.num_computed: seq.num_computed + chunk],
                "start_pos": seq.num_computed,
            })
            seq.num_computed += chunk
            budget -= chunk

            if seq.num_computed >= len(seq.prompt_ids):
                # Prefill complete; move to decode
                self.decoding.append(seq)
            else:
                next_prefilling.append(seq)

        self.prefilling = next_prefilling

        # Then, admit new requests if budget remains and decode pool not full
        while (
            self.waiting
            and budget >= self.chunk_size
            and len(self.decoding) < self.max_decode_seqs
        ):
            seq = self.waiting.popleft()
            chunk = min(len(seq.prompt_ids), self.chunk_size, budget)
            batch["prefill_chunks"].append({
                "seq_id": seq.seq_id,
                "token_ids": seq.prompt_ids[:chunk],
                "start_pos": 0,
            })
            seq.num_computed = chunk
            budget -= chunk

            if seq.num_computed >= len(seq.prompt_ids):
                self.decoding.append(seq)
            else:
                self.prefilling.append(seq)

        return batch

    def mark_decode_finished(self, seq_id: int):
        self.decoding = [s for s in self.decoding if s.seq_id != seq_id]

Comparative Analysis: Chunked vs. Disaggregated¶

Both chunked prefill and full disaggregation solve the interference problem, but with different tradeoffs:

A key point: these techniques are not mutually exclusive. A disaggregated system can also use chunked prefill within each prefill worker to avoid wasting memory bandwidth on very long sequences when the prefill batch is small.

Practical Deployment Guidance¶

When to Enable Chunked Prefill¶

Enable chunked prefill whenever:

• You observe bursty P99 ITL under mixed workloads (common in production chatbots and RAG pipelines).
• Your system receives occasional long prompts mixed with ongoing conversations.
• You are using vLLM or SGLang and have not already enabled it — it is low-risk and usually improves P99 latency.

A good starting value: max_num_batched_tokens = 2048 with chunk_size = 512. Profile your P99 ITL and TTFT with a production traffic replay, then tune.

When to Build a Disaggregated System¶

Disaggregation is warranted when:

• Prompt lengths are consistently long (>8K tokens — document processing, code understanding, long-context RAG).
• You have strict P99 ITL SLOs (< 50 ms for real-time applications) that chunked prefill alone cannot meet.
• You operate a large enough cluster that dedicating separate GPU pools is economically justifiable.
• You can place prefill and decode workers on the same NVSwitch fabric (same node or adjacent nodes) to keep KV transfer below ~5 ms.

Monitoring Key Metrics¶

# Example metrics to track for a disaggregated system
# (pseudocode — plug into your Prometheus/OpenTelemetry stack)

METRICS = {
    # Latency
    "ttft_p50_ms": "Time to first token, 50th percentile",
    "ttft_p99_ms": "Time to first token, 99th percentile",
    "itl_p50_ms":  "Inter-token latency, 50th percentile",
    "itl_p99_ms":  "Inter-token latency, 99th percentile",

    # KV transfer (disaggregated only)
    "kv_transfer_latency_p99_ms": "P99 KV cache transfer time prefill->decode",
    "kv_transfer_bytes_per_sec":  "KV transfer throughput (capacity planning)",
    "kv_transfer_queue_depth":    "Number of KV caches awaiting transfer",

    # Pool utilization
    "prefill_worker_gpu_util_pct": "GPU utilization on prefill pool",
    "decode_worker_gpu_util_pct":  "GPU utilization on decode pool",
    "decode_kv_cache_fill_pct":    "Fraction of paged KV cache in use on decode workers",

    # Scheduler health
    "prefill_queue_depth":   "Requests waiting for prefill start",
    "chunked_prefill_iters": "Average iterations to complete one prefill",
}

Connections to the Broader Inference Stack¶

Disaggregated prefill/decode does not exist in isolation. Several adjacent technologies interact with it:

Prefix Caching (see Prefix Caching & KV-Cache Reuse): If two requests share a common prefix, the prefill worker can skip recomputing that portion, and the decode worker receives a smaller KV cache. Disaggregation amplifies the value of prefix caching because the KV transfer cost is proportional to the unique (non-cached) portion of the KV cache.

Speculative Decoding (see Speculative Decoding: Draft Models, Medusa, EAGLE & Lookahead): Speculative decoding generates draft tokens on the decode worker and verifies them in a batched forward pass. This verification pass looks like a short prefill — under disaggregation, it stays on the decode worker (it is short enough not to cause interference) rather than being sent to the prefill pool.

Multi-GPU Inference (see Multi-GPU & Multi-Node Inference): Tensor parallelism and pipeline parallelism within each pool interact with KV transfer. For a tensor-parallel model (e.g., 4-way TP), each GPU holds \(1/4\) of each KV head — so the KV cache transfer is split across 4 GPUs, and all four must synchronize with the corresponding 4 decode GPUs. This requires careful collective communication design.

GPU Architecture (see GPU Architecture & The Memory Hierarchy): The fundamental reason prefill and decode prefer different hardware is rooted in the roofline model. The compute roof matters for prefill; the memory bandwidth wall matters for decode. Choosing hardware with the right roof height for each pool is a direct application of the roofline analysis.

Further Reading¶

• Zhong et al., “DistServe: Disaggregating Prefill and Decoding for Goodput-Optimized Large Language Model Serving,” OSDI 2024. The foundational paper for disaggregated serving; includes the formal goodput optimization and KV transfer analysis.
• Patel et al., “Splitwise: Efficient Generative LLM Inference Using Phase Splitting,” ISCA 2024 (Microsoft Research). Formalizes the heterogeneous hardware argument and coins “prompt phase” / “token phase” terminology.
• Agrawal et al., “Sarathi-Serve: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills,” OSDI 2024. Demonstrates chunked prefill in a production-oriented scheduler with careful measurement of the TTFT-ITL tradeoff.
• Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention,” SOSP 2023. The vLLM paper; PagedAttention is the KV-cache memory management substrate on which chunked prefill and disaggregation are built.
• vLLM documentation: “Chunked Prefill” and “Disaggregated Prefill.” The vLLM project documentation covers both features with configuration examples and benchmark guidance.
• SGLang GitHub repository (lm-sys/sglang). SGLang’s scheduler source code is an excellent reference for how chunked prefill interacts with RadixAttention prefix caching in a production system.

Exercises¶

1. (Conceptual.) A colleague proposes “fixing” prefill-decode interference by simply giving decode requests strict priority: always run every in-flight decode step first, and only run a prefill once no decode is pending. Explain why, in a naive continuous-batching engine, this does not solve the interference problem the chapter describes, and name the two mechanisms — one from the disaggregation section and one from the chunked-prefill section — that actually do.

2. (Quantitative — KV-cache transfer.) You are sizing the interconnect for a disaggregated deployment of a model with \(L = 40\) layers, GQA with \(n_\text{kv} = 8\) KV heads, \(d_\text{head} = 128\), serving prompts of \(T_p = 2048\) tokens, KV stored in BF16 (2 bytes/element). Using the chapter’s KV-size formula:

• (a) Compute the KV-cache size in bytes for one request, and express it in MB.
• (b) Compute the transfer time over NVLink4 (\(\approx 900\) GB/s) and over InfiniBand HDR (\(\approx 12.5\) GB/s).
• © The prefill itself for this prompt takes roughly 120 ms. On which of the two interconnects does the KV transfer become a first-order term in end-to-end latency, and what does the chapter recommend doing about it?

3. (Quantitative — choosing the chunk size.) A 13B model on one A100 prefills an 8,192-token prompt as a single pass in about 800 ms (as in the chapter’s worked example), and prefill time scales roughly linearly in the number of tokens processed. Each decode iteration adds about 15 ms of its own decode compute. Your SLO is P99 ITL \(\le 100\) ms. For chunk sizes \(C = 256\) and \(C = 1024\), compute: (a) the number of chunks, (b) the per-chunk prefill time, © the resulting per-iteration ITL seen by in-flight decode sequences, and (d) which value(s) of \(C\) satisfy the SLO. (e) Briefly, why not just pick the smallest \(C\) possible?

4. (Implementation — dynamic chunk sizing.) The ChunkedPrefillScheduler in the chapter uses a fixed self.chunk_size. Modify it so the effective chunk size is recomputed every schedule() call from the current decode load, using the chapter’s adaptive_chunk_size policy (large chunks when no decode is active, small chunks under decode pressure). Show the changed schedule() code, and explain what behavior a caller would observe as the decode pool fills up.

def schedule(self) -> dict:
    """
    Produce a batch descriptor for the next forward pass, using a
    chunk size chosen dynamically from current decode load.
    """
    budget = self.max_batched_tokens
    batch = {"decode": [], "prefill_chunks": []}

    # --- Dynamic chunk size for THIS iteration ---
    eff_chunk = adaptive_chunk_size(
        decode_queue_depth=len(self.decoding),
        prefill_queue_depth=len(self.waiting),
        base_chunk_size=self.chunk_size,
    )

    # --- Step 1: decode sequences (highest priority) ---
    for seq in self.decoding:
        if not seq.finished and budget >= 1:
            batch["decode"].append(seq.seq_id)
            budget -= 1

    # --- Step 2: continue partially prefilled sequences ---
    next_prefilling = deque()
    for seq in self.prefilling:
        remaining = len(seq.prompt_ids) - seq.num_computed
        chunk = min(remaining, eff_chunk, budget)   # was self.chunk_size
        if chunk <= 0:
            next_prefilling.append(seq)
            continue
        batch["prefill_chunks"].append({
            "seq_id": seq.seq_id,
            "token_ids": seq.prompt_ids[seq.num_computed: seq.num_computed + chunk],
            "start_pos": seq.num_computed,
        })
        seq.num_computed += chunk
        budget -= chunk
        if seq.num_computed >= len(seq.prompt_ids):
            self.decoding.append(seq)
        else:
            next_prefilling.append(seq)
    self.prefilling = next_prefilling

    # --- Step 3: admit new requests ---
    while (
        self.waiting
        and budget >= eff_chunk                      # was self.chunk_size
        and len(self.decoding) < self.max_decode_seqs
    ):
        seq = self.waiting.popleft()
        chunk = min(len(seq.prompt_ids), eff_chunk, budget)  # was self.chunk_size
        batch["prefill_chunks"].append({
            "seq_id": seq.seq_id,
            "token_ids": seq.prompt_ids[:chunk],
            "start_pos": 0,
        })
        seq.num_computed = chunk
        budget -= chunk
        if seq.num_computed >= len(seq.prompt_ids):
            self.decoding.append(seq)
        else:
            self.prefilling.append(seq)

    return batch

5. (Implementation — GQA in chunked-prefill attention.) The chapter’s chunked_prefill_attention assumes n_heads == n_kv_heads and builds its causal mask with a Python for loop. Rewrite the function to (a) support GQA where n_kv_heads < n_heads by expanding the KV heads, and (b) replace the per-row loop with a vectorized mask. Keep the same signature and output shape [C, n_heads, d_head].

import torch
import torch.nn.functional as F

def chunked_prefill_attention(
    q_chunk: torch.Tensor,    # [C, n_heads, d_head]
    k_full: torch.Tensor,     # [T_past + C, n_kv_heads, d_head]
    v_full: torch.Tensor,     # [T_past + C, n_kv_heads, d_head]
    T_past: int,
    C: int,
    scale: float,
) -> torch.Tensor:
    n_heads  = q_chunk.shape[1]
    n_kv     = k_full.shape[1]
    T_total  = T_past + C

    # (a) GQA expansion: repeat each KV head to cover its query-head group.
    assert n_heads % n_kv == 0, "n_heads must be a multiple of n_kv_heads"
    group = n_heads // n_kv
    if group > 1:
        k_full = k_full.repeat_interleave(group, dim=1)  # [T_total, n_heads, d_head]
        v_full = v_full.repeat_interleave(group, dim=1)

    # (b) Vectorized causal mask: row i (abs pos T_past+i) sees key j <= T_past+i.
    q_pos = torch.arange(C, device=q_chunk.device) + T_past          # [C]
    k_pos = torch.arange(T_total, device=q_chunk.device)             # [T_total]
    causal_mask = k_pos[None, :] <= q_pos[:, None]                   # [C, T_total] bool

    # Scores and masked softmax.
    q = q_chunk.transpose(0, 1)                    # [n_heads, C, d_head]
    k = k_full.transpose(0, 1)                     # [n_heads, T_total, d_head]
    scores = torch.bmm(q, k.transpose(1, 2)) * scale               # [n_heads, C, T_total]
    scores = scores.masked_fill(~causal_mask.unsqueeze(0), float('-inf'))

    attn = F.softmax(scores, dim=-1)               # [n_heads, C, T_total]
    v = v_full.transpose(0, 1)                     # [n_heads, T_total, d_head]
    out = torch.bmm(attn, v)                       # [n_heads, C, d_head]
    return out.transpose(0, 1)                     # [C, n_heads, d_head]

6. (Conceptual — pool allocation and technique choice.) Two teams share the same 16-GPU cluster and total GPU budget. Team A serves a document-summarization product: prompts average 16K tokens, outputs average 200 tokens. Team B serves a chatbot: prompts average 300 tokens, outputs average 600 tokens. (a) Using DistServe’s framing, argue how the prefill:decode replica split \(r_P : r_D\) should differ between the two workloads. (b) For each team, would you reach first for chunked prefill or full disaggregation, and why? © Which team benefits more from Splitwise-style heterogeneous hardware, and which single hardware property matters most for their bottleneck phase?