[MengqingCao]

Summary

Every TP worker of an engine, and every rank of its PCP group, builds its KV pool at the block count rank 0 reports (call_func surfaces rank 0's answer). A worker whose own estimate is smaller can therefore be asked for more blocks than it budgeted.

Before sizing the pools, this takes the minimum of free and the KV budget over the TP group and, when it spans more than one rank, the PCP group. Every rank then sizes from the same inputs:

• The two checks that read this rank's memory, the KV connector reserve against free and InsufficientPoolBudget from plan_pools, raise on all ranks together or on none. A rank short of memory can no longer raise alone while its peers wait in a reduction.
• plan_pools, a pure function of those inputs, builds the same plan on each rank, so no block count has to be folded back into it afterwards.

The PCP group is gated on pcp_group_size, not prefill_context_parallel_size, because Ulysses SP runs over the same group with PCP=1. The Memory budget log line now prints the agreed free and available_for_kv.

It changes behaviour only with PP=1, where get_num_blocks had no collective at all. With PP>1 the existing reduction after plan_pools already takes the minimum block count over the default (WORLD) group, TP and PCP ranks included, so the counts already agree; there the new agreement leaves them unchanged and costs two single-int64 gloo all-reduces per group at startup, kept so there is one code path. PP stages can still raise apart, as on main.

Split out of #3147 (decode P/D admission), where it was found. It is the bottom of the stack #3171 ← #3147 ← #3173 ← #3165 and does not depend on the PRs above it.

Commit 443b280 (three CI entries, kv_block_ranks, check_kv_block_agreement.py and its workflow step) is only there to publish a multi-rank run. It will be dropped before this PR lands.

Test plan

• tests/test_model_runner_staging_fence.py: a source-level check that free and the KV budget are reduced with group_min over each TP and PCP group's CPU group, the PCP group gated on pcp_group_size, before free_for_kv, with nothing in between overwriting them. It fails on MAX, the device group, a missing budget reduction, the earlier post-plan_pools placement, the agreement moved after free_for_kv or after the plan is published, and a re-read of free in between. 7 passed.
• Black and Ruff on the changed files.
• Multi-GPU startup on 443b280, which reduced the block count after plan_pools: TP4 + PCP2 (DeepSeek-V4-Pro) and TP2 (Llama-3-8B) report one count on every rank. Log lines and gsm8k in fix(runner): reconcile KV block counts across TP and PCP-group workers #3171 (comment).
• The same entries on 932aefb, with the agreement before plan_pools: ATOM Test in progress.
• SP>1: blocked on [FlyDSL] dist: public SP APIs: PCP graph_capture, head exchange, all_to_all, MegaMoE setup errors aiter#7346. Ulysses SP requires it, and the CI AITER wheel lacks it, so the SP4 entry is refused at startup.

[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 3171 --add-label <label>

[valar-cop]

a few small things..

tp>1 or sp/pcp>1 run to show the startup collective completes and all ranks agree?

ci green.

[valar-cop]

ci triage for 6d7d49a:

ATOM Test passed. https://github.com/ROCm/ATOM/actions/runs/43054467486

ATOM vLLM Test passed. https://github.com/ROCm/ATOM/actions/runs/43054467227

ATOM SGLang Test passed. https://github.com/ROCm/ATOM/actions/runs/43054467035

[yitingw1]

Reviewed 443b280f (+189/-0, 8 files), full diff. Also read get_num_blocks end to end on main, plan_pools/PoolPlan, group_agreement.py, the accuracy job in atom-test.yaml, the launch half of atom_test.sh, and capabilities.py.

Not repeating @valar-cop's two inline points — both are fixed in 169169f3/443b280f, and I confirmed the test fix has teeth (see below).

Verdict

The runner fix should go in. The CI half should be a separate PR.

The bug is real and universal, not a corner case. I ran the new checker against 25 real, unmodified, pre-fix ATOM TP8 server logs I had locally: 25/25 flagged, every one with all 8 ranks disagreeing, spread 0.57%–1.51%.

$ check_kv_block_agreement.py acc_GLM-5.2_serve.log --ranks 8
::error title=KV block counts::ranks disagree on num_kvcache_blocks:
  [29453, 29462, 29470, 29518, 29595, 29609, 29620, 29905]

The CI half (+130 lines, 3 new test_level: pr 8-GPU entries) has no stated justification and one of the three models was deliberately removed in #2135. More below.

What I verified by running it, not by reading

1. The checker works on production logs. The 25-log run above confirms three premises at once: the regex matches the real log format; all 8 ranks log and all land in the single file atom_test.sh captures (it is not rank-0-only); and the check catches what it is meant to catch.
2. The group choice exactly matches the engine's worker set. capabilities.py:128 — workers_per_engine = tp_world_size * pcp_group_size_of(config), no DP. So the README's "leave it unset with DP attention" is right, and TP-then-PCP min is the global min over the 2-D grid.
3. Reducing only the PAGE count is complete, not an omission. In plan_pools, STATE counts are max_num_seqs * entries_per_req + extra_entries — config-derived, identical on every rank. Only remaining // entry_bytes is budget-derived. paged_entries is the only quantity that can diverge.
4. if reduced != num_kvcache_blocks is load-bearing, not an optimization. A model with no PAGE class has paged_entries == 0 everywhere, and with_paged_entries raises ValueError when paged_class is None. The guard is what keeps those models out of that line.
5. The new AST test has teeth. Six mutations, all red: group_min→group_max; delete the with_paged_entries fold; get_tp_group()→get_dp_group(); gate on prefill_context_parallel_size; drop the PCP group entirely; .cpu_group→.device_group.
6. Tests pass: test_check_kv_block_agreement.py 3/3, test_model_runner_staging_fence.py 7/7.

---

🔴 1 — the new collective turns one rank's clear error into a hang on all the others

In the PR's get_num_blocks (lines 2037–2328), the new collective at 2173 has two rank-local raises in front of it:

line 2077  RuntimeError   # free - connector_bytes <= 0      -> depends on this rank's `free`
line 2146  RuntimeError   # InsufficientPoolBudget           -> depends on this rank's `available_for_kv`
line 2173  group_min(...) # new

Both depend on this rank's own mem_get_info(). That means they fire precisely in the scenario this PR exists to fix — one rank has less memory than its peers. That rank raises at 2146 and exits; the others block in the gloo all-reduce at 2173 until the PG times out (no explicit timeout is set in ATOM, so torch's gloo default of 30 min applies).

Concrete trigger: TP8 where any rank's available_for_kv falls below the STATE floor while rank 0's does not. The measured cross-rank spread is 0.57%–1.51%, so any config tuned near the hinted minimum --gpu-memory-utilization, or one GPU carrying ~1% more from another process, sits in that window. And 2146 is exactly the error that tells the user "set --gpu-memory-utilization >= X" — it becomes a silent 30-minute hang instead.

This convention is already written down twice in this same function: _agree_pipeline_checkpoint_spec's docstring ("one raising alone would leave the others waiting here for its geometry") and the checkpoint_spec comment ("raising before the collective would leave its peers waiting in it").

Note main does not have this hazard: with PP=1 there is no collective in get_num_blocks at all, so a short rank dies immediately and the launcher sees it. This PR introduces it for every TP>1 deployment, i.e. nearly all of them.

Ask (either, first preferred):

1. Move the agreement before plan_pools — group_min on available_for_kv instead of on num_kvcache_blocks. That fixes three things at once: both raises become uniform across ranks; the whole plan becomes bit-identical on every rank rather than just paged_entries; and the with_paged_entries refold is no longer needed. plan_pools documents itself as a pure function, so equal inputs give equal plans.
2. Keep the current placement but agree success/failure first with group_agree — already in the group_agreement.py you are importing from — so every rank raises the same error together.

⚠️ 2 — the new CI step sits before Check accuracy test results and skips it

The step is inserted between Run ATOM accuracy test and Dump server log with if: ${{ !cancelled() && ... }}. The two steps after it, Check accuracy test results and Check MTP acceptance rate, are both if: success().

So a KV-block disagreement makes this step red, which skips the accuracy threshold comparison and the MTP acceptance gate entirely. A PR that breaks both block agreement and accuracy only ever shows the first, and the measured accuracy is never compared against its threshold at all. That is a steep price for an incidental startup self-check.

Ask: move this step after Check MTP acceptance rate so it is the job's last gate, or restructure it so it does not affect success() for the steps behind it.

⚠️ 3 — Meta-Llama-3-8B-Instruct was deliberately removed in #2135; the PR body does not mention it

The _baseline_note says "Bar of the former TP1 entry", so you know it is former. But main has no accuracy entry for this model: 27ed7c9 (#2135, 2026-09-03) removed it, with this reason:

It has been failing on main on every recent run, and the model is no longer part of what ATOM needs to track.

It is now coming back at test_level: pr. Two things:

• Why does "failing on main on every recent run" no longer apply? If it still fails, every PR gains a permanently red job. If something fixed it since, say which change.
• ci: drop the Meta-Llama-3-8B-Instruct accuracy test #2135 also removed the atomesh standalone whitelist key pointing at it and recorded the consequence ("atomesh standalone now validates zero models on a pull request"). Re-adding the model does not restore that key; the asymmetry deserves a line in the body.

Ask: either use a model already in the PR set for the TP2 check (e.g. add kv_block_ranks to the existing gpt-oss-120b entry — zero new jobs), or address #2135's removal reason directly in the PR body.

⚠️ 4 — three new 8-GPU PR-level jobs is a lot for a startup invariant

PR-level accuracy entries go 14 → 17 (+21%), each an 8-GPU gsm8k run with up to a 60-minute timeout. The invariant itself is settled in the first seconds of startup; running a full gsm8k and then grepping the log spends ~99% of that machine time on something unrelated to the check.

The severity does not carry that price either. From a real GLM-5.2 TP8 log: available_for_kv is 41.39 vs 40.76 GB across ranks, block counts 29905 vs 29453 — the smallest rank is asked for 0.63 GiB more than it budgeted, against safety=3.84 GB. Today the bug is absorbed by the safety margin: a genuine latent inconsistency, not an active OOM source. (I have no data on whether the spread is much larger under PCP/SP, where ranks genuinely do different work — see below.)

Ask: add kv_block_ranks to existing entries (DSV4-Pro TP8 → 8, MiniMax-M3-MXFP4 TP8 → 8, Kimi-K3 DSpark PP2 → its value) to cover the TP path with zero new jobs; move the PCP2 and SP4 topologies to nightly if they must be covered. Better still, make this a launch-only smoke (start, wait for ready, grep, kill) that does not run gsm8k at all — minutes instead of an hour.

⚠️ 5 — the AST test pins which statements exist, not where the block sits (verified bypass)

I moved the whole reduction block to after self.pool_plan = plan / config.pool_entries = ... and the test stayed green. With that edit the engine publishes the unreduced plan and the bug is fully back; only the CI check would notice, and only on the three opted-in entries.

This is an inherent limit of source-level tests, not an oversight — the file docstring explains why AST is the only option here ("Importing ModelRunner pulls in the GPU attention stack"). But closing it is cheap.

Ask (either):

1. One more AST assertion: the reduction block's lineno must be less than the self.pool_plan = plan assignment's.
2. Better: extract the 12 lines into a pure function (e.g. reconcile_paged_entries(plan, groups)) next to group_agreement.py and test it behaviourally with fake groups, leaving get_num_blocks with a one-line call. The AST test then degrades to "the call is still there" and stops going falsely red on refactors.

Minor and not necessarily this PR's to fix: a KV-block test living in test_model_runner_staging_fence.py (a file about staging-buffer lifetime) is a bit misfiled.

📝 6 — under PP>1 the new collectives are pure overhead; the body slightly overstates the scope

The existing PP reduction right above uses torch.distributed.all_reduce(t, op=MIN) with no group=, i.e. the default WORLD group. So with PP>1 main already takes the global minimum, TP and PCP included, and the two new sub-group reductions cannot change the result.

"It applies to every deployment" is literally true, but the only case it actually fixes is PP=1; with PP>1 it is two redundant gloo collectives. Small related inconsistency: the existing reduction runs on GPU/nccl (device=self.device), the new one on CPU/gloo.

Ask: state the scope precisely in the body ("fixes PP=1; with PP>1 the existing WORLD reduction already covers it"), and either skip the new reduction when PP>1 or add a line saying why it is kept.

📝 7 — the adjacent PP reduction is missing the same guard (pre-existing)

The new code guards with_paged_entries behind if reduced != num_kvcache_blocks, which is what keeps PAGE-less models safe. The PP reduction above calls with_paged_entries unconditionally, so PP>1 with no PAGE class would hit ValueError("no PAGE entry class to resize"). Not introduced here.

Ask: since the new code demonstrates the guard is needed, apply the same condition above, or add a comment explaining why the PP path cannot reach that state.

📝 8 — nothing ties kv_block_ranks to extraArgs

kv_block_ranks: 8 and extraArgs: "-tp 4 -pcp 2" are two independently maintained numbers. Someone changing -tp 4 to -tp 8 without touching kv_block_ranks gets expected 8 'Memory budget' lines, found 16 rather than "you changed the topology".

Ask: either infer the rank count from the log itself (count the lines, assert only that the value is unique), or add a catalog schema validation rule that kv_block_ranks == tp * max(pcp, sp) parsed from extraArgs.

---

What I did not verify

Flagging these as my own gaps, not as process asks:

• No multi-GPU run. Your test plan's "Multi-GPU startup with TP>1 and with SP>1" is still unchecked, and my 🔴1 hang path is derived from reading the code and AST line numbers — I did not actually drive a rank into 2146 on a live TP8 server to watch the peers hang.
• The spread under PCP/SP is unknown to me. All 25 of my logs are plain TP8. Under PCP/SP, where ranks do genuinely different work, the spread could be far above 1.5%, which would invalidate my "absorbed by the safety margin" framing in ⚠️4 and raise this fix's urgency considerably. This is the one number most worth measuring — launch -tp 4 -pcp 2, grep the 8 num_kvcache_blocks lines. A few minutes.
• The three new catalog entries' own parameters (the M3 SP4 extraArgs string, the DSV4-Pro TP4+PCP2 env, whether each clears its threshold).
• matrix.kv_block_ranks != '' when the key is absent — I reasoned from the GH Actions docs (null coerces to ''), not from a run.

[valar-cop]

@yitingw1 both my inline points are resolved: group_min is a single import from atom/distributed/group_agreement.py:21, and the test now pins group_min, the tp/pcp group list, the pcp_group_size gate and the with_paged_entries fold — max or dp group would go red.

still no run in the test plan for tp>1 / sp>1. please paste the num_kvcache_blocks lines from one, that's the only thing left from me.

and the raise at atom/model_engine/model_runner.py:2146 sitting before the collective at 2173 is real, agree on available_for_kv first and it goes away.

[MengqingCao]

@valar-cop on your review above, the remaining ask was a multi-rank startup run. ATOM Test 42943244497 on 443b280 started two of the three new entries, and every rank logged the same count.

TP4 + PCP2, DeepSeek-V4-Pro (job): all 8 ranks built 86310 blocks. Their available_for_kv still differs, 124.57 to 124.92 GB, so the reduction is what made them agree. gsm8k 0.9538 (bar 0.94).

[atom 03:53:33] Memory budget: total_gpu=287.98GB, free=160.40GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.10GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.92GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 03:53:33] Memory budget: total_gpu=287.98GB, free=160.04GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.46GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.57GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 03:53:33] Memory budget: total_gpu=287.98GB, free=160.22GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.27GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.75GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 03:53:33] Memory budget: total_gpu=287.98GB, free=160.37GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.13GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.89GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 03:53:33] Memory budget: total_gpu=287.98GB, free=160.17GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.33GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.69GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 03:53:33] Memory budget: total_gpu=287.98GB, free=160.06GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.44GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.59GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 03:53:33] Memory budget: total_gpu=287.98GB, free=160.18GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.31GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.71GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 03:53:33] Memory budget: total_gpu=287.98GB, free=160.04GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.46GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.57GB, block_bytes=1399040, num_kvcache_blocks=86310

