[TRTLLM-13308][feat] Make the KVCacheManagerV2 KV pool rebalance safe under TP, CP, PP and attention DP

[thorjohnsen]

Description

Tracked by TRTLLM-13308.

The KV pool rebalance hook (enable_kv_pool_rebalance) was documented as "MVP scope:
single-GPU aggregated", but nothing enforced that: _can_pause_for_rebalance never checked
tp_size, and the only other constraint on the feature is Mamba + block reuse
(_util.py:236). So the hook was reachable — and unsafe — on every multi-GPU configuration.

This PR makes it safe, and then extends it. In one table:

Parallelism	Before	After
TP	reachable, each rank decided alone	rank 0 decides and broadcasts
CP	reachable, each rank decided alone	rank 0 of the CP group decides and broadcasts
Attention DP	reachable, each rank decided alone — crashed	OR-reduction across the TP dimension
PP	refused outright by _can_pause_for_rebalance	supported, via a drain of the microbatch ring

Plus one defect fix in the KVCacheManagerV2 auto-tuner that this PR does not cause but does
make expensive — please read that section first, it is the most surprising thing here.

Read this first: the auto-tuner is sized by warmup dummies, not by traffic

_KVCache.close() feeds every closed cache into the tuner's capacity statistics. Warmup and
CUDA-graph padding requests reserve capacity at the model's full declared context rather
than at a realistic sequence length, and _avg_sqr_capacity averages the square of capacity.
A handful of dummies therefore owns the statistic outright.

Measured on DeepSeek-V4-Flash-NVFP4 (tp4 = ep4, attention DP, 4× GB300), per rank:

sampled closes	4,934
of those, at capacity == max_position_embeddings == 1,048,576	76 — 1.5%
their share of the sum of squares	100.00%
RMS including them	130,139
RMS excluding them	313 — real sequences averaged ~530 tokens

The tuner was sizing pools for million-token sequences on a workload whose sequences average
~530, and the resulting target is a pool inversion: applying it moves 171 GB per rank.

Those caches are already marked stats-excluded at creation — the dummy path sets
is_dummy and _create_kv_cache calls mark_stats_excluded — close() simply never
consulted the flag. The fix is one condition per backend:

-if self.capacity > 0:
+if self.capacity > 0 and not manager.is_stats_excluded(self.id):

-if (mCapacity > 0)
+if (mCapacity > 0 && !mManager->isStatsExcluded(id))

Deliberately not reusing _should_record_stats(): that also ANDs in the user-facing
_stats_enabled toggle, and turning off stats reporting must not blind the auto-tuner to
real traffic.

Behaviour change reviewers should weigh. Dummies were ~90% of the sample count, not just
of the RMS, so the 2000-sample maturity gate used to open on a few hundred real sequences. It
now needs 2000 real completed sequences per rank. Measured: 1,024 prompts → 257 samples, never
triggered; 12,288 prompts → ~3,060 samples, triggered once. Low-traffic deployments will engage
the feature later than they do today.

Scope note. This bug predates the PR — the tuner and its sampling both ship on main
today. It is included here because this PR is what makes the mis-sizing fire, on every rank,
as a 171 GB move; shipping the rebalance without it would turn a latent mis-sizing into a
visible one. Happy to split it out if the KV-cache-manager owners prefer.

Why the trigger has to be agreed at all

TP ranks are kept identical — same layers, same KV geometry, same request stream — so the hook
should fire on the same iteration on every rank by construction. Auditing every input to
need_adjustment found that holds for all of them except one.

Deterministic across ranks:

Input	Where	Why
mNumSampledKvCaches	kvCacheManager.h:290, kvCache.cpp:563	plain ++; no RNG in the V2 tree
mAvgReusedLength, mAvgSqrCapacity, mAvgSqrHistoryLength	kvCache.cpp:105,561-562	request-derived
2000-sample / 100-sample gates	kvCacheManager.cpp:853,784	counter comparisons
target-ratio math, 1.25 skew threshold	kvCacheManager.cpp:781-796,838	pure float ops, same binary, same order → bit-identical

