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/6] [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/6] [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/6] [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/6] [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/6] [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/6] [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"])