TP2, Meta-Llama-3-8B-Instruct (job): both ranks built 107703 blocks. gsm8k 0.7377 (bar 0.73).

[atom 05:31:30] Memory budget: total_gpu=287.98GB, free=262.74GB, utilization=0.3, budget=86.40GB, peak_torch=10.67GB, non_torch=15.54GB, cudagraph_est=0.19GB, safety=5.76GB, available_for_kv=54.23GB, block_bytes=540672, num_kvcache_blocks=107703
[atom 05:31:30] Memory budget: total_gpu=287.98GB, free=262.74GB, utilization=0.3, budget=86.40GB, peak_torch=10.67GB, non_torch=15.54GB, cudagraph_est=0.19GB, safety=5.76GB, available_for_kv=54.23GB, block_bytes=540672, num_kvcache_blocks=107703

For comparison, a TP8 Qwen3-Next run on main before this change logged eight different counts, 2315143 to 2328796.

SP4 did not start (job). The CI AITER wheel lacks ROCm/aiter#7346, which atom/distributed/ulysses_sp.py requires for Ulysses SP, so startup is refused before get_num_blocks runs. That is independent of this change, and SP stays unverified until aiter#7346 lands.

CI on 443b280: Pre Checkin, Actionlint, ATOM vLLM Test and ATOM SGLang Test passed; 16 of 17 accuracy jobs passed, and the failure is SP4 above.

