[None][fix] Forward async transfer manager through multimodal scheduler wrapper

tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py

@@ -816,6 +816,17 @@ def schedule_request(
     def can_schedule(self, requests: RequestList) -> bool:
         return self.scheduler.can_schedule(requests)
 
+    def set_async_transfer_manager(self, mgr) -> None:
+        """Forward the AsyncTransferManager to the wrapped scheduler.
+
+        PyExecutor feature-detects this method to wire in-flight disagg KV
+        sends into the KVCacheV2Scheduler deadlock detector; without the
+        passthrough, a wrapped V2 scheduler never receives it.
+        """
+        set_mgr = getattr(self.scheduler, "set_async_transfer_manager", None)
+        if set_mgr is not None:
+            set_mgr(mgr)
+
 
 class MultimodalEagerEncoderScheduler(MultimodalScheduler):
     """Eagerly schedule encoder work for already-active MM requests.

tests/unittest/_torch/executor/test_multimodal_scheduler.py

@@ -32,6 +32,7 @@
     MultimodalScheduler,
     ScheduledRequests,
 )
+from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler
 from tensorrt_llm._torch.pyexecutor.scheduler.waiting_queue import FCFSWaitingQueue
 from tensorrt_llm.inputs.multimodal import (
     MULTIMODAL_ENCODER_ITEM_METADATA_KEY,
@@ -284,6 +285,23 @@ def test_multimodal_scheduler_preserves_non_multimodal_requests():
     assert output.context_requests == [request]
 
 
+@pytest.mark.parametrize("scheduler_cls", [MultimodalScheduler, MultimodalEagerEncoderScheduler])
+def test_multimodal_scheduler_forwards_async_transfer_manager(scheduler_cls):
+    v2_scheduler = KVCacheV2Scheduler.__new__(KVCacheV2Scheduler)
+    v2_scheduler._async_transfer_manager = None
+    transfer_manager = Mock()
+
+    scheduler_cls(v2_scheduler, max_batch_size=1, max_num_tokens=1).set_async_transfer_manager(
+        transfer_manager
+    )
+    # A wrapped scheduler without the hook is left untouched.
+    scheduler_cls(_BaseScheduler(), max_batch_size=1, max_num_tokens=1).set_async_transfer_manager(
+        transfer_manager
+    )
+
+    assert v2_scheduler._async_transfer_manager is transfer_manager
+
+
 def test_request_rejects_item_above_effective_startup_maximum():
     request = make_mm_request(1, [9])
     request.py_multimodal_data["image"] = {"pixel_values": torch.empty(1)}