feat(executorch): expose composable Edge export API

[shoumikhin]

Why this PR is needed

torch_tensorrt.save(..., output_format="executorch") is the easiest way to create an ExecuTorch .pte file. It completes the whole export pipeline immediately:

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn(1, 3, 224, 224, device="cuda")]
trt_module = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=inputs)

torch_tensorrt.save(
    trt_module,
    "model.pte",
    output_format="executorch",
)

That remains the recommended API for common one-step exports.

Some workflows need to stop before final ExecuTorch memory planning and serialization. For example, users may want to:

• inspect the delegated Edge graph;
• run custom Edge transforms;
• send operations that TensorRT does not support to another backend;
• preserve independent methods such as prefill and decode;
• add constant methods or generate an ETRecord;
• choose the final ExecuTorch backend configuration.

This PR adds torch_tensorrt.executorch.export(). It returns ExecuTorch's standard EdgeProgramManager, which is the supported boundary for inspecting and customizing an Edge program before calling to_executorch().

Which API should I use?

Use torch_tensorrt.save() when you want a .pte in one step.

Use torch_tensorrt.executorch.export() when you need control before the final .pte is created:

import torch_tensorrt.executorch

edge = torch_tensorrt.executorch.export(trt_module)

# Inspect or transform the Edge program here.
print(edge.exported_program().graph)

program = edge.to_executorch()
with open("model.pte", "wb") as output:
    program.write_to_file(output)

# Advanced callers must also persist external tensor data when present.
program.write_tensor_data_to_file(".")

save() now delegates to the same implementation, so the simple and advanced paths do not maintain separate lowering logic.

Example: TensorRT with a fallback backend

TensorRT always receives the first chance to claim its prebuilt engine nodes. Additional partitioners run afterward and can claim operations that TensorRT does not support.

Assume cuda_partitioner is an ExecuTorch partitioner configured for your target:

edge = torch_tensorrt.executorch.export(
    trt_module,
    partitioners=[cuda_partitioner],
)

program = edge.to_executorch()

This can produce one ExecuTorch program containing both TensorRT and CUDA delegates.

Example: multiple methods

Assume prefill_program and decode_program are independent, engine-bearing ExportedProgram objects, and the two spec lists are configured for their methods:

edge = torch_tensorrt.executorch.export(
    {
        "prefill": prefill_program,
        "decode": decode_program,
    },
    compile_specs={
        "prefill": prefill_specs,
        "decode": decode_specs,
    },
)

program = edge.to_executorch()

The mapping preserves method names and allows per-method partitioner and compile-spec pipelines. It does not imply that mutable state is shared between methods.

Example: transforms, metadata, and ETRecord

Assume the transform passes and Edge configuration below are provided by the caller:

edge = torch_tensorrt.executorch.export(
    trt_module,
    transform_passes=my_edge_passes,
    compile_config=edge_compile_config,
    constant_methods={"get_vocab_size": 256},
    generate_etrecord=True,
)

program = edge.to_executorch()
program.get_etrecord().save("model_etrecord.bin")

Accepted inputs

The new API accepts:

• a TensorRT-compiled torch.fx.GraphModule;
• an engine-bearing torch.export.ExportedProgram;
• a mapping from method names to independent ExportedProgram objects.

A plain torch.nn.Module must be compiled with Torch-TensorRT first.

Safety and memory behavior

Engine rewriting changes the exported graph. The new API stages an independent copy of graph structure, signatures, state containers, constants, and metadata before rewriting anything. Rewrite and lowering failures therefore do not consume or structurally modify the caller's original program.

Tensor and TensorRT engine payload storage is treated as immutable and shared with the staged structure. This avoids deep-copying potentially multi-gigabyte engines. Custom transform passes must not mutate shared payload objects.

Metadata describing symbolic shapes is also shared rather than copied. It points back into the live shape environment the exported program is guarded on, and that environment holds fake tensors which cannot be copied. Sharing it is what keeps dynamic-shape models working; copying it would also detach the copy from the shape symbols the graph depends on.

When a lifted TensorRT engine is replaced, the old placeholder, graph-signature entry, and constant are removed. Multiple calls that share one engine reuse one materialized uint8 payload buffer.

Compatibility

• Existing one-step save(..., output_format="executorch") behavior is preserved.
• .pte serialization and external .ptd persistence are preserved.
• TensorRT remains the first partitioner.
• Shared dynamic dimensions declared with Input(shared_dims=...) remain shared.
• Zero-engine methods are allowed. A later partitioner may claim them, or portable operators may remain undelegated.

Validation

Built from source with the Torch-TensorRT C++ runtime active, against TensorRT and
ExecuTorch, and exercised on both x86_64 and aarch64 CUDA GPUs.

Checked end to end with a real compiled TensorRT engine:

• static-shape export() writes a loadable .pte;
• dynamic-shape export() writes a loadable .pte of the same size the previous
  one-step save path produced;
• the caller's GraphModule still holds its original engine nodes after export();
• engine payloads are shared rather than re-serialized;
• save(..., output_format="executorch") still writes a .pte;
• save() still raises TypeError for a non-list partitioners or compile_specs.

Unit and integration tests in the ExecuTorch test directory: 76 passed.

Two GPU tests in test_cuda_partitioner_composition.py,
test_erfinv_routes_to_cuda_backend and
test_weighted_partition_persists_external_data, fail with
SpecViolationError: Node.meta lowered_module_0 is missing val field. They fail the
same way on main without this change, so they are pre-existing and not caused by
this work.

Note that these tests only run in the L2 tier, which requires the ci: run-l2 label
on a pull request.

Static checks on the changed files: import ordering, formatting, lint, and type
checking.

Related work

PR #4433 forwards additional Edge-lowering options through one-step save(). After this PR lands, #4433 can keep save() as a thin wrapper over torch_tensorrt.executorch.export() instead of extending the previous private lowering path.

[lanluo-nvidia]

I would suggest here first add the torch_tensorrt.save() as the default method for user wants to export a executorch pte file.
then add the advanced method, when user needs a customization before serialize into a pte file.
So that we explain user both scenario clearly.

[lanluo-nvidia]

Also please wait for PR: 4398 to be merged to main(I will push by the end of today.)
please rebase on top of that.