[valar-cop]

@MengqingCao already looked at 443b280, nothing new since my last review. say @valar-cop re-review if you want a fresh one.

[MengqingCao]

@yitingw1 on the CI half (⚠️ 2, ⚠️ 3, ⚠️ 4 and 📝 8): commit 443b280 (the three catalog entries, kv_block_ranks, check_kv_block_agreement.py with its test, and the workflow step) is only there to get a public multi-rank run of this change. The results are in #3171 (comment). It is not meant to merge, and I will drop it before this PR lands.

[Jasen2201]

LGTM

[valar-cop]

LGTM

ci Atomesh Accuracy Validation failed — not sure yet (low): No job log was provided under .valar/ci, so the cause cannot be read. The job only ran because this PR edits .github/benchmark/models_accuracy.json; its PR-level model whitelist (DeepSeek-R1-0528 main, gpt-oss-120b nightly) selects none of this PR's entries, so the change cannot affect what it validates.

next: Fetch the failed step's log before attributing it; dropping the CI commit (as the author plans) removes this trigger. https://github.com/ROCm/ATOM/actions/runs/42986166499

[MengqingCao]

CI on 932aefb, which agrees on free and the KV budget before plan_pools: ATOM Test 42986166648. Every multi-rank engine built one block count on all of its ranks, and the TP2 and TP4 + PCP2 entries passed.

