fix(executorch): lift nested and partitioner-excluded mutable buffers

.github/scripts/verify-executorch-reference-runner.sh

@@ -14,6 +14,10 @@ set +x
 #   First argument: path to an existing .pte model.
 #   EXECUTORCH_SOURCE_DIR=/path/to/executorch
 #
+# Optional second argument: path to a caller-owned KV-cache decode .pte (see
+#   examples/torchtrt_executorch_example/export_kv_cache_decode.py). When given,
+#   kv_cache_decode_check is built and run against it as well.
+#
 # Optional:
 #   TensorRT_ROOT=/path/to/extracted/TensorRT
 #     If unset, the script reuses Bazel's fetched TensorRT SDK when available
@@ -29,15 +33,20 @@ set +x
 repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
 cd "${repo_root}"
 
-if [[ $# -ne 1 ]]; then
-  echo "Usage: $0 PATH_TO_MODEL.pte" >&2
+if [[ $# -lt 1 || $# -gt 2 ]]; then
+  echo "Usage: $0 PATH_TO_MODEL.pte [PATH_TO_KV_CACHE_DECODE.pte]" >&2
   exit 1
 fi
 model_path="$1"
 if [[ ! -f "${model_path}" ]]; then
   echo "ExecuTorch model not found: ${model_path}" >&2
   exit 1
 fi
+kv_model_path="${2:-}"
+if [[ -n "${kv_model_path}" && ! -f "${kv_model_path}" ]]; then
+  echo "KV-cache decode model not found: ${kv_model_path}" >&2
+  exit 1
+fi
 
 python_executable="${PYTHON_EXECUTABLE:-}"
 if [[ -z "${python_executable}" ]]; then
@@ -274,6 +283,7 @@ require_tar_entry "torch_tensorrt/src/torch_tensorrt/executorch/CMakeLists.txt"
 require_tar_entry "torch_tensorrt/examples/executorch_reference_runner/CMakeLists.txt"
 require_tar_entry "torch_tensorrt/bin/example_executorch_runner"
 require_tar_entry "torch_tensorrt/lib/libextension_cuda.so"
+require_tar_entry "torch_tensorrt/examples/executorch_reference_runner/kv_cache_decode_check.cpp"
 require_tar_entry "torch_tensorrt/BUILD"
 
 export TORCH_TENSORRT_ROOT="${verify_root}/torch_tensorrt"
@@ -295,8 +305,13 @@ fi
 
 cmake "${cmake_args[@]}"
 
+build_targets=(example_executorch_runner)
+if [[ -n "${kv_model_path}" ]]; then
+  build_targets+=(kv_cache_decode_check)
+fi
+
 cmake --build "${verify_root}/build-executorch-reference-runner" \
-  --target example_executorch_runner \
+  --target "${build_targets[@]}" \
   -j"${MAX_JOBS:-$(nproc)}"
 
 runner_log="${verify_root}/my_runner.log"
@@ -452,3 +467,20 @@ for _log in "${runner_log}" "${packaged_runner_log}"; do
     fi
   done
 done
+
+if [[ -n "${kv_model_path}" ]]; then
+  # kv_cache_decode_check exits non-zero when a decode step does not observe the KV
+  # the previous step wrote; the grep additionally pins the assertion itself, so
+  # weakening the check inside the binary cannot quietly turn this into a no-op.
+  kv_check_log="${verify_root}/kv_cache_decode_check.log"
+  kv_check_path="${verify_root}/build-executorch-reference-runner/kv_cache_decode_check"
+  if command -v ldd >/dev/null 2>&1 &&
+    ldd "${kv_check_path}" |
+      grep -E "libtorch|libtorch_cpu|libtorch_cuda|libc10" >&2; then
+    echo "kv_cache_decode_check links PyTorch/libtorch shared libraries" >&2
+    exit 1
+  fi
+
+  "${kv_check_path}" --model_path="${kv_model_path}" 2>&1 | tee "${kv_check_log}"
+  grep -q "PASS: decode at pos=1 observed the KV written at pos=0" "${kv_check_log}"
+fi

.github/workflows/executorch-test-linux.yml

@@ -82,7 +82,10 @@ jobs:
         export EXECUTORCH_ROOT="${EXECUTORCH_SOURCE_DIR}"
         python examples/torchtrt_executorch_example/export_static_shape.py \
           --model_path="${RUNNER_TEMP}/torchtrt-python.pte"
+        python examples/torchtrt_executorch_example/export_kv_cache_decode.py \
+          --model_path="${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte"
         .github/scripts/verify-executorch-reference-runner.sh \
-          "${RUNNER_TEMP}/torchtrt-python.pte"
+          "${RUNNER_TEMP}/torchtrt-python.pte" \
+          "${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte"
         python examples/executorch_reference_runner/load_model.py \
           --model_path="${RUNNER_TEMP}/torchtrt-python.pte" --num_runs=1

cpp/include/torch_tensorrt/executorch/TensorRTBackend.h

@@ -5,7 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  *
  * ExecuTorch backend delegate that runs TensorRT engines serialized by
- * torch_tensorrt. The processed blob uses the standalone TR01 wire format from
+ * torch_tensorrt. The processed blob uses the standalone wire format from
  * py/torch_tensorrt/executorch/serialization.py and is parsed directly here.
  * This runtime path intentionally does not depend on the legacy
  * Torch-TensorRT C++ runtime or libtorch.
@@ -59,6 +59,20 @@ 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, which ExecuTorch's write-back copy_ then reads.
+  std::vector<int> output_aliased_input_idx;
+  // Per input binding [0..num_inputs): true if any output aliases this input, so
+  // its in-place (KV/user) update must land in the caller-owned storage. Built at
+  // init from aliased_io; execute() uses it to reject a non-device-resident
+  // aliased input instead of silently staging its update into delegate scratch.
+  std::vector<bool> input_is_alias_target;
+  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;