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.

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

Validated with Torch, TensorRT, and ExecuTorch available:

• 31 composable export API unit tests;
• 16 public API and optional-dependency tests;
• 5 one-step save compatibility tests;
• 5 GPU composition tests covering:
  • TensorRT before CUDA fallback;
  • .pte serialization;
  • non-empty .ptd persistence;
  • TensorRT-only export without .ptd;
  • independent prefill and decode methods retaining TensorRT delegates;
  • successful real exports leaving compiled GraphModule and lifted ExportedProgram sources unchanged.

Additional checks:

• Python type checking;
• formatting and lint;
• Python syntax compilation;
• git diff --check;
• independent code review.

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.