diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ee6e7e6dfecc..342dc8005986 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,13 @@ 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 + # 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 @@ -2662,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 @@ -2736,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: @@ -2967,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 @@ -4509,16 +4545,34 @@ 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. - - MVP scope: single-GPU aggregated, 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). + """Gate KV pool rebalance to the cases the hook supports. + + 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 + 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 - if self.dist.pp_size > 1: - return False if self.kv_cache_transceiver is not None: return False if self.is_warmup: @@ -4529,8 +4583,116 @@ 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. + # + # 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 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). + + 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 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. + + 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: """Drain ``previous_batch`` so its _KVCache instances are quiescent. @@ -4568,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 mgr.impl.need_adjustment: + 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: @@ -4592,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/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py b/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py index 6542c7cdfae0..fd8c592b2308 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,18 @@ 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 kv_cache_manager = executor.kv_cache_manager _introspection.force_rebalance_precondition(kv_cache_manager.impl, skew=skew) @@ -95,14 +103,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..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 @@ -43,6 +50,9 @@ 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, is_shutdown: bool = False, @@ -53,23 +63,59 @@ def _make_executor( previous_batch=None, 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. ``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, cp_size=cp_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. The throttle keys on iter_counter, the + # loop-wide iteration index, so tests drive that directly. + 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 @@ -140,9 +186,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()) @@ -377,5 +429,606 @@ 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 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. + """ + + @staticmethod + def _fired_on(exe: MagicMock, iterations: int) -> list[int]: + """Run ``iterations`` executor iterations and report which ones opened + the gate. + + Args: + exe: The PyExecutor stand-in to drive. + iterations: How many iterations to step through. + + 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 self._fired_on(exe, 5) == [0, 1, 2, 3, 4] + + def test_config_gate_suppresses_every_iteration(self) -> None: + exe = _make_executor(enable_kv_pool_rebalance=False, rebalance_check_interval=2) + 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] + + +# --------------------------------------------------------------------------- # +# 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_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, 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. + + 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() + + 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() + + +# --------------------------------------------------------------------------- # +# 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() + + +# --------------------------------------------------------------------------- # +# 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"]) 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..a6a8d307240c --- /dev/null +++ b/tests/unittest/_torch/multi_gpu/test_kv_pool_rebalance_tp.py @@ -0,0 +1,849 @@ +# 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.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: + """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, + ) + ) + + +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 +# --------------------------------------------------------------------------- + + +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]) + + 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. + assert len(set(all_decisions)) == 1, f"ranks disagreed after the collective: {all_decisions}" + + +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" + 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_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) + + 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_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_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. + + This is the deadlock guard. ``_can_pause_for_rebalance`` gates the + 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 + 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. + + 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. + 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 % 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: 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) -> 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": + 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: 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] + + 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( + "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: 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] + + 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("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( + "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. + + 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) + + 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("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: + """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) + + 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