[None][refactor] Split connector KV save out of _send_kv_async

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -7459,6 +7459,40 @@ def _recv_disagg_gen_cache(self, new_gen_reqs):
 
     @nvtx_range("_send_kv_async")
     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._save_kv_to_connector_async(scheduled_requests)
+        if self.kv_cache_transceiver:
+            self._check_disagg_ctx_cache_transfer_status(0)
+
+    def _send_disagg_ctx_kv_async(self,
+                                  scheduled_requests: List[LlmRequest]) -> None:
+        """Start async KV sends for finished context-only disagg requests."""
+        if not self.kv_cache_transceiver:
+            return
+        for req in scheduled_requests:
+            if req.is_context_only_request and (
+                    req.is_context_finished or req.is_finished_due_to_length
+            ) and not req.is_finished_due_to_cancellation:
+                # Forward is done for this request — release the
+                # IndexMapper slot so new requests can reuse it.
+                # KV blocks stay allocated for the upcoming transfer.
+                if hasattr(self.kv_cache_manager, 'release_index_slot'):
+                    self.kv_cache_manager.release_index_slot(req.py_request_id)
+                # Order is important here: we need to start the transfer before responding
+                # to make sure the blocks are stored for reuse before they are sent.
+                self.async_transfer_manager.start_transfer(req)
+                self.kv_cache_transceiver.respond_and_send_async(req)
+
+                if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None:
+                    req.py_kv_transfer_start_time = time.monotonic()
+
+    def _save_kv_to_connector_async(
+            self, scheduled_requests: List[LlmRequest]) -> None:
+        """Hand finished requests' KV blocks to the KV connector for async saving."""
+        if not self.kv_connector_manager:
+            return
 
         def kv_connector_request_finished(req: LlmRequest):
             try:
@@ -7472,37 +7506,14 @@ def kv_connector_request_finished(req: LlmRequest):
                         req, cache_block_ids):
                     self.async_transfer_manager.start_transfer(req)
 
-        if self.kv_cache_transceiver:
-            for req in scheduled_requests:
-                if req.is_context_only_request and (
-                        req.is_context_finished or req.is_finished_due_to_length
-                ) and not req.is_finished_due_to_cancellation:
-                    # Forward is done for this request — release the
-                    # IndexMapper slot so new requests can reuse it.
-                    # KV blocks stay allocated for the upcoming transfer.
-                    if hasattr(self.kv_cache_manager, 'release_index_slot'):
-                        self.kv_cache_manager.release_index_slot(
-                            req.py_request_id)
-                    # Order is important here: we need to start the transfer before responding
-                    # to make sure the blocks are stored for reuse before they are sent.
-                    self.async_transfer_manager.start_transfer(req)
-                    self.kv_cache_transceiver.respond_and_send_async(req)
-
-                    if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None:
-                        req.py_kv_transfer_start_time = time.monotonic()
-
-        if self.kv_connector_manager:
-            if not self.disable_overlap_scheduler:
-                requests = self.previous_batch.scheduled_requests.all_requests(
-                ) if self.previous_batch is not None else []
-            else:
-                requests = scheduled_requests
-            for req in requests:
-                if req.is_finished:
-                    kv_connector_request_finished(req)
-
-        if self.kv_cache_transceiver:
-            self._check_disagg_ctx_cache_transfer_status(0)
+        if not self.disable_overlap_scheduler:
+            requests = self.previous_batch.scheduled_requests.all_requests(
+            ) if self.previous_batch is not None else []
+        else:
+            requests = scheduled_requests
+        for req in requests:
+            if req.is_finished:
+                kv_connector_request_finished(req)
 
     @staticmethod
     def _request_vote_id(request: LlmRequest) -> int:

tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py

@@ -67,9 +67,10 @@ def create_mock_resource_manager(kv_cache_manager=None, seq_slot_manager=None):
     return resource_manager
 
 
-class _FakeExecutor:
-    """Minimal stand-in for PyExecutor so we can call the unbound
-    `_send_kv_async` method without constructing the real object."""
+class _FakeExecutor(PyExecutor):
+    """Minimal PyExecutor stand-in: inherits the real methods (so it keeps
+    working as `_send_kv_async` internals evolve) but skips
+    `PyExecutor.__init__` so no engine or distributed wiring is required."""
 
     def __init__(self, kv_cache_manager, async_transfer_manager, kv_cache_transceiver):
         self.kv_cache_manager = kv_cache_manager

tests/unittest/_torch/executor/test_send_kv_async_split.py

