feat(executorch): copy-back for non-KV mutable buffers in the TensorRT delegate (#4459)

Co-authored-by: Anthony Shoumikhin <shoumikhin@meta.com>
Co-authored-by: guac e2e <guac@localhost>
Co-authored-by: x <x>

Author: 3 people

py/torch_tensorrt/_compile.py

@@ -1148,6 +1148,9 @@ def _extract_tensor(obj: Any) -> Any:
             # whichever exporter produced the program. aot_inductor is left
             # undeclared: whether an aliased in-place mutation survives
             # functionalization under inductor is unverified.
+            #
+            # Copy-back is gated per-branch below: the legacy exporter declares it
+            # itself and consumes the trailing values while doing so.
             if not retrace:
                 from torch_tensorrt.dynamo._exporter import export
 
@@ -1173,6 +1176,22 @@ def _extract_tensor(obj: Any) -> Any:
                     _declare_aliased_kv_mutations_on_ep,
                 )
 
+                # On this path the legacy exporter is what declares copy-back, except
+                # under output_format="executorch", where torch_tensorrt.executorch
+                # .export() declares it for whatever source shape it is handed. The
+                # remaining combination leaves it undeclared.
+                if (
+                    not _use_legacy
+                    and output_format != "executorch"
+                    and module.meta.get("_copyback_mutation_buffers")
+                ):
+                    logger.warning(
+                        "Module has non-KV mutable buffer(s) needing copy-back, but "
+                        "retrace=False with use_legacy_exporter=False does not declare "
+                        "them. The saved program's signature will not reflect those "
+                        "updates. Use the legacy exporter, or retrace=True."
+                    )
+
                 if output_format == "exported_program":
                     # Must precede normalization, which rewrites the engine constants
                     # this pass reads aliased_io from.
@@ -1235,6 +1254,7 @@ def _extract_tensor(obj: Any) -> Any:
                     if node.op == "placeholder" and "val" in node.meta
                     for dim in getattr(node.meta["val"], "shape", [])
                 )
+                _use_legacy = False
                 if has_symbolic_metadata and dynamic_shapes is not None:
                     from torch_tensorrt.dynamo._exporter import export
 
@@ -1283,6 +1303,16 @@ def _extract_tensor(obj: Any) -> Any:
                     _declare_aliased_kv_mutations_on_ep,
                 )
 
