[feat][PCP] Enable PCP MoE-Merge+TBO for DeepSeek V4 by yitingw1 · Pull Request #1553 · ROCm/ATOM

yitingw1 · 2026-07-10T06:27:05Z

Motivation

PCP (Prefill Context Parallel, #1220) splits the prefill token sequence across the
PCP group but left MoE running per-rank on each 1/pcp shard, so MoE weights were replicated across PCP ranks. This PR adds a MoE-merge mode that folds the PCP dimension into the MoE tp/ep sharding, and enables PCP + TBO to overlap the extra MoE communication that introduces during prefill.

ATOM_PCP_MOE_MERGE=1 (default): all-gather hidden 1/pcp → full before MoE and slice back after, so MoE weights are sharded across PCP ranks (lower per-GPU MoE memory) at the cost of one hidden gather/scatter per MoE layer.
--enable-tbo (prefill) overlaps that gather/scatter with compute via two-batch overlap, using a request-boundary micro-batch split.

Known Limitations

PCP+TBO requires -tp 1; with tp > 1 it may hang under certain conditions (long sequence + high concurrency), so it is not yet supported.
TBO requires ATOM_PCP_MOE_MERGE=1; auto-disabled with a warning when merge is off.
Prefill TBO only — decode TBO (--enable-tbo all) and PCP + DP-attention are unsupported.
DeepSeek-V4 only; multi-node not yet validated.

Test Plan

Serve DeepSeek-V4 on 8 GPUs and compare long-prefill TTFT / throughput:

# baseline PCP (no TBO)
python -m atom.entrypoints.openai_server --model deepseek-ai/DeepSeek-V4 \
    -tp 4 -pcp 2 --kv_cache_dtype fp8

# PCP + TBO (requires ATOM_PCP_MOE_MERGE=1(default) and -tp 1)
python -m atom.entrypoints.openai_server \
    --model deepseek-ai/DeepSeek-V4 -tp 1 -pcp 8 --enable-tbo --kv_cache_dtype fp8

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

Test Result

Correctness: ATOM_PCP_MOE_MERGE=0/1 produce consistent outputs; PCP+TBO
prefill matches non-TBO PCP.

MoE-merge performance (`-tp 4 -pcp 2`, chunked-prefill OFF, output 1k)

benchmark_serving random dataset (range_ratio 0.8, ignore-eos), 8×MI308. Cell =
Mean TTFT(ms) / Total Tput(tok/s); parentheses = improvement vs the TP8 baseline
(TTFT ↓ reduction, Tput ↑ gain).

Input-length sweep (concurrency 8):

Config	8k	16k	32k	64k	128k
TP8 (baseline)	3588 / 2742	8012 / 3879	14908 / 4700	48550 / 4714	126147 / 3801
TP4PCP2 `ATOM_PCP_MOE_MERGE=0`	2307 (↓35.7%) / 2832 (↑3.3%)	4920 (↓38.6%) / 4410 (↑13.7%)	9347 (↓37.3%) / 6082 (↑29.4%)	28527 (↓41.2%) / 6939 (↑47.2%)	71656 (↓43.2%) / 6310 (↑66.0%)
TP4PCP2 `ATOM_PCP_MOE_MERGE=1`	2126 (↓40.7%) / 2892 (↑5.5%)	4175 (↓47.9%) / 4311 (↑11.1%)	10631 (↓28.7%) / 5894 (↑25.4%)	31513 (↓35.1%) / 6546 (↑38.9%)	76748 (↓39.2%) / 5969 (↑57.0%)

Concurrency sweep (input 8k):

Config	c8	c16	c32	c64
TP8 (baseline)	3588 / 2742	7075 / 4119	7931 / 5525	15145 / 6583
TP4PCP2 `ATOM_PCP_MOE_MERGE=0`	2307 (↓35.7%) / 2832 (↑3.3%)	4100 (↓42.0%) / 4217 (↑2.4%)	6881 (↓13.2%) / 6313 (↑14.3%)	11285 (↓25.5%) / 8037 (↑22.1%)
TP4PCP2 `ATOM_PCP_MOE_MERGE=1`	2126 (↓40.7%) / 2892 (↑5.5%)	3546 (↓49.9%) / 4364 (↑5.9%)	6006 (↓24.3%) / 6270 (↑13.5%)	11145 (↓26.4%) / 7843 (↑19.1%)

Both PCP modes beat the TP8 baseline across the board, with the gap widening as sequence length grows. ATOM_PCP_MOE_MERGE=1 shards experts 8-way (fits large models single-node) while staying close to =0; they converge at long sequences where attention O(n²) dominates the extra per-layer all-gather.

Heavy-prefill + PCP-degree sweep (input 64k / output 1k, c32, 8×MI308)

mbt=131072, compile + CUDAGraph, chunked-prefill OFF, util 0.8, TP MoE. Only the
-tp 1 config runs TBO (TBO requires -tp 1).

Config	TTFT (ms)	TPOT (ms)	Tput (tok/s)
TP8 (baseline)	81379	248.8	5905
DP8 (baseline)	63640	149.7	9043
DP8 + TBO	76301	152.2	8427
TP4PCP2 `ATOM_PCP_MOE_MERGE=0`	50246	157.3	9334
TP4PCP2 `ATOM_PCP_MOE_MERGE=1`	54557	173.8	8493
TP2PCP4 `ATOM_PCP_MOE_MERGE=1`	45747	154.9	9678
TP1PCP8 `ATOM_PCP_MOE_MERGE=1`	36132	130.9	11572
TP1PCP8 `ATOM_PCP_MOE_MERGE=1` + TBO	36653	131.4	11509

PCP degree drives the win (from ATOM_PCP_MOE_MERGE=1, not TBO): TP8 → TP1PCP8 gives TTFT ↓56% (81379 → 36132) and Tput ↑96% (5905 → 11572). Higher PCP degree is better under heavy prefill.
TBO gives no gain on MI308: TP1PCP8 with vs without TBO is 36653 vs 36132 (TTFT slightly worse, Tput flat). The overlap-able compute per rank shrinks as PCP degree rises, so PCP collectives are too light for TBO to hide.

Decode-length sweep (`-tp 4 -pcp 2`, input 32k, c8, 8×MI308)

Fixed input 32768 / concurrency 8 / util 0.8 / mbt 131072 / chunked-prefill OFF,
sweeping output length. Cell = Total Tput(tok/s) / Mean TPOT(ms).

Config	out 2k	out 4k	out 8k	out 16k
TP8 (baseline)	3768 / 27.78	2608 / 22.47	1795 / 19.78	1173 / 18.40
TP4PCP2 `ATOM_PCP_MOE_MERGE=0`	4412 / 25.48	2767 / 22.20	1787 / 20.57	1113 / 19.78
TP4PCP2 `ATOM_PCP_MOE_MERGE=1`	4314 / 25.62	2744 / 22.17	1802 / 20.24	1136 / 19.28

End-to-end latency (E2E = TTFT + (out−1)×Mean TPOT, seconds):

Config	out 2k	out 4k	out 8k	out 16k
TP8 (baseline)	71.2	105.9	176.3	315.7
TP4PCP2 `ATOM_PCP_MOE_MERGE=0`	60.9	99.4	177.2	332.6
TP4PCP2 `ATOM_PCP_MOE_MERGE=1`	62.2	100.2	175.5	325.5

Heavy decode should not use PCP. PCP can improve and may hurt TPOT.

PCP+TBO (prefill overlap)

TBO benefit is hardware-dependent: on MI308 it gives almost no gain (the PCP collectives are too light for the overlap to hide), while on MI355 (stronger compute → prefill compute is shorter → PCP communication is a larger share → more for TBO to overlap) it delivers a clear speedup.

MI355, shipped config -tp 1 -pcp 8 (64k/1k, c32, GPU_MAX_HW_QUEUES=5) — cell =
Mean TTFT / Median TTFT / P99 TTFT(ms) / TPOT(ms) / Tput(tok/s):

Config	Mean TTFT	Median TTFT	P99 TTFT	TPOT	Tput
TP1PCP8 `ATOM_PCP_MOE_MERGE=1`	6005	2867	22211	42.52	39545
TP1PCP8 `ATOM_PCP_MOE_MERGE=1` + TBO	5376 (−10.5%)	2415 (−15.8%)	17194 (−22.6%)	40.12 (−5.6%)	42174 (+6.6%)

Appendix. Technical Details

atom/utils/envs.py: new ATOM_PCP_MOE_MERGE (default 1); only active when pcp > 1.
atom/models/deepseek_v4.py, atom/model_ops/moe.py: MoE-merge path — all-gather hidden to full sequence before MoE, scatter back to 1/pcp after.
atom/model_engine/model_runner.py, atom/model_ops/attentions/deepseek_v4_attn.py,
atom/utils/tbo/ubatch_wrapper.py: PCP+TBO prefill grouping and attention metadata. PCP+TBO uses ATOM_TBO_PREFILL_TOKEN_SPLIT=0 (request-boundary split — each micro-batch is a whole subset of requests), not the default token-midpoint split used without PCP. Reason: PCP round-robin-shards tokens within each request across ranks, so a token-midpoint split would cut across the PCP sharding and desync the micro-batches; splitting on request boundaries keeps each request's tokens together. ATOM_TBO_PREFILL_TOKEN_SPLIT therefore has no effect under PCP.
atom/model_engine/llm_engine.py: validation — TBO requires ATOM_PCP_MOE_MERGE=1; auto-disable + warn when merge is off; decode TBO (--enable-tbo all) and DP-attention remain unsupported with PCP.
Limit PCP+TBO to -tp 1: with tp > 1 PCP+TBO hangs, so TBO is only allowed when tp == 1 (all GPUs go to PCP, e.g. -tp 1 -pcp 8).
docs/context_parallel_guide.md: usage, constraints, and the TP=1 / hardware notes.

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

P0: Relax PCP+TBO raise in llm_engine.py — only block TBO-decode+PCP; prefill-only TBO now allowed. P1 (coordinated split, b1 approach): - pcp_utils.py: add num_ubatches param to pcp_pad_len (N·pcp alignment) - model_runner.py: _preprocess overrides TBO eligibility with 1/pcp local token count; Option A gate (n_prefill % 2*pcp == 0) avoids dummy tokens; prepare_inputs builds ubatch_slices in 1/pcp space with correct request boundaries; run_model does PCP split before UBatchWrapper, updates context.positions and forward_vars["positions"/"cu_seqlens_q"]; out-of- graph pcp_allgather_rerange after UBatchWrapper. - ubatch_wrapper.py: _make_ubatch_context carries input_ids per-ubatch slice for hash-MoE (moe_pcp_merge mode B). - deepseek_v4.py: ForCausalLM.forward skips PCP split and final pcp_allgather_rerange when tbo_active(); per-ubatch input_ids allgather for hash-MoE. - deepseek_v4_attn.py: build_ubatch_prefill_metadata uses real_num_tokens for _attach_v4_per_fwd_meta; Option B (arbitrary length) planned as per-ubatch reindex — see PCP_TBO.md §11. P2: moe_pcp_merge_forward inserts TBO yield/switch around pcp_allgather_rankmajor and pcp_reduce_scatter. Attention-internal PCP collectives (compressor kv_score, kv, positions allgather) wrapped with tbo_yield_and_switch / tbo_switch_to_compute_sync to serialize RCCL submissions across TBO ubatch threads. Add ATOM_TBO_DEBUG env var and per-collective debug logging in pcp_utils.py for locating RCCL hang root cause (see PCP_TBO.md §9.5). Verified: warmup ✓, gsm8k=0.86 (Option A mixed path), server stable. Option B (arbitrary length) deferred — requires per-ubatch reindex. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

…erge_TBO

… into enable_deepseek_PCP_merge_TBO

…erge_TBO

github-actions · 2026-07-10T06:27:29Z

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

✅ Run after the PR is approved and Pre Checkin passes
✅ Run immediately when an approval review is submitted
✅ Can be requested before approval with labels

Label	Tests
`ci:full`	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
`ci:atom`	Run native ATOM model accuracy tests
`ci:vllm`	Run ATOM vLLM OOT model accuracy tests
`ci:sglang`	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1553 --add-label <label>

ZhangLirong-amd · 2026-07-10T07:07:39Z

+            and not batch.is_dummy_run
+            and getattr(self, "_pcp_tbo_balanced_active", False)
+        )
+        if _pcp_tbo_balanced:


can we make prepare_inputs more clean? Line 1897~1938 PCP+tbo related in a new function?

yitingw1 and others added 18 commits July 1, 2026 08:12

[Fea] add moe_pcp_merge

69737e2

fix moe_pcp_merge

97c87ad

Enable PCP+TBO

b851744

Merge remote-tracking branch 'origin/main' into enable_deepseek_PCP_m…

8385d8f

…erge_TBO

Merge remote-tracking branch 'origin/main' into enable_deepseek_PCP_m…

f4ad844

…erge_TBO

Using ATOM_PCP_MOE_MERGE

74227ef

Remove debug log

eb01b8b

Refactor comments

ff2da6f

Merge remote-tracking branch 'origin/main' into enable_deepseek_PCP_m…

9059e86

…erge_TBO

Refactor comments 2

b5b629f

Add PCP doc

d7fec57

Merge remote-tracking branch 'origin/main' into enable_deepseek_PCP_m…

87291f5

…erge_TBO

Merge remote-tracking branch 'yitingw1/enable_deepseek_PCP_merge_TBO'…

a1bc75a

… into enable_deepseek_PCP_merge_TBO

Merge remote-tracking branch 'origin/main' into enable_deepseek_PCP_m…

d74a390

…erge_TBO

Limit PCP+TBO only when TP=1

42870d4

Update docs

5b98657

Update docs

94a3b22

yitingw1 added 2 commits July 10, 2026 06:36

Fix format

461bb5e

Fix format

eaad013

yitingw1 requested review from ZhangLirong-amd and valarLip July 10, 2026 06:50

ZhangLirong-amd reviewed Jul 10, 2026

View reviewed changes

Refactor model_runner

de6595b

yitingw1 requested a review from junhaha666 July 10, 2026 08:11

Update docs

272505e

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[feat][PCP] Enable PCP MoE-Merge+TBO for DeepSeek V4#1553

[feat][PCP] Enable PCP MoE-Merge+TBO for DeepSeek V4#1553
yitingw1 wants to merge 22 commits into
ROCm:mainfrom
yitingw1:enable_deepseek_PCP_merge_TBO

yitingw1 commented Jul 10, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 10, 2026

Uh oh!

ZhangLirong-amd Jul 10, 2026

Uh oh!

yitingw1 Jul 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

yitingw1 commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Test Plan

Test Result

MoE-merge performance (-tp 4 -pcp 2, chunked-prefill OFF, output 1k)

Heavy-prefill + PCP-degree sweep (input 64k / output 1k, c32, 8×MI308)

Decode-length sweep (-tp 4 -pcp 2, input 32k, c8, 8×MI308)

PCP+TBO (prefill overlap)

Appendix. Technical Details

Submission Checklist

Uh oh!

github-actions Bot commented Jul 10, 2026

🏷️ CI Guide

Uh oh!

ZhangLirong-amd Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

yitingw1 Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yitingw1 commented Jul 10, 2026 •

edited

Loading

MoE-merge performance (`-tp 4 -pcp 2`, chunked-prefill OFF, output 1k)

Decode-length sweep (`-tp 4 -pcp 2`, input 32k, c8, 8×MI308)