Request Lifecycle & I/O Formats

Input & Configuration

Output → Populated DAG

Sets request.prefill_replica_id and request.decode_replica_id, then builds the DAG. The exact DAG shape depends on whether the prefill and decode replicas are the same or different (code in splitwise_global_scheduler.py):

PromptTask ──→ TokenTask
  # linked via prefill_task.chain = [decode_task]
  # no KV transfer needed

PromptTask ──→ KVCacheTransferFlow ──→ TokenTask
  # edges in request.dag (nx.DiGraph)
  # KV transfer bridges the two replicas

DAG Node Structures

# ── Base Node (all DAG nodes inherit from this) ──
Node(BaseEntity):
  node_id: int                # unique ID within this request's DAG
  state: NodeState             # NONE=0 → QUEUED=1 → RUNNING=2 → COMPLETED=4
                               #   (or BLOCKED=3 if preempted, ABORTED=5)
  request: Request             # back-reference to owning request
  chain: list[Node]            # consecutive nodes (used in same-replica case)
  num_preemptions: int = 0

# ── PromptTask (prefill stage) ──
PromptTask(Task → Node):
  task_type: TaskType = PROMPT # = 1
  prompt_size: int             # = request.num_prefill_tokens
  tokens_per_iteration: int    # = prompt_size (process all at once)
  processing_tokens: int = 0   # tokens being processed this iteration
  processed_tokens: int = 0    # total tokens processed
  generating_tokens: int = 0   # tokens being generated this iteration
  generated_tokens: int = 0    # total generated (max 1 for prefill)
  is_prefill_complete: bool = False
  cleanup_memory: bool = False # keep KV cache for decode
  batch_size: int = 1
  duration: float = 0.0       # set by execution time predictor

  # memory = 2 × num_tokens × mlp_hidden_dim × num_layers × bytes_per_token

# ── TokenTask (decode stage) ──
TokenTask(Task → Node):
  task_type: TaskType = TOKEN  # = 2
  token_size: int              # = request.num_decode_tokens − 1
                               #   (1 token already generated in prefill)
  tokens_per_iteration: int = 1 # auto-regressive: 1 token at a time
  is_prefill_complete: bool = True
  # same processing/generated fields as PromptTask

# ── KVCacheTransferFlow (PD transfer) ──
KVCacheTransferFlow(Flow → Node):
  flow_type: FlowType = KVCacheTransfer  # = 1
  src: Replica                 # prefill replica
  dest: Replica                # decode replica
  size: float                  # KV cache size in bytes
                               #   = estimate_kv_cache_size(prefill_tokens, replica)
  duration: float = 0.0        # set after transfer simulation
  notify: bool = False
  batch_size: int = 1

After DAG construction, the global scheduler emits a GlobalScheduleEvent which triggers ReplicaScheduleEvent(s) for the assigned replica(s).

Backend A: vidur (RandomForest on profiling CSVs)

Trains 11 separate RandomForest models from pre-collected GPU profiling data. After training, predictions are pre-generated for all possible token counts and stored as O(1) lookup tables (Python dicts).

# Flow:
Batch(batch_size, num_tokens_per_request)
  → extract features (num_tokens, kv_cache_size)
  → lookup pre-computed dict: {(num_tokens,): time_ms}
  → sum across 11 sub-models per layer
  → execution_time

Input File 1: Compute Profiling CSVs

Pre-collected by running each model operator on a real GPU at various token counts. The RandomForest learns f(num_tokens) → time_ms for each operator.

# File: data/profiling/compute/{DEVICE}/{MODEL}/mlp.csv
# Example: data/profiling/compute/h100/meta-llama/Llama-2-7b-hf/mlp.csv
#
# Feature columns (input to RandomForest):
num_tokens                          # int — tokens in this profiling run
num_tensor_parallel_workers         # int — TP degree used during profiling
n_head, n_kv_head, n_embd           # model architecture params
n_expanded_embd, vocab_size         # model architecture params
use_gated_mlp                       # bool

# Target columns (what RandomForest predicts) — each with min/max/mean/median/std:
time_stats.emb.*                    # embedding layer time
time_stats.input_layernorm.*        # input LayerNorm time
time_stats.attn_pre_proj.*          # QKV projection time
time_stats.attn_rope.*              # rotary position encoding time
time_stats.attn_post_proj.*         # output projection time
time_stats.post_attention_layernorm.*  # post-attention LayerNorm time
time_stats.mlp_up_proj.*            # MLP up-projection time
time_stats.mlp_act.*                # activation function time
time_stats.mlp_down_proj.*          # MLP down-projection time
time_stats.add.*                    # residual add time

# File: data/profiling/compute/{DEVICE}/{MODEL}/attention.csv
#
# Additional feature columns:
num_tokens, kv_cache_size           # input features

# Target columns:
time_stats.attn_prefill.*           # prefill attention kernel time
time_stats.attn_decode.*            # decode attention kernel time
time_stats.attn_kv_cache_save.*     # KV cache write time

Input File 2: Network Profiling CSVs

Pre-collected by running collective operations at various message sizes. Also trained via RandomForest: f(message_size) → time_ms.

# File: data/profiling/network/{NETWORK_DEVICE}/all_reduce.csv
# Example: data/profiling/network/h100_pairwise_nvlink/all_reduce.csv
#
# Columns:
size                                # int — message size in bytes
num_workers                         # int — number of GPUs
rank                                # int — GPU rank
collective                          # str — operation type ("all_reduce")
devices_per_node                    # int — GPUs per machine
max_devices_per_node                # int
time_stats.all_reduce.min           # float — min observed time (ms)
time_stats.all_reduce.max           # float
time_stats.all_reduce.mean          # float — training target
time_stats.all_reduce.median        # float
time_stats.all_reduce.std           # float

# File: data/profiling/network/{NETWORK_DEVICE}/send_recv.csv
# Same structure but for point-to-point send/recv (used for PP communication)