TP4 + PCP2, DeepSeek-V4-Pro (job): all 8 ranks built 86310 blocks from one agreed available_for_kv, 124.57 GB. On 443b280 their own values ranged from 124.57 to 124.92 GB before the count was reduced. gsm8k 0.9484 (bar 0.94).

[atom 09:45:42] Memory budget: total_gpu=287.98GB, free=160.04GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.33GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.57GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 09:45:42] Memory budget: total_gpu=287.98GB, free=160.04GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.46GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.57GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 09:45:42] Memory budget: total_gpu=287.98GB, free=160.04GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.31GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.57GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 09:45:42] Memory budget: total_gpu=287.98GB, free=160.04GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.46GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.57GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 09:45:42] Memory budget: total_gpu=287.98GB, free=160.04GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.10GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.57GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 09:45:42] Memory budget: total_gpu=287.98GB, free=160.04GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.13GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.57GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 09:45:42] Memory budget: total_gpu=287.98GB, free=160.04GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.44GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.57GB, block_bytes=1399040, num_kvcache_blocks=86310
[atom 09:45:42] Memory budget: total_gpu=287.98GB, free=160.04GB, utilization=0.9, budget=259.19GB, peak_torch=114.06GB, non_torch=14.27GB, cudagraph_est=0.34GB, safety=5.76GB, available_for_kv=124.57GB, block_bytes=1399040, num_kvcache_blocks=86310