Not deterministic — the 120 s cooldown:

double now = nowSeconds();                    // kvCacheManager.cpp:855 — steady_clock, per rank
if (now - mLastAdjustmentTime < 120.0) return false;

mLastAdjustmentTime is stamped per rank, in the constructor (:126) and at the end of
adjust() (:873). Nothing ties those readings together, so two ranks can straddle the
boundary on different iterations and rebalance one iteration apart. In that window they hold
different pool geometry, and _prepare_and_schedule_batch (py_executor.py:3618) runs
_schedule() on every rank independently with no broadcast — so they can admit different
requests and issue mismatched collectives. The window is small, but it gets a fresh roll every
120 s.

Fix. TP rank 0 decides and broadcasts (_agreed_need_adjustment). Only the trigger is
agreed — the ratios are still computed independently on every rank and never exchanged,
because they are pure functions of statistics that are already identical. Two properties make
that safe: needAdjustment() is const (kvCacheManager.h:238), so only rank 0's value
mattering is harmless; and tryUpdateTargetRatios() runs from the KvCache close path
(kvCache.cpp:564), not from this read, so every rank keeps maintaining its own ratios.

To keep the broadcast off the hot path the check is throttled to once every
KV_POOL_REBALANCE_CHECK_INTERVAL (10) iterations. Rebalance is already rate-limited to once
per 120 s, so this costs a fraction of a second of latency while cutting the collective rate by
an order of magnitude. Every condition in _can_pause_for_rebalance is rank-uniform, which
keeps the throttle counter — and therefore the iteration the collective runs on — identical
across ranks.

Attention DP: an OR-reduction, not an exclusion

An earlier revision of this PR excluded ADP from the agreement, on the argument that ADP
ranks own independent request streams and independent KV caches and so legitimately want
different ratios at different times. That argument holds for the ratios but not for the
timing, because the path the decision gates is itself collective under ADP:
_consume_previous_batch_for_rebalance → _flush_pending_transfer_responses →
_enqueue_responses, which — as its own docstring requires — every DP rank must enter even
with an empty response list, and which runs a tp_gather.

A rank rebalancing alone therefore joined a collective its peers were not in. Live on
Qwen3-Next-80B-A3B (tp4 = ep4, ADP, 4× H100) with only rank 0's need_adjustment true, that
gather paired against the tp_allgather of batch sizes in _can_queue:

TypeError: cannot unpack non-iterable int object
HangDetector: propagating hard-kill to all ranks via MPI_Abort

The TP hop now runs under ADP as any(tp_allgather(need)). any() rather than rank 0's
broadcast preserves the anti-starvation property the exclusion was reaching for: a rank that
needs a rebalance gets one without waiting for rank 0 to want the same thing. Same
configuration after the change: exit 0.

Two consequences worth weighing:

• The check now performs a collective on the ADP path, so _can_pause_for_rebalance must stay
  rank-uniform under ADP as well as under plain TP — the same requirement the non-ADP broadcast
  already relies on, now extended.
• adjust() stamps _last_adjustment_time unconditionally, so a rank pulled into a peer's
  rebalance has its own 120 s cooldown reset by it. The group's cadence self-aligns, at the
  price of deferring a need that arises immediately afterwards by up to one cooldown.

Cost on that model (2 pool groups, 17 active requests per rank): 21.1 ms for the rank that
needed it, 8.1 ms for a rank pulled in — of which ~7.3 ms is the pre-existing
torch.cuda.current_stream().synchronize() and 0.03 ms is the no-op adjust(). A check that
decides not to rebalance costs 0.13-1.46 ms, against 0.09-1.68 ms for the tp_broadcast
already on the non-ADP path.

CP

Orthogonal, and must not be skipped. Within one DP replica the CP ranks split the same
request along the sequence dimension, so they must rebalance together — including under ADP,
where the CP hop stays even though the TP hop changes shape.

PP: supported, by draining the ring rather than inline

_executor_loop_pp had no rebalance hook and _can_pause_for_rebalance rejected pp_size > 1
outright. The other two loops rebalance inline — at the top of an iteration at most one batch
is in flight and _consume_previous_batch_for_rebalance retires it on the spot. The PP loop
keeps up to num_micro_batches batches in a ring, retired only as their sample state travels
that ring, so there is no inline call that can drain it, and adjust() cannot run while any
are outstanding because it moves KV pages underneath whatever holds them.

So: stop feeding the ring and let the loop drain itself. While a rebalance is pending the loop
forces can_queue to False — its existing "skip this microbatch slot" path, the same one every
idle iteration already takes. Each iteration retires one slot and queues nothing new, so
num_micro_batches iterations empty the ring, after which a new Stage 3.4 suspends, adjusts and
resumes.

Every rank must do this on the same iteration — a rank that drained alone would desynchronize
the send/recv chain and hang the pipeline, a stronger requirement than the TP case where
divergence merely risks mismatched scheduling. Three things provide it: _agreed_need_adjustment
gains a pp_broadcast hop; the drain countdown is a rank-independent constant rather than a
per-rank condition; and the ring state the quiescence test reads is already symmetric.

The PP hop is not suppressed under ADP, unlike the TP hop: ADP replicates along the TP
dimension, so a replica's pipeline stages all serve that replica's one request stream and must
drain together, and pp_group holds exactly those stages.

If the ring is still busy when the countdown expires the rebalance is skipped with a warning
rather than forced — adjusting underneath a live _KVCache would corrupt it, while skipping
costs only a delay, since the tuner asks again on the next check interval.

Test Coverage

tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py (new) — 9 tests / 31
parameterized cases on real MPI ranks with a real MPIDist, following the test_allgather.py
harness. Needs no model weights and no KV cache.

Test	Cases	Pins down
test_tp_ranks_agree_on_rebalance_trigger	tp:2, tp:4 × 4 scenarios	rank 0's decision wins on every rank
test_cp_ranks_agree_on_rebalance_trigger	cp:2, cp:4 × 2	the CP hop, in isolation
test_cp_and_tp_ranks_agree_on_rebalance_trigger	world:4-cp:2	the TP hop consumes the CP hop's result
test_pp_ranks_agree_on_rebalance_trigger	pp:2, pp:4 × 4 scenarios	the PP hop, in isolation
test_tp_and_pp_ranks_agree_on_rebalance_trigger	world:4-tp:2	the PP hop consumes the TP hop's result
test_attention_dp_still_agrees_across_pp_stages	world:4-tp:2	expects [True, False, True, False] — the PP hop ran and the TP hop did not
test_attention_dp_ranks_or_couple	tp:2, tp:4	the OR-reduction: any rank's need pulls the group in
test_attention_dp_agrees_over_cp_then_ors_across_tp	world:4-cp:2 × 2	CP agreement survives inside an ADP replica
test_rebalance_check_stays_in_lockstep_across_ranks	tp:2, tp:4 × interval 1, 8	ranks fire the check on identical iterations

The scenario that matters is only_rank0_false: every follower's local clock says "rebalance
now" but rank 0's does not. The lockstep test drives the real collective inside the loop, so
a cadence divergence hangs rather than passing quietly.

tests/integration/test_lists/test-db/l0_dgx_h100.yml — unittest/_torch/multi_gpu was
collected by exactly one test-db entry, a 2-GPU stage, so 18 of the 31 cases could never
run in CI on any run: tp:4, pp:4, CP×TP, TP×PP and both ADP compositions — precisely the
newest logic on this branch. They all reported need 4 GPUs, have 2. The file is now named in
the 4-GPU PyTorch block as well.

