From de080e435dd883f678952ad6b434f9ad4f59ac42 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Tue, 6 Oct 2026 08:42:17 -0700
Subject: [PATCH] [None][fix] forward async transfer manager through multimodal
 scheduler wrapper

PyExecutor wires the AsyncTransferManager into the KVCacheV2Scheduler
deadlock detector via hasattr(self.scheduler, "set_async_transfer_manager").
With MM item scheduling enabled, the scheduler is wrapped in
MultimodalScheduler / MultimodalEagerEncoderScheduler, which did not expose
the hook, so the wrapped V2 scheduler never received the manager and could
count finished-but-unreaped disagg KV sends as a stall (false deadlock).

Add an explicit passthrough on MultimodalScheduler that forwards to the
wrapped scheduler when it supports the hook.

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/pyexecutor/scheduler/scheduler.py   | 11 +++++++++++
 .../executor/test_multimodal_scheduler.py      | 18 ++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
index ff180cca75fb..faa56ef1b10f 100644
--- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
+++ b/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.
diff --git a/tests/unittest/_torch/executor/test_multimodal_scheduler.py b/tests/unittest/_torch/executor/test_multimodal_scheduler.py
index f67ad06095f4..80ebc01aa945 100644
--- a/tests/unittest/_torch/executor/test_multimodal_scheduler.py
+++ b/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)}