From 7f870dd3746a71b7c9107d85acda264852011fe1 Mon Sep 17 00:00:00 2001
From: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
Date: Wed, 2 Sep 2026 14:43:31 -0700
Subject: [PATCH] [https://nvbugs/6683840][fix] Join the encoder launch before
 the decoder forward

torch.fx's tracing flag and its patch of nn.Module.__call__ are process-global,
so a dynamo compile on the encoder-launch thread captures the decoder thread's
concurrent module calls and dynamo refuses to re-enter. At tp_size 1 the encoder
future is now joined before the decoder forward runs, while ready_event stays
unsynchronized so encoder kernels still overlap decoder work on their own stream.

Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/py_executor.py      | 14 ++++++++++++--
 tests/unittest/_torch/executor/test_py_executor.py |  8 +++++---
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index ad23219f7ebf..232cec47f1a8 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -6567,7 +6567,7 @@ def _warmup_encoder_cuda_graphs_enc_dec(self) -> None:
             warmup(self.resource_manager)
 
     def _submit_encoder_step(self, encoder_requests: List[LlmRequest]) -> None:
-        """Queue encoder work, serializing it with decoder work under TP."""
+        """Queue encoder work, joining the launch before the caller runs the decoder."""
         executor = self.encoder_launch_executor
         if executor is None:
             raise RuntimeError("Encoder launch executor is unavailable.")
@@ -6603,8 +6603,18 @@ def _submit_encoder_step(self, encoder_requests: List[LlmRequest]) -> None:
                 self.inflight_req_ids.erase(request.request_id)
             return
 
+        # torch.fx's tracing state is process-global, so a dynamo compile on
+        # either thread captures the other's concurrent module calls
+        # (https://nvbugs/6683840). Join the launch, but leave ready_event
+        # unsynchronized so encoder kernels still overlap decoder work.
+        try:
+            result = future.result()
+        except Exception as e:
+            self._finish_failed_encoder_step(requests, e)
+            return
+
         self.pending_encoder_steps.append(
-            PendingEncoderStep(requests=requests, future=future))
+            PendingEncoderStep(requests=requests, future=future, result=result))
 
     @nvtx_range("_poll_encoder_steps")
     def _poll_encoder_steps(self) -> None:
diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py
index d5ce44a3bfce..b20af1b9d0a7 100644
--- a/tests/unittest/_torch/executor/test_py_executor.py
+++ b/tests/unittest/_torch/executor/test_py_executor.py
@@ -453,7 +453,6 @@ def test_async_encoder_step_lifecycle():
         ready_event=ready_event,
     )
     future = Mock()
-    future.done.side_effect = [False, True]
     future.result.return_value = result
     executor = _make_async_encoder_executor(future)
     active_request = types.SimpleNamespace(
@@ -473,10 +472,13 @@ def test_async_encoder_step_lifecycle():
         )
     )
 
+    # The launch is joined inline (https://nvbugs/6683840); publication stays
+    # asynchronous, waiting on ready_event rather than the future.
     executor._submit_encoder_step(requests)
-    executor._poll_encoder_steps()
 
-    future.result.assert_not_called()
+    future.result.assert_called_once_with()
+    future.done.assert_not_called()
+    ready_event.query.assert_not_called()
     executor._publish_encoder_step.assert_not_called()
     assert executor.inflight_req_ids.ids == {11, 12}
     assert len(executor.pending_encoder_steps) == 1