From 20cd6d32d0005c2449f8ee4c63c9a9cb886b57f9 Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Wed, 1 Apr 2026 07:14:08 -0700 Subject: [PATCH] [https://nvbugs/6029220][fix] Disable multi-stream in maybe_execute_in_parallel under torch.compile PyTorch 2.11 has a bug (pytorch/pytorch#176486) where dynamo captures CUDA stream/event operations and converts them into torch.ops.streams.* nodes. At runtime these nodes create events with uninitialized device type (CPU), causing "Event device type CPU does not match recording stream's device type CUDA" errors. Add a `disable_on_compile` flag to `maybe_execute_in_parallel`. Callers outside custom ops (model-level MoE, MTP norm, mamba mixer, etc.) set it to True so stream/event ops are not captured by dynamo. Callers inside custom ops (fused MoE, attention) leave it False since custom ops are opaque to the compiler. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_deepseekv3.py | 8 ++++++-- tensorrt_llm/_torch/models/modeling_exaone_moe.py | 1 + tensorrt_llm/_torch/models/modeling_glm.py | 2 ++ tensorrt_llm/_torch/models/modeling_llama.py | 7 ++++++- .../_torch/models/modeling_llama_min_latency.py | 7 ++++++- tensorrt_llm/_torch/models/modeling_nemotron_h.py | 1 + tensorrt_llm/_torch/models/modeling_qwen3_next.py | 2 ++ tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py | 6 ++++-- tensorrt_llm/_torch/modules/multi_stream_utils.py | 11 +++++++++-- 9 files changed, 37 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 138c6b8af6d7..8b2da5404489 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_exaone_moe.py b/tensorrt_llm/_torch/models/modeling_exaone_moe.py index fe420178558e..5a796f926c4d 100644 --- a/tensorrt_llm/_torch/models/modeling_exaone_moe.py +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_glm.py b/tensorrt_llm/_torch/models/modeling_glm.py index 3698c273afd6..c6a635d66b3e 100644 --- a/tensorrt_llm/_torch/models/modeling_glm.py +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 743e0b8ef502..aab9d3992c66 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/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] diff --git a/tensorrt_llm/_torch/models/modeling_llama_min_latency.py b/tensorrt_llm/_torch/models/modeling_llama_min_latency.py index ae3d1601fbe5..a0c45ca11277 100644 --- a/tensorrt_llm/_torch/models/modeling_llama_min_latency.py +++ b/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' diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index ca19d1da4fd3..e416edd450c9 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 68e9c5a70f5f..cf17233cf075 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/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 diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index 42c1df6a5d4f..e0028af9b26f 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/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, diff --git a/tensorrt_llm/_torch/modules/multi_stream_utils.py b/tensorrt_llm/_torch/modules/multi_stream_utils.py index c7b58c0896bb..f3e60a208646 100644 --- a/tensorrt_llm/_torch/modules/multi_stream_utils.py +++ b/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()