From 9236b8f906383d883fd3c1ac8f011db348c4527c Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Fri, 7 Aug 2026 01:58:47 +0000 Subject: [PATCH 1/9] [None][fix] Agree the KVCacheManagerV2 rebalance trigger across TP ranks 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> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 67 ++++- .../test_kv_pool_rebalance_accuracy.py | 33 +++ .../_torch/executor/test_kv_pool_rebalance.py | 125 +++++++- .../multi_gpu/test_kv_pool_rebalance_tp.py | 276 ++++++++++++++++++ 4 files changed, 498 insertions(+), 3 deletions(-) create mode 100644 tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ee6e7e6dfecc..3749ebf7f65a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -157,6 +157,14 @@ class _SleepWakeupAction(StrEnum): _SLEEP_WAKEUP_ACK_TIMEOUT_S = 30.0 _SLEEP_WAKEUP_ACK_POLL_INTERVAL_S = 0.01 +# How many executor iterations between KV pool rebalance checks. The V2 +# auto-tuner rate-limits itself to one adjustment per 120s, so a check every +# iteration is pure overhead -- and under TP each check costs a broadcast (see +# PyExecutor._agreed_need_adjustment). At typical iteration times this adds a +# fraction of a second of latency to a rebalance that happens at most twice a +# minute, while cutting the collective rate by an order of magnitude. +KV_POOL_REBALANCE_CHECK_INTERVAL = 10 + def _sleep_wakeup_ack_ready(comm, source: int, tag: _SleepWakeupTag) -> bool: """Return whether an ACK is ready without blocking on recv.""" @@ -596,6 +604,10 @@ def __init__( self.guided_decoder = guided_decoder self.disable_overlap_scheduler = disable_overlap_scheduler self.enable_kv_pool_rebalance = enable_kv_pool_rebalance + # Iteration throttle for the KV pool rebalance check. See + # _can_pause_for_rebalance / _agreed_need_adjustment. + self._rebalance_check_interval = KV_POOL_REBALANCE_CHECK_INTERVAL + self._rebalance_check_counter = 0 self.enable_early_first_token_response = enable_early_first_token_response self.virtual_memory_pools = virtual_memory_pools @@ -4511,9 +4523,17 @@ def _sync_and_process_resource_governor_queue(self): def _can_pause_for_rebalance(self) -> bool: """Gate KV pool rebalance to the cases the v1 hook supports. - MVP scope: single-GPU aggregated, no in-flight disagg transfer, + MVP scope: no pipeline parallelism, no in-flight disagg transfer, no beam search, no drafter, not during warmup or shutdown. Honors the ``enable_kv_pool_rebalance`` opt-in flag (default off). + + Every condition below must be rank-uniform, because a ``True`` here is + what puts a TP rank into ``_agreed_need_adjustment``'s collective, and + ranks entering that collective on different iterations would deadlock. + The config checks are identical by construction; ``is_warmup`` and + ``is_shutdown`` are phase flags the executor loop already has to keep in + lockstep for the many other per-iteration collectives it runs, so + relying on them here adds no new requirement. """ if not self.enable_kv_pool_rebalance: return False @@ -4529,8 +4549,51 @@ def _can_pause_for_rebalance(self) -> bool: return False if self.drafter is not None: return False + + # Throttle the check itself. Rebalance is rate-limited to once per + # 120s by the V2 auto-tuner's own cooldown, so polling every iteration + # buys nothing and costs a collective per iteration in the TP case + # (see _agreed_need_adjustment). At typical iteration times this + # delays a rebalance by a couple of seconds at most. + self._rebalance_check_counter += 1 + if self._rebalance_check_counter < self._rebalance_check_interval: + return False + self._rebalance_check_counter = 0 return True + def _agreed_need_adjustment(self) -> bool: + """Decide whether to rebalance, identically on every TP rank. + + Every input to ``need_adjustment`` is a deterministic function of the + request stream -- the sample counters and the moving averages that + produce the target ratios are all fed from request-derived values with + no randomness -- **except** the 120s cooldown, which compares against a + per-rank ``steady_clock`` reading (``kvCacheManager.cpp:855``, stamped + per rank at ``:126`` and ``: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: ``_prepare_and_schedule_batch`` runs ``_schedule()`` on every + rank independently with no broadcast, so ranks holding different pool + geometry can admit different requests and issue mismatched collectives. + + So rank 0 of the TP group decides and broadcasts. Only the *trigger* + needs agreement -- the resulting ratios do not, because they are pure + functions of statistics that are already identical across TP ranks, and + those statistics keep being maintained on every rank regardless (they + are updated from the KvCache close path, not from this read). + + Attention DP is excluded for the same reason the scheduler's own + ``tp_broadcast`` excludes it: ADP ranks own independent request streams + and independent KV caches, so they legitimately need different pool + ratios at different times. Forcing rank 0's decision on them would + starve a rank that needs to rebalance when rank 0 does not. + """ + need = self.kv_cache_manager.impl.need_adjustment + if self.dist.tp_size == 1 or self.enable_attention_dp: + return need + return self.dist.tp_broadcast(need, root=0) + def _consume_previous_batch_for_rebalance(self) -> None: """Drain ``previous_batch`` so its _KVCache instances are quiescent. @@ -4570,7 +4633,7 @@ def _maybe_rebalance_kv_pools(self) -> None: uses today after eviction. """ mgr = self.kv_cache_manager - if not mgr.impl.need_adjustment: + if not self._agreed_need_adjustment(): return torch.cuda.current_stream().synchronize() diff --git a/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py b/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py index 6542c7cdfae0..4ffb6fe1e52c 100644 --- a/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py +++ b/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py @@ -43,10 +43,19 @@ def _inject_pool_ratio_mismatch(llm: LLM, *, skew: float = 2.0) -> None: past the auto-tuner's adjustment threshold. The hook requires a model with >=2 pool groups (e.g. Gemma-3-1B with VSWA) and raises otherwise, so a future model change can't silently turn this test into a no-op. + + Also drops the executor's rebalance-check throttle to every iteration, so + the test does not depend on how ``KV_POOL_REBALANCE_CHECK_INTERVAL`` compares + to the number of iterations this short prompt set happens to run. Raise that + interval above the iteration count and the hook would never fire, leaving the + token comparison below to pass vacuously; the ratio assertion in + ``_generate_tokens`` is the backstop that would catch it. """ from tensorrt_llm.runtime.kv_cache_manager_v2 import _introspection executor = llm._executor.engine + executor._rebalance_check_interval = 1 + executor._rebalance_check_counter = 0 kv_cache_manager = executor.kv_cache_manager _introspection.force_rebalance_precondition(kv_cache_manager.impl, skew=skew) @@ -95,14 +104,38 @@ def _generate_tokens(*, model_path: str, disable_overlap: bool, enable_rebalance responsible for setting that env var (via monkeypatch or otherwise) before invoking this helper. """ + from tensorrt_llm.runtime.kv_cache_manager_v2 import _introspection + with LLM( model_path, disable_overlap_scheduler=disable_overlap, kv_cache_config=_vswa_kv_cache_config(enable_rebalance=enable_rebalance), ) as llm: + impl = llm._executor.engine.kv_cache_manager.impl if enable_rebalance: _inject_pool_ratio_mismatch(llm) + ratio_before = list(_introspection.current_gpu_ratio(impl)) outputs = llm.generate(_PROMPTS, _SAMPLING) + ratio_after = list(_introspection.current_gpu_ratio(impl)) + + # Guard against a vacuous pass. Token equality between the rebalance + # and no-rebalance arms proves nothing if adjust() never ran, and + # nothing in the run logs at info level to tell us it did. The pool + # ratio moving is the observable signature that it happened. + if enable_rebalance: + assert ratio_after != ratio_before, ( + "rebalance never fired: GPU pool ratio unchanged at " + f"{ratio_before}. The token comparison would pass vacuously. " + "Check the executor's rebalance-check throttle and the V2 " + "auto-tuner's sample-count / cooldown gates." + ) + else: + assert ratio_after == ratio_before, ( + "pool ratio moved with enable_kv_pool_rebalance=False " + f"({ratio_before} -> {ratio_after}); the baseline arm is " + "supposed to hold pool ratios fixed." + ) + return [list(o.outputs[0].token_ids) for o in outputs] diff --git a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py index d2f12754754e..de7d66515e1e 100644 --- a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py +++ b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py @@ -43,6 +43,8 @@ def _make_executor( *, enable_kv_pool_rebalance: bool = True, pp_size: int = 1, + tp_size: int = 1, + enable_attention_dp: bool = False, kv_cache_transceiver=None, is_warmup: bool = False, is_shutdown: bool = False, @@ -53,23 +55,37 @@ def _make_executor( previous_batch=None, padding_dummies=None, has_cuda_graph_runner: bool = True, + rebalance_check_interval: int = 1, ) -> MagicMock: """Construct a MagicMock shaped like PyExecutor with exactly the attributes the rebalance code path reads. ``padding_dummies`` is the runner's ``{draft_len: dummy}`` map; its dummies count as active on GPU, like real pre-allocated ones. + + ``rebalance_check_interval`` defaults to 1 so the iteration throttle is + transparent for tests that are not about the throttle itself; the throttle + gets its own coverage in ``TestRebalanceCheckThrottle``. """ exe = MagicMock(spec=PyExecutor) # Gate inputs. exe.enable_kv_pool_rebalance = enable_kv_pool_rebalance - exe.dist = MagicMock(pp_size=pp_size) + exe.dist = MagicMock(pp_size=pp_size, tp_size=tp_size) + exe.enable_attention_dp = enable_attention_dp exe.kv_cache_transceiver = kv_cache_transceiver exe.is_warmup = is_warmup exe.is_shutdown = is_shutdown exe.drafter = drafter + # Iteration throttle state. + exe._rebalance_check_interval = rebalance_check_interval + exe._rebalance_check_counter = 0 + + # Bind the real agreement helper so _maybe_rebalance_kv_pools exercises it + # rather than getting a truthy MagicMock back. + exe._agreed_need_adjustment = lambda: PyExecutor._agreed_need_adjustment(exe) + # KV cache manager (resource-manager wrapper). exe.kv_cache_manager = MagicMock() exe.kv_cache_manager.max_beam_width = max_beam_width @@ -377,5 +393,112 @@ def test_consumes_and_clears(self): assert exe.previous_batch is None +# --------------------------------------------------------------------------- # +# Iteration throttle +# --------------------------------------------------------------------------- # + + +class TestRebalanceCheckThrottle: + """``_can_pause_for_rebalance`` only lets a check through every N iterations. + + The throttle must be driven purely by rank-uniform state: all ranks have to + reach the agreement collective on the same iteration or they deadlock. + """ + + def test_returns_false_until_interval_reached(self): + exe = _make_executor(rebalance_check_interval=4) + results = [PyExecutor._can_pause_for_rebalance(exe) for _ in range(4)] + assert results == [False, False, False, True] + + def test_counter_resets_after_firing(self): + exe = _make_executor(rebalance_check_interval=3) + results = [PyExecutor._can_pause_for_rebalance(exe) for _ in range(6)] + assert results == [False, False, True, False, False, True] + + def test_interval_of_one_checks_every_iteration(self): + exe = _make_executor(rebalance_check_interval=1) + assert all(PyExecutor._can_pause_for_rebalance(exe) for _ in range(5)) + + def test_counter_not_advanced_when_a_gate_rejects(self): + # The counter sits behind every gate, so an executor that cannot + # rebalance accumulates no throttle credit. Every gate is rank-uniform, + # which is what keeps the counter -- and therefore the iteration the + # agreement collective runs on -- identical across ranks. + exe = _make_executor(enable_kv_pool_rebalance=False, rebalance_check_interval=2) + for _ in range(10): + assert PyExecutor._can_pause_for_rebalance(exe) is False + assert exe._rebalance_check_counter == 0 + + def test_counter_not_advanced_during_warmup(self): + exe = _make_executor(is_warmup=True, rebalance_check_interval=2) + for _ in range(10): + assert PyExecutor._can_pause_for_rebalance(exe) is False + assert exe._rebalance_check_counter == 0 + + # Once warmup clears, the throttle starts counting from scratch. + exe.is_warmup = False + assert [PyExecutor._can_pause_for_rebalance(exe) for _ in range(2)] == [False, True] + + +# --------------------------------------------------------------------------- # +# Cross-rank agreement on the rebalance trigger +# --------------------------------------------------------------------------- # + + +class TestAgreedNeedAdjustment: + """Rank 0 of the TP group decides; the decision is broadcast. + + Rationale: every input to ``need_adjustment`` is deterministic given the + request stream except the 120s cooldown, which reads a per-rank + ``steady_clock``. Two TP ranks straddling that boundary would rebalance on + different iterations, and TP ranks schedule independently, so their batches + could then diverge. + """ + + def test_single_rank_reads_locally_without_broadcast(self): + exe = _make_executor(tp_size=1, need_adjustment=True) + assert PyExecutor._agreed_need_adjustment(exe) is True + exe.dist.tp_broadcast.assert_not_called() + + def test_tp_broadcasts_rank0_decision(self): + exe = _make_executor(tp_size=4, need_adjustment=True) + exe.dist.tp_broadcast.return_value = True + + assert PyExecutor._agreed_need_adjustment(exe) is True + exe.dist.tp_broadcast.assert_called_once_with(True, root=0) + + def test_tp_rank_follows_broadcast_over_its_own_reading(self): + # This is the case the whole mechanism exists for: the local read says + # "rebalance" but rank 0 says no, so this rank must not rebalance. + exe = _make_executor(tp_size=2, need_adjustment=True) + exe.dist.tp_broadcast.return_value = False + + assert PyExecutor._agreed_need_adjustment(exe) is False + + def test_tp_rank_follows_broadcast_when_local_says_no(self): + exe = _make_executor(tp_size=2, need_adjustment=False) + exe.dist.tp_broadcast.return_value = True + + assert PyExecutor._agreed_need_adjustment(exe) is True + + def test_attention_dp_decides_independently(self): + # 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. + exe = _make_executor(tp_size=4, enable_attention_dp=True, need_adjustment=True) + assert PyExecutor._agreed_need_adjustment(exe) is True + exe.dist.tp_broadcast.assert_not_called() + + def test_no_rebalance_when_agreement_says_no(self, monkeypatch): + exe = _make_executor(tp_size=2, need_adjustment=True, active_requests=[_make_request(1)]) + exe.dist.tp_broadcast.return_value = False + monkeypatch.setattr("torch.cuda.current_stream", MagicMock()) + + PyExecutor._maybe_rebalance_kv_pools(exe) + + exe.kv_cache_manager.impl.adjust.assert_not_called() + exe.kv_cache_manager.suspend_request.assert_not_called() + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py new file mode 100644 index 000000000000..21d44228bcd5 --- /dev/null +++ b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py @@ -0,0 +1,276 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Multi-rank TP tests for the KVCacheManagerV2 rebalance trigger. + +``tests/unittest/_torch/executor/test_kv_pool_rebalance.py`` covers the same +logic against a mocked ``dist``; these tests run it across real MPI ranks with +a real ``MPIDist``, which is the only way to show the collective actually +agrees (and that ranks reach it in lockstep rather than deadlocking). + +What is being protected +----------------------- +Every input to ``need_adjustment`` is a deterministic function of the request +stream -- the sample counters and the moving averages behind the target ratios +are all request-derived, with no randomness -- **except** the 120s cooldown, +which compares against a per-rank ``steady_clock`` reading +(``kvCacheManager.cpp:855``). Two TP ranks can therefore straddle that boundary +on different iterations and rebalance one iteration apart. TP ranks compute +``_schedule()`` independently with no broadcast, so for that one iteration they +could admit different requests and issue mismatched collectives. + +``_agreed_need_adjustment`` closes this by letting TP rank 0 decide and +broadcasting. The tests below inject exactly the skew the mechanism exists to +absorb: ranks whose *local* readings disagree. + +These tests need no model weights and no KV cache -- they exercise the +agreement and throttle logic directly. +""" + +import pickle +import sys +import traceback +from unittest.mock import MagicMock + +import cloudpickle +import pytest +import torch +from mpi4py import MPI +from mpi4py.futures import MPIPoolExecutor + +import tensorrt_llm +from tensorrt_llm._torch.distributed.communicator import MPIDist +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm.mapping import Mapping + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__( + cloudpickle.dumps, + cloudpickle.loads, + pickle.HIGHEST_PROTOCOL, +) + +# MPIPoolExecutor leaks a worker thread on first use; keep CI green. +pytestmark = pytest.mark.threadleak(enabled=False) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def run_single_rank(tensor_parallel_size, single_rank_forward_func, *args): + """Wrapper used by MPIPoolExecutor; matches test_allgather.py.""" + rank = tensorrt_llm.mpi_rank() + torch.cuda.set_device(rank) + try: + single_rank_forward_func(tensor_parallel_size, rank, *args) + except Exception: + traceback.print_exc() + raise + return True + + +def _make_executor( + dist, + *, + need_adjustment: bool, + enable_attention_dp: bool = False, + rebalance_check_interval: int = 1, +): + """A PyExecutor stand-in carrying a *real* MPIDist. + + Only the attributes the rebalance gate and the agreement helper read are + populated; everything else stays a MagicMock so an accidental new read + shows up as a test failure rather than silently passing. + """ + exe = MagicMock(spec=PyExecutor) + exe.dist = dist + exe.enable_attention_dp = enable_attention_dp + exe.enable_kv_pool_rebalance = True + exe.kv_cache_transceiver = None + exe.is_warmup = False + exe.is_shutdown = False + exe.drafter = None + exe.kv_cache_manager = MagicMock() + exe.kv_cache_manager.max_beam_width = 1 + exe.kv_cache_manager.impl = MagicMock() + exe.kv_cache_manager.impl.need_adjustment = need_adjustment + exe._rebalance_check_interval = rebalance_check_interval + exe._rebalance_check_counter = 0 + return exe + + +def _tp_dist(world_size: int, rank: int) -> MPIDist: + return MPIDist(Mapping(world_size=world_size, rank=rank, tp_size=world_size)) + + +# --------------------------------------------------------------------------- +# Per-rank work +# --------------------------------------------------------------------------- + + +def run_agreement(world_size, rank, local_flags, expected): + """Every TP rank must end up with rank 0's decision, not its own.""" + dist = _tp_dist(world_size, rank) + exe = _make_executor(dist, need_adjustment=local_flags[rank]) + + got = PyExecutor._agreed_need_adjustment(exe) + + assert got == expected, ( + f"rank {rank}: local reading was {local_flags[rank]}, rank 0 said " + f"{local_flags[0]}, so the agreed decision should be {expected}, got {got}" + ) + + # And every rank must have reached the same answer, which is the property + # that actually keeps the schedulers from diverging. + all_decisions = dist.allgather(got) + assert len(set(all_decisions)) == 1, f"ranks disagreed after the collective: {all_decisions}" + + +def run_attention_dp_independence(world_size, rank, local_flags): + """Under attention DP each rank keeps its own decision. + + 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. + """ + dist = _tp_dist(world_size, rank) + exe = _make_executor(dist, need_adjustment=local_flags[rank], enable_attention_dp=True) + + got = PyExecutor._agreed_need_adjustment(exe) + + assert got == local_flags[rank], ( + f"rank {rank}: attention DP must decide locally; expected {local_flags[rank]}, got {got}" + ) + + +def run_throttle_lockstep(world_size, rank, interval, iterations): + """Ranks must reach the agreement collective on the *same* iterations. + + This is the deadlock guard. ``_can_pause_for_rebalance`` gates the + collective, so ranks whose throttle counters drifted apart would enter + ``tp_broadcast`` on different iterations. Driving the real collective + inside the loop means such a divergence hangs here rather than passing + silently; the allgather afterwards pins down that the firing iterations + were in fact identical. + + Ranks are deliberately given *different* local ``need_adjustment`` readings, + so it is the throttle cadence -- not agreement on the value -- that is under + test here. Everything ``_can_pause_for_rebalance`` reads is rank-uniform, + which is precisely the property being confirmed. + """ + dist = _tp_dist(world_size, rank) + exe = _make_executor(dist, need_adjustment=(rank % 2 == 0), rebalance_check_interval=interval) + + fired_on = [] + for i in range(iterations): + if PyExecutor._can_pause_for_rebalance(exe): + fired_on.append(i) + # Real collective -- diverging ranks hang instead of passing. + PyExecutor._agreed_need_adjustment(exe) + + all_fired = dist.allgather(fired_on) + assert all(f == all_fired[0] for f in all_fired), ( + f"ranks fired the rebalance check on different iterations: {all_fired}" + ) + + # Sanity: the throttle actually throttled, and it fired when expected. + expected = [i for i in range(iterations) if (i + 1) % interval == 0] + assert fired_on == expected, ( + f"throttle fired on {fired_on}, expected {expected} for interval " + f"{interval} over {iterations} iterations" + ) + assert fired_on, "throttle never fired; the test would be vacuous" + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def _skip_if_not_enough_gpus(world_size): + if torch.cuda.device_count() < world_size: + pytest.skip(f"need {world_size} GPUs, have {torch.cuda.device_count()}") + + +def _flags_for(case: str, world_size: int): + """Local per-rank ``need_adjustment`` readings for each skew scenario.""" + if case == "all_true": + return [True] * world_size + if case == "all_false": + return [False] * world_size + if case == "only_rank0_true": + # Rank 0 wants to rebalance, nobody else does -> all must follow it. + return [i == 0 for i in range(world_size)] + if case == "only_rank0_false": + # The dangerous one: every follower's clock says "rebalance now" but + # rank 0's does not. Without agreement the followers would resize and + # rank 0 would not. + return [i != 0 for i in range(world_size)] + raise ValueError(case) + + +@pytest.mark.parametrize( + "case", + ["all_true", "all_false", "only_rank0_true", "only_rank0_false"], +) +@pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"tp:{x}") +def test_tp_ranks_agree_on_rebalance_trigger(world_size, case): + """Rank 0's decision wins on every TP rank, whatever the local readings.""" + _skip_if_not_enough_gpus(world_size) + flags = _flags_for(case, world_size) + expected = flags[0] + + with MPIPoolExecutor(max_workers=world_size) as ex: + results = ex.map( + run_single_rank, + *zip(*[(world_size, run_agreement, flags, expected)] * world_size), + ) + for r in results: + assert r is True + + +@pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"tp:{x}") +def test_attention_dp_ranks_decide_independently(world_size): + """Attention DP opts out of the agreement: each rank keeps its own answer.""" + _skip_if_not_enough_gpus(world_size) + flags = _flags_for("only_rank0_false", world_size) + + with MPIPoolExecutor(max_workers=world_size) as ex: + results = ex.map( + run_single_rank, + *zip(*[(world_size, run_attention_dp_independence, flags)] * world_size), + ) + for r in results: + assert r is True + + +@pytest.mark.parametrize("interval", [1, 8], ids=lambda x: f"interval:{x}") +@pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"tp:{x}") +def test_rebalance_check_stays_in_lockstep_across_ranks(world_size, interval): + """The throttle must fire on identical iterations on every rank. + + A drift here would put ranks into ``tp_broadcast`` on different iterations, + which deadlocks -- so this test hanging is itself the failure signal. + """ + _skip_if_not_enough_gpus(world_size) + + with MPIPoolExecutor(max_workers=world_size) as ex: + results = ex.map( + run_single_rank, + *zip(*[(world_size, run_throttle_lockstep, interval, 40)] * world_size), + ) + for r in results: + assert r is True From f688ad7e50d8a39e97f55d1d2098379b63f6614a Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Fri, 7 Aug 2026 02:28:32 +0000 Subject: [PATCH 2/9] [None][fix] Extend the rebalance trigger agreement to context parallelism 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> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 23 +++++++- .../_torch/executor/test_kv_pool_rebalance.py | 44 +++++++++++++- .../multi_gpu/test_kv_pool_rebalance_tp.py | 58 ++++++++++++++++++- 3 files changed, 120 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 3749ebf7f65a..bedeb3ebac60 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4583,6 +4583,21 @@ def _agreed_need_adjustment(self) -> bool: those statistics keep being maintained on every rank regardless (they are updated from the KvCache close path, not from this read). + Context parallelism needs the same treatment. A request is split across + CP ranks, so they must admit it together, and CP runs on the same + executor loops as TP -- the loop choice keys only on ``pp_size`` -- which + means CP ranks also compute ``_schedule()`` independently. Note that + ``dist.tp_size`` is ``mapping.tp_size``, which is 1 for a pure-CP job, so + keying on it alone would silently leave CP unsynchronized. Under + Ulysses, CP is folded into attention TP (``attn_tp_size = tp_size * + cp_size``), making those ranks TP-shaped for KV purposes. + + 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 + ``V(tp_rank 0, cp_rank 0)``. Dedicated sub-communicators are used rather + than the global one, which carries regular executor traffic. + Attention DP is excluded for the same reason the scheduler's own ``tp_broadcast`` excludes it: ADP ranks own independent request streams and independent KV caches, so they legitimately need different pool @@ -4590,9 +4605,13 @@ def _agreed_need_adjustment(self) -> bool: starve a rank that needs to rebalance when rank 0 does not. """ need = self.kv_cache_manager.impl.need_adjustment - if self.dist.tp_size == 1 or self.enable_attention_dp: + if self.enable_attention_dp: return need - return self.dist.tp_broadcast(need, root=0) + if self.dist.cp_size > 1: + need = self.dist.cp_broadcast(need, root=0) + if self.dist.tp_size > 1: + need = self.dist.tp_broadcast(need, root=0) + return need def _consume_previous_batch_for_rebalance(self) -> None: """Drain ``previous_batch`` so its _KVCache instances are quiescent. diff --git a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py index de7d66515e1e..22f16b809cce 100644 --- a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py +++ b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py @@ -44,6 +44,7 @@ def _make_executor( enable_kv_pool_rebalance: bool = True, pp_size: int = 1, tp_size: int = 1, + cp_size: int = 1, enable_attention_dp: bool = False, kv_cache_transceiver=None, is_warmup: bool = False, @@ -71,7 +72,7 @@ def _make_executor( # Gate inputs. exe.enable_kv_pool_rebalance = enable_kv_pool_rebalance - exe.dist = MagicMock(pp_size=pp_size, tp_size=tp_size) + exe.dist = MagicMock(pp_size=pp_size, tp_size=tp_size, cp_size=cp_size) exe.enable_attention_dp = enable_attention_dp exe.kv_cache_transceiver = kv_cache_transceiver exe.is_warmup = is_warmup @@ -481,6 +482,47 @@ def test_tp_rank_follows_broadcast_when_local_says_no(self): assert PyExecutor._agreed_need_adjustment(exe) is True + def test_cp_broadcasts_even_when_tp_size_is_one(self): + # A pure-CP job has mapping.tp_size == 1, so keying the agreement on + # tp_size alone would leave CP unsynchronized -- yet a request is split + # across CP ranks, and CP runs on the same executor loops as TP (loop + # choice keys only on pp_size), so CP ranks also schedule independently. + exe = _make_executor(tp_size=1, cp_size=4, need_adjustment=True) + exe.dist.cp_broadcast.return_value = False + + assert PyExecutor._agreed_need_adjustment(exe) is False + exe.dist.cp_broadcast.assert_called_once_with(True, root=0) + exe.dist.tp_broadcast.assert_not_called() + + def test_tp_and_cp_chain_propagates_global_rank0(self): + # CP first, then TP: after the CP step a rank holds V(its tp_rank, cp0), + # and the TP step replaces it with V(tp0, cp0) -- global rank 0's value. + # + # The local reading and the CP result are deliberately *different* here. + # If they matched, the final assertion would pass whether or not the TP + # step actually consumes the CP step's result, and the chaining -- the + # whole point of this test -- would go unverified. + exe = _make_executor(tp_size=2, cp_size=2, need_adjustment=True) + exe.dist.cp_broadcast.return_value = False + exe.dist.tp_broadcast.return_value = False + + assert PyExecutor._agreed_need_adjustment(exe) is False + exe.dist.cp_broadcast.assert_called_once_with(True, root=0) + # Called with the CP result (False), not the local reading (True). + exe.dist.tp_broadcast.assert_called_once_with(False, root=0) + + def test_single_rank_touches_no_collective(self): + exe = _make_executor(tp_size=1, cp_size=1, need_adjustment=True) + assert PyExecutor._agreed_need_adjustment(exe) is True + exe.dist.cp_broadcast.assert_not_called() + exe.dist.tp_broadcast.assert_not_called() + + def test_attention_dp_skips_cp_broadcast_too(self): + exe = _make_executor(tp_size=2, cp_size=2, enable_attention_dp=True, need_adjustment=True) + assert PyExecutor._agreed_need_adjustment(exe) is True + exe.dist.cp_broadcast.assert_not_called() + exe.dist.tp_broadcast.assert_not_called() + def test_attention_dp_decides_independently(self): # ADP ranks own independent request streams and independent KV caches, # so forcing rank 0's decision on them would starve a rank that needs diff --git a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py index 21d44228bcd5..7d29b5a8789d 100644 --- a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py +++ b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py @@ -116,6 +116,12 @@ def _tp_dist(world_size: int, rank: int) -> MPIDist: return MPIDist(Mapping(world_size=world_size, rank=rank, tp_size=world_size)) +def _cp_dist(world_size: int, rank: int) -> MPIDist: + """A pure-CP mapping: ``tp_size == 1``, so ``cp_broadcast`` carries the + agreement and keying on ``tp_size`` alone would miss it entirely.""" + return MPIDist(Mapping(world_size=world_size, rank=rank, cp_size=world_size)) + + # --------------------------------------------------------------------------- # Per-rank work # --------------------------------------------------------------------------- @@ -128,17 +134,45 @@ def run_agreement(world_size, rank, local_flags, expected): got = PyExecutor._agreed_need_adjustment(exe) + # Collect *before* asserting. Every rank has to reach this collective: if + # some ranks bailed out on a failed assert first, the ranks that passed + # would block here until the pytest timeout, turning a clean failure into a + # five-minute hang. + all_decisions = dist.allgather(got) + assert got == expected, ( f"rank {rank}: local reading was {local_flags[rank]}, rank 0 said " f"{local_flags[0]}, so the agreed decision should be {expected}, got {got}" ) - # And every rank must have reached the same answer, which is the property # that actually keeps the schedulers from diverging. - all_decisions = dist.allgather(got) assert len(set(all_decisions)) == 1, f"ranks disagreed after the collective: {all_decisions}" +def run_cp_agreement(world_size, rank, local_flags, expected): + """Same guarantee as TP, but over a pure-CP mapping (``tp_size == 1``). + + A request is split across CP ranks, so they must admit it together; and CP + runs on the same executor loops as TP, which compute ``_schedule()`` per + rank with no broadcast. Keying the agreement on ``tp_size`` alone would + leave this configuration unsynchronized. + """ + dist = _cp_dist(world_size, rank) + assert dist.tp_size == 1, "this test is meaningless unless tp_size is 1" + exe = _make_executor(dist, need_adjustment=local_flags[rank]) + + got = PyExecutor._agreed_need_adjustment(exe) + + # Collect before asserting -- see run_agreement for why the order matters. + all_decisions = dist.allgather(got) + + assert got == expected, ( + f"CP rank {rank}: local reading was {local_flags[rank]}, rank 0 said " + f"{local_flags[0]}, so the agreed decision should be {expected}, got {got}" + ) + assert len(set(all_decisions)) == 1, f"CP ranks disagreed: {all_decisions}" + + def run_attention_dp_independence(world_size, rank, local_flags): """Under attention DP each rank keeps its own decision. @@ -242,6 +276,26 @@ def test_tp_ranks_agree_on_rebalance_trigger(world_size, case): assert r is True +@pytest.mark.parametrize( + "case", + ["only_rank0_true", "only_rank0_false"], +) +@pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"cp:{x}") +def test_cp_ranks_agree_on_rebalance_trigger(world_size, case): + """Pure CP (``tp_size == 1``) must agree just as TP does.""" + _skip_if_not_enough_gpus(world_size) + flags = _flags_for(case, world_size) + expected = flags[0] + + with MPIPoolExecutor(max_workers=world_size) as ex: + results = ex.map( + run_single_rank, + *zip(*[(world_size, run_cp_agreement, flags, expected)] * world_size), + ) + for r in results: + assert r is True + + @pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"tp:{x}") def test_attention_dp_ranks_decide_independently(world_size): """Attention DP opts out of the agreement: each rank keeps its own answer.""" From 29417dc805660796057c1f63f243059513e80585 Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:13:20 +0000 Subject: [PATCH 3/9] [None][fix] Address review on the rebalance trigger agreement 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> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 33 ++- .../test_kv_pool_rebalance_accuracy.py | 1 - .../_torch/executor/test_kv_pool_rebalance.py | 104 +++++--- .../multi_gpu/test_kv_pool_rebalance_tp.py | 234 ++++++++++++++++-- 4 files changed, 306 insertions(+), 66 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index bedeb3ebac60..07338e70f1bd 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -607,7 +607,6 @@ def __init__( # Iteration throttle for the KV pool rebalance check. See # _can_pause_for_rebalance / _agreed_need_adjustment. self._rebalance_check_interval = KV_POOL_REBALANCE_CHECK_INTERVAL - self._rebalance_check_counter = 0 self.enable_early_first_token_response = enable_early_first_token_response self.virtual_memory_pools = virtual_memory_pools @@ -4527,13 +4526,16 @@ def _can_pause_for_rebalance(self) -> bool: no beam search, no drafter, not during warmup or shutdown. Honors the ``enable_kv_pool_rebalance`` opt-in flag (default off). - Every condition below must be rank-uniform, because a ``True`` here is - what puts a TP rank into ``_agreed_need_adjustment``'s collective, and - ranks entering that collective on different iterations would deadlock. - The config checks are identical by construction; ``is_warmup`` and - ``is_shutdown`` are phase flags the executor loop already has to keep in - lockstep for the many other per-iteration collectives it runs, so - relying on them here adds no new requirement. + A ``True`` here is what puts a TP rank into + ``_agreed_need_adjustment``'s collective, so ranks that disagree on this + predicate *on the same iteration* would enter that collective in + different numbers. The config checks are identical across ranks by + construction, and ``is_warmup`` / ``is_shutdown`` are phase flags the + executor loop already has to keep in lockstep for the many other + per-iteration collectives it runs, so relying on them adds no new + requirement. Note that the throttle below reads ``iter_counter`` + instead of keeping its own counter precisely so that a rank which does + return early here cannot carry a lasting cadence offset out of it. """ if not self.enable_kv_pool_rebalance: return False @@ -4555,10 +4557,19 @@ def _can_pause_for_rebalance(self) -> bool: # buys nothing and costs a collective per iteration in the TP case # (see _agreed_need_adjustment). At typical iteration times this # delays a rebalance by a couple of seconds at most. - self._rebalance_check_counter += 1 - if self._rebalance_check_counter < self._rebalance_check_interval: + # + # Throttling on iter_counter rather than on a counter of our own is + # deliberate. A private counter would only advance on iterations where + # every gate above already passed, so a single iteration on which one + # rank returned early would leave that rank's counter permanently offset + # from its peers -- and the ranks would then reach the collective in + # _agreed_need_adjustment on different iterations from then on. + # iter_counter is bumped unconditionally at the end of every executor + # loop iteration, so being a pure function of the iteration index it + # cannot accumulate that offset. This file already throttles the same + # way for iter stats (see _kv_iter_stats_interval). + if self.iter_counter % self._rebalance_check_interval != 0: return False - self._rebalance_check_counter = 0 return True def _agreed_need_adjustment(self) -> bool: diff --git a/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py b/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py index 4ffb6fe1e52c..fd8c592b2308 100644 --- a/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py +++ b/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py @@ -55,7 +55,6 @@ def _inject_pool_ratio_mismatch(llm: LLM, *, skew: float = 2.0) -> None: executor = llm._executor.engine executor._rebalance_check_interval = 1 - executor._rebalance_check_counter = 0 kv_cache_manager = executor.kv_cache_manager _introspection.force_rebalance_precondition(kv_cache_manager.impl, skew=skew) diff --git a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py index 22f16b809cce..045c529f663f 100644 --- a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py +++ b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py @@ -79,9 +79,10 @@ def _make_executor( exe.is_shutdown = is_shutdown exe.drafter = drafter - # Iteration throttle state. + # Iteration throttle state. The throttle keys on iter_counter, the + # loop-wide iteration index, so tests drive that directly. exe._rebalance_check_interval = rebalance_check_interval - exe._rebalance_check_counter = 0 + exe.iter_counter = 0 # Bind the real agreement helper so _maybe_rebalance_kv_pools exercises it # rather than getting a truthy MagicMock back. @@ -402,43 +403,82 @@ def test_consumes_and_clears(self): class TestRebalanceCheckThrottle: """``_can_pause_for_rebalance`` only lets a check through every N iterations. - The throttle must be driven purely by rank-uniform state: all ranks have to - reach the agreement collective on the same iteration or they deadlock. + The throttle keys on ``iter_counter``, the loop-wide iteration index, rather + than on a counter of its own. That makes the set of iterations a rank fires + the agreement collective on a pure function of the iteration number, so it + cannot drift between ranks. """ - def test_returns_false_until_interval_reached(self): - exe = _make_executor(rebalance_check_interval=4) - results = [PyExecutor._can_pause_for_rebalance(exe) for _ in range(4)] - assert results == [False, False, False, True] + @staticmethod + def _fired_on(exe: MagicMock, iterations: int) -> list[int]: + """Run ``iterations`` executor iterations and report which ones opened + the gate. - def test_counter_resets_after_firing(self): - exe = _make_executor(rebalance_check_interval=3) - results = [PyExecutor._can_pause_for_rebalance(exe) for _ in range(6)] - assert results == [False, False, True, False, False, True] + Args: + exe: The PyExecutor stand-in to drive. + iterations: How many iterations to step through. - def test_interval_of_one_checks_every_iteration(self): + Returns: + The ``iter_counter`` values on which the gate returned ``True``. + """ + fired = [] + for i in range(iterations): + exe.iter_counter = i + if PyExecutor._can_pause_for_rebalance(exe): + fired.append(i) + return fired + + def test_fires_once_per_interval(self) -> None: + exe = _make_executor(rebalance_check_interval=4) + assert self._fired_on(exe, 12) == [0, 4, 8] + + def test_interval_of_one_checks_every_iteration(self) -> None: exe = _make_executor(rebalance_check_interval=1) - assert all(PyExecutor._can_pause_for_rebalance(exe) for _ in range(5)) + assert self._fired_on(exe, 5) == [0, 1, 2, 3, 4] - def test_counter_not_advanced_when_a_gate_rejects(self): - # The counter sits behind every gate, so an executor that cannot - # rebalance accumulates no throttle credit. Every gate is rank-uniform, - # which is what keeps the counter -- and therefore the iteration the - # agreement collective runs on -- identical across ranks. + def test_config_gate_suppresses_every_iteration(self) -> None: exe = _make_executor(enable_kv_pool_rebalance=False, rebalance_check_interval=2) - for _ in range(10): - assert PyExecutor._can_pause_for_rebalance(exe) is False - assert exe._rebalance_check_counter == 0 - - def test_counter_not_advanced_during_warmup(self): - exe = _make_executor(is_warmup=True, rebalance_check_interval=2) - for _ in range(10): - assert PyExecutor._can_pause_for_rebalance(exe) is False - assert exe._rebalance_check_counter == 0 - - # Once warmup clears, the throttle starts counting from scratch. - exe.is_warmup = False - assert [PyExecutor._can_pause_for_rebalance(exe) for _ in range(2)] == [False, True] + assert self._fired_on(exe, 10) == [] + + def test_gate_rejection_does_not_shift_the_schedule(self) -> None: + """Regression guard for the cadence-drift deadlock. + + A throttle counting its own eligible iterations would only advance on + iterations that cleared every gate, so a rank that bailed out early even + once would fire on a different set of iterations than its peers from + then on -- and the two would meet ``_agreed_need_adjustment``'s + broadcast out of step. Keying on ``iter_counter`` makes the schedule + independent of that history: a rank that skips a check rejoins the + common cadence instead of being permanently offset from it. + """ + interval = 4 + iterations = 16 + # A peer rank that never bails out: the cadence everyone must share. + reference = self._fired_on(_make_executor(rebalance_check_interval=interval), iterations) + + # This rank alone is briefly shut down, over a window that straddles one + # of its firing iterations. + window = (3, 4, 5) + exe = _make_executor(rebalance_check_interval=interval) + fired = [] + for i in range(iterations): + exe.iter_counter = i + exe.is_shutdown = i in window + if PyExecutor._can_pause_for_rebalance(exe): + fired.append(i) + + # The property that matters: this rank may *miss* checks, but it must + # never fire on an iteration its peers do not. Anything else means it + # carried a lasting offset out of the window and would meet them at + # _agreed_need_adjustment's broadcast out of step. + assert set(fired) <= set(reference), ( + f"rank drifted off the shared cadence: fired on {fired}, peers fire " + f"on {reference}; the extra iterations {sorted(set(fired) - set(reference))} " + "would enter the agreement collective alone" + ) + # And it really did skip the check inside the window (not vacuous). + assert [i for i in reference if i not in fired] == [4] + assert fired == [0, 8, 12] # --------------------------------------------------------------------------- # diff --git a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py index 7d29b5a8789d..3f7a1efcac5c 100644 --- a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py +++ b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py @@ -108,27 +108,78 @@ def _make_executor( exe.kv_cache_manager.impl = MagicMock() exe.kv_cache_manager.impl.need_adjustment = need_adjustment exe._rebalance_check_interval = rebalance_check_interval - exe._rebalance_check_counter = 0 + exe.iter_counter = 0 return exe def _tp_dist(world_size: int, rank: int) -> MPIDist: + """Build a pure-TP communicator for this rank. + + Args: + world_size: Total number of ranks, all of them TP ranks. + rank: This process's global rank. + + Returns: + An ``MPIDist`` whose mapping has ``tp_size == world_size``. + """ return MPIDist(Mapping(world_size=world_size, rank=rank, tp_size=world_size)) def _cp_dist(world_size: int, rank: int) -> MPIDist: - """A pure-CP mapping: ``tp_size == 1``, so ``cp_broadcast`` carries the - agreement and keying on ``tp_size`` alone would miss it entirely.""" + """Build a pure-CP communicator for this rank. + + ``tp_size`` is 1 here, so ``cp_broadcast`` alone carries the agreement and + keying the mechanism on ``tp_size`` would miss this configuration entirely. + + Args: + world_size: Total number of ranks, all of them CP ranks. + rank: This process's global rank. + + Returns: + An ``MPIDist`` whose mapping has ``cp_size == world_size`` and + ``tp_size == 1``. + """ return MPIDist(Mapping(world_size=world_size, rank=rank, cp_size=world_size)) +def _cp_tp_dist(cp_size: int, tp_size: int, rank: int) -> MPIDist: + """Build a combined CP x TP communicator for this rank. + + This is the only mapping that makes ``_agreed_need_adjustment`` run *both* + broadcasts, which is the shape the production chain actually takes. + + Args: + cp_size: Context-parallel width. + tp_size: Tensor-parallel width. + rank: This process's global rank. + + Returns: + An ``MPIDist`` over ``cp_size * tp_size`` ranks with both dimensions > 1. + """ + return MPIDist( + Mapping( + world_size=cp_size * tp_size, + rank=rank, + tp_size=tp_size, + cp_size=cp_size, + ) + ) + + # --------------------------------------------------------------------------- # Per-rank work # --------------------------------------------------------------------------- -def run_agreement(world_size, rank, local_flags, expected): - """Every TP rank must end up with rank 0's decision, not its own.""" +def run_agreement(world_size: int, rank: int, local_flags: list[bool], expected: bool) -> None: + """Every TP rank must end up with rank 0's decision, not its own. + + Args: + world_size: Number of TP ranks. + rank: This process's global rank. + local_flags: Per-rank local ``need_adjustment`` readings. + expected: The decision every rank should agree on. + """ dist = _tp_dist(world_size, rank) exe = _make_executor(dist, need_adjustment=local_flags[rank]) @@ -149,13 +200,19 @@ def run_agreement(world_size, rank, local_flags, expected): assert len(set(all_decisions)) == 1, f"ranks disagreed after the collective: {all_decisions}" -def run_cp_agreement(world_size, rank, local_flags, expected): +def run_cp_agreement(world_size: int, rank: int, local_flags: list[bool], expected: bool) -> None: """Same guarantee as TP, but over a pure-CP mapping (``tp_size == 1``). A request is split across CP ranks, so they must admit it together; and CP runs on the same executor loops as TP, which compute ``_schedule()`` per rank with no broadcast. Keying the agreement on ``tp_size`` alone would leave this configuration unsynchronized. + + Args: + world_size: Number of CP ranks. + rank: This process's global rank. + local_flags: Per-rank local ``need_adjustment`` readings. + expected: The decision every rank should agree on. """ dist = _cp_dist(world_size, rank) assert dist.tp_size == 1, "this test is meaningless unless tp_size is 1" @@ -173,12 +230,63 @@ def run_cp_agreement(world_size, rank, local_flags, expected): assert len(set(all_decisions)) == 1, f"CP ranks disagreed: {all_decisions}" -def run_attention_dp_independence(world_size, rank, local_flags): +def run_cp_tp_agreement( + world_size: int, rank: int, cp_size: int, local_flags: list[bool], expected: bool +) -> None: + """Global rank 0's decision must reach every rank through *both* broadcasts. + + This is the only case that exercises the production chain end to end: + ``cp_broadcast`` and then ``tp_broadcast``. The pure-TP test has + ``cp_size == 1`` and the pure-CP test has ``tp_size == 1``, so each skips one + of the two steps; neither can show that the TP step consumes the *result* of + the CP step rather than the rank's own local reading. + + With ``local_flags`` true only at global rank 0, a rank whose ``tp_rank`` is + non-zero can only end up ``True`` if its CP root picked up rank 0's decision + in the first step and then passed it on in the second. + + Args: + world_size: Total number of ranks, equal to ``cp_size * tp_size``. + rank: This process's global rank. + cp_size: Context-parallel width; TP width is derived from it. + local_flags: Per-rank local ``need_adjustment`` readings. + expected: The decision every rank should agree on. + """ + tp_size = world_size // cp_size + dist = _cp_tp_dist(cp_size, tp_size, rank) + # Both dimensions must really be >1, or this degenerates into one of the + # single-dimension tests and stops covering the chain. + assert dist.cp_size > 1 and dist.tp_size > 1, ( + f"this test only means something when both dimensions are > 1, got " + f"cp_size={dist.cp_size} tp_size={dist.tp_size}" + ) + exe = _make_executor(dist, need_adjustment=local_flags[rank]) + + got = PyExecutor._agreed_need_adjustment(exe) + + # Collect before asserting -- see run_agreement for why the order matters. + all_decisions = dist.allgather(got) + + assert got == expected, ( + f"rank {rank} (tp_rank={dist.mapping.tp_rank}, " + f"cp_rank={dist.mapping.cp_rank}): local reading was {local_flags[rank]}, " + f"global rank 0 said {local_flags[0]}, so the agreed decision should be " + f"{expected}, got {got}" + ) + assert len(set(all_decisions)) == 1, f"CPxTP ranks disagreed: {all_decisions}" + + +def run_attention_dp_independence(world_size: int, rank: int, local_flags: list[bool]) -> None: """Under attention DP each rank keeps its own decision. 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. + + Args: + world_size: Number of TP ranks. + rank: This process's global rank. + local_flags: Per-rank local ``need_adjustment`` readings. """ dist = _tp_dist(world_size, rank) exe = _make_executor(dist, need_adjustment=local_flags[rank], enable_attention_dp=True) @@ -190,11 +298,11 @@ def run_attention_dp_independence(world_size, rank, local_flags): ) -def run_throttle_lockstep(world_size, rank, interval, iterations): +def run_throttle_lockstep(world_size: int, rank: int, interval: int, iterations: int) -> None: """Ranks must reach the agreement collective on the *same* iterations. This is the deadlock guard. ``_can_pause_for_rebalance`` gates the - collective, so ranks whose throttle counters drifted apart would enter + collective, so ranks whose throttle cadence drifted apart would enter ``tp_broadcast`` on different iterations. Driving the real collective inside the loop means such a divergence hangs here rather than passing silently; the allgather afterwards pins down that the firing iterations @@ -202,14 +310,27 @@ def run_throttle_lockstep(world_size, rank, interval, iterations): Ranks are deliberately given *different* local ``need_adjustment`` readings, so it is the throttle cadence -- not agreement on the value -- that is under - test here. Everything ``_can_pause_for_rebalance`` reads is rank-uniform, - which is precisely the property being confirmed. + test here. + + Note the cadence is a pure function of ``iter_counter``, so drift can only + come from ranks disagreeing about the iteration index itself. The related + hazard -- a rank skipping a check because of a rank-local gate -- cannot be + covered here, because by construction it would leave the other ranks + blocked in a collective that rank never enters; that one is pinned down + without a real collective in ``TestRebalanceCheckThrottle``. + + Args: + world_size: Number of TP ranks. + rank: This process's global rank. + interval: Throttle interval to configure. + iterations: How many executor iterations to simulate. """ dist = _tp_dist(world_size, rank) exe = _make_executor(dist, need_adjustment=(rank % 2 == 0), rebalance_check_interval=interval) fired_on = [] for i in range(iterations): + exe.iter_counter = i if PyExecutor._can_pause_for_rebalance(exe): fired_on.append(i) # Real collective -- diverging ranks hang instead of passing. @@ -221,7 +342,7 @@ def run_throttle_lockstep(world_size, rank, interval, iterations): ) # Sanity: the throttle actually throttled, and it fired when expected. - expected = [i for i in range(iterations) if (i + 1) % interval == 0] + expected = [i for i in range(iterations) if i % interval == 0] assert fired_on == expected, ( f"throttle fired on {fired_on}, expected {expected} for interval " f"{interval} over {iterations} iterations" @@ -234,13 +355,29 @@ def run_throttle_lockstep(world_size, rank, interval, iterations): # --------------------------------------------------------------------------- -def _skip_if_not_enough_gpus(world_size): +def _skip_if_not_enough_gpus(world_size: int) -> None: + """Skip the calling test unless ``world_size`` GPUs are visible. + + Args: + world_size: Number of ranks the test needs. + """ if torch.cuda.device_count() < world_size: pytest.skip(f"need {world_size} GPUs, have {torch.cuda.device_count()}") -def _flags_for(case: str, world_size: int): - """Local per-rank ``need_adjustment`` readings for each skew scenario.""" +def _flags_for(case: str, world_size: int) -> list[bool]: + """Local per-rank ``need_adjustment`` readings for each skew scenario. + + Args: + case: Name of the skew scenario. + world_size: Number of ranks to produce readings for. + + Returns: + One local reading per rank. + + Raises: + ValueError: If ``case`` is not a known scenario. + """ if case == "all_true": return [True] * world_size if case == "all_false": @@ -261,8 +398,13 @@ def _flags_for(case: str, world_size: int): ["all_true", "all_false", "only_rank0_true", "only_rank0_false"], ) @pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"tp:{x}") -def test_tp_ranks_agree_on_rebalance_trigger(world_size, case): - """Rank 0's decision wins on every TP rank, whatever the local readings.""" +def test_tp_ranks_agree_on_rebalance_trigger(world_size: int, case: str) -> None: + """Rank 0's decision wins on every TP rank, whatever the local readings. + + Args: + world_size: Number of TP ranks to run on. + case: Skew scenario name, resolved by ``_flags_for``. + """ _skip_if_not_enough_gpus(world_size) flags = _flags_for(case, world_size) expected = flags[0] @@ -281,8 +423,13 @@ def test_tp_ranks_agree_on_rebalance_trigger(world_size, case): ["only_rank0_true", "only_rank0_false"], ) @pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"cp:{x}") -def test_cp_ranks_agree_on_rebalance_trigger(world_size, case): - """Pure CP (``tp_size == 1``) must agree just as TP does.""" +def test_cp_ranks_agree_on_rebalance_trigger(world_size: int, case: str) -> None: + """Pure CP (``tp_size == 1``) must agree just as TP does. + + Args: + world_size: Number of CP ranks to run on. + case: Skew scenario name, resolved by ``_flags_for``. + """ _skip_if_not_enough_gpus(world_size) flags = _flags_for(case, world_size) expected = flags[0] @@ -296,9 +443,48 @@ def test_cp_ranks_agree_on_rebalance_trigger(world_size, case): assert r is True +@pytest.mark.parametrize("cp_size", [2], ids=lambda x: f"cp:{x}") +@pytest.mark.parametrize("world_size", [4], ids=lambda x: f"world:{x}") +def test_cp_and_tp_ranks_agree_on_rebalance_trigger(world_size: int, cp_size: int) -> None: + """Both broadcasts chained: CP x TP propagates global rank 0's decision. + + The pure-TP and pure-CP cases each collapse one dimension to 1, so between + them they never run ``cp_broadcast`` and ``tp_broadcast`` back to back. This + case does, on a 2x2 topology. + + ``only_rank0_true`` is the discriminating scenario: every rank ends ``True`` + only if the TP step consumed the CP step's *result*. Had it broadcast each + rank's own local reading instead, the ranks in the second TP group would + come back ``False``. + + Args: + world_size: Total ranks; must equal ``cp_size * tp_size``. + cp_size: Context-parallel width; TP width is ``world_size // cp_size``. + """ + _skip_if_not_enough_gpus(world_size) + flags = _flags_for("only_rank0_true", world_size) + expected = flags[0] + assert flags == [True, False, False, False], ( + "this case only discriminates the chained broadcast when rank 0 alone " + f"reads True, got {flags}" + ) + + with MPIPoolExecutor(max_workers=world_size) as ex: + results = ex.map( + run_single_rank, + *zip(*[(world_size, run_cp_tp_agreement, cp_size, flags, expected)] * world_size), + ) + for r in results: + assert r is True + + @pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"tp:{x}") -def test_attention_dp_ranks_decide_independently(world_size): - """Attention DP opts out of the agreement: each rank keeps its own answer.""" +def test_attention_dp_ranks_decide_independently(world_size: int) -> None: + """Attention DP opts out of the agreement: each rank keeps its own answer. + + Args: + world_size: Number of TP ranks to run on. + """ _skip_if_not_enough_gpus(world_size) flags = _flags_for("only_rank0_false", world_size) @@ -313,11 +499,15 @@ def test_attention_dp_ranks_decide_independently(world_size): @pytest.mark.parametrize("interval", [1, 8], ids=lambda x: f"interval:{x}") @pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"tp:{x}") -def test_rebalance_check_stays_in_lockstep_across_ranks(world_size, interval): +def test_rebalance_check_stays_in_lockstep_across_ranks(world_size: int, interval: int) -> None: """The throttle must fire on identical iterations on every rank. A drift here would put ranks into ``tp_broadcast`` on different iterations, which deadlocks -- so this test hanging is itself the failure signal. + + Args: + world_size: Number of TP ranks to run on. + interval: Throttle interval to configure on every rank. """ _skip_if_not_enough_gpus(world_size) From 1f1cb52086fe1ab4e172f97fab1c5818764f4431 Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:04:36 +0000 Subject: [PATCH 4/9] [None][fix] Keep the CP broadcast under attention DP _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> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 25 ++++-- .../_torch/executor/test_kv_pool_rebalance.py | 23 +++++- .../multi_gpu/test_kv_pool_rebalance_tp.py | 76 +++++++++++++++++++ 3 files changed, 116 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 07338e70f1bd..e92623e791f1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4609,18 +4609,29 @@ def _agreed_need_adjustment(self) -> bool: ``V(tp_rank 0, cp_rank 0)``. Dedicated sub-communicators are used rather than the global one, which carries regular executor traffic. - Attention DP is excluded for the same reason the scheduler's own - ``tp_broadcast`` excludes it: ADP ranks own independent request streams - and independent KV caches, so they legitimately need different pool - ratios at different times. Forcing rank 0's decision on them would + Attention DP suppresses the **TP** hop only, matching the scheduler's + own propagation (``py_executor.py:2470-2477``), which likewise gates its + ``tp_broadcast`` on ``not enable_attention_dp`` while running its + ``cp_broadcast`` unconditionally. 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 independent KV caches and legitimately need different pool + ratios at different times; forcing rank 0's decision on them would starve a rank that needs to rebalance when rank 0 does not. + + CP is orthogonal to that and must **not** be skipped. Within a single + DP replica the CP ranks still split the same request along the sequence + dimension, so they must admit it together for exactly the reason given + above for pure CP. Nothing in ``Mapping`` or ``LlmArgs`` rejects + ``enable_attention_dp`` with ``cp_size > 1``, so skipping the CP hop + under ADP would reintroduce this PR's divergence inside every replica. + The net effect is that each DP replica decides on its own ``cp_rank``-0 + reading, and its CP ranks follow it. """ need = self.kv_cache_manager.impl.need_adjustment - if self.enable_attention_dp: - return need if self.dist.cp_size > 1: need = self.dist.cp_broadcast(need, root=0) - if self.dist.tp_size > 1: + if self.dist.tp_size > 1 and not self.enable_attention_dp: need = self.dist.tp_broadcast(need, root=0) return need diff --git a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py index 045c529f663f..8db7d274e9f1 100644 --- a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py +++ b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py @@ -557,8 +557,29 @@ def test_single_rank_touches_no_collective(self): exe.dist.cp_broadcast.assert_not_called() exe.dist.tp_broadcast.assert_not_called() - def test_attention_dp_skips_cp_broadcast_too(self): + def test_attention_dp_still_broadcasts_over_cp(self): + """ADP suppresses the TP hop only -- CP ranks must still agree. + + Under ADP the TP dimension is the DP dimension, so those ranks own + independent request streams and decide independently. But CP is + orthogonal: inside one DP replica the CP ranks split the *same* request + along the sequence dimension, so they have to admit it together. + Skipping the CP hop here would reintroduce the divergence this whole + mechanism exists to remove, once per replica. This mirrors the + scheduler's own propagation, which gates ``tp_broadcast`` on + ``not enable_attention_dp`` but runs ``cp_broadcast`` unconditionally. + """ exe = _make_executor(tp_size=2, cp_size=2, enable_attention_dp=True, need_adjustment=True) + exe.dist.cp_broadcast.return_value = False + + # The CP result wins over this rank's own reading... + assert PyExecutor._agreed_need_adjustment(exe) is False + exe.dist.cp_broadcast.assert_called_once_with(True, root=0) + # ...but the TP hop stays suppressed, so replicas remain independent. + exe.dist.tp_broadcast.assert_not_called() + + def test_attention_dp_without_cp_touches_no_collective(self): + exe = _make_executor(tp_size=2, cp_size=1, enable_attention_dp=True, need_adjustment=True) assert PyExecutor._agreed_need_adjustment(exe) is True exe.dist.cp_broadcast.assert_not_called() exe.dist.tp_broadcast.assert_not_called() diff --git a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py index 3f7a1efcac5c..48340ee97343 100644 --- a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py +++ b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py @@ -298,6 +298,45 @@ def run_attention_dp_independence(world_size: int, rank: int, local_flags: list[ ) +def run_adp_cp_agreement( + world_size: int, rank: int, cp_size: int, local_flags: list[bool], expected: list[bool] +) -> None: + """Under ADP the TP hop is suppressed but the CP hop must still run. + + ADP makes the TP dimension the DP dimension, so those ranks decide + independently. CP is orthogonal: inside one DP replica the CP ranks split + the *same* request along the sequence dimension and must admit it together. + Skipping the CP broadcast here would reintroduce this mechanism's own + divergence once per replica. + + Args: + world_size: Total ranks, equal to ``cp_size * tp_size``. + rank: This process's global rank. + cp_size: Context-parallel width; TP width is derived from it. + local_flags: Per-rank local ``need_adjustment`` readings. + expected: Per-rank agreed decision after the CP-only broadcast. + """ + tp_size = world_size // cp_size + dist = _cp_tp_dist(cp_size, tp_size, rank) + exe = _make_executor(dist, need_adjustment=local_flags[rank], enable_attention_dp=True) + + got = PyExecutor._agreed_need_adjustment(exe) + + # Collect before asserting -- see run_agreement for why the order matters. + all_decisions = dist.allgather(got) + + assert got == expected[rank], ( + f"rank {rank} (tp_rank={dist.mapping.tp_rank}, " + f"cp_rank={dist.mapping.cp_rank}): local reading was {local_flags[rank]}, " + f"its CP root read {expected[rank]}, so the agreed decision should be " + f"{expected[rank]}, got {got}" + ) + assert all_decisions == expected, ( + f"ADP+CP decisions were {all_decisions}, expected {expected}: CP ranks " + "must follow their replica's root while replicas stay independent" + ) + + def run_throttle_lockstep(world_size: int, rank: int, interval: int, iterations: int) -> None: """Ranks must reach the agreement collective on the *same* iterations. @@ -497,6 +536,43 @@ def test_attention_dp_ranks_decide_independently(world_size: int) -> None: assert r is True +@pytest.mark.parametrize("cp_size", [2], ids=lambda x: f"cp:{x}") +@pytest.mark.parametrize("world_size", [4], ids=lambda x: f"world:{x}") +def test_attention_dp_agrees_over_cp_but_not_tp(world_size: int, cp_size: int) -> None: + """ADP + CP: CP ranks follow their replica's root, replicas stay independent. + + ADP suppresses only the TP hop. Nothing in ``Mapping`` or ``LlmArgs`` + rejects ``enable_attention_dp`` with ``cp_size > 1``, so this topology is + reachable, and skipping the CP hop in it would leave each replica's CP ranks + deciding independently -- the divergence class this mechanism removes. + + The flags are chosen so the case fails 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). + + Args: + world_size: Total ranks; must equal ``cp_size * tp_size``. + cp_size: Context-parallel width; TP width is ``world_size // cp_size``. + """ + _skip_if_not_enough_gpus(world_size) + + # cp_groups are consecutive ranks within a TP slice: [[0, 1], [2, 3]]. + flags = [True, False, False, True] + expected = [flags[(r // cp_size) * cp_size] for r in range(world_size)] + + # Non-vacuity, both directions. + assert expected != flags, "flags must force the CP broadcast to change some rank" + assert len(set(expected)) > 1, "replicas must end up disagreeing, or the TP hop is untested" + + with MPIPoolExecutor(max_workers=world_size) as ex: + results = ex.map( + run_single_rank, + *zip(*[(world_size, run_adp_cp_agreement, cp_size, flags, expected)] * world_size), + ) + for r in results: + assert r is True + + @pytest.mark.parametrize("interval", [1, 8], ids=lambda x: f"interval:{x}") @pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"tp:{x}") def test_rebalance_check_stays_in_lockstep_across_ranks(world_size: int, interval: int) -> None: From ef96782327a801bc57a26d6436847616ae0b94bc Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Wed, 12 Aug 2026 21:46:12 +0000 Subject: [PATCH 5/9] [None][feat] Extend the KVCacheManagerV2 rebalance to pipeline parallelism _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> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 164 +++++++++++- .../_torch/executor/test_kv_pool_rebalance.py | 235 +++++++++++++++- .../multi_gpu/test_kv_pool_rebalance_tp.py | 253 ++++++++++++++++++ 3 files changed, 642 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e92623e791f1..342dc8005986 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -607,6 +607,10 @@ def __init__( # Iteration throttle for the KV pool rebalance check. See # _can_pause_for_rebalance / _agreed_need_adjustment. self._rebalance_check_interval = KV_POOL_REBALANCE_CHECK_INTERVAL + # Countdown of drain iterations remaining before a pipeline-parallel + # rebalance can run; None when no rebalance is pending. Only + # _executor_loop_pp uses it -- see _start_pp_rebalance_drain. + self._pp_rebalance_drain_iters: Optional[int] = None self.enable_early_first_token_response = enable_early_first_token_response self.virtual_memory_pools = virtual_memory_pools @@ -2673,6 +2677,16 @@ def _executor_loop_pp(self): if self.enable_iter_perf_stats: iter_start_time = time.time() + # A rebalance cannot run inline here the way it does in the + # other two loops -- the ring has to drain first -- so this + # only starts the drain. Skipped while one is already + # pending so the decision is not retaken mid-drain. + if (self._uses_kv_manager_v2() + and self._pp_rebalance_drain_iters is None + and self._can_pause_for_rebalance() + and self._agreed_need_adjustment()): + self._start_pp_rebalance_drain() + self._handle_disagg_cache_errors_synced() # Fetch new requests from request queue @@ -2747,6 +2761,12 @@ def _executor_loop_pp(self): self._run_encoder_step(scheduled_batch.encoder_requests) can_queue, _ = self._can_queue(scheduled_batch) + if self._pp_rebalance_drain_iters is not None: + # Draining for a KV pool rebalance: stop feeding the ring + # so it can empty out. Every rank starts and ends the + # drain on the same iteration, so suppressing the queue + # here keeps them in lockstep rather than breaking it. + can_queue = False if not can_queue: self._revert_gen_alloc(scheduled_batch) if not can_queue: @@ -2978,6 +2998,11 @@ def handle_executed_batches(executed_batch_num: int): # Stage 3.3: Handle executed batches. handle_executed_batches(executed_batch_num) + # Stage 3.4: Rebalance the KV pools once the drain started at + # the top of some earlier iteration has emptied the ring. + if self._uses_kv_manager_v2(): + self._maybe_finish_pp_rebalance() + # Stage 4: March forward in microbatch slots microbatch_id = (microbatch_id + 1) % self.num_micro_batches self.iter_counter += 1 @@ -4520,11 +4545,20 @@ def _sync_and_process_resource_governor_queue(self): raise ValueError(f"Invalid request type: {type(request)}.") def _can_pause_for_rebalance(self) -> bool: - """Gate KV pool rebalance to the cases the v1 hook supports. + """Gate KV pool rebalance to the cases the hook supports. - MVP scope: no pipeline parallelism, no in-flight disagg transfer, - no beam search, no drafter, not during warmup or shutdown. - Honors the ``enable_kv_pool_rebalance`` opt-in flag (default off). + Scope: no in-flight disagg transfer, no beam search, no drafter, not + during warmup or shutdown. Honors the ``enable_kv_pool_rebalance`` + opt-in flag (default off). + + Pipeline parallelism *is* supported, but not in the same shape as the + other two loops. ``_executor_loop`` and ``_executor_loop_overlap`` + rebalance inline, because at the top of an iteration at most one + in-flight batch exists (``previous_batch``) and it can be consumed on + the spot. ``_executor_loop_pp`` keeps up to ``num_micro_batches`` + batches in flight in a ring, so a ``True`` here only *starts* a drain; + the rebalance itself happens once the ring is empty. See + ``_start_pp_rebalance_drain``. A ``True`` here is what puts a TP rank into ``_agreed_need_adjustment``'s collective, so ranks that disagree on this @@ -4539,8 +4573,6 @@ def _can_pause_for_rebalance(self) -> bool: """ if not self.enable_kv_pool_rebalance: return False - if self.dist.pp_size > 1: - return False if self.kv_cache_transceiver is not None: return False if self.is_warmup: @@ -4627,12 +4659,38 @@ def _agreed_need_adjustment(self) -> bool: under ADP would reintroduce this PR's divergence inside every replica. The net effect is that each DP replica decides on its own ``cp_rank``-0 reading, and its CP ranks follow it. + + Pipeline parallelism needs agreement too, for a different reason. PP + ranks do *not* schedule independently -- the first PP rank schedules and + propagates the result (``_pp_schedule_and_propagate``) -- so the + divergent-``_schedule()`` argument above does not apply. What does + apply is the drain: rebalancing under PP requires every rank to stop + feeding the microbatch ring **on the same iteration** + (``_start_pp_rebalance_drain``). A rank that drained alone while its + peers kept queueing microbatches would desynchronize the per-iteration + send/recv chain -- the sample-state relay and + ``ring_broadcast_executed_batch_num`` -- and hang the pipeline. Each PP + rank also holds a *different* slice of layers, hence its own pools, its + own ``need_adjustment`` reading and its own cooldown clock, so left + alone they would not agree by construction. + + Unlike the TP hop, the PP hop must **not** be suppressed under attention + DP. ADP replicates along the TP dimension, so a replica's pipeline + stages all serve that replica's own request stream and must drain + together; ``pp_group`` for a given ``tp_rank`` holds exactly those + stages, so broadcasting over it propagates the replica's first-stage + decision to its own stages and nothing wider. + + Chaining CP, then TP, then PP leaves every rank holding global rank 0's + value (or, under ADP, its own replica's first-stage value). """ 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) + if self.dist.pp_size > 1: + need = self.dist.pp_broadcast(need, root=0) return need def _consume_previous_batch_for_rebalance(self) -> None: @@ -4672,13 +4730,28 @@ def _maybe_rebalance_kv_pools(self) -> None: scheduler reactivates them through prepare_context / try_allocate_generation on the next iteration, the same path it uses today after eviction. + + This is the inline path used by ``_executor_loop`` and + ``_executor_loop_overlap``, where consuming ``previous_batch`` is + enough to reach quiescence. ``_executor_loop_pp`` cannot rebalance + inline and instead drains its microbatch ring first, then calls + ``_rebalance_kv_pools_now`` directly. """ - mgr = self.kv_cache_manager if not self._agreed_need_adjustment(): return torch.cuda.current_stream().synchronize() self._consume_previous_batch_for_rebalance() + self._rebalance_kv_pools_now() + + def _rebalance_kv_pools_now(self) -> None: + """Suspend every active request, ``adjust()``, resume. + + The caller guarantees quiescence -- no forward may be in flight and + no _KVCache may be mid-update -- because ``adjust()`` moves pages + underneath whatever holds them. + """ + mgr = self.kv_cache_manager paused: List[LlmRequest] = [] for req in self.active_requests: @@ -4696,6 +4769,83 @@ def _maybe_rebalance_kv_pools(self) -> None: mgr.resume_request(req) self._resume_padding_dummies_after_rebalance(mgr, paused_dummies) + def _start_pp_rebalance_drain(self) -> None: + """Begin emptying the microbatch ring ahead of a PP rebalance. + + ``adjust()`` needs every _KVCache quiescent, but ``_executor_loop_pp`` + keeps up to ``num_micro_batches`` batches in flight across the pipeline + at once. Those cannot be consumed on the spot the way the overlap + loop's single ``previous_batch`` can: a microbatch is completed by the + sample-state relay travelling around the PP ring, which only advances + as iterations run. + + So instead of draining inline we stop *feeding* the ring and let the + loop drain it: while a drain is pending, ``_executor_loop_pp`` forces + ``can_queue`` to False, which is the loop's existing "skip this + microbatch slot" path -- the same one an empty batch takes on an idle + server (``_can_queue`` is just ``batch_size > 0``). Each iteration + retires one slot and queues nothing new, so ``num_micro_batches`` + iterations empty the ring. + + The countdown is a rank-independent constant and every rank starts it + on the same iteration (``_agreed_need_adjustment`` broadcasts over the + PP group), so all ranks stop feeding, reach quiescence, and rebalance + together. That is what keeps the per-iteration send/recv chain in + lockstep through the drain. + """ + self._pp_rebalance_drain_iters = self.num_micro_batches + + def _pp_ring_is_quiescent(self) -> bool: + """True when no microbatch is in flight anywhere in the ring. + + Both halves matter: ``micro_batches`` covers batches this rank has + queued but not yet retired, and ``unhandled_batch_counter`` covers + sample states that have been handed to the relay but whose results + have not been applied to the requests yet. + """ + return (self._pp_rebalance_drain_iters is not None + and self.unhandled_batch_counter == 0 + and all(mb is None for mb in self.micro_batches)) + + def _maybe_finish_pp_rebalance(self) -> None: + """Run the pending PP rebalance once the ring has drained. + + Called at the end of every ``_executor_loop_pp`` iteration. Counts + down the drain and, when it expires, rebalances. Both the countdown + and the quiescence test are pure functions of state that the loop keeps + symmetric across ranks, so every rank takes the same branch on the same + iteration. + + If the ring is somehow still busy when the countdown expires the + rebalance is skipped rather than forced: adjusting underneath a live + _KVCache would corrupt it, and skipping costs only a delay -- the + auto-tuner still wants the adjustment and will ask again on the next + check interval. + """ + if self._pp_rebalance_drain_iters is None: + return + self._pp_rebalance_drain_iters -= 1 + if self._pp_rebalance_drain_iters > 0: + return + + if not self._pp_ring_is_quiescent(): + logger.warning( + "KV pool rebalance skipped: the PP microbatch ring was still " + f"busy after {self.num_micro_batches} drain iterations " + f"(unhandled_batch_counter={self.unhandled_batch_counter}, " + f"in-flight slots=" + f"{sum(mb is not None for mb in self.micro_batches)}).") + self._pp_rebalance_drain_iters = None + return + + torch.cuda.current_stream().synchronize() + if self.pp_multi_stream_sample: + # Sampling for the last microbatch runs on its own stream; its + # writes must land before pages move underneath them. + self.sample_stream.synchronize() + self._rebalance_kv_pools_now() + self._pp_rebalance_drain_iters = None + def _suspend_padding_dummies_for_rebalance( self, mgr: KVCacheManagerV2) -> List[Tuple[int, LlmRequest]]: """Suspend the CUDA-graph padding dummies before ``adjust()``. diff --git a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py index 8db7d274e9f1..e4ff1c4706c6 100644 --- a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py +++ b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py @@ -57,6 +57,11 @@ def _make_executor( padding_dummies=None, has_cuda_graph_runner: bool = True, rebalance_check_interval: int = 1, + num_micro_batches: int = 1, + micro_batches=None, + unhandled_batch_counter: int = 0, + pp_rebalance_drain_iters=None, + pp_multi_stream_sample: bool = False, ) -> MagicMock: """Construct a MagicMock shaped like PyExecutor with exactly the attributes the rebalance code path reads. @@ -84,10 +89,26 @@ def _make_executor( exe._rebalance_check_interval = rebalance_check_interval exe.iter_counter = 0 + # Pipeline-parallel microbatch ring. Only _executor_loop_pp's drain path + # reads these; the default is a quiescent depth-1 ring, i.e. the shape the + # non-PP loops present. + exe.num_micro_batches = num_micro_batches + exe.micro_batches = ( + list(micro_batches) if micro_batches is not None else [None] * num_micro_batches + ) + exe.unhandled_batch_counter = unhandled_batch_counter + exe._pp_rebalance_drain_iters = pp_rebalance_drain_iters + exe.pp_multi_stream_sample = pp_multi_stream_sample + # Bind the real agreement helper so _maybe_rebalance_kv_pools exercises it # rather than getting a truthy MagicMock back. exe._agreed_need_adjustment = lambda: PyExecutor._agreed_need_adjustment(exe) + # Same for the suspend/adjust/resume core, which both the inline path and + # the PP drain delegate to. + exe._rebalance_kv_pools_now = lambda: PyExecutor._rebalance_kv_pools_now(exe) + exe._pp_ring_is_quiescent = lambda: PyExecutor._pp_ring_is_quiescent(exe) + # KV cache manager (resource-manager wrapper). exe.kv_cache_manager = MagicMock() exe.kv_cache_manager.max_beam_width = max_beam_width @@ -158,9 +179,15 @@ def test_flag_off_returns_false(self): exe = _make_executor(enable_kv_pool_rebalance=False) assert PyExecutor._can_pause_for_rebalance(exe) is False - def test_pp_size_gt_one_returns_false(self): + def test_pp_size_gt_one_is_allowed(self): + """PP no longer short-circuits the gate. + + Under PP a ``True`` here starts a ring drain rather than rebalancing + inline (see ``_start_pp_rebalance_drain``), but the gate itself must + let PP through for that to happen at all. + """ exe = _make_executor(pp_size=2) - assert PyExecutor._can_pause_for_rebalance(exe) is False + assert PyExecutor._can_pause_for_rebalance(exe) is True def test_transceiver_present_returns_false(self): exe = _make_executor(kv_cache_transceiver=MagicMock()) @@ -552,10 +579,53 @@ def test_tp_and_cp_chain_propagates_global_rank0(self): exe.dist.tp_broadcast.assert_called_once_with(False, root=0) def test_single_rank_touches_no_collective(self): - exe = _make_executor(tp_size=1, cp_size=1, need_adjustment=True) + exe = _make_executor(tp_size=1, cp_size=1, pp_size=1, need_adjustment=True) assert PyExecutor._agreed_need_adjustment(exe) is True exe.dist.cp_broadcast.assert_not_called() exe.dist.tp_broadcast.assert_not_called() + exe.dist.pp_broadcast.assert_not_called() + + def test_pp_broadcasts_first_stage_decision(self): + # PP ranks hold different layers, hence different pools, different + # need_adjustment readings and independent cooldown clocks. They must + # still start the ring drain on the same iteration or the pipeline's + # per-iteration send/recv chain desynchronizes. + exe = _make_executor(pp_size=4, need_adjustment=True) + exe.dist.pp_broadcast.return_value = False + + assert PyExecutor._agreed_need_adjustment(exe) is False + exe.dist.pp_broadcast.assert_called_once_with(True, root=0) + + def test_pp_hop_is_last_so_global_rank0_wins(self): + # CP, then TP, then PP: each hop must consume the previous hop's result, + # leaving every rank holding global rank 0's value. The three return + # values differ from the local reading and from each other's inputs so + # that a dropped link in the chain fails the assertion. + exe = _make_executor(tp_size=2, cp_size=2, pp_size=2, need_adjustment=True) + exe.dist.cp_broadcast.return_value = False + exe.dist.tp_broadcast.return_value = True + exe.dist.pp_broadcast.return_value = False + + assert PyExecutor._agreed_need_adjustment(exe) is False + exe.dist.cp_broadcast.assert_called_once_with(True, root=0) + exe.dist.tp_broadcast.assert_called_once_with(False, root=0) + exe.dist.pp_broadcast.assert_called_once_with(True, root=0) + + def test_pp_hop_survives_attention_dp(self): + """ADP suppresses the TP hop but must never suppress the PP hop. + + ADP replicates along the TP dimension, so a replica's pipeline stages + all serve that replica's own request stream and have to drain together. + ``pp_group`` for a given tp_rank holds exactly those stages, so the + broadcast propagates the replica's first-stage decision and nothing + wider. + """ + exe = _make_executor(tp_size=2, pp_size=2, enable_attention_dp=True, need_adjustment=True) + exe.dist.pp_broadcast.return_value = False + + assert PyExecutor._agreed_need_adjustment(exe) is False + exe.dist.tp_broadcast.assert_not_called() + exe.dist.pp_broadcast.assert_called_once_with(True, root=0) def test_attention_dp_still_broadcasts_over_cp(self): """ADP suppresses the TP hop only -- CP ranks must still agree. @@ -603,5 +673,164 @@ def test_no_rebalance_when_agreement_says_no(self, monkeypatch): exe.kv_cache_manager.suspend_request.assert_not_called() +# --------------------------------------------------------------------------- # +# Pipeline-parallel drain +# --------------------------------------------------------------------------- # + + +def _batch_state() -> MagicMock: + """Stand-in for an in-flight BatchStatePP occupying a ring slot.""" + return MagicMock() + + +class TestPpRingIsQuiescent: + """``adjust()`` may only run when nothing is in flight in the PP ring.""" + + def test_empty_ring_with_drain_pending_is_quiescent(self): + exe = _make_executor(pp_size=2, num_micro_batches=2, pp_rebalance_drain_iters=1) + assert PyExecutor._pp_ring_is_quiescent(exe) is True + + def test_occupied_slot_is_not_quiescent(self): + exe = _make_executor( + pp_size=2, + num_micro_batches=2, + micro_batches=[None, _batch_state()], + pp_rebalance_drain_iters=1, + ) + assert PyExecutor._pp_ring_is_quiescent(exe) is False + + def test_unhandled_sample_state_is_not_quiescent(self): + # The slot is already cleared, but the sample state is still travelling + # the relay and has not been applied to its requests yet. + exe = _make_executor( + pp_size=2, num_micro_batches=2, unhandled_batch_counter=1, pp_rebalance_drain_iters=1 + ) + assert PyExecutor._pp_ring_is_quiescent(exe) is False + + def test_no_drain_pending_is_not_quiescent(self): + exe = _make_executor(pp_size=2, num_micro_batches=2, pp_rebalance_drain_iters=None) + assert PyExecutor._pp_ring_is_quiescent(exe) is False + + +class TestStartPpRebalanceDrain: + def test_countdown_covers_the_whole_ring(self): + # One slot retires per iteration while nothing new is queued, so a full + # ring's worth of iterations empties it. + exe = _make_executor(pp_size=4, num_micro_batches=4) + PyExecutor._start_pp_rebalance_drain(exe) + assert exe._pp_rebalance_drain_iters == 4 + + +class TestMaybeFinishPpRebalance: + """The drain counts down, then rebalances exactly once.""" + + @staticmethod + def _exe(monkeypatch, **kwargs): + exe = _make_executor(pp_size=2, **kwargs) + monkeypatch.setattr("torch.cuda.current_stream", MagicMock()) + return exe + + def test_no_op_when_no_drain_pending(self, monkeypatch): + exe = self._exe(monkeypatch, num_micro_batches=2, pp_rebalance_drain_iters=None) + PyExecutor._maybe_finish_pp_rebalance(exe) + exe.kv_cache_manager.impl.adjust.assert_not_called() + + def test_does_not_rebalance_before_the_ring_drains(self, monkeypatch): + exe = self._exe(monkeypatch, num_micro_batches=4, pp_rebalance_drain_iters=4) + PyExecutor._maybe_finish_pp_rebalance(exe) + + assert exe._pp_rebalance_drain_iters == 3 + exe.kv_cache_manager.impl.adjust.assert_not_called() + + def test_rebalances_when_the_countdown_expires(self, monkeypatch): + req = _make_request(1) + exe = self._exe( + monkeypatch, num_micro_batches=2, pp_rebalance_drain_iters=1, active_requests=[req] + ) + + PyExecutor._maybe_finish_pp_rebalance(exe) + + exe.kv_cache_manager.suspend_request.assert_called_once_with(req) + exe.kv_cache_manager.impl.adjust.assert_called_once() + exe.kv_cache_manager.resume_request.assert_called_once_with(req) + # Cleared, so the next check interval starts a fresh drain. + assert exe._pp_rebalance_drain_iters is None + + def test_full_drain_rebalances_exactly_once(self, monkeypatch): + exe = self._exe(monkeypatch, num_micro_batches=3, active_requests=[_make_request(1)]) + PyExecutor._start_pp_rebalance_drain(exe) + + for _ in range(3): + PyExecutor._maybe_finish_pp_rebalance(exe) + # Extra iterations after the drain must not rebalance again. + for _ in range(3): + PyExecutor._maybe_finish_pp_rebalance(exe) + + exe.kv_cache_manager.impl.adjust.assert_called_once() + + def test_busy_ring_skips_rather_than_corrupting(self, monkeypatch): + # Adjusting underneath a live _KVCache would move pages out from under + # it. Skipping only costs a delay: the auto-tuner still wants the + # adjustment and asks again on the next check interval. + exe = self._exe( + monkeypatch, + num_micro_batches=2, + micro_batches=[None, _batch_state()], + pp_rebalance_drain_iters=1, + active_requests=[_make_request(1)], + ) + + PyExecutor._maybe_finish_pp_rebalance(exe) + + exe.kv_cache_manager.impl.adjust.assert_not_called() + exe.kv_cache_manager.suspend_request.assert_not_called() + assert exe._pp_rebalance_drain_iters is None + + def test_sample_stream_is_synced_when_multi_stream_sampling(self, monkeypatch): + # Sampling for the last microbatch runs on its own stream; its writes + # must land before adjust() moves pages underneath them. + exe = self._exe( + monkeypatch, + num_micro_batches=2, + pp_rebalance_drain_iters=1, + pp_multi_stream_sample=True, + ) + exe.sample_stream = MagicMock() + + PyExecutor._maybe_finish_pp_rebalance(exe) + + exe.sample_stream.synchronize.assert_called_once() + + def test_sample_stream_untouched_without_multi_stream_sampling(self, monkeypatch): + exe = self._exe( + monkeypatch, + num_micro_batches=2, + pp_rebalance_drain_iters=1, + pp_multi_stream_sample=False, + ) + exe.sample_stream = MagicMock() + + PyExecutor._maybe_finish_pp_rebalance(exe) + + exe.sample_stream.synchronize.assert_not_called() + + def test_previous_batch_is_not_consumed(self, monkeypatch): + """The PP path must not run the overlap loop's drain helper. + + Under PP a microbatch is retired by the sample-state relay, not by + ``_consume_previous_batch_for_rebalance``; ``previous_batch`` is kept + only to synchronize its sampler event. Running the overlap helper on it + would double-handle a batch the ring already completed. + """ + exe = self._exe( + monkeypatch, num_micro_batches=2, pp_rebalance_drain_iters=1, previous_batch=MagicMock() + ) + + PyExecutor._maybe_finish_pp_rebalance(exe) + + exe._consume_previous_batch_for_rebalance.assert_not_called() + exe.kv_cache_manager.impl.adjust.assert_called_once() + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py index 48340ee97343..a6a8d307240c 100644 --- a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py +++ b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py @@ -166,6 +166,48 @@ def _cp_tp_dist(cp_size: int, tp_size: int, rank: int) -> MPIDist: ) +def _pp_dist(world_size: int, rank: int) -> MPIDist: + """Build a pure-PP communicator for this rank. + + ``tp_size`` and ``cp_size`` are both 1, so ``pp_broadcast`` alone carries + the agreement. PP ranks hold different layers and therefore different + pools, so without this hop they would never agree by construction. + + Args: + world_size: Total number of ranks, all of them PP ranks. + rank: This process's global rank. + + Returns: + An ``MPIDist`` whose mapping has ``pp_size == world_size``. + """ + return MPIDist(Mapping(world_size=world_size, rank=rank, pp_size=world_size)) + + +def _tp_pp_dist(tp_size: int, pp_size: int, rank: int) -> MPIDist: + """Build a combined TP x PP communicator for this rank. + + Ranks are laid out ``rank == pp_rank * tp_size + tp_rank``, so with + ``tp_size == pp_size == 2`` the TP groups are ``[0, 1]`` and ``[2, 3]`` + while the PP groups are ``[0, 2]`` and ``[1, 3]``. + + Args: + tp_size: Tensor-parallel width. + pp_size: Pipeline-parallel width. + rank: This process's global rank. + + Returns: + An ``MPIDist`` over ``tp_size * pp_size`` ranks with both dimensions > 1. + """ + return MPIDist( + Mapping( + world_size=tp_size * pp_size, + rank=rank, + tp_size=tp_size, + pp_size=pp_size, + ) + ) + + # --------------------------------------------------------------------------- # Per-rank work # --------------------------------------------------------------------------- @@ -337,6 +379,135 @@ def run_adp_cp_agreement( ) +def run_pp_agreement(world_size: int, rank: int, local_flags: list[bool], expected: bool) -> None: + """Every PP rank must end up with the first stage's decision, not its own. + + PP ranks hold different layers, so each has its own pools, its own + ``need_adjustment`` reading and its own cooldown clock -- left alone they + would not agree by construction. They must agree because rebalancing under + PP means draining the microbatch ring, and a rank that stopped feeding the + ring while its peers kept going would desynchronize the per-iteration + send/recv chain and hang the pipeline. + + Args: + world_size: Number of PP ranks. + rank: This process's global rank. + local_flags: Per-rank local ``need_adjustment`` readings. + expected: The decision every rank should agree on. + """ + dist = _pp_dist(world_size, rank) + # A pure-PP mapping collapses TP and CP, so pp_broadcast alone carries this. + assert dist.pp_size > 1 and dist.tp_size == 1 and dist.cp_size == 1, ( + f"expected a pure-PP mapping, got pp_size={dist.pp_size} " + f"tp_size={dist.tp_size} cp_size={dist.cp_size}" + ) + exe = _make_executor(dist, need_adjustment=local_flags[rank]) + + got = PyExecutor._agreed_need_adjustment(exe) + + # Collect before asserting -- see run_agreement for why the order matters. + all_decisions = dist.allgather(got) + + assert got == expected, ( + f"rank {rank} (pp_rank={dist.mapping.pp_rank}): local reading was " + f"{local_flags[rank]}, the first PP stage said {local_flags[0]}, so the " + f"agreed decision should be {expected}, got {got}" + ) + assert len(set(all_decisions)) == 1, f"PP ranks disagreed: {all_decisions}" + + +def run_tp_pp_agreement( + world_size: int, rank: int, tp_size: int, local_flags: list[bool], expected: bool +) -> None: + """The TP and PP hops must chain, leaving global rank 0's decision everywhere. + + Ranks are laid out ``rank == pp_rank * tp_size + tp_rank``, so on a 2x2 the + TP groups are ``[0, 1]`` and ``[2, 3]`` and the PP groups are ``[0, 2]`` and + ``[1, 3]``. With only global rank 0 reading ``True``: + + * the TP step gives ranks 0 and 1 rank 0's ``True``, and ranks 2 and 3 rank + 2's ``False``; + * the PP step then broadcasts rank 0's ``True`` to rank 2, and rank 1's + *post-TP* ``True`` to rank 3. + + So every rank ends ``True`` -- but only if the PP hop consumes the TP hop's + result. Had it broadcast each rank's own local reading, ranks 1 and 3 would + come back ``False``, which is exactly what makes this case discriminating. + + Args: + world_size: Total ranks, equal to ``tp_size * pp_size``. + rank: This process's global rank. + tp_size: Tensor-parallel width; PP width is derived from it. + local_flags: Per-rank local ``need_adjustment`` readings. + expected: The decision every rank should agree on. + """ + pp_size = world_size // tp_size + dist = _tp_pp_dist(tp_size, pp_size, rank) + assert dist.pp_size > 1 and dist.tp_size > 1, ( + f"this test only means something when both dimensions are > 1, got " + f"pp_size={dist.pp_size} tp_size={dist.tp_size}" + ) + exe = _make_executor(dist, need_adjustment=local_flags[rank]) + + got = PyExecutor._agreed_need_adjustment(exe) + + # Collect before asserting -- see run_agreement for why the order matters. + all_decisions = dist.allgather(got) + + assert got == expected, ( + f"rank {rank} (tp_rank={dist.mapping.tp_rank}, " + f"pp_rank={dist.mapping.pp_rank}): local reading was {local_flags[rank]}, " + f"global rank 0 said {local_flags[0]}, so the agreed decision should be " + f"{expected}, got {got}" + ) + assert len(set(all_decisions)) == 1, f"TPxPP ranks disagreed: {all_decisions}" + + +def run_adp_pp_agreement( + world_size: int, rank: int, tp_size: int, local_flags: list[bool], expected: list[bool] +) -> None: + """Under ADP the TP hop is suppressed but the PP hop must still run. + + ADP replicates along the TP dimension, so each ``tp_rank`` owns an + independent request stream -- but a replica's *pipeline stages* all serve + that one stream and must drain together or the replica's pipeline hangs. + ``pp_group`` for a given ``tp_rank`` holds exactly those stages. + + On a 2x2 with only global rank 0 reading ``True``, the expected outcome is + ``[True, False, True, False]``: ranks 0 and 2 form one replica and follow + rank 0, ranks 1 and 3 form the other and follow rank 1. That single vector + pins down both halves at once -- rank 2 being ``True`` shows the PP hop ran, + and rank 1 staying ``False`` shows the TP hop did not. + + Args: + world_size: Total ranks, equal to ``tp_size * pp_size``. + rank: This process's global rank. + tp_size: Tensor-parallel width; PP width is derived from it. + local_flags: Per-rank local ``need_adjustment`` readings. + expected: Per-rank agreed decision after the PP-only broadcast. + """ + pp_size = world_size // tp_size + dist = _tp_pp_dist(tp_size, pp_size, rank) + exe = _make_executor(dist, need_adjustment=local_flags[rank], enable_attention_dp=True) + + got = PyExecutor._agreed_need_adjustment(exe) + + # Collect before asserting -- see run_agreement for why the order matters. + all_decisions = dist.allgather(got) + + assert got == expected[rank], ( + f"rank {rank} (tp_rank={dist.mapping.tp_rank}, " + f"pp_rank={dist.mapping.pp_rank}): local reading was {local_flags[rank]}, " + f"its replica's first stage read {expected[rank]}, so the agreed decision " + f"should be {expected[rank]}, got {got}" + ) + assert all_decisions == expected, ( + f"ADP+PP decisions were {all_decisions}, expected {expected}: pipeline " + "stages must follow their replica's first stage while replicas stay " + "independent" + ) + + def run_throttle_lockstep(world_size: int, rank: int, interval: int, iterations: int) -> None: """Ranks must reach the agreement collective on the *same* iterations. @@ -517,6 +688,88 @@ def test_cp_and_tp_ranks_agree_on_rebalance_trigger(world_size: int, cp_size: in assert r is True +@pytest.mark.parametrize( + "case", + ["all_true", "all_false", "only_rank0_true", "only_rank0_false"], +) +@pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"pp:{x}") +def test_pp_ranks_agree_on_rebalance_trigger(world_size: int, case: str) -> None: + """The first PP stage's decision wins on every stage of the pipeline. + + Args: + world_size: Number of PP ranks to run on. + case: Skew scenario name, resolved by ``_flags_for``. + """ + _skip_if_not_enough_gpus(world_size) + flags = _flags_for(case, world_size) + expected = flags[0] + + with MPIPoolExecutor(max_workers=world_size) as ex: + results = ex.map( + run_single_rank, + *zip(*[(world_size, run_pp_agreement, flags, expected)] * world_size), + ) + for r in results: + assert r is True + + +@pytest.mark.parametrize("tp_size", [2], ids=lambda x: f"tp:{x}") +@pytest.mark.parametrize("world_size", [4], ids=lambda x: f"world:{x}") +def test_tp_and_pp_ranks_agree_on_rebalance_trigger(world_size: int, tp_size: int) -> None: + """TP then PP chained: the pair propagates global rank 0's decision. + + ``only_rank0_true`` is the discriminating scenario -- see + ``run_tp_pp_agreement`` for the rank-by-rank walk-through. + + Args: + world_size: Total ranks; must equal ``tp_size * pp_size``. + tp_size: Tensor-parallel width; PP width is ``world_size // tp_size``. + """ + _skip_if_not_enough_gpus(world_size) + flags = _flags_for("only_rank0_true", world_size) + expected = flags[0] + assert flags == [True, False, False, False], ( + "this case only discriminates the chained broadcast when rank 0 alone " + f"reads True, got {flags}" + ) + + with MPIPoolExecutor(max_workers=world_size) as ex: + results = ex.map( + run_single_rank, + *zip(*[(world_size, run_tp_pp_agreement, tp_size, flags, expected)] * world_size), + ) + for r in results: + assert r is True + + +@pytest.mark.parametrize("tp_size", [2], ids=lambda x: f"tp:{x}") +@pytest.mark.parametrize("world_size", [4], ids=lambda x: f"world:{x}") +def test_attention_dp_still_agrees_across_pp_stages(world_size: int, tp_size: int) -> None: + """ADP suppresses the TP hop but must never suppress the PP hop. + + Each DP replica is itself a pipeline whose stages have to drain together; + suppressing the PP hop would hang the replica rather than merely let it + drift. + + Args: + world_size: Total ranks; must equal ``tp_size * pp_size``. + tp_size: Tensor-parallel width; PP width is ``world_size // tp_size``. + """ + _skip_if_not_enough_gpus(world_size) + flags = _flags_for("only_rank0_true", world_size) + # rank = pp_rank * tp_size + tp_rank, so the replicas are {0, 2} and {1, 3} + # and each follows its own first stage: rank 0 reads True, rank 1 False. + expected = [True, False, True, False] + + with MPIPoolExecutor(max_workers=world_size) as ex: + results = ex.map( + run_single_rank, + *zip(*[(world_size, run_adp_pp_agreement, tp_size, flags, expected)] * world_size), + ) + for r in results: + assert r is True + + @pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"tp:{x}") def test_attention_dp_ranks_decide_independently(world_size: int) -> None: """Attention DP opts out of the agreement: each rank keeps its own answer. From c21efaa5d54a427baa84be54ddd17e28e8d1aad4 Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Wed, 12 Aug 2026 22:10:24 +0000 Subject: [PATCH 6/9] [None][test] Unit-test the pipeline-parallel rebalance drain wiring 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> --- .../_torch/executor/test_kv_pool_rebalance.py | 198 ++++++++++++++++++ 1 file changed, 198 insertions(+) diff --git a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py index e4ff1c4706c6..f8e14ca2ca2a 100644 --- a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py +++ b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py @@ -25,6 +25,13 @@ The accuracy of pool rebalancing itself (i.e., that suspend/adjust/resume preserves generated tokens) is covered by the integration accuracy test; here we only verify the call chain and gate logic. + +``TestPpLoopDrainWiring`` is the one exception to the "no real executor" +rule. The pipeline-parallel drain is not a method that can be called in +isolation -- it is spread across three points inside +``_executor_loop_pp`` -- so that class drives the real loop over an +``object.__new__(PyExecutor)`` instance, the same way +``test_py_executor.py`` does for the PP scheduling path. """ from unittest.mock import MagicMock @@ -832,5 +839,196 @@ def test_previous_batch_is_not_consumed(self, monkeypatch): exe.kv_cache_manager.impl.adjust.assert_called_once() +# --------------------------------------------------------------------------- # +# Pipeline-parallel loop wiring +# --------------------------------------------------------------------------- # + + +class _ReachedForward(RuntimeError): + """Raised from the mocked forward to stop the loop and prove it queued.""" + + +def _make_pp_loop_executor(monkeypatch, *, num_micro_batches=2, agreement=(True, False)): + """A PyExecutor real enough to run ``_executor_loop_pp``. + + The drain is not reachable through a single method: it is started at the + top of the iteration, enforced where ``can_queue`` is computed, and + completed in Stage 3.4. Only the real loop exercises all three, so this + builds a bare instance and fills in exactly the attributes one iteration + touches -- the same approach ``test_py_executor.py`` takes for the PP + scheduling path. + + ``_can_queue`` deliberately returns True. Every iteration therefore + *wants* to queue, so a slot left empty can only be the drain's doing. + """ + exe = object.__new__(PyExecutor) + + profiler = MagicMock() + profiler.__enter__.return_value = MagicMock() + exe._profiler = MagicMock(return_value=profiler) + exe.hang_detector = MagicMock() + exe.device_id = 0 + exe.enable_iter_perf_stats = False + exe.iter_counter = 0 + + # Rebalance state. _uses_kv_manager_v2() reads the explicit flag first. + exe._is_kv_manager_v2 = True + exe._pp_rebalance_drain_iters = None + exe.pp_multi_stream_sample = False + exe._can_pause_for_rebalance = MagicMock(return_value=True) + exe._agreed_need_adjustment = MagicMock(side_effect=list(agreement)) + # The suspend/adjust/resume core is covered by its own tests; here we only + # care that the loop reaches it, and on which iteration. + exe._rebalance_kv_pools_now = MagicMock() + + # Per-iteration no-ops. + exe._handle_disagg_cache_errors_synced = MagicMock() + exe._handle_control_request = MagicMock() + exe._pad_attention_dp_dummy_request = MagicMock() + exe._revert_gen_alloc = MagicMock() + exe._add_inflight_ids = MagicMock() + exe._handle_dynamic_draft_len = MagicMock() + exe.resource_manager = MagicMock() + exe.wait_on_pp_send_handles = MagicMock() + exe.kv_cache_transceiver = None + + # Loop termination: the loop breaks when should_stop_processing goes true, + # which is a property over these three. + exe.is_shutdown = False + exe.active_requests = [] + exe.waiting_queue = [] + + # Safety net. These tests normally end when the forward raises, but a + # drain that never completes would queue nothing, never reach the forward, + # and spin forever -- turning a clean assertion failure into a hung test. + # Shutting the loop down after a generous bound keeps such a regression a + # *failure* rather than a hang. + iterations = [] + + def _fetch_new_requests(): + iterations.append(1) + if len(iterations) > 8 * num_micro_batches: + exe.is_shutdown = True + return [] + + exe._fetch_and_activate_new_requests = MagicMock(side_effect=_fetch_new_requests) + + scheduled_batch = MagicMock() + scheduled_batch.batch_size = 1 + scheduled_batch.num_encoder_requests = 0 + scheduled_batch.num_context_requests = 1 + scheduled_batch.num_generation_requests = 0 + scheduled_batch.encoder_requests = [] + scheduled_batch.generation_requests = [] + exe._pp_schedule_and_propagate = MagicMock(return_value=(scheduled_batch, [], 0, False)) + exe._can_queue = MagicMock(return_value=(True, True)) + + # First PP rank of a two-stage pipeline, so the loop takes the + # inter-stage forward and the first-rank branch of the ring broadcast. + exe.dist = MagicMock( + rank=0, + pp_rank=0, + pp_size=2, + tp_size=1, + cp_size=1, + is_first_pp_rank=True, + is_last_pp_rank=False, + ) + + # Microbatch ring, empty and quiescent to begin with. + exe.num_micro_batches = num_micro_batches + exe.micro_batches = [None] * num_micro_batches + exe.send_handles = [None] * num_micro_batches + exe.send_expected_batch_num_handles = [None] * num_micro_batches + exe.unhandled_batch_counter = 0 + exe.previous_batch = None + exe.pp_async_broadcast_sample_state = False + exe.executed_batch_response_queue = MagicMock() + exe.executed_batch_response_queue.empty.return_value = True + + # Reaching the forward means this iteration queued a microbatch, which is + # exactly what the drain must prevent. Raising there both records the + # fact and ends the loop. + exe._forward_step_inter_pp = MagicMock(side_effect=_ReachedForward) + + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.set_device", MagicMock() + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.cudart.cudaSetDevice", MagicMock() + ) + monkeypatch.setattr("tensorrt_llm._torch.pyexecutor.py_executor.CUASSERT", MagicMock()) + return exe + + +class TestPpLoopDrainWiring: + """The three points in ``_executor_loop_pp`` that make the drain work.""" + + def test_loop_queues_normally_when_no_rebalance_is_pending(self, monkeypatch): + """Control case: without a drain the loop reaches the forward. + + Without this the drain test below would pass even if ``can_queue`` were + never honored -- e.g. if the loop simply never queued anything. + """ + exe = _make_pp_loop_executor(monkeypatch, agreement=(False,)) + + with pytest.raises(_ReachedForward): + PyExecutor._executor_loop_pp(exe) + + exe._forward_step_inter_pp.assert_called_once() + exe._rebalance_kv_pools_now.assert_not_called() + + def test_drain_suppresses_queueing_then_rebalances_and_resumes(self, monkeypatch): + """The whole wiring in one pass, over a two-slot ring. + + Iterations 1 and 2 must queue nothing even though ``_can_queue`` says + they could; the rebalance must land at the end of iteration 2, once the + ring has had a full cycle to empty; and iteration 3 must go back to + queueing, which the forward's exception reports. + """ + exe = _make_pp_loop_executor(monkeypatch, num_micro_batches=2, agreement=(True, False)) + + with pytest.raises(_ReachedForward): + PyExecutor._executor_loop_pp(exe) + + # Two drain iterations queued nothing, the third one did. + exe._forward_step_inter_pp.assert_called_once() + assert exe._revert_gen_alloc.call_count == 2 + assert exe.micro_batches == [None, None] + + # And the rebalance ran exactly once, at the end of the drain. + exe._rebalance_kv_pools_now.assert_called_once() + assert exe._pp_rebalance_drain_iters is None + + def test_drain_length_follows_the_ring_depth(self, monkeypatch): + """A deeper ring drains for longer before rebalancing. + + Pins the countdown to ``num_micro_batches`` rather than to a constant: + with four slots the loop must skip four iterations, not two. + """ + exe = _make_pp_loop_executor(monkeypatch, num_micro_batches=4, agreement=(True, False)) + + with pytest.raises(_ReachedForward): + PyExecutor._executor_loop_pp(exe) + + assert exe._revert_gen_alloc.call_count == 4 + exe._rebalance_kv_pools_now.assert_called_once() + + def test_no_drain_is_started_while_one_is_pending(self, monkeypatch): + """The decision is taken once per rebalance, not once per iteration. + + ``_agreed_need_adjustment`` runs a collective, so re-entering it + mid-drain would put this rank into a broadcast its peers are not in. + """ + exe = _make_pp_loop_executor(monkeypatch, num_micro_batches=4, agreement=(True, False)) + + with pytest.raises(_ReachedForward): + PyExecutor._executor_loop_pp(exe) + + # Once to start the drain, once on the iteration after it finished -- + # never on the three iterations in between. + assert exe._agreed_need_adjustment.call_count == 2 + + if __name__ == "__main__": pytest.main([__file__, "-v"]) From 70dd51f8b4cb31079e7b39903acfebe7344ec630 Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Thu, 13 Aug 2026 20:00:03 +0000 Subject: [PATCH 7/9] [None][fix] Agree the rebalance trigger across attention-DP ranks Under attention DP the rebalance decision was rank-local: _agreed_need_adjustment skipped the TP hop entirely, on the argument that ADP ranks own independent request streams and independent KV caches and so legitimately want different pool 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 calls _flush_pending_transfer_responses, which -- as its own docstring requires -- enters _enqueue_responses on every DP rank even with an empty response list, and that 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, attention DP, 4xH100), with only rank 0's need_adjustment true, the gather paired against the tp_allgather of batch sizes in _can_queue and the executor loop died: 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 an OR-reduction instead of being skipped, so the ranks agree on *when* to rebalance while each still computes its own ratios. any() rather than rank 0's broadcast keeps the anti-starvation property the suppression 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, with ranks 1-3 rebalancing off rank 0's need without having been injected. Measured cost on that model (2 pool groups, 17 active requests per rank): - full hook, rank that needed it 21.1 ms - full hook, rank pulled in by a peer 8.1 ms, 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, so the added collective is not measurably more expensive than the one it joins - adjust() itself is bandwidth-bound at ~290 GB/s aggregate per node, so synchronizing the ranks does not reduce total stall -- it concentrates it into one hiccup instead of N. Correctness, not performance, is the reason for this change. Two behavioural consequences reviewers should weigh. The check now performs a collective on the ADP path, so _can_pause_for_rebalance has to stay rank-uniform under ADP as well as under plain TP -- the same requirement the non-ADP broadcast already relies on, now extended. And adjust() stamps _last_adjustment_time unconditionally, so a rank pulled into a peer's rebalance has its own 120s cooldown reset by it: the group's cadence self-aligns, at the price of deferring a need that arises immediately afterwards. The unsafe call site predates this branch, but main's rebalance is documented as single-GPU-only and ADP is reachable only in multi-GPU configurations, which is what this PR exists to support. Tests updated to the new semantics, including the two that asserted the old one by name: attention_dp_ranks_decide_independently and attention_dp_without_cp_touches_no_collective -- the latter asserted precisely the property that made the crash possible. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 59 +++++- .../_torch/executor/test_kv_pool_rebalance.py | 59 ++++-- .../multi_gpu/test_kv_pool_rebalance_tp.py | 174 ++++++++++++------ 3 files changed, 221 insertions(+), 71 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 342dc8005986..4f4d057b479f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4648,8 +4648,56 @@ def _agreed_need_adjustment(self) -> bool: DP dimension (``mapping.py``: ``dp_size = tp_size if enable_attention_dp else 1``), so those ranks own independent request streams and independent KV caches and legitimately need different pool - ratios at different times; forcing rank 0's decision on them would - starve a rank that needs to rebalance when rank 0 does not. + ratios at different times. + + That independence covers the *ratios*, but not the *timing*: the ranks + must still enter the rebalance together, so under ADP the TP hop + becomes an OR-reduction instead of being skipped. The reason is that + the rebalance path is itself collective under ADP. + ``_consume_previous_batch_for_rebalance`` calls + ``_flush_pending_transfer_responses``, which -- as its own docstring + requires -- enters ``_enqueue_responses`` on every DP rank even with an + empty response list, and that performs a ``tp_gather``. A rank + rebalancing alone therefore joins a collective its peers are not in, and + pairs up against whatever they happen to be running (in a live 4-rank + Qwen3-Next run: the ``tp_allgather`` of batch sizes in ``_can_queue``, + which returned ints where responses were expected and crashed the + executor loop). + + ``any()`` rather than rank 0's reading is what keeps the starvation + argument intact: a rank that needs to rebalance still gets one instead + of waiting for rank 0 to want the same thing. + + The cost to a rank that did *not* need one is not zero. It runs the + whole hook: the stream sync, ``_consume_previous_batch_for_rebalance`` + (which in the overlap loop collapses that iteration's overlap), then + suspend-all / ``adjust()`` / resume-all. Measured end to end on + Qwen3-Next-80B (tp4 = ep4, ADP, 2 pool groups, 17 active requests) with + one rank needing the rebalance and three pulled in by it: + + rank that needed it 21.1 ms = 0.4 drain + 7.3 sync/agree + 13.5 rebalance + ranks pulled in 8.1 ms = 0.3 drain + 7.3 sync/agree + 0.5 rebalance + + So an uninvolved rank pays ~8 ms, not the 0.04 ms its no-op ``adjust()`` + suggests, and the dominant term is the pre-existing + ``torch.cuda.current_stream().synchronize()`` below rather than anything + this agreement adds -- a check that decides *not* to rebalance, which + skips that sync, costs 0.1-1.5 ms. + + Note the cooldown couples too: ``adjust()`` stamps + ``_last_adjustment_time`` unconditionally, so a rank dragged into a peer's + rebalance has its own 120s clock reset by it. That keeps the group's + cadence aligned rather than letting N ranks each rebalance on their own + schedule, at the price of deferring a need that arises immediately + afterwards by up to one cooldown. + + The extra collective costs one allgather per check, and the check is + already throttled to one iteration in + ``KV_POOL_REBALANCE_CHECK_INTERVAL``, against the ``tp_allgather`` that + ``_can_queue`` runs every iteration under ADP anyway. It does add a + lockstep requirement that ADP did not have before: every TP rank must + reach this on the same iteration, so the ``_can_pause_for_rebalance`` + gates above must stay rank-uniform under ADP as well as under plain TP. CP is orthogonal to that and must **not** be skipped. Within a single DP replica the CP ranks still split the same request along the sequence @@ -4687,8 +4735,11 @@ def _agreed_need_adjustment(self) -> bool: 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) + if self.dist.tp_size > 1: + if self.enable_attention_dp: + need = any(self.dist.tp_allgather(need)) + else: + need = self.dist.tp_broadcast(need, root=0) if self.dist.pp_size > 1: need = self.dist.pp_broadcast(need, root=0) return need diff --git a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py index f8e14ca2ca2a..43dbc9c8806c 100644 --- a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py +++ b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py @@ -619,19 +619,23 @@ def test_pp_hop_is_last_so_global_rank0_wins(self): exe.dist.pp_broadcast.assert_called_once_with(True, root=0) def test_pp_hop_survives_attention_dp(self): - """ADP suppresses the TP hop but must never suppress the PP hop. - - ADP replicates along the TP dimension, so a replica's pipeline stages - all serve that replica's own request stream and have to drain together. - ``pp_group`` for a given tp_rank holds exactly those stages, so the - broadcast propagates the replica's first-stage decision and nothing - wider. + """ADP must never suppress the PP hop, and the TP hop now OR-reduces. + + A replica's pipeline stages all serve that replica's own request stream + and have to drain together, so the PP broadcast is unchanged. The TP + hop is no longer skipped under ADP -- it becomes an allgather+OR, + because the rebalance path itself is collective across TP (see + ``test_attention_dp_or_couples_across_ranks``). The chain is therefore + OR over TP, then broadcast over PP. """ exe = _make_executor(tp_size=2, pp_size=2, enable_attention_dp=True, need_adjustment=True) + exe.dist.tp_allgather.return_value = [True, False] exe.dist.pp_broadcast.return_value = False assert PyExecutor._agreed_need_adjustment(exe) is False exe.dist.tp_broadcast.assert_not_called() + exe.dist.tp_allgather.assert_called_once_with(True) + # The OR of [True, False] is what the PP hop must be handed. exe.dist.pp_broadcast.assert_called_once_with(True, root=0) def test_attention_dp_still_broadcasts_over_cp(self): @@ -655,19 +659,48 @@ def test_attention_dp_still_broadcasts_over_cp(self): # ...but the TP hop stays suppressed, so replicas remain independent. exe.dist.tp_broadcast.assert_not_called() - def test_attention_dp_without_cp_touches_no_collective(self): + def test_attention_dp_without_cp_uses_the_tp_allgather_only(self): + """No CP hop when cp_size == 1, but the TP hop still runs under ADP. + + This case used to assert that ADP touched *no* collective at all. That + was the bug: the rebalance path it gates reaches + ``_flush_pending_transfer_responses`` -> ``_enqueue_responses``, which + under ADP runs a ``tp_gather`` on every DP rank. Deciding locally meant + entering that gather alone. + """ exe = _make_executor(tp_size=2, cp_size=1, enable_attention_dp=True, need_adjustment=True) + exe.dist.tp_allgather.return_value = [True, True] + assert PyExecutor._agreed_need_adjustment(exe) is True exe.dist.cp_broadcast.assert_not_called() exe.dist.tp_broadcast.assert_not_called() + exe.dist.tp_allgather.assert_called_once_with(True) + + def test_attention_dp_or_couples_across_ranks(self): + """ADP agrees on *when* to rebalance while keeping per-rank ratios. + + ``any()`` rather than a rank-0 broadcast is deliberate: it preserves the + anti-starvation property that motivated suppressing the hop in the first + place -- a rank that needs a rebalance gets one without waiting for rank + 0 to want the same thing -- while still bringing every rank into the + collective path together. A rank that did not need one still runs the + whole hook; only its ``adjust()`` is a no-op (0.04 ms on + Qwen3-Next-80B), not the surrounding drain and suspend/resume. + """ + exe = _make_executor(tp_size=4, enable_attention_dp=True, need_adjustment=False) + exe.dist.tp_allgather.return_value = [False, False, True, False] - def test_attention_dp_decides_independently(self): - # 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. - exe = _make_executor(tp_size=4, enable_attention_dp=True, need_adjustment=True) + # This rank read False; a peer read True, so it joins the rebalance. assert PyExecutor._agreed_need_adjustment(exe) is True exe.dist.tp_broadcast.assert_not_called() + exe.dist.tp_allgather.assert_called_once_with(False) + + def test_attention_dp_stays_put_when_no_rank_needs_it(self): + """The OR must not manufacture a rebalance nobody asked for.""" + exe = _make_executor(tp_size=4, enable_attention_dp=True, need_adjustment=False) + exe.dist.tp_allgather.return_value = [False, False, False, False] + + assert PyExecutor._agreed_need_adjustment(exe) is False def test_no_rebalance_when_agreement_says_no(self, monkeypatch): exe = _make_executor(tp_size=2, need_adjustment=True, active_requests=[_make_request(1)]) diff --git a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py index a6a8d307240c..0289b2656d7d 100644 --- a/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py +++ b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py @@ -318,12 +318,28 @@ def run_cp_tp_agreement( assert len(set(all_decisions)) == 1, f"CPxTP ranks disagreed: {all_decisions}" -def run_attention_dp_independence(world_size: int, rank: int, local_flags: list[bool]) -> None: - """Under attention DP each rank keeps its own decision. - - 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. +def run_attention_dp_or_coupling(world_size: int, rank: int, local_flags: list[bool]) -> None: + """Under attention DP the ranks OR their readings rather than decide alone. + + ADP ranks own independent request streams and independent KV caches, so the + resulting *ratios* are legitimately per-rank -- but the *timing* cannot be. + The rebalance path is collective under ADP: + ``_consume_previous_batch_for_rebalance`` calls + ``_flush_pending_transfer_responses``, which enters ``_enqueue_responses`` + on every DP rank even with an empty response list (its own docstring + requires that), and that runs a ``tp_gather``. A rank that rebalanced alone + would join a collective its peers are not in. Observed on a live 4-rank + Qwen3-Next run before this was fixed: the gather paired against the + ``tp_allgather`` of batch sizes in ``_can_queue`` and the executor loop died + with ``cannot unpack non-iterable int object``. + + ``any()`` rather than rank 0's broadcast is what preserves the + anti-starvation property the old suppression was reaching for: a rank that + needs a rebalance still gets one instead of waiting for rank 0 to want the + same thing. A rank that does not need one is not free, though -- it runs + the whole hook, and only the ``adjust()`` within it is a no-op (0.04 ms + measured, inside ~0.5 ms of suspend/resume; the preceding drain is + unmeasured). Args: world_size: Number of TP ranks. @@ -335,28 +351,37 @@ def run_attention_dp_independence(world_size: int, rank: int, local_flags: list[ got = PyExecutor._agreed_need_adjustment(exe) - assert got == local_flags[rank], ( - f"rank {rank}: attention DP must decide locally; expected {local_flags[rank]}, got {got}" + expected = any(local_flags) + assert got == expected, ( + f"rank {rank}: attention DP must OR the readings so every rank enters the " + f"collective rebalance path together; local reading was {local_flags[rank]}, " + f"the group read {local_flags}, expected {expected}, got {got}" ) def run_adp_cp_agreement( world_size: int, rank: int, cp_size: int, local_flags: list[bool], expected: list[bool] ) -> None: - """Under ADP the TP hop is suppressed but the CP hop must still run. + """Under ADP the CP hop broadcasts and the TP hop then OR-reduces. - ADP makes the TP dimension the DP dimension, so those ranks decide - independently. CP is orthogonal: inside one DP replica the CP ranks split - the *same* request along the sequence dimension and must admit it together. - Skipping the CP broadcast here would reintroduce this mechanism's own - divergence once per replica. + Two different mechanisms, chained, and this case pins down both. CP first: + inside one DP replica the CP ranks split the *same* request along the + sequence dimension and must admit it together, so they take their CP root's + reading. Then TP: under ADP the replicas must still enter the rebalance + together, because the path is collective (see + ``run_attention_dp_or_coupling``), so the post-CP values are OR-ed across + the TP dimension. + + The observable difference from the old suppressed-TP behaviour is that the + replicas no longer end on different answers. That was never safe -- the + ``tp_gather`` in ``_enqueue_responses`` spans the replicas. Args: world_size: Total ranks, equal to ``cp_size * tp_size``. rank: This process's global rank. cp_size: Context-parallel width; TP width is derived from it. local_flags: Per-rank local ``need_adjustment`` readings. - expected: Per-rank agreed decision after the CP-only broadcast. + expected: Per-rank agreed decision after the CP broadcast and TP OR. """ tp_size = world_size // cp_size dist = _cp_tp_dist(cp_size, tp_size, rank) @@ -375,6 +400,7 @@ def run_adp_cp_agreement( ) assert all_decisions == expected, ( f"ADP+CP decisions were {all_decisions}, expected {expected}: CP ranks " + "must follow their CP root and the replicas must then OR together; " "must follow their replica's root while replicas stay independent" ) @@ -466,18 +492,20 @@ def run_tp_pp_agreement( def run_adp_pp_agreement( world_size: int, rank: int, tp_size: int, local_flags: list[bool], expected: list[bool] ) -> None: - """Under ADP the TP hop is suppressed but the PP hop must still run. + """Under ADP both hops run: TP OR-reduces, then PP broadcasts. - ADP replicates along the TP dimension, so each ``tp_rank`` owns an - independent request stream -- but a replica's *pipeline stages* all serve - that one stream and must drain together or the replica's pipeline hangs. - ``pp_group`` for a given ``tp_rank`` holds exactly those stages. + A replica's *pipeline stages* all serve that replica's one request stream + and must drain together or the replica's pipeline hangs -- that is the PP + hop, and it is unchanged. What changed is the TP hop: it no longer skips + under ADP, because the rebalance path is collective across the TP dimension + (see ``run_attention_dp_or_coupling``), so the replicas must enter together + too. - On a 2x2 with only global rank 0 reading ``True``, the expected outcome is - ``[True, False, True, False]``: ranks 0 and 2 form one replica and follow - rank 0, ranks 1 and 3 form the other and follow rank 1. That single vector - pins down both halves at once -- rank 2 being ``True`` shows the PP hop ran, - and rank 1 staying ``False`` shows the TP hop did not. + On a 2x2 with only global rank 0 reading ``True`` the outcome is therefore + all ``True``: the TP OR carries rank 0's need to rank 1 within stage 0, and + the PP broadcast then carries stage 0's decision to stage 1. The previous + expectation here was ``[True, False, True, False]``, which encoded the + replica independence a live 4-rank Qwen3-Next run disproved. Args: world_size: Total ranks, equal to ``tp_size * pp_size``. @@ -498,13 +526,13 @@ def run_adp_pp_agreement( assert got == expected[rank], ( f"rank {rank} (tp_rank={dist.mapping.tp_rank}, " f"pp_rank={dist.mapping.pp_rank}): local reading was {local_flags[rank]}, " - f"its replica's first stage read {expected[rank]}, so the agreed decision " - f"should be {expected[rank]}, got {got}" + f"its first stage agreed on {expected[rank]} after the TP OR, so the " + f"agreed decision should be {expected[rank]}, got {got}" ) assert all_decisions == expected, ( - f"ADP+PP decisions were {all_decisions}, expected {expected}: pipeline " - "stages must follow their replica's first stage while replicas stay " - "independent" + f"ADP+PP decisions were {all_decisions}, expected {expected}: the TP OR " + "must couple the replicas within a stage and the PP broadcast must then " + "carry the first stage's decision to the rest" ) @@ -745,11 +773,12 @@ def test_tp_and_pp_ranks_agree_on_rebalance_trigger(world_size: int, tp_size: in @pytest.mark.parametrize("tp_size", [2], ids=lambda x: f"tp:{x}") @pytest.mark.parametrize("world_size", [4], ids=lambda x: f"world:{x}") def test_attention_dp_still_agrees_across_pp_stages(world_size: int, tp_size: int) -> None: - """ADP suppresses the TP hop but must never suppress the PP hop. + """ADP must never suppress the PP hop, and no longer suppresses the TP hop. Each DP replica is itself a pipeline whose stages have to drain together; suppressing the PP hop would hang the replica rather than merely let it - drift. + drift. The TP hop now OR-reduces instead of being skipped, so a need + anywhere in a stage reaches every rank in it. Args: world_size: Total ranks; must equal ``tp_size * pp_size``. @@ -757,9 +786,9 @@ def test_attention_dp_still_agrees_across_pp_stages(world_size: int, tp_size: in """ _skip_if_not_enough_gpus(world_size) flags = _flags_for("only_rank0_true", world_size) - # rank = pp_rank * tp_size + tp_rank, so the replicas are {0, 2} and {1, 3} - # and each follows its own first stage: rank 0 reads True, rank 1 False. - expected = [True, False, True, False] + # rank = pp_rank * tp_size + tp_rank. Stage 0 is {0, 1} and ORs rank 0's + # True across itself; the PP hop then carries that to stage 1 = {2, 3}. + expected = [True] * world_size with MPIPoolExecutor(max_workers=world_size) as ex: results = ex.map( @@ -771,8 +800,13 @@ def test_attention_dp_still_agrees_across_pp_stages(world_size: int, tp_size: in @pytest.mark.parametrize("world_size", [2, 4], ids=lambda x: f"tp:{x}") -def test_attention_dp_ranks_decide_independently(world_size: int) -> None: - """Attention DP opts out of the agreement: each rank keeps its own answer. +def test_attention_dp_ranks_or_couple(world_size: int) -> None: + """Attention DP agrees on *when*, not on *what*: the readings are OR-ed. + + ``only_rank0_false`` is the discriminating vector: rank 0 reads False while + its peers read True. A broadcast from rank 0 would give everyone False and + starve the peers that need the rebalance; the old suppression would leave + rank 0 out of a collective its peers enter. Only the OR gives all True. Args: world_size: Number of TP ranks to run on. @@ -780,42 +814,74 @@ def test_attention_dp_ranks_decide_independently(world_size: int) -> None: _skip_if_not_enough_gpus(world_size) flags = _flags_for("only_rank0_false", world_size) + # Non-vacuity: the vector must distinguish OR from broadcast-from-rank-0. + assert any(flags) != flags[0], ( + "flags must differ between rank 0 and the OR, or the test cannot tell " + "an OR-reduction from a rank-0 broadcast" + ) + with MPIPoolExecutor(max_workers=world_size) as ex: results = ex.map( run_single_rank, - *zip(*[(world_size, run_attention_dp_independence, flags)] * world_size), + *zip(*[(world_size, run_attention_dp_or_coupling, flags)] * world_size), ) for r in results: assert r is True +@pytest.mark.parametrize( + "flags", + [[False, True, False, False], [True, False, False, False]], + ids=["cp_root_wins", "or_across_replicas"], +) @pytest.mark.parametrize("cp_size", [2], ids=lambda x: f"cp:{x}") @pytest.mark.parametrize("world_size", [4], ids=lambda x: f"world:{x}") -def test_attention_dp_agrees_over_cp_but_not_tp(world_size: int, cp_size: int) -> None: - """ADP + CP: CP ranks follow their replica's root, replicas stay independent. +def test_attention_dp_agrees_over_cp_then_ors_across_tp( + world_size: int, cp_size: int, flags: list[bool] +) -> None: + """ADP + CP: CP ranks follow their root, then the replicas OR together. + + Nothing in ``Mapping`` or ``LlmArgs`` rejects ``enable_attention_dp`` with + ``cp_size > 1``, so this topology is reachable and both hops have to be + pinned down in it. - ADP suppresses only the TP hop. Nothing in ``Mapping`` or ``LlmArgs`` - rejects ``enable_attention_dp`` with ``cp_size > 1``, so this topology is - reachable, and skipping the CP hop in it would leave each replica's CP ranks - deciding independently -- the divergence class this mechanism removes. + Two flag vectors, because one cannot pin both properties: - The flags are chosen so the case fails 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). + * ``cp_root_wins`` puts the only True on a *non-root* CP rank. The CP + broadcast discards it, so the OR that follows finds nothing and every rank + ends False. A plain global OR of the local readings would give True + everywhere, so this is what proves the CP hop runs *before* the TP hop and + is not merely folded into it. + * ``or_across_replicas`` puts the only True on a CP root in one replica. It + survives the broadcast and the OR then carries it to the other replica, so + every rank ends True -- the coupling the collective rebalance path needs. Args: world_size: Total ranks; must equal ``cp_size * tp_size``. cp_size: Context-parallel width; TP width is ``world_size // cp_size``. + flags: Per-rank local ``need_adjustment`` readings for this case. """ _skip_if_not_enough_gpus(world_size) - # cp_groups are consecutive ranks within a TP slice: [[0, 1], [2, 3]]. - flags = [True, False, False, True] - expected = [flags[(r // cp_size) * cp_size] for r in range(world_size)] - - # Non-vacuity, both directions. - assert expected != flags, "flags must force the CP broadcast to change some rank" - assert len(set(expected)) > 1, "replicas must end up disagreeing, or the TP hop is untested" + # rank = tp_rank * cp_size + cp_rank, so cp_groups are consecutive ranks + # ([[0, 1], [2, 3]]) and tp_groups stride by cp_size ([[0, 2], [1, 3]]). + after_cp = [flags[(r // cp_size) * cp_size] for r in range(world_size)] + expected = [ + any(after_cp[x] for x in range(world_size) if x % cp_size == r % cp_size) + for r in range(world_size) + ] + + # Non-vacuity: the chain must not collapse to a plain global OR of `flags`, + # which is exactly the mistake this pair of vectors exists to catch. + assert expected != flags, "flags must force some hop to change some rank" + if not any(after_cp): + assert any(flags), ( + "the cp_root_wins vector must start with a True somewhere, or it " + "proves nothing about ordering" + ) + assert expected == [False] * world_size, ( + "a True on a non-root CP rank must be discarded by the CP broadcast" + ) with MPIPoolExecutor(max_workers=world_size) as ex: results = ex.map( From 990a8776a779a7a79ecc559f43625c99e6e9d0ec Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Fri, 14 Aug 2026 00:03:56 +0000 Subject: [PATCH 8/9] [None][test] Run the rebalance agreement test on a 4-GPU CI stage tests/unittest/_torch/multi_gpu is collected by exactly one test-db entry, the 2-GPU H100 pre-merge stage. Every 4-rank case of the KV-pool rebalance agreement test therefore skips in CI with "need 4 GPUs, have 2" -- 17 of the 30 cases in the last full run (pipeline 53654), including all of the composition coverage added by the recent work: tp4, pp4, CP x TP, TP x PP, and the attention-DP paths. The newest logic on this branch is exactly what CI could not reach. Name the file in the 4-GPU pytorch/others block so those cases execute. The 2-GPU stage keeps running the file as before via the directory entry; the 2-rank cases are cheap and the overlap is not worth a fragile -k filter. Verified on 4x GB300: all 31 cases pass in 5m22s. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_dgx_h100.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 395038927cf0..c8a0432a658d 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -98,6 +98,11 @@ l0_dgx_h100: tests: # ------------- PyTorch tests --------------- - unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu4" + # unittest/_torch/multi_gpu is otherwise only collected by the 2-GPU stage, so + # the 4-rank cases of the KV-pool rebalance agreement test (tp4, pp4, CP x TP, + # TP x PP, attention-DP over CP) always skip with "need 4 GPUs, have 2". Name + # the file here so those compositions actually run. + - unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py # ------------- Model specific tests --------------- - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=False-attn_backend=TRTLLM-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_4gpus[tp4-fp8kv=True-attn_backend=TRTLLM-torch_compile=True] From f95c99cb552df1d4c55ab499485ee89383ee1e9d Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Fri, 14 Aug 2026 04:07:37 +0000 Subject: [PATCH 9/9] [None][fix] Exclude dummy KV caches from the V2 auto-tuner's statistics Warmup and CUDA-graph padding requests reserve capacity at the model's full declared context rather than at a realistic sequence length, and the tuner averages the *square* of capacity. A handful of them therefore dominates the statistic outright and the pools get sized for sequences that never arrive. Measured on DeepSeek-V4-Flash (tp4=ep4, attention DP, 4x GB300), per rank: 76 of 4,934 sampled closes -- 1.5% -- supplied 100.00% of the sum of squares, every one at capacity 1,048,576 == max_position_embeddings. RMS with them 130,139; without them 313. avg_capacity read 131,703 against real sequences averaging ~530 tokens, and the resulting target ratio was computed for million-token sequences. Those caches are already marked stats-excluded at creation -- the dummy path sets is_dummy and _create_kv_cache calls mark_stats_excluded -- so close() only has to honour the flag that is already there. Deliberately not reusing _should_record_stats(), which also ANDs in the user-facing _stats_enabled toggle; disabling stats reporting must not blind the auto-tuner to real traffic. After the fix, on the same model and workload, avg_capacity reads 531-534 against ~530-token sequences and the tuner triggers organically. Behaviour change worth noting for reviewers: dummies were ~90% of the sample count, so the 2000-sample maturity gate used to open on a few hundred real sequences. It now needs 2000 real completed sequences per rank -- 1,024 prompts reached 257 samples and never triggered, 12,288 reached ~3,060 and triggered once. Low-traffic deployments will engage the feature later than they do today. Test asserts on need_adjustment rather than the counters, since neither has a getter on the C++ backend, 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 the guard fails the test. Python backend: 130 tests OK (skipped=13). The C++ half of this change is NOT yet compiled or exercised; it needs a build before the default backend is fixed. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> --- .../kv_cache_manager_v2/kvCache.cpp | 6 +- .../kv_cache_manager_v2/_core/_kv_cache.py | 7 ++- .../test_kv_cache_manager_v2.py | 57 +++++++++++++++++++ 3 files changed, 68 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index c6a1c51c5388..63b2b6b1cbb7 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -560,7 +560,11 @@ void KvCache::close() stopCommitting(); TLLM_CHECK_DEBUG(_checkSanity()); - if (mCapacity > 0) + // Dummy/warmup caches are reserved at the model's full declared context, not at a realistic + // sequence length, and mAvgSqrCapacity is an RMS -- so a handful of them dominates the + // statistic outright and the tuner sizes pools for sequences that never arrive. They are + // already tracked as stats-excluded at creation; honour that here too. + if (mCapacity > 0 && !mManager->isStatsExcluded(id)) { mAvgCapacity.update(static_cast(mCapacity)); mManager->updateAvgSqrCapacity(mAvgCapacity.value() * mAvgCapacity.value()); diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 4061af81297f..8c71a29468f2 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -583,7 +583,12 @@ def close(self) -> None: self.stop_committing() assert NDEBUG or self._check_sanity() manager = self.manager - if self.capacity > 0: + # Dummy/warmup caches are reserved at the model's full declared context, + # not at a realistic sequence length, and _avg_sqr_capacity is an RMS -- + # so a handful of them dominates the statistic outright and the tuner + # sizes pools for sequences that never arrive. They are already tracked + # as stats-excluded at creation; honour that here too. + if self.capacity > 0 and not manager.is_stats_excluded(self.id): self._avg_capacity.update(self.capacity) manager._avg_sqr_capacity.update(self._avg_capacity.value**2) manager._avg_sqr_history_length.update(self._avg_history_length.value**2) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 4d0900017593..6d42bbf0d9b2 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -4447,6 +4447,63 @@ def test_adjust_resizes_pool_groups(self) -> None: self.assertGreater(slots_after[0], slots_before[0], f"{slots_before} -> {slots_after}") self.assertLess(slots_after[1], slots_before[1], f"{slots_before} -> {slots_after}") + def _close_one_cache(self, capacity: int, *, dummy: bool, cache_id: int) -> None: + """Create, size and close one KV cache, optionally marked as a dummy. + + Only the close matters here: it is where the auto-tuner samples + capacity and history length. + """ + prompt = [self.next_token() for _ in range(self._TOKENS_PER_BLOCK)] + kv_cache = self.manager.create_kv_cache(ReuseScope(), prompt, id=cache_id) + if dummy: + self.manager.mark_stats_excluded(cache_id) + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(capacity, capacity - 1)) + s.take_finish_event().synchronize() + kv_cache.close() + + def test_dummy_kv_caches_do_not_feed_the_tuner(self) -> None: + """Stats-excluded caches must not move the target pool ratio. + + Warmup and CUDA-graph padding requests reserve capacity at the model's + full declared context rather than at a realistic sequence length, and + the tuner averages the *square* of capacity -- so a handful of them + dominates the statistic outright and the pools get sized for sequences + that never arrive. Those caches are already marked stats-excluded at + creation; ``_KVCache.close`` has to honour that. + """ + self.prepare_two_pool_groups() + + # Open the sample-count and cooldown gates but leave the target ratio + # alone -- unlike force_rebalance_precondition, which also skews it. + # need_adjustment is then driven purely by whether a close moved the + # target away from the current ratio. + _introspection.set_num_sampled_kv_caches(self.manager, 2001) + _introspection.set_last_adjustment_time(self.manager, 0.0) + self.assertFalse( + self.manager.need_adjustment, + "target should still equal the current ratio before any close", + ) + + big = 64 * self._TOKENS_PER_BLOCK + self._close_one_cache(big, dummy=True, cache_id=1) + self.assertFalse( + self.manager.need_adjustment, + "a stats-excluded cache moved the target ratio; dummy requests are " + "feeding the auto-tuner", + ) + + # Control: the same close *without* the exclusion must move the target, + # otherwise the assertion above passes for the wrong reason. + self._close_one_cache(big, dummy=False, cache_id=2) + self.assertTrue( + self.manager.need_adjustment, + "a real close of the same shape did not move the target either, so " + "the check above is vacuous", + ) + def test_kv_survives_adjust(self) -> None: """Committed blocks must still verify after pages migrate between slots.""" self.prepare_two_pool_groups()