[https://nvbugs/6029220][fix] Disable multi-stream in maybe_execute_i…

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -1159,9 +1159,12 @@ def _compute_routed_output():
         # NOTE: define compiled helpers at module scope to avoid defining decorators inside compiled frames
 
         routed_output, shared_output = maybe_execute_in_parallel(
-            _compute_routed_output, _compute_shared_output,
+            _compute_routed_output,
+            _compute_shared_output,
             self.event_dict[EventType.Main],
-            self.event_dict[EventType.MoeShared], self.aux_stream)
+            self.event_dict[EventType.MoeShared],
+            self.aux_stream,
+            disable_on_compile=True)
 
         if not do_finalize:
             return [shared_output, *routed_output]
@@ -1644,6 +1647,7 @@ def norm_hidden():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream,
+            disable_on_compile=True,
         )
         hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1)
         # Split hidden_states columnwise based on TP

tensorrt_llm/_torch/models/modeling_exaone_moe.py

@@ -542,6 +542,7 @@ def norm_hidden():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream,
+            disable_on_compile=True,
         )
         hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1)
         # Split hidden_states columnwise based on TP

tensorrt_llm/_torch/models/modeling_glm.py

@@ -432,6 +432,7 @@ def _compute_routed_output():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream,
+            disable_on_compile=True,
         )
 
         if not do_finalize:
@@ -864,6 +865,7 @@ def norm_hidden():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream,
+            disable_on_compile=True,
         )
         hidden_states = torch.concat([inputs_embeds, hidden_states], dim=-1)
         # Split hidden_states columnwise based on TP

tensorrt_llm/_torch/models/modeling_llama.py

@@ -338,7 +338,12 @@ def forward(
         fn1 = lambda: self.compute_routed_output(
             hidden_states, all_rank_num_tokens, cutlass_min_latency_mode)
         shared_output, routed_output = maybe_execute_in_parallel(
-            fn0, fn1, self.moe_event[0], self.moe_event[1], self.aux_stream)
+            fn0,
+            fn1,
+            self.moe_event[0],
+            self.moe_event[1],
+            self.aux_stream,
+            disable_on_compile=True)
         if cutlass_min_latency_mode:
             return [shared_output, *routed_output]
 

tensorrt_llm/_torch/models/modeling_llama_min_latency.py

@@ -612,7 +612,12 @@ def forward(
         fn1 = lambda: self.compute_routed_output(
             hidden_states, all_rank_num_tokens, hidden_states_high)
         shared_output, routed_output = maybe_execute_in_parallel(
-            fn0, fn1, self.moe_event[0], self.moe_event[1], self.aux_stream)
+            fn0,
+            fn1,
+            self.moe_event[0],
+            self.moe_event[1],
+            self.aux_stream,
+            disable_on_compile=True)
 
         assert shared_output.size() == routed_output.size(
         ), f'unmatched tensor shape'

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -322,6 +322,7 @@ def _compute_routed_output():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream_shared,
+            disable_on_compile=True,
         )
 
         final_hidden_states = shared_output + routed_output

tensorrt_llm/_torch/models/modeling_qwen3_next.py

@@ -254,6 +254,7 @@ def _compute_shared_output():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.MoeShared],
             self.aux_stream,
+            disable_on_compile=True,
         )
         if not do_finalize:
             return final_hidden_states
@@ -808,6 +809,7 @@ def _compute_projected_states_ba():
             self.event_dict[EventType.Main],
             self.event_dict[EventType.Attention],
             self.aux_stream,
+            disable_on_compile=True,
         )
 
         # Use fused kernel when possible to avoid elementwise ops

tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py

@@ -421,10 +421,12 @@ def convert_dt():
             # convert will go second and we lose PDL, but we're using cuda
             # graphs for low latency so that seems ok.
             # If any of the contiguous calls below actually fire, that also breaks PDL.
-            xbc_d, dt_d = maybe_execute_in_parallel(conv1d, convert_dt,
+            xbc_d, dt_d = maybe_execute_in_parallel(conv1d,
+                                                    convert_dt,
                                                     self.events[0],
                                                     self.events[1],
-                                                    self.aux_steram)
+                                                    self.aux_steram,
+                                                    disable_on_compile=True)
 
             x_d, B_d, C_d = torch.split(
                 xbc_d,

tensorrt_llm/_torch/modules/multi_stream_utils.py

@@ -37,7 +37,8 @@ def maybe_execute_in_parallel(
         fn1: Callable,
         event0: torch.cuda.Event,
         event1: torch.cuda.Event,
-        aux_stream: Optional[torch.cuda.Stream] = None) -> tuple[Any, Any]:
+        aux_stream: Optional[torch.cuda.Stream] = None,
+        disable_on_compile: bool = False) -> tuple[Any, Any]:
     """Utility function to run two functions in two cuda streams in parallel. Multi-stream is
     only enabled when cuda graph is turned on because switch stream has extra host overhead.
 
@@ -52,12 +53,18 @@ def maybe_execute_in_parallel(
         event1 (torch.cuda.Event): cuda event for fn1
         aux_stream (Optional[torch.cuda.Stream]): the second cuda stream for fn1.
             Multi-stream is disabled when aux_stream is None.
+        disable_on_compile (bool): if True, disable multi-stream when
+            torch.compile is tracing. Callers that are not inside a custom op
+            should set this to True so that stream/event ops are not captured
+            by dynamo. Callers inside custom ops (e.g. attention, MoE) should
+            leave this as False since custom ops are opaque to the compiler.
 
     Returns:
         tuple[Any, Any]: the return values of fn0() and fn1()
     """
 
-    multi_stream = do_multi_stream() and aux_stream is not None
+    multi_stream = (do_multi_stream() and aux_stream is not None and
+                    not (disable_on_compile and torch.compiler.is_compiling()))
 
     if multi_stream:
         event0.record()