Prefill/Decode Disaggregation Simulation

1. What is PD Disaggregation?

In standard LLM serving, every GPU replica runs both prefill (processing the full prompt in parallel) and decode (generating tokens one-by-one autoregressively). These two phases have fundamentally different computational profiles: prefill is compute-bound with high arithmetic intensity and long per-request latency, while decode is memory-bandwidth-bound with low arithmetic intensity but extreme latency sensitivity since each token delays the user-visible response.

When co-located on the same GPU, prefill and decode contend for the same resources. A long prefill batch can stall ongoing decode iterations, causing Time-Between-Tokens (TBT) spikes. Conversely, interleaving small decode batches with prefill wastes GPU compute cycles. The disaggregation approach -- pioneered by the SplitWise and DistServe papers -- separates prefill and decode onto dedicated GPU pools (P-nodes and D-nodes), allowing each pool to be independently optimized for its workload profile.

The key tradeoff is straightforward: disaggregation introduces a KV cache transfer cost between the P-node (which generates the KV cache during prefill) and the D-node (which needs it for decode). If the transfer bandwidth is sufficient and the model's KV cache per token is small, disaggregation can dramatically improve both TTFT (Time-To-First-Token) and TBT. If the transfer is slow or the KV cache is large, the overhead may negate the benefits.

2. How SimAI Models PD Disaggregation

SimAI extends the original Vidur simulator with four tightly integrated additions: (1) a ReplicaType classification that tags each replica as PREFILL, DECODE, or MIXED; (2) a request DAG architecture that models the KV cache transfer as an explicit graph edge; (3) the SplitwiseGlobalScheduler that routes requests to the correct pool; and (4) KV cache size estimation and transfer flow modeling. Together, these enable full PD disaggregation simulation without any GPU hardware.

2.1 ReplicaType Classification

The foundation of PD disaggregation is the ReplicaType enum. In original Vidur, all replicas are implicitly MIXED -- they handle both prefill and decode. SimAI adds explicit P/D typing so each replica knows its role.

from enum import IntEnum

class ReplicaType(IntEnum):
    MIXED   = 0   # Co-located: both prefill + decode (original Vidur)
    PREFILL = 1   # Dedicated prefill nodes (P-nodes)
    DECODE  = 2   # Dedicated decode nodes (D-nodes)

2.2 Request DAG Architecture

SimAI represents each request as a directed acyclic graph (DAG) using nx.DiGraph(). In the co-located case, the DAG is simple: PromptTask -> TokenTask. In the disaggregated case, a KVCacheTransferFlow node is inserted between them, modeling the data movement cost.

class Request(BaseEntity):
    def __init__(self, ...):
        # DAG for task dependency modeling
        self.dag = nx.DiGraph()
        self.root_node = None
        self.flow_node = None

        # PD disaggregation tracking
        self.prefill_replica_id = None
        self.decode_replica_id = None

        # Timing fields
        self.prefill_arrived_at = arrived_at
        self.decode_arrived_at = float('inf')
        self.decode_time = float('inf')

        # KV cache transfer metadata
        self.pd_p2p_comm_size = float('inf')
        self.pd_p2p_comm_time = float('inf')
        self.pd_p2p_bytes_per_token = None
        self.pd_p2p_comm_dtype = None

2.3 SplitwiseGlobalScheduler

The SplitwiseGlobalScheduler is the central orchestrator for PD disaggregation. During initialization, it splits the replica pool into prefill and decode subsets based on pd_node_ratio, then creates separate sub-schedulers for each pool. When a request arrives, it creates a PromptTask and TokenTask, routes the PromptTask to a P-node and the TokenTask to a D-node, and inserts the KV cache transfer flow between them.

class SplitwiseGlobalScheduler(BaseGlobalScheduler):
    def __init__(self, config: SimulationConfig, replicas: Dict[int, Replica]):
        super().__init__(config, replicas)

        # P:D ratio from config (e.g., 0.5 = 1P:1D, 0.33 = 1P:2D)
        self.pd_node_ratio = self._replicas[0].pd_node_ratio

        # Split replicas into P-pool and D-pool
        self._num_prefill_nodes = self._num_replicas * self.pd_node_ratio
        self._num_decode_nodes  = self._num_replicas - self._num_prefill_nodes

        # Build prefill replica dict (IDs 0..N_p-1)
        self.prefill_replicas = {}
        for replica_id, replica in self._replicas.items():
            if replica_id < self._num_prefill_nodes:
                self.prefill_replicas[replica_id] = replica
                replica.replica_type = ReplicaType.PREFILL

        # Build decode replica dict (IDs N_p..N-1)
        self.decode_replicas = {}
        for replica_id, replica in self._replicas.items():
            if replica_id >= self._num_prefill_nodes:
                self.decode_replicas[replica_id] = replica
                replica.replica_type = ReplicaType.DECODE

        # Create sub-schedulers for each pool (default: LOR)
        self.prefill_scheduler = self.get_global_scheduler(self.prefill_replicas)
        self.decode_scheduler  = self.get_global_scheduler(self.decode_replicas)

        # Round-robin counters
        self.p_request_counter = 0
        self.d_request_counter = 0

The schedule() method processes each request in three steps:

2.3.1 PD Node Selection Logic — Source Code Deep Dive

Strategy A: Round-Robin (ACTIVE in schedule())

The schedule() method uses simple modular counters to distribute requests across pools. This is the code path that actually executes:

for request in self._request_queue:
    # --- Prefill node selection: simple round-robin ---
    # (LOR code is commented out on line 351)
    # replica_id = min(pending_prefill_requests_map.items(), key=lambda x: x[1])[0]

    replica_id = self.p_request_counter % len(self.prefill_replicas)  # Round-robin
    self.p_request_counter += 1

    prefill_replica = self.prefill_replicas[replica_id]
    request.prefill_replica_id = replica_id

    # --- Decode node selection: round-robin with offset ---
    # (LOR code is commented out on lines 364-366)
    # replica_id = min(pending_decode_requests_map.items(), key=lambda x: x[1])[0]

    replica_id = (self.d_request_counter % len(self.decode_replicas)) \
                 + len(self.prefill_replicas)  # Offset by N_prefill
    self.d_request_counter += 1

    decode_replica = self.decode_replicas[replica_id]
    request.decode_replica_id = replica_id

Strategy B: Load-Aware Selection (DEFINED but NOT CALLED)

The class defines two load-aware methods — find_best_prefill_replica() and find_best_decode_replica() — that implement smarter selection based on pending token count and memory utilization. However, neither method is ever called by schedule() in the current codebase:

def find_best_prefill_replica(self, prefill_replicas, prefill_task):
    """Select prefill replica with minimum pending tokens, if queue not overloaded"""
    if len(prefill_replicas) == 0:
        return None
    prefill_replica = min(prefill_replicas,
                          key=lambda replica: replica.sched_pending_tokens)  # Load-aware!
    if self.is_queue_long(prefill_replica, prefill_task):
        return None  # Back-pressure: reject if queue too long
    return prefill_replica

def find_best_decode_replica(self, decode_replicas, prefill_task, decode_task):
    """Select decode replica with minimum pending tokens, if memory sufficient"""
    if len(decode_replicas) == 0:
        return None
    decode_replica = min(decode_replicas,
                          key=lambda replica: replica.sched_pending_tokens)  # Load-aware!
    if self.is_memory_loaded(decode_replica, [prefill_task, decode_task]):
        return None  # Back-pressure: reject if OOM risk
    return decode_replica

Also: Sub-Scheduler Initialization Mismatch