+                # create_trt_exp_program already declares the copy-back mutations and
+                # consumes their trailing outputs, so what trails the graph on that path
+                # is a genuine user output. Only the torch.export paths leave the values
+                # to be reclassified.
+                _copyback_bufs = (
+                    []
+                    if _use_legacy
+                    else module.meta.get("_copyback_mutation_buffers", [])
+                )
+
                 if output_format == "aot_inductor" and any(
                     getattr(sub, "aliased_io", None)
                     for _sub_name, sub in module.named_modules()
@@ -1297,7 +1327,9 @@ def _extract_tensor(obj: Any) -> Any:
                 if output_format == "exported_program":
                     # Must precede normalization, which rewrites the engine constants
                     # this pass reads aliased_io from.
-                    exp_program = _declare_aliased_kv_mutations_on_ep(exp_program)
+                    exp_program = _declare_aliased_kv_mutations_on_ep(
+                        exp_program, copyback_buffers=_copyback_bufs
+                    )
                     _normalize_engine_constants_to_python(exp_program)
                     function_overload_with_kwargs(
                         torch.export.save,
@@ -1318,7 +1350,9 @@ def _extract_tensor(obj: Any) -> Any:
                         package_path=file_path,
                     )
                 elif output_format == "executorch":
-                    exp_program = _declare_aliased_kv_mutations_on_ep(exp_program)
+                    exp_program = _declare_aliased_kv_mutations_on_ep(
+                        exp_program, copyback_buffers=_copyback_bufs
+                    )
                     _save_as_executorch(
                         exp_program,
                         file_path,

py/torch_tensorrt/dynamo/_compiler.py

@@ -45,6 +45,10 @@
     pre_export_lowering,
 )
 from torch_tensorrt.dynamo.lowering._buffer_lifting import (
+    aliased_input_bindings,
+    assert_no_kv_alias_markers_survived,
+    assert_predicted_kv_aliased,
+    hide_copyback_outputs,
     inline_lifted_buffers_into_gm,
     lift_mutated_buffers,
 )
@@ -815,7 +819,11 @@ def compile(
     # prerequisite for IKVCacheUpdateLayer / aliased I/O to fire on a
     # module-held cache). Returns a fresh GraphModule whose forward signature
     # reflects the new placeholders.
-    gm, lifted_buffers = lift_mutated_buffers(gm)
+    gm, lifted_buffers = lift_mutated_buffers(gm, settings)
+    _copyback_mutation_buffers = gm.meta.get("_copyback_mutation_buffers", [])
+    _copyback_bindings = gm.meta.get("_copyback_bindings", [])
+    _predicted_kv_bindings = gm.meta.get("_predicted_kv_bindings", [])
+    _no_kv_alias_writes = gm.meta.get("_no_kv_alias_writes", [])
     if lifted_buffers:
         # Append each lifted buffer as an engine input AFTER the user inputs.
         # Buffer tensors live on the gm's state; prepare an Input spec for
@@ -835,6 +843,12 @@ def compile(
     logger.debug(f"CPU memory usage after post_lowering: {get_cpu_memory_usage()} MB")
     logger.debug("Lowered Input graph: " + str(gm.graph))
 
+    # The marker is what keeps a re-routed write out of the engine's aliased_io, and
+    # it rides on node.meta through every pass above. Read it back before conversion
+    # so a pass that dropped it is named here rather than surfacing as a module that
+    # raises on every call.
+    assert_no_kv_alias_markers_survived(gm, _no_kv_alias_writes)
+
     # Move the weights in the state_dict to CPU
     if offload_module_to_cpu:
         deallocate_module(gm)
@@ -856,6 +870,23 @@ def compile(
         engine_cache,
         graph_signature=exported_program.graph_signature,
     )
+    if _copyback_mutation_buffers:
+        trt_gm.meta["_copyback_mutation_buffers"] = _copyback_mutation_buffers
+    # Ground-truth check on both directions of the classification: every write lift
+    # called KV (engine-aliased, so its copy_ was dropped) must appear in a compiled
+    # engine's aliased_io or its write-back is silently lost, and no write lift
+    # called copy-back may appear there or the trailing output it added is one the
+    # runtime truncates while the graph still reads it. Fail loudly for either. A
+    # dryrun returns before conversion, so there is no ground truth to check against.
+    assert_predicted_kv_aliased(
+        aliased_input_bindings(
+            getattr(sub, "aliased_io", None) for _name, sub in trt_gm.named_children()
+        ),
+        _predicted_kv_bindings,
+        settings,
+        engines_built=not settings.dryrun,
+        copyback_bindings=_copyback_bindings,
+    )
     if lifted_buffers:
         # Inline buffers into the compiled gm as get_attr nodes + registered
         # buffers. The resulting gm's forward takes only user inputs; buffers
@@ -864,6 +895,12 @@ def compile(
         # serializable by torch_tensorrt.save / torch.export (no external
         # Python wrapper that would be lost on a round-trip).
         trt_gm = inline_lifted_buffers_into_gm(trt_gm, lifted_buffers)
+    # A copy-back value stays on the output node, where the exporters read it and
+    # where dead-code elimination cannot reach it, but nothing between here and the
+    # ExecuTorch runtime writes it into the buffer. Returning it would report a
+    # mutation the caller cannot act on, so the compiled module keeps the arity of
+    # the model it was compiled from.
+    hide_copyback_outputs(trt_gm, len(_copyback_mutation_buffers))
     return trt_gm
 
 
@@ -1803,6 +1840,12 @@ def convert_exported_program_to_serialized_trt_engine(
     automatically; this lower-level entry point exposes the same machinery
     for callers that want to manage the bindings themselves.
 
+    Only writes the engine can alias in place are supported here. A mutation the
+    engine cannot alias needs its new value copied back into the buffer after the
+    call, and this entry point reports neither which output carries which buffer nor
+    performs the copy, so it raises rather than returning an engine whose buffer
+    would never update. Use :func:`torch_tensorrt.dynamo.compile` for those models.
+
     Arguments:
         exported_program (torch.export.ExportedProgram): Source module, running torch.export on a ``torch.nn.Module``
         inputs (Optional[Sequence[Sequence[Any]]]): List of specifications of input shape, dtype and memory layout for inputs to the module. This argument is required. Input Sizes can be specified as torch sizes, tuples or lists. dtypes can be specified using
@@ -1886,6 +1929,11 @@ def convert_exported_program_to_serialized_trt_engine(
         **kwargs: Any,
     Returns:
         bytes: Serialized TensorRT engine, can either be saved to a file or deserialized via TensorRT APIs
+    Raises:
+        RuntimeError: if ``lift_mutable_buffers=True`` and the model mutates a buffer
+            the engine cannot alias in place, or if a write predicted to be aliased is
+            absent from the built engine's ``aliased_io``. Either way the buffer would
+            silently never update.
     """
 
     if kwargs.get("debug", False):
@@ -2074,8 +2122,26 @@ def convert_exported_program_to_serialized_trt_engine(
     # resulting bindings at runtime — they are appended after the user inputs
     # in the order returned here.
     lifted_buffers: List[Tuple[str, str, torch.Tensor]] = []
+    predicted_kv_bindings: List[str] = []
     if lift_mutable_buffers:
-        gm, lifted_buffers = lift_mutated_buffers(gm)
+        gm, lifted_buffers = lift_mutated_buffers(gm, settings)
+        # Read before lowering: `gm` is replaced below and the meta does not follow it.
+        predicted_kv_bindings = gm.meta.get("_predicted_kv_bindings", [])
+        # A write the engine cannot alias in place is classified as copy-back: its
+        # new value is appended as an extra engine output, and whatever loads the
+        # serialized program copies that output into the buffer afterwards. This
+        # entry point returns engine bytes and no program, so it neither reports
+        # which output carries which buffer nor gives the copy anywhere to be
+        # declared, and the caller has no way to complete it.
+        copyback_buffers = gm.meta.get("_copyback_mutation_buffers", [])
+        if copyback_buffers:
+            raise RuntimeError(
+                "convert_exported_program_to_serialized_trt_engine cannot express the "
+                f"write-back for mutable buffer(s) {copyback_buffers}: the engine "
+                "cannot alias them in place, so the buffers would never update. Use "
+                "torch_tensorrt.dynamo.compile and save the result, which declares "
+                "the buffer mutation for the runtime to apply."
+            )
         if lifted_buffers:
             buffer_tensors = [t for _, _, t in lifted_buffers]
             buffer_inputs = prepare_inputs(buffer_tensors)
@@ -2176,6 +2242,12 @@ def convert_exported_program_to_serialized_trt_engine(
         )
         raise RuntimeError(f"While interpreting the module got an error: {e}") from e
 
+    assert_predicted_kv_aliased(
+        aliased_input_bindings([interpreter_result.aliased_io]),
+        predicted_kv_bindings,
+        settings,
+    )
+
     serialized_engine: bytes = interpreter_result.serialized_engine
     return serialized_engine