TP2, Meta-Llama-3-8B-Instruct (job): both ranks built 107631 blocks. gsm8k 0.7309 (bar 0.73).

[atom 10:52:12] Memory budget: total_gpu=287.69GB, free=262.49GB, utilization=0.3, budget=86.31GB, peak_torch=10.67GB, non_torch=15.49GB, cudagraph_est=0.19GB, safety=5.75GB, available_for_kv=54.20GB, block_bytes=540672, num_kvcache_blocks=107631
[atom 10:52:12] Memory budget: total_gpu=287.69GB, free=262.49GB, utilization=0.3, budget=86.31GB, peak_torch=10.67GB, non_torch=15.49GB, cudagraph_est=0.19GB, safety=5.75GB, available_for_kv=54.20GB, block_bytes=540672, num_kvcache_blocks=107631

SP4 (job) was refused at startup again: the CI AITER wheel still lacks ROCm/aiter#7346.

Other accuracy jobs: 13 of the other 14 passed. GLM-5.3-Flash-kpool-16shot (job) failed. Its 8 ranks built one count, 453540 blocks, but gsm8k stalled at 742/1319 with client timeouts and ran out of its 30-minute budget. I don't see a link to this change:

• Between this run and the one on 443b280, where the job passed, vllm-atom: GLM-5.3-Flash prefix caching and MTP on the plugin backend #3137 on main switched that entry from bf16 to fp8 KV, and its block bytes went from 836352 to 430848.
• The same job passed on two other PRs in the same hour.