The constructor sets self._sub_scheduler = "lor" and creates LOR (Least Outstanding Requests) sub-schedulers for both pools via get_global_scheduler(). These sub-schedulers are stored as self.prefill_scheduler and self.decode_scheduler, but neither is ever invoked — the schedule() method implements its own routing logic instead of delegating to them.

Why? Evidence of Work-in-Progress, Not a Design Choice

The source code contains clear evidence that this is a work-in-progress codebase where load-aware routing was planned but not yet completed, rather than a deliberate design decision to use round-robin. Four pieces of evidence:

# TODO: > > 参考 sw写的；但也很多区别；换一个名字；类似pd分离的其他名字；不严格是sw了
# TODO: Refer to Splitwise implementation; but many differences; need a new name;
#       similar to pd separation; not strictly Splitwise anymore
class SplitwiseGlobalScheduler(BaseGlobalScheduler):
    def __init__(self, config, replicas):
        # ① Config-driven sub-scheduler was planned but commented out:
        # self._sub_scheduler = self._config.splitwise_scheduler_sub_scheduler

        # ② TODO explicitly marks this as incomplete:
        # TODO > improve _sub_scheduler flexible choice

        # ③ Round-robin was tried, then switched to LOR:
        # self._sub_scheduler = "round_robin"
        self._sub_scheduler = "lor"  # ← Creates LOR sub-schedulers

        # ④ But schedule() bypasses sub-schedulers entirely,
        #    using inline round-robin counters instead (lines 354, 369)

The development trajectory is clear from these comments:

2.4 KV Cache Transfer Flow

The add_kv_cache_transfer() method is the heart of the DAG transformation. It estimates the KV cache size for the given prompt, creates a FlowType.KVCacheTransfer flow node, removes the direct PromptTask->TokenTask edge, and inserts the transfer flow as an intermediary.

def add_kv_cache_transfer(self, request, src_replica, dest_replica, bandwidth):
    """Transform DAG: prompt->token into prompt->kv_transfer->token."""
    prefill_task = request.root_node            # Root = PromptTask
    decode_task  = next(request.successors(prefill_task))  # Next = TokenTask

    # Estimate KV cache size based on prompt length and model config
    flow_size = request.estimate_kv_cache_size(
        num_tokens=prefill_task.prompt_size,
        replica=src_replica)

    # Create KVCacheTransfer flow node
    kv_transfer_flow = request.create_flow(
        FlowType.KVCacheTransfer,
        size=flow_size,
        src=src_replica,
        dest=dest_replica)
    kv_transfer_flow.notify = True

    # Rewire the DAG
    request.flow_node = kv_transfer_flow
    request.dag.remove_edge(prefill_task, decode_task)
    request.dag.add_edge(prefill_task, kv_transfer_flow)
    request.dag.add_edge(kv_transfer_flow, decode_task)

    # Assign tasks to replicas
    prefill_task.instance = src_replica
    decode_task.instance  = dest_replica

    # Simulate transfer latency with configurable bandwidth
    kv_transfer_flow.link = DummyLink(
        name="DummyLink",
        bandwidth=bandwidth)

The KV cache size is computed by the estimate_kv_cache_size() method on the Request object. The formula accounts for both the key and value tensors across all layers:

kv_size = 2 * num_tokens * hidden_dim * num_layers * dtype_bytes

def estimate_kv_cache_size(self, num_tokens=None, replica=None):
    """Returns the KV-cache size (in bytes) after generating num_tokens."""
    # Determine bytes per element based on dtype
    if replica.pd_p2p_comm_dtype == 'float16':
        pd_p2p_bytes_per_token = 2
    elif replica.pd_p2p_comm_dtype == 'float32':
        pd_p2p_bytes_per_token = 4
    elif replica.pd_p2p_comm_dtype == 'bfloat16':
        pd_p2p_bytes_per_token = 2
    elif replica.pd_p2p_comm_dtype == 'int8':
        pd_p2p_bytes_per_token = 1
    # ... (int16=2, int32=4, int64=8, float64=8)

    self.pd_p2p_bytes_per_token = pd_p2p_bytes_per_token
    self.pd_p2p_comm_dtype = replica.pd_p2p_comm_dtype

    # KV = 2 (K+V) * seq_len * hidden_dim * layers * dtype_bytes
    return 2 * num_tokens * replica.mlp_hidden_dim \
            * replica.num_layers * pd_p2p_bytes_per_token