tests/unittest/_torch/executor/test_kv_pool_rebalance.py — 21 → 57 tests, covering the
throttle, the agreement helper, the PP drain countdown, the quiescence predicate, the busy-ring
skip and the sample-stream sync.

tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py —
test_dummy_kv_caches_do_not_feed_the_tuner. Asserts on need_adjustment rather than the
counters (neither has a getter on the C++ backend), opens the gates with the raw setters rather
than force_rebalance_precondition (which would also skew the target and mask the thing under
test), and carries a control arm that closes an identically shaped cache without the
exclusion — without it the check would pass for the wrong reason. Mutation-tested: reverting
either guard fails it with the intended message.

tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py — the throttle would
otherwise have made this test vacuous: it runs fewer iterations than a large interval, so
adjust() would never fire and the token comparison would pass while testing nothing. It now
un-throttles explicitly, and asserts the GPU pool ratio actually moved in the rebalance arm
and stayed fixed in the baseline arm. Verified by mutation:

E  AssertionError: rebalance never fired: GPU pool ratio unchanged at
   [0.5000054240226746, 0.49999457597732544]. The token comparison would pass vacuously.

Results

Suite	Where	Result
test_kv_pool_rebalance_tp.py	4× GB300	31 cases passed, 5m22s
test_kv_pool_rebalance_tp.py	CI, DGX_H100-2_GPUs	13 passed, 17 skipped (the 4-GPU gap above)
test_kv_pool_rebalance.py	local	57 passed
test_kv_pool_rebalance_accuracy.py	8× H100	2 passed, 105 s
kv_cache_manager_v2_tests	Python backend	130 passed, OK (skipped=13)
kv_cache_manager_v2_tests	C++ backend (default), 4x GB300	130 passed, OK (skipped=13), 261 s
Full pre-merge L0_MergeRequest_PR #53654	on c21efaa	SUCCESS, incl. 24 Multi-GPU stages, 1906 passed / 0 failed

Measurements backing the cost claims

adjust() is bandwidth-bound; its cost is set by bytes in motion, not by pool-group count.
Two points 31× apart on the same model and hardware (DeepSeek-V4-Flash, 3 pool groups, 4×
GB300): 5.50 GB in 91 ms and 171 GB in 2.15 s, both at 60-80 GB/s per rank. On Qwen3-Next-80B
(2 groups, 8× H100): 3.6 GB in 48 ms at ~75 GB/s.

One fully organic end-to-end rebalance was captured with no injection anywhere
(DeepSeek-V4, 12,288 prompts, ~3,060 real samples/rank, 477 s): it fired once across 376
hook calls, the ADP OR-reduction pulled all four ranks in on the same iteration, all four landed
on the same target, moved 169.36 GB in ~2.55 s, then converged and stopped. sync+agree was
0.5-1.6 ms.

Known open items

• adjust() re-checks neither the maturity gate nor the cooldown (kvCacheManager.cpp:862
  vs :851), so once the agreement returns True every rank applies whatever target it is
  carrying, including a rank whose own maturity gate had not opened. Measurement shows no harm
  in the case that occurs — in the organic run every rank was mature and landed on the same
  target — so this is a theoretical concern about ADP replicas drifting apart, not a
  demonstrated one. Flagged for a codeowner's eye rather than fixed here.
• Backends order pool groups differently — Python pool group 1 is C++ pool group 0.
  Confirmed, not fixed, and orthogonal to this PR.
• "Wants to rebalance" is not yet "is faster for it." Nothing here measures throughput
  across the move. At the observed sequence shape the startup layout caps concurrency at ~67
  sequences while one pool group holds slots for 47,841, and the tuner's target balances all
  three at ~1,466 — but that is slot arithmetic, not a measured throughput result, and is not
  quoted as one.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

🤖 Generated with Claude Code

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

