[None][fix] Agree the KVCacheManagerV2 rebalance trigger across TP ranks - #17391
[None][fix] Agree the KVCacheManagerV2 rebalance trigger across TP ranks#17391thorjohnsen wants to merge 6 commits into
Conversation
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe 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. ChangesKV pool rebalance coordination
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
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/py_executor.pytests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.pytests/unittest/_torch/executor/test_kv_pool_rebalance.pytests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py
There was a problem hiding this comment.
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
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_kv_pool_rebalance.pytests/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
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
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.
|
|
||
| Broadcasting over the CP group and then the TP group propagates global | ||
| rank 0's decision to everyone: after the CP step each rank holds | ||
| ``V(its tp_rank, cp_rank 0)``, and the TP step then replaces that with |
There was a problem hiding this comment.
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 needIf 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.
There was a problem hiding this comment.
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 needThe 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 thattp_broadcaststays uncalled).- Added
test_attention_dp_without_cp_touches_no_collectiveso 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
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
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.
|
/bot run --disable-fail-fast |
|
PR_Github #65044 [ run ] triggered by Bot. Commit: |
|
PR_Github #65044 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65127 [ run ] triggered by Bot. Commit: |
|
PR_Github #65127 [ run ] completed with state
|
|
@mikeiovine @cascade812 @Tabrizian @QiJune — could one of you approve? Need a @liji-nv @SimengLiu-nv — could one of you approve? Need a @tburt-nv @dpitman-nvda @QiJune @BowenFu — could one of you add the |
|
Removed the "ci: full pre-merge approved" label because @BowenFu could not be verified as an active member of NVIDIA/trt-llm-ci-approvers. Ask a member of that team to apply it. |
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
TP ranks are deliberately kept identical -- same layers, same KV cache geometry, same request stream -- so the rebalance hook should fire on the same iteration on every rank and compute the same pool ratios without any coordination. Auditing that property found it holds for every input to need_adjustment except one. The sample counters and the moving averages behind the target ratios are all fed from request-derived values (kvCache.cpp:105,561-563) with no randomness anywhere in the V2 tree, and the ratio math is pure float arithmetic on those values, so it is bit-identical across ranks. The 120s cooldown is not: it compares against a per-rank steady_clock reading (kvCacheManager.cpp:855), stamped per rank at construction (:126) and at the end of adjust() (:873). Two TP ranks can therefore straddle the cooldown boundary on different iterations and rebalance one iteration apart. That window is enough to break TP, because _prepare_and_schedule_batch runs _schedule() on every rank independently with no broadcast: ranks holding different pool geometry can admit different requests and issue mismatched collectives. The window is small but it gets a fresh roll every 120s. Note this was reachable in practice. _can_pause_for_rebalance's docstring said "MVP scope: single-GPU aggregated", but the predicate never checked tp_size, and no other guard did either -- the only additional constraint on enable_kv_pool_rebalance is Mamba + block reuse. 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. That is safe because needAdjustment() is const (kvCacheManager.h:238) and tryUpdateTargetRatios() runs from the KvCache close path (kvCache.cpp:564), not from this read, so every rank keeps maintaining its ratios regardless of which rank decides. Attention DP is excluded, matching the predicate the scheduler's own tp_broadcast already uses: ADP ranks own independent request streams and independent KV caches, so forcing rank 0's decision on them would starve a rank that needs to rebalance when rank 0 does not. To keep the broadcast off the hot path, the check is throttled to once every KV_POOL_REBALANCE_CHECK_INTERVAL iterations. Rebalance is already rate-limited to once per 120s, 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. Tests: - test_kv_pool_rebalance_tp.py (new): multi-rank TP coverage on real MPI ranks with a real MPIDist, at TP=2 and TP=4. Injects the skew the mechanism exists to absorb, including the case where 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. Needs no model weights. - test_kv_pool_rebalance.py: coverage for the throttle and the agreement helper (15 -> 26 tests). - test_kv_pool_rebalance_accuracy.py: the throttle would otherwise have made this test vacuous, since it runs fewer iterations than a large interval. It now un-throttles explicitly and asserts the GPU pool ratio actually moved in the rebalance arm and stayed fixed in the baseline arm, so a future change that stops adjust() from firing fails instead of silently passing. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
…lism The agreement added in the previous commit keyed on dist.tp_size, which is mapping.tp_size. A pure context-parallel job has tp_size == 1, so it took the local-read path and did no collective at all, leaving CP ranks with exactly the per-rank clock race the mechanism was built to eliminate. CP meets both conditions that make agreement necessary: - A request is split across CP ranks, so they must admit it together. - CP ranks schedule independently. Executor loop selection keys only on pp_size, so a CP job with pp_size == 1 runs _executor_loop / _executor_loop_overlap, whose _prepare_and_schedule_batch calls _schedule() on every rank with no broadcast. Under Ulysses this is doubly true: attn_tp_size = tp_size * cp_size and attn_cp_size = 1, so CP ranks are TP-shaped for KV purposes while mapping.tp_size is still 1. Broadcast over the CP group and then the TP group. With pipeline parallelism already gated off, world_size == tp_size * cp_size, so a rank is identified by (tp_rank, cp_rank): after the CP step a rank holds V(its tp_rank, 0), and the TP step replaces that with V(0, 0), which is global rank 0's decision. Each step is skipped when its dimension is 1, so a pure-TP job still pays exactly one collective and a single-rank job pays none. The dedicated sub-communicators are used rather than the global one, which carries regular executor traffic including request broadcasting -- the same consideration that makes the sleep/wakeup protocol duplicate its own communicator. Tests: - Multi-rank coverage over a pure-CP mapping at CP=2 and CP=4, asserting rank 0's decision wins. Verified by mutation: restoring the tp_size-only predicate fails all four cases with "CP rank 1: local reading was False, rank 0 said True, so the agreed decision should be True, got False". - Unit coverage for the CP-only path, the CP-then-TP chain, and that attention DP skips both collectives. - The multi-rank workers now allgather before asserting. Previously a rank that failed its assert never reached the collective, so the ranks that passed blocked until the pytest timeout; a clean failure was reported as a five-minute hang. The mutation run above now fails in 59s. - The CP-then-TP chain test deliberately gives the local reading and the CP result different values. With both set to True it passed whether or not the TP step consumed the CP step's result, leaving the chaining unverified. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
Throttle the rebalance check on iter_counter instead of a private counter. A private counter only advanced on iterations where every gate in _can_pause_for_rebalance had already passed, so one iteration on which a single rank returned early -- is_shutdown is set per rank on the fatal error path -- left that rank's counter permanently offset from its peers. The ranks would then reach _agreed_need_adjustment's broadcast on different iterations from then on. iter_counter is bumped unconditionally at the end of every executor loop iteration, so the firing schedule becomes a pure function of the iteration index and cannot accumulate an offset. This matches how iter stats are already throttled in the same file. _rebalance_check_counter is now dead and removed. Add a CP x TP multi-rank case. The pure-TP test has cp_size == 1 and the pure-CP test has tp_size == 1, so between them the chained cp_broadcast -> tp_broadcast the production path actually takes was never executed; neither could show that the TP step consumes the CP step's result rather than the rank's own local reading. Verified by mutation: broadcasting the local reading in the TP step fails the new case on the cp_rank == 1 ranks while both single-dimension tests stay green. Rework TestRebalanceCheckThrottle for the new cadence and add a regression test asserting a rank that skips checks never fires on an iteration its peers do not. Verified by mutation: restoring the private counter fails it with the drifted schedule in the message. Annotate and document the multi-GPU test helpers per the coding guidelines. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
_agreed_need_adjustment returned early on enable_attention_dp, which skipped the CP broadcast as well as the TP one. Only the TP hop should be suppressed. Under ADP the TP dimension is the DP dimension (mapping.py: dp_size = tp_size if enable_attention_dp else 1), so those ranks own independent request streams and KV caches and must decide independently. CP is orthogonal: inside a single DP replica the CP ranks still split the same request along the sequence dimension and have to admit it together, which is the same argument this mechanism already makes for pure CP. Nothing in Mapping or LlmArgs rejects enable_attention_dp with cp_size > 1, so the early return left every replica's CP ranks deciding the trigger independently -- the divergence class this change exists to remove. This also aligns with the scheduler's own propagation at py_executor.py:2470-2477, which gates its tp_broadcast on not enable_attention_dp but runs its cp_broadcast unconditionally. The original rationale cited that tp_broadcast as precedent for excluding ADP outright and missed the cp_broadcast directly below it. Invert test_attention_dp_skips_cp_broadcast_too, which asserted the old behaviour, into test_attention_dp_still_broadcasts_over_cp, and add test_attention_dp_without_cp_touches_no_collective to keep the no-collective case covered. Add a real-MPI case, test_attention_dp_agrees_over_cp_but_not_tp (world_size=4, cp_size=2, tp_size=2), with flags chosen to fail in both directions: two ranks 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. Verified by mutation: restoring the early return yields [True, False, False, True] against the expected [True, True, False, False] while the pure-TP, pure-CP and CP-x-TP cases all stay green. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
ce68152 to
1f1cb52
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #65686 [ run ] triggered by Bot. Commit: |
…elism _executor_loop_pp had no rebalance hook and _can_pause_for_rebalance rejected pp_size > 1 outright, so pipeline-parallel jobs never rebalanced their KV pools. The other two loops rebalance inline: at the top of an iteration at most one batch is in flight (previous_batch) and _consume_previous_batch_for_ rebalance retires it on the spot. The PP loop keeps up to num_micro_batches batches in flight in a ring, and a microbatch is retired only when its sample state finishes travelling that ring, which advances just as iterations run. There is no inline call that can drain it, and adjust() cannot run while any of them is outstanding because it moves KV pages underneath whatever holds them. So rather than drain inline, stop feeding the ring and let the loop drain itself. While a rebalance is pending the loop forces can_queue to False, which is its existing "skip this microbatch slot" path -- the same one every idle iteration already takes, since _can_queue is just batch_size > 0. Each iteration then 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 while its peers kept queueing would desynchronize the per-iteration send/recv chain and hang the pipeline, so this is 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 so the decision is shared; 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 across ranks. The PP hop must not be suppressed under attention DP, 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; pp_group holds exactly those stages. If the ring is somehow 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 auto-tuner asks again on the next check interval. The suspend/adjust/resume core is split out as _rebalance_kv_pools_now so both paths share it. The PP path must not run _consume_previous_batch_for_ rebalance: under PP a microbatch is retired by the relay and previous_batch is kept only to synchronize its sampler event, so running the overlap loop's helper on it would double-handle a batch the ring already completed. Tests: 16 new unit tests covering the drain countdown, the quiescence predicate, the busy-ring skip and the sample-stream sync, plus 10 new multi-rank cases over real MPI communicators -- pure PP at widths 2 and 4, a TP2xPP2 chain that only passes if the PP hop consumes the TP hop's result, and a TP2xPP2 ADP case expecting [True, False, True, False], which pins down both that the PP hop ran and that the TP hop did not. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
|
/bot kill |
|
PR_Github #65693 [ kill ] triggered by Bot. Commit: |
|
PR_Github #65686 [ run ] completed with state |
|
PR_Github #65693 [ kill ] completed with state |
The drain's pieces were unit-tested but its wiring was not. Starting the drain, forcing can_queue to False while it runs, and completing it in Stage 3.4 all live inside _executor_loop_pp, which no unit test drove, so only live multi-GPU runs covered them and a regression would have reached CI unnoticed. TestPpLoopDrainWiring drives the real loop over an object.__new__ (PyExecutor) instance populated with just what one iteration touches -- the approach test_py_executor.py already uses for the PP scheduling path. Two choices make it discriminating: _can_queue returns True, so every iteration wants to queue and an empty slot can only be the drain's doing; and the mocked forward raises, so reaching it both records that the iteration queued and ends the loop. The main case walks a two-slot ring end to end -- iterations 1 and 2 queue nothing despite _can_queue saying they could, the rebalance lands at the end of iteration 2, and iteration 3 returns to queueing. The others cover the control case (no drain pending reaches the forward), the countdown tracking ring depth rather than a constant, and the decision being taken once per rebalance instead of once per iteration -- which matters because _agreed_need_adjustment runs a collective, so re-entering it mid-drain would put a rank into a broadcast its peers are not in. The fixture shuts the loop down after a bounded number of iterations. Without that, dropping the Stage 3.4 call makes the drain never complete, so nothing is queued, the forward is never reached and the loop spins forever -- which turns a clean assertion failure into a hung stage. Each of the wiring mutations now fails in well under a second instead. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #65701 [ run ] triggered by Bot. Commit: |
|
PR_Github #65701 [ run ] completed with state
|
|
[by Codex] @lowsfer Could you review this PR? Thanks! |
6 participants
Description
TP ranks are deliberately kept identical — same layers, same KV cache geometry, same request stream — so the KV pool rebalance hook should fire on the same iteration on every rank and compute the same pool ratios by construction, without coordination. Auditing that property found it holds for every input to
need_adjustmentexcept one.Deterministic (identical across TP ranks given identical requests):
mNumSampledKvCacheskvCacheManager.h:290, incremented atkvCache.cpp:563++; no RNG anywhere in the V2 treemAvgReusedLength,mAvgSqrCapacity,mAvgSqrHistoryLengthkvCache.cpp:105,561-562kvCacheManager.cpp:853,784kvCacheManager.cpp:781-796,838Not deterministic — the 120 s cooldown:
mLastAdjustmentTimeis stamped per rank, in the constructor (:126) and at the end ofadjust()(:873). Nothing ties those readings together.Why one iteration of skew is enough to break TP
Two ranks can straddle the cooldown 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. (_pp_schedule_and_propagateat:2436does broadcast, which is why PP is structurally immune.) The window is small, but it gets a fresh roll every 120 s.This was reachable in practice:
_can_pause_for_rebalance's docstring said "MVP scope: single-GPU aggregated", but the predicate never checkedtp_size, and no other guard did either — the only additional constraint onenable_kv_pool_rebalanceis Mamba + block reuse (_util.py:236).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()isconst(kvCacheManager.h:238) — no side effects, so only rank 0's value mattering is harmless.tryUpdateTargetRatios()runs from the KvCache close path (kvCache.cpp:564), not from this read, so every rank keeps maintaining its ratios regardless of who decides.Attention DP is excluded, matching the predicate the scheduler's own
tp_broadcastalready uses (py_executor.py:2459): ADP ranks own independent request streams and independent KV caches, so forcing rank 0's decision on them would starve a rank that needs to rebalance when rank 0 does not.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_rebalanceis rank-uniform, which keeps the throttle counter — and therefore the iteration the collective runs on — identical across ranks.For context on cost: the rebalance stall was measured at 28–45 ms, ~99 % of it inside
adjust()itself (drain/suspend/resume are sub-millisecond).Test Coverage
tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py(new) — 14 cases on real MPI ranks with a realMPIDist, following thetest_allgather.pyharness. Needs no model weights or KV cache. Picked up automatically byl0_dgx_h100.yml(unittest/_torch/multi_gpu); max world size 4.test_tp_ranks_agree_on_rebalance_triggertest_attention_dp_ranks_decide_independentlytest_rebalance_check_stays_in_lockstep_across_ranksThe scenario that matters is
only_rank0_false: every follower's local clock says "rebalance now" but rank 0's does not. Without agreement the followers would resize while rank 0 did not. The lockstep test drives the real collective inside the loop, so a cadence divergence hangs rather than passing quietly.tests/unittest/_torch/executor/test_kv_pool_rebalance.py— 15 → 26 tests, covering the throttle and the agreement helper.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, soadjust()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 — removing the un-throttle while keeping the assertion produces:Results (8×H100, sm90)
test_kv_pool_rebalance_tp.pytest_kv_pool_rebalance.pytest_kv_pool_rebalance_accuracy.pykv_cache_manager_v2_tests(C++ backend, regression check)OK (skipped=13)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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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
Dev Engineer Review
QA Engineer Review
Added test coverage includes:
No
tests/integration/test_lists/,test-db/, orqa/files were modified. Test-list coverage for the added test functions is not provided. Reported results include 14 TP tests, 26 executor tests, 2 accuracy tests, and 127 C++ backend regression tests passing.Verdict: needs follow-up.