From 7e66d75c649fcd81dffb670491686c8d91e0c2cb Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:15:01 +0800 Subject: [PATCH 1/3] [None][refactor] Add DisaggTransferCoordinator skeleton and loop transcript tests Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- .../disaggregation/executor/coordinator.py | 129 ++++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 122 +++---- .../executor/test_disagg_coordinator.py | 77 +++++ .../executor/test_disagg_loop_transcript.py | 305 ++++++++++++++++++ .../executor/test_send_kv_async_split.py | 2 +- 5 files changed, 580 insertions(+), 55 deletions(-) create mode 100644 tensorrt_llm/_torch/disaggregation/executor/coordinator.py create mode 100644 tests/unittest/_torch/executor/test_disagg_coordinator.py create mode 100644 tests/unittest/_torch/executor/test_disagg_loop_transcript.py diff --git a/tensorrt_llm/_torch/disaggregation/executor/coordinator.py b/tensorrt_llm/_torch/disaggregation/executor/coordinator.py new file mode 100644 index 000000000000..a2a2b9918901 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/executor/coordinator.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Executor-facing entry points for disaggregated KV transfer. + +The executor loops call the disagg state machine only through a +``DisaggTransferCoordinator``. This module must not import ``PyExecutor`` or +hold a reference to it: everything it needs is injected as callables. +""" + +from dataclasses import dataclass, fields +from typing import TYPE_CHECKING, Callable, List, Tuple + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + +if TYPE_CHECKING: + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ScheduledRequests + + +@dataclass(frozen=True) +class DisaggLoopDelegates: + """Executor callables the coordinator forwards to. + + Transitional: each field is removed once the corresponding logic moves + into the coordinator. + """ + + handle_errors_synced: Callable[[], None] + prepare_context_schedulable: Callable[[List[LlmRequest]], None] + poll_gen_transfers: Callable[[], None] + check_transfer_timeouts: Callable[[], None] + admit: Callable[[List[LlmRequest]], Tuple[List[LlmRequest], bool]] + revert_deferred_gen_init: Callable[[List[LlmRequest], List[LlmRequest]], None] + receive_gen_init: Callable[[List[LlmRequest]], None] + poll_progress_when_idle: Callable[[], None] + prepare_transmission_completed: Callable[["ScheduledRequests"], None] + send_completed_context: Callable[[List[LlmRequest]], None] + reap_context_sends: Callable[[int], None] + pace_idle: Callable[[], None] + + +class DisaggTransferCoordinator: + """Disagg transfer entry points used by every executor loop variant. + + Several methods perform rank-consensus collectives inside the transceiver + or over ``dist``; every rank must call them the same number of times per + iteration. The loops therefore call them unconditionally and rely on the + coordinator (or ``NoopDisaggCoordinator``) to be a no-op when + disaggregation is off. + """ + + def __init__(self, delegates: DisaggLoopDelegates) -> None: + self._d = delegates + + # -- loop head ----------------------------------------------------------- + + def handle_errors_synced(self) -> None: + """Fail requests whose transfer errored; rank-synchronized.""" + self._d.handle_errors_synced() + + def prepare_context_schedulable(self, new_requests: List[LlmRequest]) -> None: + """Let the transceiver gate generation-first context requests.""" + self._d.prepare_context_schedulable(new_requests) + + def poll_gen_transfers(self) -> None: + """Poll receive-side transfers and their timeouts; rank-synchronized.""" + self._d.poll_gen_transfers() + + def check_transfer_timeouts(self) -> None: + """Flag transfers that exceeded ``kv_transfer_timeout_ms``.""" + self._d.check_transfer_timeouts() + + # -- scheduling ---------------------------------------------------------- + + def admit(self, fitting_gen_init: List[LlmRequest]) -> Tuple[List[LlmRequest], bool]: + """Select the gen-init requests that may start receiving this iteration. + + Returns ``(admitted, blocked_by_active_transfers)``. + """ + return self._d.admit(fitting_gen_init) + + def revert_deferred_gen_init( + self, candidates: List[LlmRequest], admitted: List[LlmRequest] + ) -> None: + """Release KV allocated for candidates that were not admitted.""" + self._d.revert_deferred_gen_init(candidates, admitted) + + def receive_gen_init(self, admitted: List[LlmRequest]) -> None: + """Prepare resources and start the KV receive for admitted requests.""" + self._d.receive_gen_init(admitted) + + def poll_progress_when_idle(self) -> None: + """Reap completed context sends; rank-symmetric.""" + self._d.poll_progress_when_idle() + + # -- batch execution ----------------------------------------------------- + + def prepare_transmission_completed(self, scheduled_batch: "ScheduledRequests") -> None: + """Turn gen requests whose receive completed into running requests.""" + self._d.prepare_transmission_completed(scheduled_batch) + + def send_completed_context(self, requests: List[LlmRequest]) -> None: + """Start async KV sends for finished context-only requests.""" + self._d.send_completed_context(requests) + + def reap_context_sends(self, at_least: int = 0) -> None: + """Poll send-side transfers and release settled requests.""" + self._d.reap_context_sends(at_least) + + # -- loop tail ----------------------------------------------------------- + + def pace_idle(self) -> None: + """Sleep briefly when only a transfer completing can make progress.""" + self._d.pace_idle() + + +class NoopDisaggCoordinator(DisaggTransferCoordinator): + """Coordinator used when the executor has no KV cache transceiver.""" + + def __init__(self) -> None: + super().__init__( + DisaggLoopDelegates(**{f.name: _noop for f in fields(DisaggLoopDelegates)}) + ) + + def admit(self, fitting_gen_init: List[LlmRequest]) -> Tuple[List[LlmRequest], bool]: + return fitting_gen_init, False + + +def _noop(*_args, **_kwargs) -> None: + return None diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 7375046296ac..97ec7483ca6e 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -53,6 +53,9 @@ from ..disaggregation.executor.admission import \ DisaggTransferAdmissionController +from ..disaggregation.executor.coordinator import (DisaggLoopDelegates, + DisaggTransferCoordinator, + NoopDisaggCoordinator) from ..disaggregation.executor.pp_termination import DisaggPPTerminationHandler from ..disaggregation.executor.transfer_manager import AsyncTransferManager from ..distributed import Distributed @@ -2498,10 +2501,8 @@ def _pp_schedule_and_propagate(self, microbatch_id: int): and is_dp_broadcast): scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs = self._schedule( ) - if self.kv_cache_transceiver: - fitting_disagg_gen_init_requests, wait_for_disagg_gen_transfer_progress = ( - self._apply_disagg_transfer_admission( - fitting_disagg_gen_init_requests)) + fitting_disagg_gen_init_requests, wait_for_disagg_gen_transfer_progress = ( + self.disagg.admit(fitting_disagg_gen_init_requests)) serializable_schedule = SerializableSchedulerOutput.from_scheduler_result( scheduled_batch, fitting_disagg_gen_init_requests, num_fitting_reqs, wait_for_disagg_gen_transfer_progress) @@ -2576,8 +2577,8 @@ def _pp_retry_until_can_schedule(self, scheduled_batch): ) # Let cache transceiver finish at least one cache transmission and release requests' KV cache resources - self._check_disagg_ctx_cache_transfer_status(1) - self._check_kv_transfer_timeout() + self.disagg.reap_context_sends(1) + self.disagg.check_transfer_timeouts() else: raise RuntimeError( f"Reach maximum PP retry count ({self.pp_scheduler_max_retry_count}) but still cannot run first PP's schedule result. Please consider increasing the KV cache size by setting `free_gpu_memory_fraction` to a larger value. Or you can set `TLLM_PP_SCHEDULER_MAX_RETRY_COUNT` to a larger value to allow more retries." @@ -2608,7 +2609,7 @@ def _executor_loop_pp(self): and self._agreed_need_adjustment()): self._start_pp_rebalance_drain() - self._handle_disagg_cache_errors_synced() + self.disagg.handle_errors_synced() # Fetch new requests from request queue new_requests = self._fetch_and_activate_new_requests() @@ -2618,9 +2619,8 @@ def _executor_loop_pp(self): self._handle_control_request() - if self.kv_cache_transceiver: - self._check_disagg_ctx_schedulable_status(new_requests) - self._check_disagg_gen_transfer_status() + self.disagg.prepare_context_schedulable(new_requests) + self.disagg.poll_gen_transfers() if self.enable_iter_perf_stats: iter_stats = self._get_init_iter_stats( @@ -2642,13 +2642,12 @@ def _executor_loop_pp(self): self.active_requests) local_scheduler_output = self.scheduler.schedule_request( self.active_requests, self.inflight_req_ids) - if self.kv_cache_transceiver: - local_disagg_candidates = getattr( - local_scheduler_output, - "fitting_disagg_gen_init_requests", []) - self._revert_deferred_disagg_gen_init_alloc( - local_disagg_candidates, - fitting_disagg_gen_init_requests) + local_disagg_candidates = getattr( + local_scheduler_output, + "fitting_disagg_gen_init_requests", []) + self.disagg.revert_deferred_gen_init( + local_disagg_candidates, + fitting_disagg_gen_init_requests) if (self._mm_encoder_item_scheduling_enabled and scheduled_batch.scheduled_mm_encoder_items): @@ -2659,11 +2658,8 @@ def _executor_loop_pp(self): self._pause_recompute_paused_requests(scheduled_batch) # For requests that are fitting disagg gen init, also prepare resources for KV cache manager - if self.kv_cache_transceiver: - self._prepare_disagg_gen_init( - fitting_disagg_gen_init_requests) - - self._check_disagg_transfer_progress_when_idle() + self.disagg.receive_gen_init(fitting_disagg_gen_init_requests) + self.disagg.poll_progress_when_idle() self.num_scheduled_requests = scheduled_batch.batch_size @@ -2704,10 +2700,7 @@ def _executor_loop_pp(self): self._add_inflight_ids(scheduled_batch) - if self.kv_cache_transceiver: - # For generation requests which have completed KV cache transfer - self._prepare_disagg_gen_transmission_complete( - scheduled_batch) + self.disagg.prepare_transmission_completed(scheduled_batch) self._handle_dynamic_draft_len(scheduled_batch) @@ -2924,7 +2917,7 @@ def handle_executed_batches(executed_batch_num: int): self._maybe_finish_pp_rebalance() if not can_queue and self._pp_ring_is_drained(): - self._pace_idle_disagg_loop() + self.disagg.pace_idle() # Stage 4: March forward in microbatch slots microbatch_id = (microbatch_id + 1) % self.num_micro_batches @@ -3263,8 +3256,7 @@ def _handle_executed_batch(self, finished_requests = self._handle_responses() # Complete ctx send sessions AFTER responses are created so # _handle_responses sees the request before it is terminated. - if self.kv_cache_transceiver: - self._check_disagg_ctx_cache_transfer_status(0) + self.disagg.reap_context_sends(0) sample_state_scheduled_requests = executed_batch.scheduled_requests attn_metadata = getattr(self.model_engine, 'attn_metadata', None) @@ -3506,6 +3498,38 @@ def _commit_kv_cache_stats(self, self.kv_cache_manager.commit_scheduled_kv_cache_stats( scheduled_batch) + @property + def disagg(self) -> DisaggTransferCoordinator: + """Disagg transfer entry points; built on first use.""" + coordinator = self.__dict__.get("_disagg_coordinator") + if coordinator is None: + coordinator = self._build_disagg_coordinator() + self._disagg_coordinator = coordinator + return coordinator + + def _build_disagg_coordinator(self) -> DisaggTransferCoordinator: + if getattr(self, "kv_cache_transceiver", None) is None: + return NoopDisaggCoordinator() + return DisaggTransferCoordinator( + DisaggLoopDelegates( + handle_errors_synced=self._handle_disagg_cache_errors_synced, + prepare_context_schedulable=self. + _check_disagg_ctx_schedulable_status, + poll_gen_transfers=self._check_disagg_gen_transfer_status, + check_transfer_timeouts=self._check_kv_transfer_timeout, + admit=self._apply_disagg_transfer_admission, + revert_deferred_gen_init=self. + _revert_deferred_disagg_gen_init_alloc, + receive_gen_init=self._prepare_disagg_gen_init, + poll_progress_when_idle=self. + _check_disagg_transfer_progress_when_idle, + prepare_transmission_completed=self. + _prepare_disagg_gen_transmission_complete, + send_completed_context=self._send_disagg_ctx_kv_async, + reap_context_sends=self._check_disagg_ctx_cache_transfer_status, + pace_idle=self._pace_idle_disagg_loop, + )) + def _get_disagg_transfer_admission_controller( self) -> DisaggTransferAdmissionController: controller = getattr(self, "_disagg_transfer_admission_controller", @@ -3739,10 +3763,9 @@ def _prepare_and_schedule_batch(self): self._handle_control_request() - if self.kv_cache_transceiver: - self._check_disagg_ctx_schedulable_status(new_requests) - self._check_disagg_gen_transfer_status() - self._check_kv_transfer_timeout() + self.disagg.prepare_context_schedulable(new_requests) + self.disagg.poll_gen_transfers() + self.disagg.check_transfer_timeouts() iter_stats = None if self.enable_iter_perf_stats: @@ -3837,13 +3860,12 @@ def _prepare_and_schedule_batch(self): if self.kv_cache_transceiver: wait_for_disagg_gen_transfer_progress = False admitted_disagg_gen_init_requests, wait_for_disagg_gen_transfer_progress = ( - self._apply_disagg_transfer_admission( - scheduler_fitting_disagg_gen_init_requests)) + self.disagg.admit(scheduler_fitting_disagg_gen_init_requests)) # Prepare KV cache manager resources only for requests admitted # into the transfer window this iteration. - self._prepare_disagg_gen_init(admitted_disagg_gen_init_requests) + self.disagg.receive_gen_init(admitted_disagg_gen_init_requests) - self._check_disagg_transfer_progress_when_idle() + self.disagg.poll_progress_when_idle() # In gen-only benchmark mode, all requests must fit in KV cache # simultaneously. If some requests are stuck in INIT state and the @@ -4194,7 +4216,7 @@ def _executor_loop(self): if self._is_kv_manager_v2 and self._can_pause_for_rebalance(): self._maybe_rebalance_kv_pools() - self._handle_disagg_cache_errors_synced() + self.disagg.handle_errors_synced() scheduled_batch, iter_stats = self._prepare_and_schedule_batch() @@ -4248,11 +4270,8 @@ def _executor_loop(self): can_queue, _ = self._can_queue(scheduled_batch) if can_queue: + self.disagg.prepare_transmission_completed(scheduled_batch) if self.kv_cache_transceiver: - # For generation requests which have completed KV cache transfer - self._prepare_disagg_gen_transmission_complete( - scheduled_batch) - # Return the first token to the client self._handle_first_token_response(scheduled_batch) @@ -4386,8 +4405,7 @@ def _executor_loop(self): finished_requests = self._handle_responses() # Complete ctx send sessions AFTER responses are created so # _handle_responses sees the request before it is terminated. - if self.kv_cache_transceiver: - self._check_disagg_ctx_cache_transfer_status(0) + self.disagg.reap_context_sends(0) # Compute GPU times after _handle_responses creates metric entries # (safe in non-overlap mode: no next iteration to overwrite events) self.perf_manager.compute_batch_gpu_times( @@ -4440,7 +4458,7 @@ def _executor_loop(self): self._flush_iter_stats_synced() if not can_queue: - self._pace_idle_disagg_loop() + self.disagg.pace_idle() self.iter_counter += 1 @@ -5025,7 +5043,7 @@ def _executor_loop_overlap(self): if self._is_kv_manager_v2 and self._can_pause_for_rebalance(): self._maybe_rebalance_kv_pools() - self._handle_disagg_cache_errors_synced() + self.disagg.handle_errors_synced() # Need to wait for the copy of previous iteration before # modifying any host memory copied to GPU. Scheduler V2 @@ -5068,10 +5086,7 @@ def _executor_loop_overlap(self): scheduled_batch) if can_queue: - if self.kv_cache_transceiver: - # For generation requests which have completed KV cache transfer - self._prepare_disagg_gen_transmission_complete( - scheduled_batch) + self.disagg.prepare_transmission_completed(scheduled_batch) has_draft_batch = self.drafter is not None and self.previous_batch is not None and self.use_spec_decode and self.drafter.should_forward_draft_model( scheduled_batch) @@ -5318,7 +5333,7 @@ def _executor_loop_overlap(self): self._kv_connector_terminate_requests() if not can_queue: - self._pace_idle_disagg_loop() + self.disagg.pace_idle() self.iter_counter += 1 @@ -7509,10 +7524,9 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): def _send_kv_async(self, scheduled_requests: List[LlmRequest]): # Order matters: reaping before the connector registers its transfer # can release a request the connector still needs. - self._send_disagg_ctx_kv_async(scheduled_requests) + self.disagg.send_completed_context(scheduled_requests) self._save_kv_to_connector_async(scheduled_requests) - if self.kv_cache_transceiver: - self._check_disagg_ctx_cache_transfer_status(0) + self.disagg.reap_context_sends(0) def _send_disagg_ctx_kv_async(self, scheduled_requests: List[LlmRequest]) -> None: diff --git a/tests/unittest/_torch/executor/test_disagg_coordinator.py b/tests/unittest/_torch/executor/test_disagg_coordinator.py new file mode 100644 index 000000000000..d04a874d9409 --- /dev/null +++ b/tests/unittest/_torch/executor/test_disagg_coordinator.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the DisaggTransferCoordinator skeleton.""" + +import ast +import inspect +from dataclasses import fields +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.disaggregation.executor import coordinator as coordinator_module +from tensorrt_llm._torch.disaggregation.executor.coordinator import ( + DisaggLoopDelegates, + DisaggTransferCoordinator, + NoopDisaggCoordinator, +) + +pytestmark = pytest.mark.cpu_only + + +def _public_methods(cls) -> set: + return { + name + for name, member in inspect.getmembers(cls, predicate=inspect.isfunction) + if not name.startswith("_") + } + + +def test_coordinator_module_does_not_depend_on_py_executor() -> None: + """The coordinator must be constructible and testable without PyExecutor.""" + tree = ast.parse(inspect.getsource(coordinator_module)) + imported = set() + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + imported.add(node.module or "") + imported.update(alias.name for alias in node.names) + elif isinstance(node, ast.Import): + imported.update(alias.name for alias in node.names) + assert not any("py_executor" in name or name == "PyExecutor" for name in imported) + + +def test_every_coordinator_method_has_a_delegate() -> None: + """Every entry point must be backed by a delegate: a silently no-op method on + the real coordinator would drop a rank-consensus collective and hang peers.""" + delegate_names = {f.name for f in fields(DisaggLoopDelegates)} + assert _public_methods(DisaggTransferCoordinator) == delegate_names + + +def test_real_coordinator_forwards_arguments_and_results() -> None: + delegates = DisaggLoopDelegates(**{f.name: Mock() for f in fields(DisaggLoopDelegates)}) + delegates.admit.return_value = (["admitted"], True) + coordinator = DisaggTransferCoordinator(delegates) + + assert coordinator.admit(["fitting"]) == (["admitted"], True) + coordinator.reap_context_sends(1) + coordinator.revert_deferred_gen_init(["a"], ["b"]) + + delegates.admit.assert_called_once_with(["fitting"]) + delegates.reap_context_sends.assert_called_once_with(1) + delegates.revert_deferred_gen_init.assert_called_once_with(["a"], ["b"]) + + +def test_noop_coordinator_admits_everything_unchanged() -> None: + """Without a transceiver, scheduler-fitting gen-init requests must pass + through unfiltered and never report a transfer-budget block.""" + fitting = [object(), object()] + assert NoopDisaggCoordinator().admit(fitting) == (fitting, False) + + +def test_noop_coordinator_accepts_every_loop_call() -> None: + """Loops call the coordinator unconditionally, so the no-op variant must + accept every call the real one does.""" + noop = NoopDisaggCoordinator() + for name in _public_methods(DisaggTransferCoordinator) - {"admit"}: + params = inspect.signature(getattr(noop, name)).parameters + getattr(noop, name)(*[Mock() for _ in params]) diff --git a/tests/unittest/_torch/executor/test_disagg_loop_transcript.py b/tests/unittest/_torch/executor/test_disagg_loop_transcript.py new file mode 100644 index 000000000000..adcac7621331 --- /dev/null +++ b/tests/unittest/_torch/executor/test_disagg_loop_transcript.py @@ -0,0 +1,305 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-iteration transcript of communication boundaries in each executor loop. + +Each test drives one idle iteration (nothing schedulable) followed by shutdown +through a real loop body on a single rank and pins the sequence of coordinator +calls plus the executor-owned ADP-synchronized flushes. This protects the call +points against being dropped or reordered while disagg logic moves into the +coordinator; a changed sequence is a review signal, not necessarily a bug. + +Rank symmetry is checked only in the narrow form that fits one process: the +first and a non-first PP rank must issue the same collective-sensitive calls in +the same order during an idle iteration. Real multi-rank blocking semantics are +covered elsewhere (Gloo tests; FakeDist arrives with CS-2). Regular disagg PP +termination advances from executed-batch handling; a recompute-pause fallback +can call the same termination handler from an idle iteration. Neither path is +covered here (nothing is pending in these iterations); both belong to the +executed-batch/lifecycle transcripts of PR-5. +""" + +import inspect +import queue +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock + +import pytest + +from tensorrt_llm._torch.disaggregation.executor.coordinator import ( + DisaggTransferCoordinator, + NoopDisaggCoordinator, +) +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import ( + ScheduledRequests, + SerializableSchedulerOutput, +) + +pytestmark = pytest.mark.cpu_only + +# Coordinator entry points whose delegates run a rank-consensus collective. +# Derived from the delegate targets in PyExecutor._build_disagg_coordinator; +# update alongside them. +_COLLECTIVE_COORDINATOR_CALLS = { + "handle_errors_synced", # dist.allreduce / tp_allgather under ADP + "prepare_context_schedulable", # transceiver.prepare_context_requests consensus + "poll_gen_transfers", # gen transfer status consensus + "poll_progress_when_idle", # ctx transfer status consensus + "receive_gen_init", # async receive polls gen status consensus + "reap_context_sends", # ctx transfer status consensus +} +# Executor-owned per-iteration collectives that must stay in lockstep with the +# coordinator calls under ADP. +_COLLECTIVE_EXECUTOR_EVENTS = { + "flush_pending_transfer_responses", # tp_gather in _enqueue_responses + "handle_kv_transfer_timeouts_synced", # tp_allgather +} +_COLLECTIVE_SENSITIVE = _COLLECTIVE_COORDINATOR_CALLS | _COLLECTIVE_EXECUTOR_EVENTS + + +def _entry_points() -> list: + return sorted( + name + for name, member in inspect.getmembers( + DisaggTransferCoordinator, predicate=inspect.isfunction + ) + if not name.startswith("_") + ) + + +def _recording_coordinator(calls: list) -> DisaggTransferCoordinator: + coordinator = NoopDisaggCoordinator() + for name in _entry_points(): + + def record(*args, _name=name): + calls.append((_name, *args)) + + setattr(coordinator, name, record) + + def admit(fitting): + calls.append(("admit", fitting)) + return fitting, False + + coordinator.admit = admit + return coordinator + + +def _collective_calls(calls: list) -> list: + return [call for call in calls if call[0] in _COLLECTIVE_SENSITIVE] + + +def _idle_executor(monkeypatch, calls: list) -> PyExecutor: + """Bare executor whose first iteration schedules nothing and whose second + iteration observes shutdown. Executor-owned helpers are stubbed; the loop + bodies and the communication boundaries are real call sites.""" + for target in ("torch.cuda.set_device", "cudart.cudaSetDevice", "CUASSERT"): + monkeypatch.setattr(f"tensorrt_llm._torch.pyexecutor.py_executor.{target}", Mock()) + + executor = object.__new__(PyExecutor) + executor._disagg_coordinator = _recording_coordinator(calls) + executor._flush_pending_transfer_responses = lambda: calls.append( + ("flush_pending_transfer_responses",) + ) + executor._handle_kv_transfer_timeouts_synced = lambda: calls.append( + ("handle_kv_transfer_timeouts_synced",) + ) + executor.kv_cache_transceiver = Mock() + executor.async_transfer_manager = Mock() + executor.async_transfer_manager.has_any_inflight_requests.return_value = False + executor.kv_connector_manager = None + + executor.device_id = 0 + profiler = MagicMock() + profiler.__enter__.return_value = Mock() + executor._profiler = Mock(return_value=profiler) + executor.hang_detector = MagicMock() + executor.enable_iter_perf_stats = False + executor.enable_attention_dp = False + executor.is_benchmark_disagg = False + executor.is_warmup = False + executor.iter_counter = 0 + executor._resource_governor_enabled = False + executor._is_kv_manager_v2 = False + executor._mm_encoder_item_scheduling_enabled = False + executor.enable_early_first_token_response = False + executor.drafter = None + executor.model_engine = None + executor.previous_batch = None + executor.active_requests = [] + executor.waiting_queue = [] + executor.inflight_req_ids = set() + executor.is_shutdown = False + + def fetch_new_requests(): + if fetch_new_requests.called: + executor.is_shutdown = True + fetch_new_requests.called = True + return [] + + fetch_new_requests.called = False + executor._fetch_and_activate_new_requests = fetch_new_requests + executor._schedule = Mock(return_value=(ScheduledRequests(), [], 0)) + executor._can_queue = Mock(return_value=(False, False)) + executor._check_benchmark_disagg_gate = Mock(return_value=(True, False)) + executor._sync_gen_only_benchmark_has_insufficient_kv = Mock(return_value=False) + for name in ( + "_poll_encoder_steps", + "_handle_control_request", + "_pad_attention_dp_dummy_request", + "_prefetch_for_context_requests", + "_pad_empty_attention_dp_batch", + "_terminate_requests", + "_pause_requests", + "_revert_gen_alloc", + "_finalize_adp_dummy_allocation", + "_wait_for_model_engine_input_copy", + "_enqueue_responses", + ): + setattr(executor, name, Mock()) + return executor + + +def _pp_executor(monkeypatch, calls: list, *, rank: int) -> PyExecutor: + """Idle executor on a two-stage pipeline; rank 0 schedules, rank 1 receives + the schedule from its predecessor and re-runs the scheduler locally.""" + executor = _idle_executor(monkeypatch, calls) + executor.dist = Mock( + rank=rank, + pp_rank=rank, + pp_size=2, + tp_size=1, + cp_size=1, + world_size=2, + is_first_pp_rank=rank == 0, + is_last_pp_rank=rank == 1, + next_pp_rank=1 - rank, + prev_pp_rank=1 - rank, + ) + executor.num_micro_batches = 1 + executor.micro_batches = [None] + executor.send_handles = [None] + executor.send_schedule_handles = [None] + executor.send_expected_batch_num_handles = [None] + executor.wait_on_pp_send_handles = Mock() + executor.executed_batch_response_queue = queue.Queue() + executor.unhandled_batch_counter = 0 + executor.pp_async_broadcast_sample_state = True + executor._pp_rebalance_drain_iters = None + executor._progress_recompute_pause_termination_if_idle = Mock() + if rank != 0: + empty_schedule = SerializableSchedulerOutput.from_scheduler_result( + ScheduledRequests(), [], 0 + ) + # First recv is the schedule, second is the executed-batch count. + executor.dist.recv_object = Mock(side_effect=[empty_schedule, 0]) + executor.scheduler = Mock() + executor.scheduler.can_schedule.return_value = True + executor.scheduler.schedule_request.return_value = SimpleNamespace( + fitting_disagg_gen_init_requests=[] + ) + executor.kv_cache_manager = Mock(spec=[]) + return executor + + +_SCHEDULE_HEAD = [ + ("handle_errors_synced",), + ("prepare_context_schedulable", []), + ("poll_gen_transfers",), + ("check_transfer_timeouts",), + ("admit", []), + ("receive_gen_init", []), + ("poll_progress_when_idle",), +] +# Non-PP loops flush at loop exit; the PP loop does not. +_SHUTDOWN_PASS = [("handle_errors_synced",), ("flush_pending_transfer_responses",)] +_PP_SHUTDOWN_PASS = [("handle_errors_synced",)] + + +def test_executor_loop_transcript(monkeypatch) -> None: + """The timeout-consensus drain precedes the response flush in this loop.""" + calls = [] + executor = _idle_executor(monkeypatch, calls) + executor.dist = Mock(tp_size=1, world_size=1) + + PyExecutor._executor_loop(executor) + + idle_pass = _SCHEDULE_HEAD + [ + ("handle_kv_transfer_timeouts_synced",), + ("flush_pending_transfer_responses",), + ("pace_idle",), + ] + assert calls == idle_pass + _SHUTDOWN_PASS + + +def test_executor_loop_overlap_transcript(monkeypatch) -> None: + """The overlap loop flushes responses before the timeout-consensus drain, + the reverse of the non-overlap loop; each order is consistent across ranks + on its own.""" + calls = [] + executor = _idle_executor(monkeypatch, calls) + executor.dist = Mock(tp_size=1, world_size=1) + + PyExecutor._executor_loop_overlap(executor) + + idle_pass = _SCHEDULE_HEAD + [ + ("flush_pending_transfer_responses",), + ("handle_kv_transfer_timeouts_synced",), + ("pace_idle",), + ] + assert calls == idle_pass + _SHUTDOWN_PASS + + +def test_executor_loop_pp_transcript_on_first_rank(monkeypatch) -> None: + """The PP loop admits inside schedule propagation, checks transfer timeouts + only on the retry and executed-batch paths, and flushes responses only from + executed-batch handling, so an idle iteration has none of those.""" + calls = [] + PyExecutor._executor_loop_pp(_pp_executor(monkeypatch, calls, rank=0)) + + assert ( + calls + == [ + ("handle_errors_synced",), + ("prepare_context_schedulable", []), + ("poll_gen_transfers",), + ("admit", []), + ("receive_gen_init", []), + ("poll_progress_when_idle",), + ("pace_idle",), + ] + + _PP_SHUTDOWN_PASS + ) + + +def test_executor_loop_pp_transcript_on_non_first_rank(monkeypatch) -> None: + """A non-first rank does not admit; it reverts KV for candidates its local + scheduler picked but the first rank did not admit.""" + calls = [] + PyExecutor._executor_loop_pp(_pp_executor(monkeypatch, calls, rank=1)) + + assert ( + calls + == [ + ("handle_errors_synced",), + ("prepare_context_schedulable", []), + ("poll_gen_transfers",), + ("revert_deferred_gen_init", [], []), + ("receive_gen_init", []), + ("poll_progress_when_idle",), + ("pace_idle",), + ] + + _PP_SHUTDOWN_PASS + ) + + +def test_pp_ranks_issue_the_same_collective_sensitive_calls(monkeypatch) -> None: + """First and non-first PP ranks take different local paths, but every + collective-sensitive boundary must be reached the same number of times in + the same order or the consensus inside it deadlocks.""" + assert _COLLECTIVE_COORDINATOR_CALLS <= set(_entry_points()) + first, other = [], [] + PyExecutor._executor_loop_pp(_pp_executor(monkeypatch, first, rank=0)) + PyExecutor._executor_loop_pp(_pp_executor(monkeypatch, other, rank=1)) + + assert _collective_calls(first) == _collective_calls(other) + assert _collective_calls(first) # the comparison is not vacuous diff --git a/tests/unittest/_torch/executor/test_send_kv_async_split.py b/tests/unittest/_torch/executor/test_send_kv_async_split.py index a257e5ec5c08..62f36ccead9c 100644 --- a/tests/unittest/_torch/executor/test_send_kv_async_split.py +++ b/tests/unittest/_torch/executor/test_send_kv_async_split.py @@ -54,7 +54,7 @@ def test_wrapper_keeps_connector_leg_without_transceiver() -> None: PyExecutor._send_kv_async(executor, []) - assert calls == ["disagg_send", "connector_save"] + assert calls == ["connector_save"] def test_disagg_send_leg_is_noop_without_transceiver() -> None: From 76518c4c6338c6335b32ee1f31d2dd8ac9259093 Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:32:47 +0800 Subject: [PATCH 2/3] [None][test] Do not touch the is_warmup setter in the transcript fixture Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- tests/unittest/_torch/executor/test_disagg_loop_transcript.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unittest/_torch/executor/test_disagg_loop_transcript.py b/tests/unittest/_torch/executor/test_disagg_loop_transcript.py index adcac7621331..29cafc578760 100644 --- a/tests/unittest/_torch/executor/test_disagg_loop_transcript.py +++ b/tests/unittest/_torch/executor/test_disagg_loop_transcript.py @@ -116,7 +116,6 @@ def _idle_executor(monkeypatch, calls: list) -> PyExecutor: executor.enable_iter_perf_stats = False executor.enable_attention_dp = False executor.is_benchmark_disagg = False - executor.is_warmup = False executor.iter_counter = 0 executor._resource_governor_enabled = False executor._is_kv_manager_v2 = False From 08a99c155ffd34dea5121675ced5028fab6ae299 Mon Sep 17 00:00:00 2001 From: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> Date: Thu, 3 Sep 2026 10:19:48 +0800 Subject: [PATCH 3/3] [None][test] Parametrize coordinator forwarding test over every delegate Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com> --- .../executor/test_disagg_coordinator.py | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/tests/unittest/_torch/executor/test_disagg_coordinator.py b/tests/unittest/_torch/executor/test_disagg_coordinator.py index d04a874d9409..35d69425753e 100644 --- a/tests/unittest/_torch/executor/test_disagg_coordinator.py +++ b/tests/unittest/_torch/executor/test_disagg_coordinator.py @@ -47,18 +47,24 @@ def test_every_coordinator_method_has_a_delegate() -> None: assert _public_methods(DisaggTransferCoordinator) == delegate_names -def test_real_coordinator_forwards_arguments_and_results() -> None: +@pytest.mark.parametrize("name", [f.name for f in fields(DisaggLoopDelegates)]) +def test_real_coordinator_forwards_each_method_to_its_delegate(name: str) -> None: + """Each entry point must reach exactly its own delegate with the arguments + unchanged; a cross-wired or dropped call changes loop behavior and may break + rank symmetry for collective-sensitive entry points.""" delegates = DisaggLoopDelegates(**{f.name: Mock() for f in fields(DisaggLoopDelegates)}) - delegates.admit.return_value = (["admitted"], True) coordinator = DisaggTransferCoordinator(delegates) + method = getattr(coordinator, name) + args = [object() for _ in inspect.signature(method).parameters] - assert coordinator.admit(["fitting"]) == (["admitted"], True) - coordinator.reap_context_sends(1) - coordinator.revert_deferred_gen_init(["a"], ["b"]) + result = method(*args) - delegates.admit.assert_called_once_with(["fitting"]) - delegates.reap_context_sends.assert_called_once_with(1) - delegates.revert_deferred_gen_init.assert_called_once_with(["a"], ["b"]) + getattr(delegates, name).assert_called_once_with(*args) + for other in fields(DisaggLoopDelegates): + if other.name != name: + getattr(delegates, other.name).assert_not_called() + expected = getattr(delegates, name).return_value if name == "admit" else None + assert result is expected def test_noop_coordinator_admits_everything_unchanged() -> None: