Graph break overhead #3946

cehongwang · 2025-12-03T19:18:02Z

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
This change requires a documentation update

Checklist:

My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

core/runtime/TRTEngine.cpp

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

github-actions

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/_compiler.py	2025-12-12 00:21:37.519598+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/_compiler.py	2025-12-12 00:22:13.058405+00:00
@@ -947,11 +947,11 @@
    # Here we delete the frozen parameters from the graph module. Note this does not affect the submodules. We are going to delete the frozen parameters from the submodules in the convert_module function.
    # This is done to release CPU memory.
    for attr in dir(gm):
        if attr.startswith("_frozen_param"):
            delattr(gm, attr)
-            
+
    trt_module = None
    for name, _ in partitioned_module.named_children():
        submodule = getattr(partitioned_module, name)
        # filter on the GraphModule
        if not isinstance(submodule, torch.fx.graph_module.GraphModule):

narendasan · 2025-12-16T21:08:34Z

core/runtime/execute_engine.cpp

-      TORCHTRT_CHECK(
-          compiled_engine->exec_ctx->setInputShape(name.c_str(), dims), "Error while setting the input shape");
-
+      if (shape_changed) {


Can't we use the same shape keys we used for cudagraphs? Why are we implementing another system?

narendasan · 2025-12-16T21:10:33Z

core/runtime/TRTEngine.cpp

    num_io = std::make_pair(inputs_size, outputs);
+
+    this->current_device_id = at::cuda::current_device();
+    this->stream = c10::cuda::getCurrentCUDAStream(this->current_device_id);


Is this PR supposed to implement the stream operator?

narendasan · 2025-12-16T21:13:24Z

py/torch_tensorrt/dynamo/_compiler.py

    for attr in dir(gm):
        if attr.startswith("_frozen_param"):
            delattr(gm, attr)
+    trt_module = None


What is this? Why does this variable need to be outside the of scope of iterating through the module?

py/torch_tensorrt/dynamo/_compiler.py

setup.py

py/torch_tensorrt/dynamo/runtime/_stream_handler.py

core/runtime/execute_engine.cpp

narendasan · 2025-12-29T18:01:09Z

core/runtime/execute_engine.cpp

+      setup_input_tensors(inputs, compiled_engine, cudagraphs_enabled, shape_changed);
      // Check if input shapes can be inferred.
-      int32_t const io_size{compiled_engine->cuda_engine->getNbIOTensors()};
+      int32_t const io_size{compiled_engine->io_size};


Are things like this actually faster? Like have we looked at the implementation to see its not cached at the TRT level?

Yes they do have small difference. I am trying to minize repeating code and this repeatly call this outbound function in every iteration

core/runtime/execute_engine.cpp

narendasan · 2025-12-29T18:05:35Z

core/runtime/TRTEngine.h

  std::string shape_key = "None";
  bool use_pre_allocated_outputs = false;
  std::vector<at::Tensor> pre_allocated_outputs;
+  std::vector<at::Tensor> allocated_outputs;


Is there a way to reuse pre_allocated_outputs?

@cehongwang we need to figure out how to merge these two implementations, we cant just keep adding new code paths for every situation. Either the output allocator needs to change to leverage the new system or the new system needs to leverage what the output allocator uses

narendasan · 2025-12-31T17:21:31Z

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

It seems like you are just renaming variables here? Why do we need to make that change?

the pre_allocated_outputs -> allocated_outputs change

I also changed self.output_tensors = self.create_output_tensors() to self.allocated_outputs = self.create_output_tensors(). Basically merged two variable together. Isn't that what you wanted me to do?

narendasan · 2025-12-31T17:24:12Z

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py

            ):
                if can_use_pre_allocated_outputs:
-                    outputs = self.pre_allocated_outputs
+                    outputs = self.allocated_outputs


meta-cla bot added the cla signed label Dec 3, 2025

github-actions bot added component: core Issues re: The core compiler component: build system Issues re: Build system component: api [Python] Issues re: Python API component: runtime component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Dec 3, 2025

github-actions bot requested a review from narendasan December 3, 2025 19:18

cehongwang force-pushed the graph-break-overhead branch from 8c9284a to 46829e6 Compare December 5, 2025 01:01

narendasan reviewed Dec 10, 2025

View reviewed changes

core/runtime/TRTEngine.cpp Outdated Show resolved Hide resolved

narendasan reviewed Dec 10, 2025

View reviewed changes

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py Outdated Show resolved Hide resolved

narendasan reviewed Dec 10, 2025

View reviewed changes

py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py Outdated Show resolved Hide resolved

cehongwang force-pushed the graph-break-overhead branch 2 times, most recently from 7541aca to 10662c0 Compare December 12, 2025 00:21

github-actions bot requested changes Dec 12, 2025

View reviewed changes

cehongwang force-pushed the graph-break-overhead branch 4 times, most recently from 4813804 to 4615705 Compare December 16, 2025 20:43