The executor now throttles KV pool rebalance checks and synchronizes tensor-parallel and context-parallel decisions. Unit, accuracy, and real-MPI tests validate interval behavior, rank agreement, attention-DP independence, and pool-ratio changes.

Changes

KV pool rebalance coordination

Layer / File(s)	Summary
Executor rebalance coordination tensorrt_llm/_torch/pyexecutor/py_executor.py	Adds a rebalance-check interval and counter. Throttles eligibility checks. Broadcasts rank-0 decisions for tensor- and context-parallel execution while preserving local decisions for single-rank and attention-DP paths.
Local and accuracy validation tests/unittest/_torch/executor/test_kv_pool_rebalance.py, tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py	Tests interval and counter behavior, distributed decision modes, rejected agreement, and pool-ratio changes with and without rebalance.
Distributed MPI validation tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py	Adds GPU-backed MPI tests for tensor-parallel, context-parallel, combined CP×TP, attention-DP, and lockstep throttle behavior.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PyExecutor
  participant KVCacheManager
  participant MPIDist
  PyExecutor->>KVCacheManager: read local adjustment state
  PyExecutor->>MPIDist: broadcast rank-0 decision for tensor or context parallelism
  MPIDist-->>PyExecutor: return agreed decision
  PyExecutor->>KVCacheManager: rebalance when the decision is true

Loading

Suggested labels: ci: full pre-merge approved

Suggested reviewers: qijune, shixiaowei02, tabrizian

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and concisely identifies the feature and its safety improvements across KV-cache parallelism modes.
Description check	✅ Passed	The description clearly explains the problem, solution, scope, risks, test coverage, results, and checklist status.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py`:
- Around line 77-81: Remove the broad try/except block from run_single_rank and
invoke single_rank_forward_func directly, allowing MPIPoolExecutor to propagate
the original failure unchanged; only add a handler if a documented specific
exception requires handling.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a9e5c753-c65b-457e-ae26-c5efdc6a2333

📥 Commits

Reviewing files that changed from the base of the PR and between 5cc76ef and ac3e393.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py
• tests/unittest/_torch/executor/test_kv_pool_rebalance.py
• tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py`:
- Around line 119-122: Update
tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py at lines 119-122
for _cp_dist: add Google-style Args and Returns sections while preserving its
existing behavior; at lines 152-159 annotate every parameter and the return type
as None, then add a Google-style docstring; at lines 284-285 annotate world_size
and case plus -> None, then add a Google-style docstring.
- Around line 279-296: Add a real-MPI CP-then-TP agreement case alongside
test_cp_ranks_agree_on_rebalance_trigger using world_size=4, cp_size=2, and
tp_size=2, with divergent flags such as [True, False, False, False]. Exercise
the production collective chain through cp_broadcast() followed by
tp_broadcast(), and assert every rank receives rank 0’s decision; retain the
existing pure-CP coverage.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8c9b6071-a5f0-4f25-b6b0-9e9d6d2a759b

📥 Commits

Reviewing files that changed from the base of the PR and between ac3e393 and 43f8135.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tests/unittest/_torch/executor/test_kv_pool_rebalance.py
• tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py

🚧 Files skipped from review as they are similar to previous changes (2)

• tests/unittest/_torch/executor/test_kv_pool_rebalance.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py

[brnguyen2]

Thorough work — the PR description alone settles most questions a reviewer would have, and the anti-vacuity assertions in the accuracy test plus the real-MPI lockstep tests are exactly the right coverage for this kind of collective.

One correctness concern (inline at [py_executor.py:4461](https://github.com/NVIDIA/TensorRT-LLM/pull/17391/files#diff-f0b4c3c02708916fd189f863c48b982a55f34ae94996089b23aa0a6f0571fe19R4461)): the attention-DP early-return skips the CP broadcast as well as the TP one, but the scheduler propagation cited as precedent gates only the TP hop on ADP — cp_broadcast at [py_executor.py:2475](https://github.com/NVIDIA/TensorRT-LLM/pull/17391/files#diff-f0b4c3c02708916fd189f863c48b982a55f34ae94996089b23aa0a6f0571fe19R2475) runs unconditionally. Under ADP + CP, CP ranks share a request stream, so they need the trigger agreement even though the TP-dimension exclusion is correct. Details and a suggested fix inline.

Also verified while reviewing (no action needed): adjust() re-checks only the deterministic per-level ratio-skew test, not the cooldown or the 2000-sample gate (kvCacheManager.cpp:862-874), so a follower whose local cooldown hadn't expired but is forced True by rank 0 performs the same _adjustLevel work rather than no-opping — and re-stamps mLastAdjustmentTime, which conveniently re-syncs the per-rank clocks after the first agreed adjustment. That property is what makes trigger-only agreement sufficient; might be worth a sentence in the _agreed_need_adjustment docstring, since the mechanism silently depends on it.

On the [None] tag: this is a real bug fix (cross-rank divergence reachable in any TP job with the feature enabled). Worth filing a tracking ticket and referencing it in the title, per the usual convention for fixes.

[brnguyen2]

This early return skips the CP broadcast too, but the scheduler propagation cited as the precedent gates only the TP hop on attention DP — cp_broadcast at [py_executor.py:2475](https://github.com/NVIDIA/TensorRT-LLM/pull/17391/files#diff-f0b4c3c02708916fd189f863c48b982a55f34ae94996089b23aa0a6f0571fe19R2475) runs unconditionally, ADP or not. That's because ADP gives each TP rank its own request stream, but CP ranks within a DP group still split the same requests and must admit them together. Under ADP + CP (nothing in Mapping or llm_args rejects the combination), CP ranks would decide the rebalance trigger independently here — the same divergence class this PR fixes.

Suggested structure, mirroring [py_executor.py:2471](https://github.com/NVIDIA/TensorRT-LLM/pull/17391/files#diff-f0b4c3c02708916fd189f863c48b982a55f34ae94996089b23aa0a6f0571fe19R2471)-2478:

need = self.kv_cache_manager.impl.need_adjustment
if self.dist.cp_size > 1:
    need = self.dist.cp_broadcast(need, root=0)
if self.dist.tp_size > 1 and not self.enable_attention_dp:
    need = self.dist.tp_broadcast(need, root=0)
return need

If ADP + CP + rebalance is instead considered unreachable today, a comment saying so (and why) would do — but test_attention_dp_skips_cp_broadcast_too and test_attention_dp_ranks_decide_independently currently bake the skip in as intended behavior, so either the code or the tests should change.

[thorjohnsen]

Good catch — you're right, and I had the precedent backwards. Fixed in ba6fff1.

I cited the scheduler's tp_broadcast as justification for excluding ADP outright and missed that the cp_broadcast immediately below it (py_executor.py:2474) runs unconditionally. That asymmetry is the whole answer, and the reasoning you give for it is the same one this PR already makes for pure CP — I just failed to carry it into the ADP branch.

Confirmed reachable: nothing in Mapping or LlmArgs rejects enable_attention_dp with cp_size > 1. Mapping.__init__'s only ADP-related assert is enable_lm_head_tp_in_adp requires enable_attention_dp (mapping.py:190), and dp_size = tp_size if enable_attention_dp else 1 (:296) makes CP orthogonal to the DP dimension rather than folded into it.

Took your structure verbatim:

need = self.kv_cache_manager.impl.need_adjustment
if self.dist.cp_size > 1:
    need = self.dist.cp_broadcast(need, root=0)
if self.dist.tp_size > 1 and not self.enable_attention_dp:
    need = self.dist.tp_broadcast(need, root=0)
return need

The CP-then-TP order differs from the scheduler's TP-then-CP, but the two agree in both regimes: without ADP, TP-then-CP gives rank(t,c) <- V(0,c) <- V(0,0) and CP-then-TP gives rank(t,c) <- V(t,0) <- V(0,0); with the TP hop suppressed both reduce to rank(t,c) <- V(t,0), i.e. each replica decides on its own cp_rank-0 reading.

On tests — you were right that they baked the skip in, so I inverted rather than extended:

• test_attention_dp_skips_cp_broadcast_too -> test_attention_dp_still_broadcasts_over_cp (asserts the CP result overrides the local reading and that tp_broadcast stays uncalled).
• Added test_attention_dp_without_cp_touches_no_collective so the genuine no-collective case keeps coverage.
• Added a real-MPI case, test_attention_dp_agrees_over_cp_but_not_tp (world_size=4, cp_size=2, tp_size=2). Flags are [True, False, False, True] so it fails in both directions: ranks 1 and 3 are overridden by their CP root (proving the CP hop ran) while the two replicas end on different answers (proving the TP hop did not).

Mutation-verified — restoring the early return gives:

AssertionError: ADP+CP decisions were [True, False, False, True],
expected [True, True, False, False]: CP ranks must follow their replica's
root while replicas stay independent

which is exactly the raw local flags, i.e. no broadcast at all. The pure-TP, pure-CP and CP-x-TP cases all stay green under that mutation, which is your point that nothing covered this topology.

One note for precision: the propagation you cite lives in _pp_schedule_and_propagate, the PP scheduling path, and rebalance is gated off for pp_size > 1. So it's a semantic precedent rather than an operative one — in the non-PP loops _schedule() runs per rank with no broadcast at all, which is this PR's premise. Doesn't change the conclusion; the ADP+CP semantics stand on their own.

Suites after the fix: 30 mock tests, 20 real-MPI tests, pre-commit clean.

[BowenFu]

Verified existing/default behavior remains unchanged behind the default-off rebalance gate; when enabled, the loop-wide cadence and CP/TP agreement are covered by focused unit and MPI tests, including ADP+CP.
The unresolved py_executor.py thread is non-blocking: ba6fff1 fixes the cadence/ADP issue it raised, and the author response documents the correction.

[thorjohnsen]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65044 [ run ] triggered by Bot. Commit: ba6fff1 Link to invocation

[tensorrt-cicd]

PR_Github #65965 [ run ] triggered by Bot. Commit: c21efaa Link to invocation

[tensorrt-cicd]

PR_Github #65965 [ run ] completed with state SUCCESS. Commit: c21efaa
/LLM/main/L0_MergeRequest_PR pipeline #53654 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

[thorjohnsen]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #66106 [ run ] triggered by Bot. Commit: 990a877 Link to invocation

[nvpohanh]

@lowsfer could you review this? thanks!

[thorjohnsen]

/bot kill

[thorjohnsen]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #66194 [ kill ] triggered by Bot. Commit: f95c99c Link to invocation

[tensorrt-cicd]

PR_Github #66196 [ run ] triggered by Bot. Commit: f95c99c Link to invocation

[tensorrt-cicd]

PR_Github #66194 [ kill ] completed with state ABORTED. Commit: f95c99c

Link to invocation

[tensorrt-cicd]

PR_Github #66106 [ run ] completed with state ABORTED. Commit: 990a877

Link to invocation

[thorjohnsen]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #66305 [ run ] triggered by Bot. Commit: f95c99c Link to invocation

[tensorrt-cicd]

PR_Github #66196 [ run ] completed with state ABORTED. Commit: f95c99c

Link to invocation

[tensorrt-cicd]

PR_Github #66305 [ run ] completed with state FAILURE. Commit: f95c99c
/LLM/main/L0_MergeRequest_PR pipeline #53960 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[thorjohnsen]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #66497 [ run ] triggered by Bot. Commit: f95c99c Link to invocation

[tensorrt-cicd]

PR_Github #66497 [ run ] completed with state FAILURE. Commit: f95c99c
/LLM/main/L0_MergeRequest_PR pipeline #54136 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation