Skip to content

feat: custom binding names for out of runtime deployment#4309

Open
narendasan wants to merge 1 commit into
mainfrom
narendasan/custom_binding_names
Open

feat: custom binding names for out of runtime deployment#4309
narendasan wants to merge 1 commit into
mainfrom
narendasan/custom_binding_names

Conversation

@narendasan

Copy link
Copy Markdown
Collaborator

Description

Allows users of the convert_exported_program_to_serialized_engine API to specify binding names for all bindings utilizing the exported programs pytrees for i/o flattening to flatten the provided names and associate names to specific io positions.

TODO: Add support for aliasing inputs and outputs through this method (see: #4251)

Fixes #4212

Type of change

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

  • New feature (non-breaking change which adds functionality)
  • 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

@meta-cla meta-cla Bot added the cla signed label May 29, 2026
@github-actions github-actions Bot added documentation Improvements or additions to documentation component: tests Issues re: Tests component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels May 29, 2026
@github-actions github-actions Bot requested a review from zewenli98 May 29, 2026 19:38

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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

--- /home/runner/work/TensorRT/TensorRT/examples/dynamo/engine_converter_binding_names.py	2026-05-29 19:38:55.472711+00:00
+++ /home/runner/work/TensorRT/TensorRT/examples/dynamo/engine_converter_binding_names.py	2026-05-29 19:39:15.963114+00:00
@@ -302,11 +302,11 @@

try:
    torch_tensorrt.dynamo.convert_exported_program_to_serialized_trt_engine(
        exported,
        arg_inputs=(x,),
-        output_binding_names=("only_one",),   # wrong arity for the 2-tuple return
+        output_binding_names=("only_one",),  # wrong arity for the 2-tuple return
        require_full_compilation=True,
        min_block_size=1,
        use_python_runtime=False,
        immutable_weights=True,
    )
--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py	2026-05-29 19:38:55.486463+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py	2026-05-29 19:39:17.826414+00:00
@@ -397,11 +397,10 @@
        run_module_start_time = datetime.now()
        super().run()
        _LOGGER.info(
            f"TRT INetwork construction elapsed time: {datetime.now() - run_module_start_time}"
        )
-

    @staticmethod
    def find_weight(
        weight_name: str,
        weight_refit_map: dict[str, Any],
--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/conversion/test_engine_converter_binding_names.py	2026-05-29 19:38:55.514192+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/conversion/test_engine_converter_binding_names.py	2026-05-29 19:39:21.059130+00:00
@@ -12,10 +12,11 @@
We ``tree_flatten`` each and verify spec equality against the exported
program's specs.  On success the flat lists concatenate (args then
kwargs) into FX's flattened placeholder order; the interpreter just
indexes positionally.
"""
+
import unittest

import tensorrt as trt
import torch
import torch_tensorrt
@@ -361,10 +362,12 @@
        (input_name,) = self._binding_names(engine, trt.TensorIOMode.INPUT)
        outs = self._run_engine(engine, {input_name: x.contiguous()})
        with torch.no_grad():
            ref = model(x)
        torch.testing.assert_close(outs["p_out"], ref["primary"], rtol=1e-2, atol=1e-2)
-        torch.testing.assert_close(outs["s_out"], ref["secondary"], rtol=1e-2, atol=1e-2)
+        torch.testing.assert_close(
+            outs["s_out"], ref["secondary"], rtol=1e-2, atol=1e-2
+        )


if __name__ == "__main__":
    run_tests()

@narendasan narendasan force-pushed the narendasan/custom_binding_names branch from 62c562d to eb1ffee Compare May 29, 2026 22:39
@narendasan narendasan force-pushed the narendasan/custom_binding_names branch from eb1ffee to 4c787ab Compare June 10, 2026 16:43
@lanluo-nvidia lanluo-nvidia self-requested a review June 10, 2026 18:47

@lanluo-nvidia lanluo-nvidia left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: conversion Issues re: Conversion stage component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: tests Issues re: Tests documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

❓ [Question] How can you set the output names when converting to serialized engine?

2 participants