A rerun would settle whether it is flaky.

Other workflows: Pre Checkin, Actionlint, ATOM vLLM Test and ATOM SGLang Test passed. Atomesh Accuracy Validation failed, as it does on every recent PR.

[valarLip]

• tests/test_model_runner_staging_fence.py:182-229 still checks source text (ast.unparse(...) == "groups = [get_tp_group()]", line-number ordering), and it is now longer than before. Any rename or reformat breaks it, and it catches no real behavior change. Please replace it with a behavioral test: monkeypatch group_min / mem_get_info so one rank reports less free memory, then assert that get_num_blocks returns the reduced block count, and that the connector-reserve RuntimeError is raised on every rank, not only on the low one.
• model_runner.py:2077-2087: config.pcp_group_size > 1 duplicates the group.world_size > 1 check. Mirror the existing loop at model_runner.py:4320-4322: for group in (get_tp_group(), get_pcp_group()): if group.world_size > 1: ....
• Trim the 6-line comment to one sentence on why: every worker builds the pool at rank 0's count, so agree before any rank-local raise.
• The if reduced != ... guard from the previous round is gone, thanks. Keep the PP reduction on the result as it is (stages differ).

(Reviewed at 6d7d49a.)

[valar-cop]

LGTM

ci ATOM Test / Accuracy (Kimi-K3 DSpark, moonshotai/Kimi-K3, --kv_cache_dtype fp8 -tp 8 ...) failed — not sure yet (low): The captured log is empty (0 bytes) so the failure cannot be read. It is a plain TP8 entry that does take the new collective, but the same code path (collective present, counts reduced to the group min) passed on the previous head 932aefb, where only GLM-5.3-Flash-kpool failed, so there is no evidence tying this to the PR.

next: Fetch the failed step's log from the run, or re-run the job, and attribute only from the log. https://github.com/ROCm/ATOM/actions/runs/43133797573/job/130908180673

ci ATOM Test / Accuracy (DeepSeek-V4-Pro DSpark, deepseek-ai/DeepSeek-V4-Pro-DSpark, --kv_cache_dtype fp8 --inde...) failed — main is already broken the same way (medium): The same job fails on main in this window (main runs at 17:21, 16:09 and 15:21 red, 14:44 and earlier green), so the failure predates and is not attributable to this PR; its log here is also empty.

next: Treat as baseline; wait for main to go green before reading this job as a signal. https://github.com/ROCm/ATOM/actions/runs/43133797573/job/130908181071

[MengqingCao]

• tests/test_model_runner_staging_fence.py:182-229 still checks source text (ast.unparse(...) == "groups = [get_tp_group()]", line-number ordering), and it is now longer than before. Any rename or reformat breaks it, and it catches no real behavior change. Please replace it with a behavioral test: monkeypatch group_min / mem_get_info so one rank reports less free memory, then assert that get_num_blocks returns the reduced block count, and that the connector-reserve RuntimeError is raised on every rank, not only on the low one.

Thanks for the review. Addressed in 869fe7a:
AST test replaced with a behavioural one. The reduction now lives in agree_kv_budget in atom/distributed/group_agreement.py, which only imports torch, so it can be tested without a GPU. tests/test_kv_budget_agreement.py runs it with a stand-in group_min: the rank with more free memory gets the group minimum, and the connector-reserve RuntimeError is raised on the larger rank too when a peer is short. The source-text assertions in test_model_runner_staging_fence.py are removed. The stand-in doesn't exercise a real collective, so the multi-rank run is still the check for that.

• model_runner.py:2077-2087: config.pcp_group_size > 1 duplicates the group.world_size > 1 check. Mirror the existing loop at model_runner.py:4320-4322: for group in (get_tp_group(), get_pcp_group()): if group.world_size > 1: ....

Gate is now world_size. get_num_blocks calls agree_kv_budget(..., (get_tp_group(), get_pcp_group())) under torch.distributed.is_initialized(), and each group is reduced only when group.world_size > 1. The pcp_group_size > 1 check is gone. SP is covered because its ranks are in the PCP group.

• Trim the 6-line comment to one sentence on why: every worker builds the pool at rank 0's count, so agree before any rank-local raise.

Comment trimmed to one sentence: "Every worker builds its pool at rank 0's count, so take the group minimum before any rank raises alone and leaves the others waiting."

• The if reduced != ... guard from the previous round is gone, thanks. Keep the PP reduction on the result as it is (stages differ).

PP reduction unchanged, as you suggested.

[Jasen2201]

LGTM

[valarLip]

⚠️ Action required before merge (3)

