Enabled Cuda Graph

Author: cehongwang

examples/apps/flux-demo.py

@@ -95,8 +95,8 @@ def load_lora(path):
 
 
 generate_image(["Test"], 2)
-load_lora("")
-generate_image(["A golden retriever holding a sign to code"], 2)
+# load_lora("")
+# generate_image(["A golden retriever holding a sign to code"], 2)
 
 # Create Gradio interface
 with gr.Blocks(title="Flux Demo with Torch-TensorRT") as demo:

py/torch_tensorrt/dynamo/runtime/_MutableTorchTensorRTModule.py

@@ -16,6 +16,7 @@
     to_torch_device,
     to_torch_tensorrt_device,
 )
+from torch_tensorrt.runtime._cudagraphs import get_cuda_graph_module
 
 logger = logging.getLogger(__name__)
 
@@ -335,14 +336,16 @@ def compile(self) -> None:
         )
         self.original_model.to("cpu")
         torch.cuda.empty_cache()
-        # torch_tensorrt.runtime.set_cudagraphs_mode(self.enable_cuda_graph)
-        # if self.enable_cuda_graph:
-        #     self.gm = torch_tensorrt.runtime.enable_cudagraphs(self.gm)
+        if self.enable_cuda_graph:
+            self._enable_cuda_graph()
         if self.enable_weight_streaming:
             self.weight_streaming_ctx = torch_tensorrt.runtime.weight_streaming(self.gm)
             requested_budget = int(16 * 2 << 20)
             self.weight_streaming_ctx.device_budget = requested_budget
 
+    def _enable_cuda_graph(self) -> None:
+        self.gm = get_cuda_graph_module(self.gm)
+
     def _validate_inputs(self, *args: Any, **kwargs: Any) -> None:
 
         if not self.arg_inputs and not self.kwarg_inputs:

py/torch_tensorrt/runtime/_cudagraphs.py

@@ -69,52 +69,58 @@ def __init__(self, compiled_module: torch.nn.Module) -> None:
         self.old_mode = _PY_RT_CUDAGRAPHS
         self.compiled_module = compiled_module
 
-    def __enter__(self) -> torch.nn.Module:
-        global _PY_RT_CUDAGRAPHS
-
-        num_torch_module = 0
-        num_trt_module = 0
-        for name, module in self.compiled_module.named_children():
-            # need to disable cudagraphs if any model requires output allocator
-            if (
-                hasattr(module, "requires_output_allocator")
-                and module.requires_output_allocator
-            ):
-                raise RuntimeError(
-                    "The model contains submodules that require a dynamic output allocator at runtime, which is incompatible with CUDA Graphs. Please disable CUDA Graphs."
-                )
-            if "_run_on_acc" in name:
-                num_trt_module += 1
-            elif "_run_on_gpu" in name:
-                num_torch_module += 1
-
-        if num_torch_module > 0:
-            # Set whole cudagraphs mode and returns wrapped module
-            _PY_RT_CUDAGRAPHS = CudaGraphsMode.WHOLE_GRAPH_CUDAGRAPHS
-            # Set new mode for C++
-            if torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime:
-                torch.ops.tensorrt.set_cudagraphs_mode(_PY_RT_CUDAGRAPHS)
-
-            logger.debug(
-                "Found pytorch subgraphs in module, wrapping module in CudaGraphsTorchTensorRTModule"
-            )
-            return CudaGraphsTorchTensorRTModule(self.compiled_module)
-        else:
-            if num_trt_module > 0:
-                logger.debug("No graph breaks detected, using runtime cudagraphs mode")
-            else:
-                logger.debug(
-                    "Please consider dynamo if there is graph breaks. Using runtime cudagraphs mode"
-                )
-            # Enable cudagraphs for TRT submodule
-            set_cudagraphs_mode(True)
-            return self.compiled_module
+    def __enter__(self) -> torch.nn.Module | torch.fx.GraphModule:
+        return get_cuda_graph_module(self.compiled_module)
 
     def __exit__(self, *args: Any) -> None:
         # Set cudagraphs back to old mode
         set_cudagraphs_mode(self.old_mode)
 
 
+def get_cuda_graph_module(
+    compiled_module: torch.fx.GraphModule,
+) -> torch.nn.Module | torch.fx.GraphModule:
+    global _PY_RT_CUDAGRAPHS
+
+    num_torch_module = 0
+    num_trt_module = 0
+    for name, module in compiled_module.named_children():
+        # need to disable cudagraphs if any model requires output allocator
+        if (
+            hasattr(module, "requires_output_allocator")
+            and module.requires_output_allocator
+        ):
+            raise RuntimeError(
+                "The model contains submodules that require a dynamic output allocator at runtime, which is incompatible with CUDA Graphs. Please disable CUDA Graphs."
+            )
+        if "_run_on_acc" in name:
+            num_trt_module += 1
+        elif "_run_on_gpu" in name:
+            num_torch_module += 1
+
+    if num_torch_module > 0:
+        # Set whole cudagraphs mode and returns wrapped module
+        _PY_RT_CUDAGRAPHS = CudaGraphsMode.WHOLE_GRAPH_CUDAGRAPHS
+        # Set new mode for C++
+        if torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime:
+            torch.ops.tensorrt.set_cudagraphs_mode(_PY_RT_CUDAGRAPHS)
+
+        logger.debug(
+            "Found pytorch subgraphs in module, wrapping module in CudaGraphsTorchTensorRTModule"
+        )
+        return CudaGraphsTorchTensorRTModule(compiled_module)
+    else:
+        if num_trt_module > 0:
+            logger.debug("No graph breaks detected, using runtime cudagraphs mode")
+        else:
+            logger.debug(
+                "Please consider dynamo if there is graph breaks. Using runtime cudagraphs mode"
+            )
+        # Enable cudagraphs for TRT submodule
+        set_cudagraphs_mode(True)
+        return compiled_module
+
+
 def enable_cudagraphs(
     compiled_module: Union[torch.fx.GraphModule, torch.nn.Module],
 ) -> _CudagraphsContextManager: