Prefill Context Parallel (PCP) Guide

For long-context serving, prefill is bottlenecked on the sequence dimension: the DSA indexer scores every query against all history, and that cost grows with sequence length and is replicated across Tensor-Parallel (TP) ranks. Plain TP shards weights/heads/experts, not tokens, so it cannot reduce this cost.

Prefill Context Parallel (PCP) is an independent parallelism dimension that splits the prefill token sequence across the PCP process group, so each GPU processes only 1/pcp of the tokens during prefill. This cuts the per-GPU prefill work (and the indexer’s sequence-length cost) to 1/pcp, lowering TTFT and raising long-prefill throughput. Decode is left unchanged. PCP composes with TP and Expert Parallelism (EP): the total world size is world = tp × pcp.

  pcp = 2, prefill tokens: 0 1 2 3 4 5
    GPU (pcp rank 0):  0   2   4      ← each GPU processes 1/pcp of the tokens
    GPU (pcp rank 1):    1   3   5
  Full KV is kept on every rank; the 1/pcp outputs are all-gathered back to the
  full sequence before the LM head. Decode runs as usual (no split).

Model support. PCP currently supports DeepSeek-V4 only. Support for more models will be added incrementally.

When to use PCP

  • Best fit: long-context / large-prompt prefill on DeepSeek-V4, where prefill TTFT dominates.

  • Combine with: --enable-tbo (prefill) to overlap the MoE communication PCP introduces (see Overlapping communication with TBO). TBO is only usable with ATOM_PCP_MOE_MERGE=1 and -tp 1 (see Constraints & Compatibility).

  • Requires: world = tp × pcp GPUs, e.g. -tp 4 -pcp 2 on 8 GPUs.

  • Little benefit / avoid: decode-heavy or short-prompt workloads. PCP only shards prefill tokens to lower TTFT on long sequences; decode is left unsharded and runs redundantly across the PCP ranks. For long-decode (large output_len) workloads PCP can hurt TPOT — do not enable it there; use TP/EP as usual.

Quick Reference

Flag / Variable

Default

Purpose

-pcp N / --prefill-context-parallel-size N

1

Enable PCP with size N (world = tp × pcp)

ATOM_PCP_MOE_MERGE

1

Whether to shard MoE across the PCP ranks too

--enable-tbo [prefill|all]

off

Overlap compute with PCP communication. With PCP, only prefill TBO is supported, and only when ATOM_PCP_MOE_MERGE=1 and -tp 1

--no-enable_chunked_prefill

chunked on

Disable chunked prefill. Recommended with PCP (see Tuning for long sequences)

--max-num-batched-tokens N

16384

Max tokens scheduled per step. Raise (e.g. 131072) for long-sequence PCP so a full long prompt is prefilled in one step

Goal (8 GPUs)

Command

Long-context prefill

-tp 4 -pcp 2

Long-context prefill + overlap

-tp 1 -pcp 8 --enable-tbo

Disable PCP (baseline)

-tp 8 -pcp 1

TBO requires -tp 1. Two-batch overlap is only supported with ATOM_PCP_MOE_MERGE=1 and TP=1 (all GPUs go to PCP, e.g. -tp 1 -pcp 8).

CLI usage

-pcp N                          # or --prefill-context-parallel-size N; world = tp × pcp
--enable-tbo                    # prefill-only TBO overlap; requires ATOM_PCP_MOE_MERGE=1 and -tp 1 (prefill only supported with PCP)

ATOM_PCP_MOE_MERGE=1            # default: shard MoE across PCP ranks (gather/scatter)
ATOM_PCP_MOE_MERGE=0            # run MoE per-rank on its 1/pcp shard, no extra MoE comm

ATOM_PCP_MOE_MERGE only has an effect when PCP is enabled (pcp > 1):

Value

MoE behaviour

When to use

1 (default, recommended)

PCP is folded into the MoE tensor/expert sharding, so MoE weights are also sharded across PCP ranks. Lowers per-GPU MoE weight memory, at the cost of one hidden gather/scatter per MoE layer (which TBO overlaps).

Most deployments

0

Each GPU runs MoE independently on its 1/pcp shard with no extra MoE communication; MoE weights are replicated across PCP ranks.

Avoid extra MoE comm and have memory headroom for replicated MoE weights

Launching server

DeepSeek-V4: TP4 + PCP2 (8 GPUs)

python -m atom.entrypoints.openai_server \
    --model deepseek-ai/DeepSeek-V4 \
    -tp 4 -pcp 2 \
    --kv_cache_dtype fp8

DeepSeek-V4: TP1 + PCP8 + prefill TBO overlap (8 GPUs)

TBO requires ATOM_PCP_MOE_MERGE=1 (the default) and -tp 1 — put every GPU into PCP.

ATOM_PCP_MOE_MERGE=1 \
python -m atom.entrypoints.openai_server \
    --model deepseek-ai/DeepSeek-V4 \
    -tp 1 -pcp 8 \
    --enable-tbo \
    --kv_cache_dtype fp8