1. Keep the KV-connector check out of group_agreement.py. atom/distributed/group_agreement.py:60-70 is the generic "verdicts every rank takes alike" module, but it now holds the ATOM_PD_MLA_STAGING / ATOM_PD_MLA_LANDING error. Its caller also takes connector_bytes off the budget (model_runner.py:2065-2070) while the helper takes it off free. Have the helper return only the agreed free and budget, then do free_for_kv, the raise and the min in get_num_blocks right after the call. Every rank still raises together, since free is already agreed. That also removes the need for the 4-field AgreedKvBudget.
2. Bring the PR body up to date with 869fe7a. The body still says the agreement is gated on pcp_group_size (the code now checks world_size per group, group_agreement.py:55). It also lists test_model_runner_staging_fence.py "7 passed" as the test, but that test was removed in favour of tests/test_kv_budget_agreement.py, and it says 443b280 "will be dropped" though it is already reverted (3d6053b). Update the test plan, and the unchecked CI box, to match the current head.
3. Make the fake group_min independent of the values passed. tests/test_kv_budget_agreement.py:30 looks up the peer by value (peers[value]). If any free happens to equal any budget, the fake returns the wrong peer without any error. Key the peer on the call (e.g. pop from a list per call, in free, budget order) instead.

Not blocking: heavy ATOM Test jobs at this head are still pending. On earlier heads, DeepSeek-V4-Pro DSpark also failed on main.

---

_{Re-reviewed at 869fe7a. Already fixed: AST source test replaced by a behavioural test; pcp_group_size gate replaced by the world_size > 1 loop (group_agreement.py:54-55); comment trimmed to one reason (model_runner.py:2071-2072); PP reduction kept as is (model_runner.py:2143-2156); agreement moved before both rank-local raises (@yitingw1 / valar-cop).}

[MengqingCao]

The DeepSeek-V4-Pro DSpark accuracy failure on 869fe7a53 is not from this change.

Accuracy (DeepSeek-V4-Pro DSpark) dies during warmup_model, not in get_num_blocks. torch.compile calls VllmBackend a second time (atom/utils/backends.py) and raises AssertionError: VllmBackend can only be called once. The runner then exits, and the engine reports Received unexpected SHUTDOWN signal from DP rank 1 during initialization. atom_test_output.txt is missing because the server never came up.

The same job fails the same way on main. The latest completed ATOM Test, 43143449277 (a11bb34c8, 2026-11-12 02:44 UTC), logs the same assertion and then a SHUTDOWN from a DP rank during initialization. This job was already red on the two completed main runs before that: 43134309257 and 43067313296.

[Jasen2201]

@zhuyuhua-v pls help merge the pr

[Jasen2201]

⚠️ Action required before merge (3)

1. Keep the KV-connector check out of group_agreement.py. atom/distributed/group_agreement.py:60-70 is the generic "verdicts every rank takes alike" module, but it now holds the ATOM_PD_MLA_STAGING / ATOM_PD_MLA_LANDING error. Its caller also takes connector_bytes off the budget (model_runner.py:2065-2070) while the helper takes it off free. Have the helper return only the agreed free and budget, then do free_for_kv, the raise and the min in get_num_blocks right after the call. Every rank still raises together, since free is already agreed. That also removes the need for the 4-field AgreedKvBudget.
2. Bring the PR body up to date with 869fe7a. The body still says the agreement is gated on pcp_group_size (the code now checks world_size per group, group_agreement.py:55). It also lists test_model_runner_staging_fence.py "7 passed" as the test, but that test was removed in favour of tests/test_kv_budget_agreement.py, and it says 443b280 "will be dropped" though it is already reverted (3d6053b). Update the test plan, and the unchecked CI box, to match the current head.
3. Make the fake group_min independent of the values passed. tests/test_kv_budget_agreement.py:30 looks up the peer by value (peers[value]). If any free happens to equal any budget, the fake returns the wrong peer without any error. Key the peer on the call (e.g. pop from a list per call, in free, budget order) instead.

Not blocking: heavy ATOM Test jobs at this head are still pending. On earlier heads, DeepSeek-V4-Pro DSpark also failed on main.

Re-reviewed at 869fe7a. Already fixed: AST source test replaced by a behavioural test; pcp_group_size gate replaced by the world_size > 1 loop (group_agreement.py:54-55); comment trimmed to one reason (model_runner.py:2071-2072); PP reduction kept as is (model_runner.py:2143-2156); agreement moved before both rank-local raises (@yitingw1 / valar-cop).

no need fix in this pr