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?

narendasan · 2025-12-16T21:16:40Z

py/torch_tensorrt/dynamo/_compiler.py

+            trt_module = getattr(partitioned_module, name)
+            trt_module.setup_engine()
+
+    if trt_module:


I feel like this is a pretty round about way to get the last engine

Also the order from items is not deterministic so I dont think this is guaranteed to give you the last module anyway. You might need to work back from outputs.

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

py/torch_tensorrt/dynamo/_compiler.py

        settings.use_fast_partitioner = True

    dryrun_stats_display(dryrun_tracker, settings.dryrun)
+    if not settings.dryrun:


Why is this conditioned on dry run? Dry run should reflect the deployment graph

narendasan · 2025-12-16T21:21:27Z

setup.py

-        cmd.append("--compilation_mode=dbg")
-    else:
-        cmd.append("--compilation_mode=opt")
+    # if develop:


narendasan · 2025-12-16T21:27:11Z

py/torch_tensorrt/dynamo/runtime/_stream_handler.py

+
+
+@torch.library.custom_op("tensorrt::enter_compute_stream", mutates_args=())
+def enter_compute_stream() -> None:


IMO these should likely be dependent on the input tensors and pull their stream information from there. You arent necessarily guaranteed for the input tensors to be on the default stream. It also means you need storage for the streams.

Also how does this work in c++?

narendasan · 2025-12-16T21:28:04Z

py/torch_tensorrt/dynamo/runtime/_stream_handler.py

+    partitioned_module: torch.fx.GraphModule,
+) -> torch.fx.GraphModule:
+    for node in partitioned_module.graph.nodes:
+        if node.op == "placeholder":


what if there are multiple placeholders or 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 from 9248896 to 1917ce3 Compare December 11, 2025 00:11

cehongwang added 2 commits December 12, 2025 00:21

Rebased the commits

4bc3ce1

Added perf script

0b42683

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

github-actions bot requested changes Dec 12, 2025

View reviewed changes

Renamed to be less confusing

29df36d

cehongwang force-pushed the graph-break-overhead branch from 10662c0 to 29df36d Compare December 12, 2025 00:24

Fixed bugs

b72a538

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

narendasan reviewed Dec 16, 2025

View reviewed changes

setup.py Outdated

cmd.append("--compilation_mode=dbg")

else:

cmd.append("--compilation_mode=opt")

# if develop:

Copy link

Collaborator

narendasan Dec 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove

narendasan reviewed Dec 16, 2025

View reviewed changes

Changed the stream handling

c611e5f

cehongwang force-pushed the graph-break-overhead branch from 4615705 to c611e5f Compare December 17, 2025 01:10

github-actions bot removed the component: build system Issues re: Build system label Dec 17, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Graph break overhead #3946

Graph break overhead #3946

Uh oh!

cehongwang commented Dec 3, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot left a comment

Uh oh!

narendasan Dec 16, 2025

Uh oh!

narendasan Dec 16, 2025

Uh oh!

narendasan Dec 16, 2025

Uh oh!

narendasan Dec 16, 2025

Uh oh!

narendasan Dec 16, 2025

Uh oh!

narendasan Dec 16, 2025

Uh oh!

narendasan Dec 16, 2025

Uh oh!

narendasan Dec 16, 2025

Uh oh!

narendasan Dec 16, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants



		@torch.library.custom_op("tensorrt::enter_compute_stream", mutates_args=())
		def enter_compute_stream() -> None:

Graph break overhead #3946

Are you sure you want to change the base?

Graph break overhead #3946

Uh oh!

Conversation

cehongwang commented Dec 3, 2025

Description

Type of change

Checklist:

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants