[yitingw1]

Motivation

Three correctness fixes on the PCP path (-pcp, measured at -tp 1 -pcp 8). Each is a place where ranks disagreed about which rows they hold, or where the code assumed they agree when they do not.

Change	Kind	Effect
Sync sampled ids outside the spec branch	Invariant repair + startup guard	Closes a declared-but-false invariant; refuses PCP + TBO + spec at startup
DSpark aux hidden states under PCP	Acceptance fix	With the id sync: DSpark acceptance under PCP 59.64 % → 64.79 % on main (TP: 64.83 %)
FP4 indexer padding rows	Crash fix (exposed by #2688)	V4 PCP prefill no longer asserts and hangs when the token count is not a multiple of pcp (FP4 is the default index cache on MI355X)

This PR originally also sharded the shared expert across PCP and summed the DSpark draft FFN across PCP. #2725 reworked the PCP MoE reduction (now inside FusedMoE, routed only, shared expert on each rank's own rows), which already takes the prefill gain and sums the draft's MoE; both were dropped here. The remaining decode/memory gain of a sharded shared expert does not fit #2725's structure and is left for a separate discussion.

Technical Details

1. Syncing the sampled ids

At tp=1 PCP ranks decode redundantly but not bit-identically: ~1 bf16 ULP flips near-tie argmaxes. On DeepSeek-V3.2, 32.8 % of decode steps sample divergent ids and 1.5 % reach the next forward's input_ids, where the decode-time MoE pcp_all_reduce assumes identical rows. The broadcast was inside the spec branch only, so PCP without a drafter and prefill steps with MTP were never synced. It now runs outside it, behind one _needs_id_sync(is_deferred_out, has_drafter) predicate shared with the TP broadcast so the two cannot drift; the verdict stays synced inside the spec branch.

PCP + TBO + spec is refused at startup: TBO stripes positions per request group while the draft split stripes the batch as one sequence, so their shards disagree.

2. DSpark aux hidden states under PCP

DSpark's aux hidden states are captured by forward hooks on middle target layers. After a PCP row-split prefill those layers hold only this rank's 1/pcp round-robin stripe, but Drafter.aux_for read them as if they were the full rows - most of what the draft consumed was the previous forward's leftovers.

• aux_for now all-gathers the buffers when the target forward row-split.
• That decision is target_pcp_split: a snapshot of feat(moe): route before the rows move; per-pass real row counts; V4.1 FP8 KV #2725's context.pcp_split - the flag the model writes where it row-splits and FusedMoE gathers on - taken by ModelRunner._settle_target_pcp_split() right after run_model, in both ModelRunner.forward and RapidServe's decode forward. The snapshot is needed because drafting clears pcp_split per pass before the draft runs. It is never derived from hook-side bookkeeping, which a cudagraph replay skips (an earlier attempt that did so desynchronised ranks and hung NCCL).
• _pcp_active_for_draft_model reads the same snapshot instead of re-deriving the split from the draft model's class name, which had sent every V4 drafter other than DeepseekV4MTP to the V2 gate.
• propose() checks that aux capture is configured instead of calling aux_for(), which would have gathered a second time.

3. FP4 indexer padding rows

Prefill FP4 scorer	Finds a row's page table by	PCP padding rows
before #2688: ..._fp4_prefill	batch_id_per_q_token[r]	bid = -1, skipped
since #2688: row-group kernel	group b → block_tables[b]	own group bs+1, but the table has bs rows → assert

pcp_owned_rows_by_request now folds padding rows into the last request's group. They sit at position 0, so they read no key - the same outcome as the old skip. Both callers (V4 and the V3.2 MLA path) share the function. FP4 is native V4's default index cache everywhere but gfx942 (config.py), so on MI355X the default V4 PCP setup hit this on any prefill whose token count is not a multiple of pcp.

Test Plan

pytest tests/test_pcp_ops.py tests/model_ops/test_indexer_rows.py tests/test_dspark.py -q

# V4-Pro-DSpark, -tp 1 -pcp 8, --method dspark k=5, kv fp8, index cache fp4,
# mbt=512 (forces middle prefill chunks), gsm8k 5-shot
bash run.sh main_padfix pr_full   # full 1319, same node

Test Result

Unit tests (with the related MoE / draft-graph suites): 293 passed, 4 skipped; ruff and black clean. Each new test fails with its fix reverted.

DSpark + PCP on main (V4-Pro-DSpark, MI355X, -tp 1 -pcp 8, gsm8k 5-shot, full 1319, same node, serial). main needs the FP4 padding fix to run at all, so its arm carries that one line:

	main + padding fix	this PR (ff59d4b93)
Acceptance	59.64 %	64.79 %	+5.15 pp
avg toks/fwd	3.98	4.24
accept-length 0/1/2/3/4/5 (%)	11.6 / 12.6 / 16.7 / 15.3 / 12.6 / 31.2	7.9 / 11.9 / 14.9 / 15.7 / 13.0 / 36.6
gsm8k strict	0.9538	0.9553
asserts	0	0

The gap is this PR as a whole (aux gather and sampled-id sync) against main; it brings PCP level with tp8 (64.83 % on an earlier base). A later refactor (target_pcp_split now snapshots pcp_split) re-smoked at 63.35 % on 200 questions, in line with earlier 200-question runs (63.14-63.53 %). The runs also cover the FP4 indexer fix through real middle prefill chunks. Token sync on DeepSeek-V3.2 (measured before #2725): pcp8+MTP 94.64 % vs tp8+MTP 94.69 % acceptance, gsm8k 0.9560 / 0.9507.

Known Limitations

1. The PCP + TBO + spec guard keys on config, so it also refuses runs that would never trigger TBO.
2. DSpark's confidence_schedule (off by default) is not PCP-safe: each rank picks its own verify length and nothing broadcasts it.
3. The PCP unit tests import aiter; a CI runner without it skips them.

Submission Checklist

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

🤖 Generated with Claude Code

[github-actions]

🏷️ 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 2125 --add-label <label>