Tips:

  • -tp 8 -pcp 1 (or omitting -pcp) disables PCP and serves as the baseline.

  • --enable-tbo overlaps the MoE communication introduced by ATOM_PCP_MOE_MERGE=1. It only helps in that mode: with ATOM_PCP_MOE_MERGE=0 there is no MoE communication to overlap, so TBO is auto-disabled (a warning is logged).

  • --enable-tbo additionally requires -tp 1; with tp > 1 it may hang under long-sequence / high-concurrency workloads (not yet supported). Give all GPUs to PCP instead, e.g. -tp 1 -pcp 8.

  • Under PCP, TBO uses a request-boundary split (each micro-batch is a whole subset of requests) — the non-default TBO split mode — instead of the token-midpoint split used without PCP. ATOM_TBO_PREFILL_TOKEN_SPLIT therefore has no effect when PCP is enabled.

  • --torch-profiler-dir ./log can be added to collect traces for performance analysis.

Tuning for long sequences

PCP targets long-prefill TTFT, but the default scheduler settings work against it. When serving long prompts:

  • Disable chunked prefill (--no-enable_chunked_prefill), or enlarge the chunk (--attn_prefill_chunk_size). It is on by default and splits a long prompt across steps, so PCP gets fewer tokens to shard per step and TBO’s request-boundary split (needs ≥2 whole sequences per step) falls back to the non-overlapped path.

  • Raise --max-num-batched-tokens to 131072 (≥ input_len; ≥ 2 × input_len for TBO’s balanced split). The default 16384 forces a long prompt across multiple steps instead of prefilling it in one.

  • Avoid PCP for long-decode workloads. PCP shards only prefill; decode runs unsharded. A large output_len is decode-dominated, where PCP adds no benefit and can raise TPOT — use plain TP/EP.

Example (long-context prefill, chunked off, larger batch budget):

ATOM_PCP_MOE_MERGE=1 \
python -m atom.entrypoints.openai_server \
    --model deepseek-ai/DeepSeek-V4 \
    -tp 1 -pcp 8 --enable-tbo \
    --no-enable_chunked_prefill \
    --max-num-batched-tokens 131072 \
    --kv_cache_dtype fp8

Performance baseline

Benchmark against a running server with a long input length (PCP targets prefill):

python -m atom.benchmarks.benchmark_serving \
  --model=deepseek-ai/DeepSeek-V4 --backend=vllm --base-url=http://localhost:7777 \
  --dataset-name=random \
  --random-input-len=32768 --random-output-len=512 \
  --num-prompts=128 --max-concurrency=64 \
  --request-rate=inf --ignore-eos

Compare -tp 4 -pcp 2 against the -tp 8 baseline and watch Mean TTFT and output throughput; the gap widens as --random-input-len grows.

PCP was introduced in ROCm/ATOM#1220, which reported, on 8×MI308 for -tp 4 -pcp 2 vs -tp 8, a 35–43% Mean-TTFT reduction and up to ~49% higher throughput on long prefill. Actual gains depend on model, sequence length, and hardware.

PCP + TBO (prefill overlap)

TBO’s benefit is hardware-dependent:

  • MI308: current testing shows PCP + TBO gives almost no gain over PCP alone — the overlap does not meaningfully hide the MoE communication on this hardware. Prefer plain PCP (-tp 4 -pcp 2) here.

  • MI355: PCP + TBO does deliver a speedup. Enable it with ATOM_PCP_MOE_MERGE=1 -tp 1 -pcp 8 --enable-tbo.

Because TBO requires -tp 1, compare TBO configs against the matching -tp 1 PCP baseline (not the -tp 4 -pcp 2 numbers above).

Constraints & Compatibility

Constraint

Notes

Models

DeepSeek-V4 only (more coming)

World size

tp × pcp ≤ 8; multi-node not yet validated

PCP + DP-attention

Not supported (raises at startup)

PCP + TBO decode (--enable-tbo all)

Not supported (raises at startup); use --enable-tbo prefill-only

ATOM_PCP_MOE_MERGE=0 + --enable-tbo

TBO auto-disabled (warning logged)

PCP + TBO with tp > 1

Not supported. May hang under long-sequence / high-concurrency workloads. TBO requires -tp 1 (all GPUs in PCP, e.g. -tp 1 -pcp 8)

PCP + TBO prefill is supported only with ATOM_PCP_MOE_MERGE=1 and -tp 1. Decode is unchanged by PCP in all configurations.

How it works

  1. At the start of the prefill forward, the token sequence is split round-robin across the PCP ranks (token i → rank i % pcp), padded so the count divides evenly.

  2. Each rank runs attention / indexer / compressor on its 1/pcp token shard. The full KV is kept on every rank (all-gathered as needed), so the attention kernels are unchanged.

  3. MoE either runs on the local 1/pcp shard (ATOM_PCP_MOE_MERGE=0) or gathers to the full sequence and scatters back (=1, default).

  4. After the final layer, the 1/pcp hidden states are all-gathered back to the full sequence, the original token order is restored, and the LM head runs.

  5. Decode is untouched: every rank keeps the full KV and runs normally, so PCP adds no decode-time cost.

Source Files

File

Description

atom/model_engine/arg_utils.py

--prefill-context-parallel-size, --enable-tbo CLI

atom/utils/envs.py

ATOM_PCP_MOE_MERGE

atom/distributed/pcp_utils.py

PCP communication and helper primitives

atom/models/deepseek_v4.py

DeepSeek-V4 PCP forward path and MoE handling

atom/model_ops/attentions/deepseek_v4_attn.py

PCP attention metadata (incl. PCP + TBO prefill)

atom/model_engine/model_runner.py

PCP token split and PCP + TBO grouping

atom/model_engine/llm_engine.py

PCP / TBO / DP-attention validation