[None][fix] Flush overlap KV cache events on every attention DP rank

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -5454,30 +5454,31 @@ def _agreed_need_adjustment(self) -> bool:
         return need
 
     def _consume_previous_batch_for_rebalance(self) -> None:
-        """Drain ``previous_batch`` so its _KVCache instances are quiescent.
-
-        No-op when ``previous_batch is None`` -- i.e., always a no-op in
-        the non-overlap loop, since that loop never sets previous_batch.
-        In the overlap loop this fires when the rebalance hook catches a
-        pending in-flight iteration; we consume it inline so suspend can
-        safely run.
-
-        Mirrors the inline sequence in ``_executor_loop_overlap`` that
-        handles ``previous_batch``.  Unlike the inline code we are not
-        guarded by ``should_process_previous_batch``: the rebalance gate
-        already excludes the multi-rank-divergence cases that flag exists
-        to handle.
+        """Drain in-flight batch work before KV pool rebalance.
+
+        In overlap mode, ADP ranks must participate in the transfer-response,
+        batch-response, and KV-event collectives even without a local batch.
+        Other modes remain a no-op when there is no previous batch.
         """
-        if self.previous_batch is None:
+        previous_batch = self.previous_batch
+        if previous_batch is None and (not self.enable_attention_dp
+                                       or self.disable_overlap_scheduler):
             return
-        self._update_requests(self.previous_batch.sample_state)
-        self._send_kv_async(
-            self.previous_batch.scheduled_requests.all_requests())
+        if previous_batch is not None:
+            self._update_requests(previous_batch.sample_state)
+            self._send_kv_async(
+                previous_batch.scheduled_requests.all_requests())
         self._flush_pending_transfer_responses()
-        self._process_previous_batch()
-        self.perf_manager.compute_batch_gpu_times(
-            self.previous_batch.scheduled_requests.all_requests())
-        self.previous_batch = None
+        if previous_batch is not None:
+            self._process_previous_batch()
+        else:
+            self._enqueue_responses([])
+        if self.enable_kv_cache_events:
+            self._add_kv_cache_events()
+        if previous_batch is not None:
+            self.perf_manager.compute_batch_gpu_times(
+                previous_batch.scheduled_requests.all_requests())
+            self.previous_batch = None
 
     def _maybe_rebalance_kv_pools(self) -> None:
         """Rebalance KV pool ratios when the V2 auto-tuner asks for it.
@@ -6003,6 +6004,12 @@ def _executor_loop_overlap(self):
                 else:
                     self._enqueue_responses([])
 
+                # KV-cache event flushing performs an attention-DP allgather.
+                # Keep it outside the rank-divergent previous-batch branch so
+                # every rank enters the collective in the same order.
+                if self.enable_kv_cache_events:
+                    self._add_kv_cache_events()
+
                 # Drain buffers from the (per-rank-divergent)
                 # _process_previous_batch above; rank-symmetric companion to
                 # the _enqueue_responses([]) call in the else branch.
@@ -6171,8 +6178,6 @@ def _process_previous_batch(self):
             attn_metadata,
             kv_cache_dtype_byte_size,
             skip_kv_cache_manager=self._is_kv_manager_v2)
-        if self.enable_kv_cache_events:
-            self._add_kv_cache_events()
 
         if self.enable_iter_perf_stats:
             self._process_iter_stats(finished_requests, self.active_requests,

tests/unittest/_torch/disaggregation/test_disagg_loop_transcript.py

@@ -38,11 +38,14 @@
     PyExecutorEffects,
     PyExecutorRequestRegistry,
 )
+from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
 from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType
 from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import (
     ScheduledRequests,
     SerializableSchedulerOutput,
 )
+from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheEventManager
 
 pytestmark = pytest.mark.cpu_only
 
@@ -120,6 +123,7 @@ def _idle_executor(monkeypatch, calls: list) -> PyExecutor:
     executor._profiler = Mock(return_value=profiler)
     executor.hang_detector = MagicMock()
     executor.enable_iter_perf_stats = False
+    executor.enable_kv_cache_events = False
     executor.enable_attention_dp = False
     executor.is_benchmark_disagg = False
     executor.iter_counter = 0
@@ -256,6 +260,108 @@ def test_executor_loop_overlap_transcript(monkeypatch) -> None:
     assert calls == idle_pass + _SHUTDOWN_PASS
 
 
+@pytest.mark.parametrize(
+    "states",
+    [("absent", "processed"), ("processed", "deferred"), ("deferred", "absent")],
+    ids=["absent-processed", "processed-deferred", "deferred-absent"],
+)
+def test_overlap_kv_events_keep_adp_collectives_aligned(
+    monkeypatch, states: tuple[str, str]
+) -> None:
+    """Buffered V2 events occupy the same collective slot despite rank-local batch state."""
+    transcripts = []
+    for rank, state in enumerate(states):
+        calls = []
+        gathered_events = []
+
+        def gather(local_events: list) -> list:
+            calls.append(("kv_event_allgather",))
+            gathered_events.append(list(local_events))
+            return [local_events, []] if rank == 0 else [[], local_events]
+
+        manager = object.__new__(KVCacheManagerV2)
+        manager.event_manager = KVCacheEventManager(
+            max_kv_event_entries=8, attention_dp_rank=rank, attention_dp_gather=gather
+        )
+        executor = _idle_executor(monkeypatch, calls)
+        executor.enable_attention_dp = True
+        executor.enable_kv_cache_events = True
+        executor._is_kv_manager_v2 = True
+        executor.kv_cache_manager = manager
+        executor.enable_kv_pool_rebalance = False
+        executor.dist = Mock(rank=rank, tp_size=2, world_size=2)
+        executor.speculation_gate = None
+        executor.use_spec_decode = False
+        executor._pause_recompute_paused_requests = Mock()
+        executor._commit_kv_cache_stats = Mock()
+        executor._update_requests = Mock()
+        executor._send_kv_async = Mock()
+        executor._handle_canceled_requests = Mock()
+        executor.perf_manager = Mock()
+        executor.resource_manager = Mock()
+        executor.resource_manager.resource_managers = {
+            ResourceManagerType.KV_CACHE_MANAGER: manager
+        }
+        executor._enqueue_responses = lambda responses: calls.append(("enqueue_responses",))
+
+        def handle_responses(*, emit_first_iter: bool) -> list:
+            executor._enqueue_responses([])
+            return []
+
+        def update_generation_resources(*args) -> None:
+            calls.append(("update_generation_resources",))
+            manager.event_manager.add_created_event([1])
+
+        def update_resources(*args, skip_kv_cache_manager: bool) -> None:
+            assert skip_kv_cache_manager is True
+            calls.append(("update_resources",))
+
+        manager.update_resources = Mock(side_effect=update_generation_resources)
+        executor._handle_responses = handle_responses
+        executor.resource_manager.update_resources.side_effect = update_resources
+        if state != "absent":
+            executor.previous_batch = SimpleNamespace(
+                scheduled_requests=ScheduledRequests(), sample_state=Mock()
+            )
+        # A locally ready rank defers its previous batch while a peer cannot queue.
+        executor._can_queue.return_value = (False, state == "deferred")
+
+        PyExecutor._executor_loop_overlap(executor)
+
+        assert calls[-len(_SHUTDOWN_PASS) :] == _SHUTDOWN_PASS
+        iteration_calls = calls[: -len(_SHUTDOWN_PASS)]
+        assert len(gathered_events) == 1
+        assert len(gathered_events[0]) == (1 if state == "processed" else 0)
+        if state == "processed":
+            manager.update_resources.assert_called_once()
+            assert iteration_calls.index(("update_generation_resources",)) < iteration_calls.index(
+                ("enqueue_responses",)
+            )
+            assert iteration_calls.index(("enqueue_responses",)) < iteration_calls.index(
+                ("update_resources",)
+            )
+            assert iteration_calls.index(("update_resources",)) < iteration_calls.index(
+                ("kv_event_allgather",)
+            )
+            executor._update_requests.assert_called_once()
+        else:
+            manager.update_resources.assert_not_called()
+            executor._update_requests.assert_not_called()
+            executor.resource_manager.update_resources.assert_not_called()
+
+        collective_names = _COLLECTIVE_SENSITIVE | {"enqueue_responses", "kv_event_allgather"}
+        transcript = [call for call in iteration_calls if call[0] in collective_names]
+        assert transcript[-4:] == [
+            ("flush_pending_transfer_responses",),
+            ("enqueue_responses",),
+            ("kv_event_allgather",),
+            ("handle_timeouts_synced",),
+        ]
+        transcripts.append(transcript)
+
+    assert transcripts[0] == transcripts[1]
+
+
 @pytest.mark.parametrize(
     "loop",
     [PyExecutor._executor_loop, PyExecutor._executor_loop_overlap],