@@ -0,0 +1,249 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Structure tests for the _send_kv_async split.
+
+``_send_kv_async`` composes two independent legs plus a reap, and the order
+is load-bearing (see the comment in the wrapper): the disagg send must
+register its transfer before the connector does, and the ctx reap must run
+last so a quickly-completed send cannot terminate a request whose connector
+transfer is not registered yet. These tests pin that structure, and pin the
+property this split exists for: the connector leg keeps running when the
+transceiver is disabled.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import Mock, call
+
+import pytest
+
+from tensorrt_llm._torch.disaggregation.executor.transfer_manager import AsyncTransferManager
+from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor
+from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType
+
+pytestmark = pytest.mark.cpu_only
+
+
+def _stub_executor() -> PyExecutor:
+    return object.__new__(PyExecutor)
+
+
+def _wrapper_calls(executor: PyExecutor) -> list:
+    calls = []
+    executor._send_disagg_ctx_kv_async = lambda reqs: calls.append("disagg_send")
+    executor._save_kv_to_connector_async = lambda reqs: calls.append("connector_save")
+    executor._check_disagg_ctx_cache_transfer_status = lambda n: calls.append(f"ctx_reap:{n}")
+    return calls
+
+
+def test_wrapper_order_disagg_then_connector_then_reap() -> None:
+    executor = _stub_executor()
+    calls = _wrapper_calls(executor)
+    executor.kv_cache_transceiver = object()
+
+    PyExecutor._send_kv_async(executor, [])
+
+    assert calls == ["disagg_send", "connector_save", "ctx_reap:0"]
+
+
+def test_wrapper_keeps_connector_leg_without_transceiver() -> None:
+    """Connector-only configs must keep working when the disagg path is off."""
+    executor = _stub_executor()
+    calls = _wrapper_calls(executor)
+    executor.kv_cache_transceiver = None
+
+    PyExecutor._send_kv_async(executor, [])
+
+    assert calls == ["disagg_send", "connector_save"]
+
+
+def test_disagg_send_leg_is_noop_without_transceiver() -> None:
+    executor = _stub_executor()
+    executor.kv_cache_transceiver = None
+    # Use a request that passes the send filter: without the guard, the loop
+    # body would hit unset executor attributes and raise.
+    PyExecutor._send_disagg_ctx_kv_async(executor, [_finished_ctx_only_request()])
+
+
+def test_connector_save_leg_is_noop_without_connector() -> None:
+    executor = _stub_executor()
+    executor.kv_connector_manager = None
+    PyExecutor._save_kv_to_connector_async(executor, [Mock()])
+
+
+def _finished_ctx_only_request(request_id: int = 1) -> SimpleNamespace:
+    return SimpleNamespace(
+        is_context_only_request=True,
+        is_context_finished=True,
+        is_finished_due_to_length=False,
+        is_finished_due_to_cancellation=False,
+        py_request_id=request_id,
+        py_kv_transfer_start_time=None,
+    )
+
+
+def test_disagg_send_leg_stores_blocks_before_sending() -> None:
+    executor = _stub_executor()
+    transceiver = Mock()
+    transceiver.kv_transfer_timeout_ms = 1000
+    executor.kv_cache_transceiver = transceiver
+    executor.kv_cache_manager = Mock(spec=[])  # no release_index_slot
+    executor.async_transfer_manager = Mock()
+    order = Mock()
+    order.attach_mock(executor.async_transfer_manager.start_transfer, "start")
+    order.attach_mock(transceiver.respond_and_send_async, "send")
+    req = _finished_ctx_only_request()
+
+    PyExecutor._send_disagg_ctx_kv_async(executor, [req])
+
+    assert order.mock_calls == [call.start(req), call.send(req)]
+    assert req.py_kv_transfer_start_time is not None
+
+
+def test_disagg_send_leg_skips_timeout_stamp_when_disabled() -> None:
+    executor = _stub_executor()
+    transceiver = Mock()
+    transceiver.kv_transfer_timeout_ms = None
+    executor.kv_cache_transceiver = transceiver
+    executor.kv_cache_manager = Mock(spec=[])
+    executor.async_transfer_manager = Mock()
+    req = _finished_ctx_only_request()
+
+    PyExecutor._send_disagg_ctx_kv_async(executor, [req])
+
+    transceiver.respond_and_send_async.assert_called_once_with(req)
+    assert req.py_kv_transfer_start_time is None
+
+
+def test_disagg_send_leg_skips_cancelled_and_unfinished_requests() -> None:
+    executor = _stub_executor()
+    transceiver = Mock()
+    executor.kv_cache_transceiver = transceiver
+    executor.kv_cache_manager = Mock(spec=[])
+    executor.async_transfer_manager = Mock()
+    cancelled = _finished_ctx_only_request(2)
+    cancelled.is_finished_due_to_cancellation = True
+    unfinished = _finished_ctx_only_request(3)
+    unfinished.is_context_finished = False
+
+    PyExecutor._send_disagg_ctx_kv_async(executor, [cancelled, unfinished])
+
+    executor.async_transfer_manager.start_transfer.assert_not_called()
+    transceiver.respond_and_send_async.assert_not_called()
+
+
+def _connector_executor() -> PyExecutor:
+    executor = _stub_executor()
+    executor.kv_connector_manager = Mock()
+    executor.kv_connector_manager.request_finished.return_value = True
+    executor.kv_cache_manager = Mock()
+    executor.kv_cache_manager.get_cache_indices.return_value = [7]
+    executor.async_transfer_manager = Mock()
+    return executor
+
+
+def test_connector_save_uses_previous_batch_with_overlap_scheduler() -> None:
+    executor = _connector_executor()
+    executor.disable_overlap_scheduler = False
+    prev_req = SimpleNamespace(is_finished=True, py_request_id=2)
+    executor.previous_batch = SimpleNamespace(
+        scheduled_requests=SimpleNamespace(all_requests=lambda: [prev_req])
+    )
+    current_req = SimpleNamespace(is_finished=True, py_request_id=3)
+
+    PyExecutor._save_kv_to_connector_async(executor, [current_req])
+
+    executor.kv_connector_manager.request_finished.assert_called_once_with(prev_req, [7])
+    executor.async_transfer_manager.start_transfer.assert_called_once_with(prev_req)
+
+
+def test_connector_save_uses_scheduled_batch_without_overlap_scheduler() -> None:
+    executor = _connector_executor()
+    executor.disable_overlap_scheduler = True
+    finished = SimpleNamespace(is_finished=True, py_request_id=4)
+    running = SimpleNamespace(is_finished=False, py_request_id=5)
+
+    PyExecutor._save_kv_to_connector_async(executor, [finished, running])
+
+    executor.kv_connector_manager.request_finished.assert_called_once_with(finished, [7])
+    executor.async_transfer_manager.start_transfer.assert_called_once_with(finished)
+
+
+def test_connector_save_skips_transfer_when_connector_declines() -> None:
+    executor = _connector_executor()
+    executor.disable_overlap_scheduler = True
+    executor.kv_connector_manager.request_finished.return_value = False
+    finished = SimpleNamespace(is_finished=True, py_request_id=6)
+
+    PyExecutor._save_kv_to_connector_async(executor, [finished])
+
+    executor.async_transfer_manager.start_transfer.assert_not_called()
+
+
+def _dual_claim_executor() -> PyExecutor:
+    """Executor running the real wrapper, real legs, and a real
+    AsyncTransferManager; only the transceiver, connector, and KV cache
+    manager boundaries are mocked."""
+    executor = _stub_executor()
+    kv_cache_manager = Mock()
+    kv_cache_manager.get_cache_indices.return_value = [7]
+    executor.kv_cache_manager = kv_cache_manager
+    resource_manager = SimpleNamespace(
+        resource_managers={ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}
+    )
+    executor.async_transfer_manager = AsyncTransferManager(resource_manager)
+    transceiver = Mock()
+    transceiver.kv_transfer_timeout_ms = None
+    executor.kv_cache_transceiver = transceiver
+    executor.kv_connector_manager = Mock()
+    executor.disable_overlap_scheduler = True
+    executor.active_requests = []
+    executor.force_terminate_ctx_for_partial_reuse = False
+    executor._disagg_timed_out_ctx_cancelled_ids = set()
+    executor._terminate_request = Mock()
+    # Make the reap's trailing _check_cache_transfer_errors a no-op.
+    executor.enable_attention_dp = True
+    executor.dist = SimpleNamespace(world_size=2)
+    return executor
+
+
+def _dual_claim_request(request_id: int) -> SimpleNamespace:
+    req = _finished_ctx_only_request(request_id)
+    req.is_finished = True  # the connector leg selects finished requests
+    req.py_kv_transfer_timed_out = False
+    req.state = None  # start_transfer overwrites
+    return req
+
+
+def test_reap_keeps_request_still_claimed_by_connector() -> None:
+    """The hazard the wrapper order exists for: a send that completes within
+    the same iteration must not release a request the connector also claimed.
+    Reaping before the connector leg would drop the transfer refcount to zero
+    and terminate the request; this test fails under that reordering."""
+    executor = _dual_claim_executor()
+    executor.kv_connector_manager.request_finished.return_value = True
+    req = _dual_claim_request(9)
+    # The send completes instantly, so the reap sees it in the same call.
+    executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([9], [])
+
+    PyExecutor._send_kv_async(executor, [req])
+
+    # The reap released only the send's claim; the connector's claim keeps
+    # the request pinned.
+    assert 9 in executor.async_transfer_manager.requests_in_transfer()
+    executor.kv_cache_manager.unpin_blocks_by_id.assert_not_called()
+    executor._terminate_request.assert_not_called()
+
+
+def test_reap_releases_request_once_connector_declines() -> None:
+    """Counterpart: with only the send's claim outstanding, the same fast
+    completion does release and terminate the request."""
+    executor = _dual_claim_executor()
+    executor.kv_connector_manager.request_finished.return_value = False
+    req = _dual_claim_request(9)
+    executor.kv_cache_transceiver.check_context_transfer_status.return_value = ([9], [])
+
+    PyExecutor._send_kv_async(executor, [req])
+
+    assert 9 not in executor.async_transfer_manager.requests_in_transfer()
+    executor.kv_cache_manager.unpin_blocks_by_id.assert_called_once()
+    executor._terminate_request.assert_called_once_with(req)