fix(executorch): support KV-cache aliased I/O in the TensorRT delegate

cpp/include/torch_tensorrt/executorch/TensorRTBackend.h

@@ -59,6 +59,16 @@ struct EngineHandle {
   std::vector<size_t> cached_output_sizes;
   size_t num_inputs = 0;
   size_t num_outputs = 0;
+  // Per output binding [0..num_outputs): index into input_binding_names of the
+  // input it aliases (in-place KV-cache / user alias), or -1 for a normal output.
+  // Built at init from the blob's aliased_io. The KV buffers are threaded by
+  // ExecuTorch as caller-owned mutable-buffer delegate args (input AND aliased
+  // output): execute() binds each aliased TRT output binding to its aliased
+  // input's caller-provided pointer (in-place) and reflects the result into the
+  // delegate output EValue (a no-op when the memory planner already aliased the
+  // two -> zero-copy).
+  std::vector<int> output_aliased_input_idx;
+  size_t num_aliased_outputs = 0;
   int device_id = 0;
   bool unified_memory = false;
   std::mutex mu;

cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h

@@ -8,13 +8,24 @@
 namespace torch_tensorrt {
 namespace executorch_backend {
 
+// One aliased output->input binding pair (KV-cache in-place update, or a
+// user-declared alias). The engine's output binding shares device memory with
+// the named input binding; the runtime binds the output to the input's tensor
+// so the update lands in-place in the caller-owned buffer.
+struct AliasedBinding {
+  std::string output; // output binding name
+  std::string input; // input binding name it aliases
+  std::string kind; // "kv_cache_update" (TRT-enforced) or "user"
+};
+
 struct TensorRTBlobHeader {
   uint32_t metadata_offset = 0;
   uint32_t metadata_size = 0;
   uint32_t engine_offset = 0;
   uint64_t engine_size = 0;
   std::vector<std::string> input_binding_names;
   std::vector<std::string> output_binding_names;
+  std::vector<AliasedBinding> aliased_io;
   bool hardware_compatible = false;
   int device_id = 0;
 

cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp

@@ -14,6 +14,7 @@
 #include <memory>
 #include <mutex>
 #include <string>
+#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -289,6 +290,64 @@ Result<DelegateHandle*> TensorRTBackend::init(
     return err;
   }
 
+  // Map each aliased output binding to the index of the input it aliases so
+  // execute() can bind it to that input's device pointer (in-place).
+  // Non-aliased models have an empty header.aliased_io -> all -1, unchanged path.
+  handle->output_aliased_input_idx.assign(handle->num_outputs, -1);
+  for (const auto& ab : header.aliased_io) {
+    int oi = -1;
+    for (size_t k = 0; k < handle->output_binding_names.size(); ++k) {
+      if (handle->output_binding_names[k] == ab.output) {
+        oi = static_cast<int>(k);
+        break;
+      }
+    }
+    int ii = -1;
+    for (size_t k = 0; k < handle->input_binding_names.size(); ++k) {
+      if (handle->input_binding_names[k] == ab.input) {
+        ii = static_cast<int>(k);
+        break;
+      }
+    }
+    if (oi < 0 || ii < 0) {
+      ET_LOG(
+          Error,
+          "TensorRTBackend::init: aliased_io names not found (output='%s', input='%s')",
+          ab.output.c_str(),
+          ab.input.c_str());
+      return Error::InvalidProgram;
+    }
+    // AliasKind::USER aliases are not shape-enforced by TensorRT (unlike
+    // kv_cache_update, which IKVCacheUpdateLayer guarantees), so confirm the
+    // aliased output and input share a shape before binding them to the same
+    // storage.
+    if (ab.kind == "user") {
+      const nvinfer1::Dims od = handle->engine->getTensorShape(ab.output.c_str());
+      const nvinfer1::Dims id = handle->engine->getTensorShape(ab.input.c_str());
+      bool compatible = od.nbDims == id.nbDims;
+      for (int d = 0; compatible && d < od.nbDims; ++d) {
+        compatible = od.d[d] == id.d[d];
+      }
+      if (!compatible) {
+        ET_LOG(
+            Error,
+            "TensorRTBackend::init: user alias output '%s' shape is incompatible with input '%s'",
+            ab.output.c_str(),
+            ab.input.c_str());
+        return Error::InvalidProgram;
+      }
+    }
+    handle->output_aliased_input_idx[static_cast<size_t>(oi)] = ii;
+    ++handle->num_aliased_outputs;
+  }
+
+  if (handle->num_aliased_outputs > 0) {
+    ET_LOG(
+        Info,
+        "TensorRTBackend::init: %zu aliased output(s) bound in-place to caller-owned inputs",
+        handle->num_aliased_outputs);
+  }
+
   err = initialize_input_profiles(*handle);
   if (err != Error::Ok) {
     return err;
@@ -325,9 +384,17 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
 
   const size_t num_inputs = engine->num_inputs;
   const size_t num_outputs = engine->num_outputs;
-  if (args.size() < num_inputs + num_outputs) {
+  // Caller-owned KV: every input is a delegate arg, and each aliased output is
+  // threaded as a delegate output arg (the caller-owned mutable buffer's mutation
+  // slot), so all engine bindings map 1:1 to delegate args.
+  const size_t num_delegate_outputs = num_outputs;
+  const size_t num_delegate_inputs = num_inputs;
+  if (args.size() < num_delegate_inputs + num_delegate_outputs) {
     ET_LOG(
-        Error, "TensorRTBackend::execute: expected at least %zu args, got %zu", num_inputs + num_outputs, args.size());
+        Error,
+        "TensorRTBackend::execute: expected at least %zu args, got %zu",
+        num_delegate_inputs + num_delegate_outputs,
+        args.size());
     return Error::InvalidArgument;
   }
 
@@ -395,16 +462,22 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
   // ------------------------------------------------------------------
   // 1. Bind input shapes and addresses
   // ------------------------------------------------------------------
+  // Device pointer each input binding was bound to; aliased outputs reuse the
+  // pointer of the input they alias so their update lands in-place.
+  std::vector<void*> input_bind_ptrs(num_inputs, nullptr);
+  size_t arg_idx = 0; // running index into delegate args
   for (size_t i = 0; i < num_inputs; ++i) {
-    EValue* arg = args[i];
-    TORCHTRT_ET_CHECK_NOT_NULL(arg, Error::InvalidArgument, "TensorRTBackend::execute: input %zu is not a tensor", i);
+    const std::string& name = engine->input_binding_names[i];
+
+    EValue* arg = args[arg_idx++];
+    TORCHTRT_ET_CHECK_NOT_NULL(
+        arg, Error::InvalidArgument, "TensorRTBackend::execute: input arg %zu is not a tensor", i);
     if (!arg->isTensor()) {
       ET_LOG(Error, "TensorRTBackend::execute: input %zu is not a tensor", i);
       return Error::InvalidArgument;
     }
 
     exec_aten::Tensor et_in = arg->toTensor();
-    const std::string& name = engine->input_binding_names[i];
     nvinfer1::Dims dims = to_trt_dims(et_in);
     if (dims.nbDims > nvinfer1::Dims::MAX_DIMS) {
       ET_LOG(Error, "TensorRTBackend::execute: input '%s' rank exceeds TensorRT limit", name.c_str());
@@ -472,6 +545,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
       }
     }
 
+    input_bind_ptrs[i] = bind_ptr;
     if (!ctx->setTensorAddress(name.c_str(), bind_ptr)) {
       ET_LOG(Error, "TensorRTBackend::execute: setTensorAddress failed for input '%s'", name.c_str());
       return Error::InvalidState;
@@ -499,17 +573,65 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
   // nbytes() and before the Python binding reads back the shape.
   // If the buffer is CPU, stage through a temporary CUDA allocation.
   // ------------------------------------------------------------------
+  // (arg index, device_src ptr) for outputs staged through a device buffer.
   std::vector<std::pair<size_t, void*>> outputs_needing_copy;
+  // Caller-owned KV: (dst = delegate output EValue ptr, src = aliased input ptr,
+  // nbytes). The engine updates the aliased input in place; reflect that into the
+  // delegate output EValue after enqueue so ExecuTorch's write-back copy_ sees the
+  // updated cache. Skipped when dst == src (memory planner aliased them: zero-copy).
+  std::vector<std::tuple<void*, void*, size_t>> aliased_reflects;
   for (size_t o = 0; o < num_outputs; ++o) {
-    EValue* arg = args[num_inputs + o];
+    const std::string& name = engine->output_binding_names[o];
+
+    // Aliased output (KV-cache / user): the engine updates the aliased input in
+    // place, so bind this output binding to the aliased input's device pointer.
+    const int alias_in = engine->output_aliased_input_idx[o];
+    if (alias_in >= 0) {
+      void* bind_ptr = input_bind_ptrs[static_cast<size_t>(alias_in)];
+      if (bind_ptr == nullptr) {
+        ET_LOG(Error, "TensorRTBackend::execute: aliased output '%s' has no bound input pointer", name.c_str());
+        return Error::InvalidState;
+      }
+      if (!ctx->setTensorAddress(name.c_str(), bind_ptr)) {
+        ET_LOG(Error, "TensorRTBackend::execute: setTensorAddress failed for aliased output '%s'", name.c_str());
+        return Error::InvalidState;
+      }
+      // The aliased output IS a delegate output arg (the caller-owned mutable
+      // buffer's mutation slot). Consume it and record a reflect so ExecuTorch's
+      // write-back copy_ sees the engine's in-place update.
+      const size_t arg_i = arg_idx++;
+      EValue* out_arg = args[arg_i];
+      TORCHTRT_ET_CHECK_NOT_NULL(
+          out_arg, Error::InvalidArgument, "TensorRTBackend::execute: aliased output %zu is not a tensor", o);
+      if (!out_arg->isTensor()) {
+        ET_LOG(Error, "TensorRTBackend::execute: aliased output %zu is not a tensor", o);
+        return Error::InvalidArgument;
+      }
+      exec_aten::Tensor et_alias_out = out_arg->toTensor();
+      nvinfer1::Dims a_dims = ctx->getTensorShape(name.c_str());
+      if (a_dims.nbDims >= 0 && a_dims.nbDims <= nvinfer1::Dims::MAX_DIMS) {
+        SizesType a_sizes[nvinfer1::Dims::MAX_DIMS];
+        for (int d = 0; d < a_dims.nbDims; ++d) {
+          a_sizes[d] = static_cast<SizesType>(a_dims.d[d]);
+        }
+        (void)executorch::runtime::resize_tensor(et_alias_out, {a_sizes, static_cast<size_t>(a_dims.nbDims)});
+      }
+      void* dst = et_alias_out.nbytes() > 0 ? et_alias_out.mutable_data_ptr() : nullptr;
+      if (dst != nullptr && dst != bind_ptr) {
+        aliased_reflects.emplace_back(dst, bind_ptr, et_alias_out.nbytes());
+      }
+      continue;
+    }
+
+    const size_t arg_i = arg_idx++; // continue the shared running arg index after the inputs
+    EValue* arg = args[arg_i];
     TORCHTRT_ET_CHECK_NOT_NULL(arg, Error::InvalidArgument, "TensorRTBackend::execute: output %zu is not a tensor", o);
     if (!arg->isTensor()) {
       ET_LOG(Error, "TensorRTBackend::execute: output %zu is not a tensor", o);
       return Error::InvalidArgument;
     }
 
     exec_aten::Tensor et_out = arg->toTensor();
-    const std::string& name = engine->output_binding_names[o];
 
     // Update the ExecuTorch tensor shape to the actual TRT output shape.
     // getTensorShape() is valid after inferShapes() has been called.
@@ -556,7 +678,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
       }
       bind_ptr = engine->cached_output_ptrs[o];
       output_staged_to_host = true;
-      outputs_needing_copy.push_back({o, bind_ptr});
+      outputs_needing_copy.push_back({arg_i, bind_ptr});
     }
 
     if (!ctx->setTensorAddress(name.c_str(), bind_ptr)) {
@@ -577,6 +699,18 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
     return Error::InvalidState;
   }
 
+  // Caller-owned KV: reflect each engine in-place update into its delegate output
+  // EValue (D2D on the same stream, after the engine work). No-op list under
+  // zero-copy (dst == src filtered out at bind time).
+  for (const auto& r : aliased_reflects) {
+    cuda_err = cudaMemcpyAsync(std::get<0>(r), std::get<1>(r), std::get<2>(r), cudaMemcpyDeviceToDevice, stream);
+    if (cuda_err != cudaSuccess) {
+      ET_LOG(
+          Error, "TensorRTBackend::execute: aliased-output reflect D2D copy failed: %s", cudaGetErrorString(cuda_err));
+      return Error::InvalidProgram;
+    }
+  }
+
   // The engine work is now in flight on `stream`. Decide whether to wait for it:
   //   must_sync = an output is staged to host (the caller reads the D2H result on
   //   return), an input was staged from host (its async H2D read the caller's host
@@ -590,7 +724,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
   const bool must_sync = output_staged_to_host || input_staged_from_host || !g_user_stream_set;
   if (must_sync) {
     for (auto& output : outputs_needing_copy) {
-      exec_aten::Tensor et_out = args[num_inputs + output.first]->toTensor();
+      exec_aten::Tensor et_out = args[output.first]->toTensor();
       cuda_err =
           cudaMemcpyAsync(et_out.mutable_data_ptr(), output.second, et_out.nbytes(), cudaMemcpyDeviceToHost, stream);
       if (cuda_err != cudaSuccess) {

cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp

@@ -136,6 +136,7 @@ bool parse_int_after_key(const std::string& json, std::size_t search_from, const
 bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) {
   out.input_binding_names.clear();
   out.output_binding_names.clear();
+  out.aliased_io.clear();
   out.hardware_compatible = false;
   out.device_id = 0;
 
@@ -229,6 +230,77 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) {
     }
   }
 
+  // Optional aliased_io array: [{"output":..,"input":..,"kind":..}, ...].
+  // Absent in older blobs -> leave empty (backward compatible). Mirrors the
+  // io_bindings walk above using the same string helpers.
+  const std::size_t alias_key = json.find("\"aliased_io\"");
+  if (alias_key != std::string::npos) {
+    std::size_t apos = json.find('[', alias_key);
+    if (apos == std::string::npos) {
+      return false;
+    }
+    ++apos;
+    while (true) {
+      apos = skip_ws(json, apos);
+      if (apos >= json.size()) {
+        return false;
+      }
+      if (json[apos] == ']') {
+        ++apos;
+        break;
+      }
+      if (json[apos] == ',') {
+        ++apos;
+        continue;
+      }
+      if (json[apos] != '{') {
+        return false;
+      }
+      ++apos;
+
+      AliasedBinding ab;
+      while (true) {
+        apos = skip_ws(json, apos);
+        if (apos >= json.size()) {
+          return false;
+        }
+        if (json[apos] == '}') {
+          ++apos;
+          break;
+        }
+        if (json[apos] == ',') {
+          ++apos;
+          continue;
+        }
+        std::string key;
+        apos = parse_string(json, apos, key);
+        if (apos == std::string::npos) {
+          return false;
+        }
+        apos = skip_ws(json, apos);
+        if (apos >= json.size() || json[apos] != ':') {
+          return false;
+        }
+        apos = skip_ws(json, apos + 1);
+        if (key == "output") {
+          apos = parse_string(json, apos, ab.output);
+        } else if (key == "input") {
+          apos = parse_string(json, apos, ab.input);
+        } else if (key == "kind") {
+          apos = parse_string(json, apos, ab.kind);
+        } else {
+          apos = skip_value(json, apos);
+        }
+        if (apos == std::string::npos) {
+          return false;
+        }
+      }
+      if (!ab.output.empty() && !ab.input.empty()) {
+        out.aliased_io.push_back(std::move(ab));
+      }
+    }
+  }
+
   return parse_bool_after_key(json, pos, "\"hardware_compatible\"", out.hardware_compatible) &&
       parse_int_after_key(json, pos, "\"device_id\"", out.device_id);
 }

py/torch_tensorrt/_compile.py

@@ -1134,6 +1134,13 @@ def _extract_tensor(obj: Any) -> Any:
                         package_path=file_path,
                     )
                 elif output_format == "executorch":
+                    from torch_tensorrt.dynamo._exporter import (
+                        _declare_aliased_kv_mutations_on_ep,
+                    )
+
+                    # retrace=True: torch.export truncates the engines' aliased KV
+                    # outputs, so declare them as buffer mutations before lowering.
+                    exp_program = _declare_aliased_kv_mutations_on_ep(exp_program)
                     _save_as_executorch(
                         exp_program,
                         file_path,