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

narendasan · 2025-12-19T21:14:51Z

core/runtime/execute_engine.cpp

      at::Tensor contig_input = inputs[i].view(shape).contiguous();
      formatted_inputs.emplace_back(std::move(contig_input));
-
+      bool need_cudagraphs_record = cudagraphs_enabled and


Just use &&, ||, !

narendasan · 2025-12-19T21:17:08Z

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 or compiled_engine->allocated_outputs.size() == 0) {


What is the allocated_outputs size test? Just that we havent allocated the outputs yet?

Can we fold some of these conditions ahead of time?

So it's like for the first time you run the model or when the input shape get changed, we need to (re)set the input shape right? I don't think we can fold it

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