3. SplitWise Replica Scheduler

At the global level, the SplitwiseGlobalScheduler uses round-robin to decide which replica handles each phase (as shown in Section 2.3.1). Once assigned, the per-replica scheduling behavior differs between prefill and decode replicas — prefill replicas optimize for throughput, while decode replicas optimize for latency. The codebase also contains unused load-aware methods that could replace round-robin in future versions.

4. KV Cache Transfer Cost Modeling

The transfer cost between P-nodes and D-nodes is the critical overhead introduced by disaggregation. SimAI models this cost in two ways: (1) an analytical model using pd_p2p_comm_bandwidth from the config, where transfer_time = kv_size / bandwidth; and (2) integration with astra-sim + NS-3 for detailed network simulation when higher fidelity is needed.

4.1 Transfer Size Across Models

The KV cache size per token varies dramatically across model architectures due to differences in attention head configuration. Models using GQA (Grouped Query Attention) or MLA (Multi-head Latent Attention) have significantly smaller KV caches than standard MHA (Multi-Head Attention).

4.2 Transfer Time Estimation

For a prompt of 2048 tokens on LLaMA-70B with float16, the KV cache transfer size is approximately 640 MB (320 KB/token * 2048 tokens). At the default 800 Gbps bandwidth, this takes about 6.4 ms -- negligible compared to the prefill compute time. However, at 100 Gbps (e.g., RoCE across racks), it rises to ~51 ms, which may become significant.

# Analytical model (DummyLink)
transfer_time = kv_cache_size_bytes / (pd_p2p_comm_bandwidth_bps / 8)

# Example: LLaMA-70B, 2048 tokens, float16, 800 Gbps
kv_size = 2 * 2048 * 8192 * 80 * 2  # = 5,368,709,120 bytes (~5.0 GB)
transfer_time = 5368709120 / (800e9 / 8)   # = 53.7 ms

# With NS-3 simulation: accounts for congestion, link utilization, etc.

5. Event Chain Comparison

The discrete-event simulation in SimAI processes different event chains depending on whether PD disaggregation is enabled. Below we compare the two event chains side-by-side.

6. Configuration Parameters

PD disaggregation in SimAI is controlled by three main parameters in the ReplicaConfig dataclass. These are set via command-line arguments or configuration files.

@dataclass
class ReplicaConfig:
    # ... other fields ...

    # PD disaggregation parameters
    pd_p2p_comm_bandwidth: int = field(
        default=800,
        metadata={"help": "PD P2P communication bandwidth (bps)"},
    )

    pd_p2p_comm_dtype: str = field(
        default='float16',
        metadata={"help": "Data type for KV cache transfer"},
    )

    pd_node_ratio: float = field(
        default=0.5,
        metadata={"help": "Ratio of P replicas to total replicas"},
    )

    # Related bandwidth parameters
    nvlink_bandwidth: int = field(
        default=1600,
        metadata={"help": "NVLink bandwidth for TP/EP (bps)"},
    )

    rdma_bandwidth: int = field(
        default=800,
        metadata={"help": "RDMA bandwidth for TP/EP (bps)"},
    )

# Balanced P:D (default)
--pd_node_ratio 0.5 --pd_p2p_comm_bandwidth 800 --pd_p2p_comm_dtype float16

# More decode capacity (long outputs, chat workload)
--pd_node_ratio 0.33 --pd_p2p_comm_bandwidth 800 --pd_p2p_comm_dtype float16

# More prefill capacity (long prompts, summarization workload)
--pd_node_ratio 0.67 --pd_p2p_comm_bandwidth 800 --pd_p2p_comm_dtype float16

# INT8 quantized KV transfer (halves transfer size)
--pd_node_ratio 0.5 --pd_p2p_comm_bandwidth 400 --pd_p2p_comm_dtype int8

7. Output Metrics

When PD disaggregation is enabled, SimAI tracks additional per-request metrics that capture the behavior of each phase independently. These metrics are recorded in the MetricsStore and exported for analysis.

class MetricConstants:
    # --- PD-specific timing metrics ---
    PREFILL_ARRIVED_AT    = "prefill_arrived_at"       # When request reached P-node
    PREFILL_COMPLETED_AT  = "prefill_completed_at"     # When prefill finished
    DECODE_ARRIVED_AT     = "decode_arrived_at"        # When KV arrived at D-node
    DECODE_TIME           = "decode_time"              # Total decode duration

    # --- PD-specific routing metrics ---
    PREFILL_REPLICA_ID    = "prefill_replica_id"       # Which P-node was used
    DECODE_REPLICA_ID     = "decode_replica_id"        # Which D-node was used

    # --- KV transfer metrics ---
    PD_P2P_COMM_SIZE      = "pd_p2p_comm_size"         # KV cache bytes transferred
    PD_P2P_COMM_TIME      = "pd_p2p_comm_time"         # Transfer latency

    # --- Derived latency metrics ---
    PREFILL_TIME_E2E      = "prefill_e2e_time"         # TTFT = prefill_completed - arrived
    TBT                   = "tbt"                      # Time Between Tokens (TPOT)

These metrics enable fine-grained analysis of PD disaggregation performance:

The metrics store also computes derived values during the on_request_end() callback:

# From vidur/metrics/metrics_store.py

# TTFT: Time from arrival to first token (prefill latency)
prefill_e2e_time = request.prefill_completed_at - request.arrived_at

# Scheduling delay for prefill
prefill_scheduling_delay = request.prefill_completed_at - request.scheduled_at

# Model execution time (excluding scheduling)
prefill_model_execution_time = (
    request.prefill_completed_at - request.scheduled_at
)

# Decode phase duration
decode_model_execution_time = (
    request.completed_at - request.prefill_completed_at
)

8. Key Insights

9. Scheduling Flow Deep Dive

To fully understand how PD disaggregation works in SimAI, it is helpful to trace a single request through the entire scheduling pipeline. The following step-flow shows every major decision point and data mutation that occurs.

10. Practical Guide: Running PD Disaggregation Simulations

This section provides practical guidance for setting up and running PD disaggregation simulations with SimAI.

10.1 Choosing the Right P:D Ratio

The P:D ratio is workload-dependent. Use the following guidelines as a starting point:

10.2 Bandwidth Sensitivity Analysis

The pd_p2p_comm_bandwidth parameter has a direct impact on the KV transfer latency. Running a sweep across bandwidths helps identify the minimum viable interconnect for a given workload. Below is the relationship between bandwidth and transfer time for a 2048-token prompt on LLaMA-70B with float16:

10.3 dtype Selection for KV Transfer

Reducing the dtype of the KV cache during transfer can halve the transfer time without significantly impacting model quality. The pd_p2p_comm_dtype parameter supports multiple options:

10.4 Comparing Co-located vs Disaggregated

To evaluate whether PD disaggregation is beneficial for your workload, run the simulation twice with the same parameters except for the global scheduler type. Compare the following metrics:

# Use round_robin or lor global scheduler
--global_scheduler_type round_robin
--num_replicas 4

# Use splitwise global scheduler
--global_scheduler_type splitwise
--num_replicas 4
--pd_node_ratio 0.5
--pd_p2p_comm_bandwidth 800

Key metrics to compare between the two runs:

Summary: Source Files for PD Disaggregation