feat(executorch): caller-owned KV-cache for the TensorRT delegate

Adds end-to-end caller-owned KV-cache support to the ExecuTorch TensorRT
delegate: the KV buffers are owned by the caller above the delegate and threaded
in as mutable-buffer delegate args, instead of being self-allocated inside a
(stateless) TensorRT engine.

Runtime + serialization (delegate):
- serialize each engine's aliased (KV-cache / in-place) I/O into the delegate blob
  (serialization.py, backend.py, TensorRTBlobHeader.{h,cpp});
- at runtime bind each aliased TRT output binding to its aliased input's
  caller-provided pointer (in-place) and reflect the result into the delegate
  output EValue -- a no-op when the memory planner already aliased the two
  (TensorRTBackend.{h,cpp}).

Export/lowering (torch_tensorrt):
- expose each engine's aliased outputs as graph-level BUFFER_MUTATIONs so
  ExecuTorch keeps the KV buffers as caller-owned mutable buffers: at transform
  time for the legacy exporter (retrace=False), and via a post-export pass
  (_declare_aliased_kv_mutations_on_ep) for torch.export (retrace=True), which
  otherwise truncates the aliased outputs at the fx boundary;
- keep delegate-mutated buffers above the delegate in TensorRTPartitioner
  (tag_constant_data would otherwise freeze them as constants).

The retrace=True pass runs for exported_program as well as executorch. The
truncation happens at the fx boundary for every output format, so declaring only
on the executorch path left an exported_program saved with the mutation absent
from its signature while the engine still updated the cache in place. It is
declared before _normalize_engine_constants_to_python, which rewrites the engine
constants the pass reads aliased_io from. retrace=False was already correct for
every format via create_trt_exp_program. aot_inductor stays undeclared and
warns: whether an aliased in-place mutation survives functionalization under
inductor is unverified.

Tests cover serialization round-trip, the exposure-flag dispatch across both
retrace modes, the buffer-mutation declaration, and the partitioner un-tagging.

Author: Conarnar

cpp/include/torch_tensorrt/executorch/TensorRTBackend.h

@@ -59,6 +59,21 @@ 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;
+  // 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;
 

cpp/src/torch_tensorrt/executorch/TensorRTBackend.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,84 @@ 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()) {
+        // A missing "kind" key means an older blob (the Python serializer omits
+        // it for KV aliases); default to the TRT-enforced kind so init()'s kind
+        // validation treats an absent key the same as the Python runtime rather
+        // than rejecting it as unknown.
+        if (ab.kind.empty()) {
+          ab.kind = "kv_cache_update";
+        }
+        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);
 }

cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp

@@ -7,6 +7,7 @@ filegroup(
     srcs = [
         "CMakeLists.txt",
         "README.md",
+        "kv_cache_decode_check.cpp",
         "main.cpp",
     ],
 )
@@ -20,3 +21,14 @@ cc_binary(
         "@executorch//:executorch_file_data_loader",
     ],
 )
+
+cc_binary(
+    name = "kv_cache_decode_check",
+    srcs = ["kv_cache_decode_check.cpp"],
+    deps = [
+        "//cpp:tensorrt_executorch_backend",
+        "@cuda//:cudart",
+        "@executorch//:executorch_core",
+        "@executorch//:executorch_file_data_loader",
+    ],
+)

examples/executorch_reference_runner/BUILD

@@ -61,3 +61,18 @@ target_link_libraries(
     executorch::extensions
     executorch::kernels
     torchtrt::executorch_backend)
+
+# Caller-owned KV-cache persistence check (see kv_cache_decode_check.cpp). It
+# cudaMalloc's the device-tagged planned arenas that hold the KV buffers and
+# copies the logits back to host, so it links the CUDA runtime directly.
+find_package(CUDAToolkit REQUIRED)
+add_executable(kv_cache_decode_check kv_cache_decode_check.cpp)
+target_link_libraries(
+  kv_cache_decode_check
+  PRIVATE
+    executorch
+    executorch::backends
+    executorch::extensions
+    executorch::kernels
+    torchtrt::executorch_backend
+    CUDA::cudart)

examples/executorch_reference_runner/CMakeLists.txt

@@ -103,3 +103,30 @@ Loading the method initializes the TensorRT ExecuTorch backend for any
 Torch-TensorRT delegate subgraphs embedded in the `.pte`. The Python
 `torch_tensorrt` package is needed when exporting the `.pte`; it is not needed
 by this native runner at inference time.
+
+## Caller-Owned KV-Cache Persistence Check
+
+`kv_cache_decode_check` is a small self-asserting runner for a caller-owned
+KV-cache decode `.pte` (its aliased KV output is bound in place to the caller's
+mutable buffer, which persists across `execute()` calls).
+
+Export a minimal single-layer decode model:
+
+```bash
+python examples/torchtrt_executorch_example/export_kv_cache_decode.py \
+  --model_path=kv_cache_decode.pte
+```
+
+The same CMake build produces the check runner (`kv_cache_decode_check`
+target). Run it:
+
+```bash
+./build-executorch-reference-runner/kv_cache_decode_check --model_path=kv_cache_decode.pte
+```
+
+It loads the method twice (each starting from a zeroed cache) and runs a decode
+at `input_pos=1` once with no prior step and once after a step at `input_pos=0`.
+Because the causal attention at position 1 covers positions 0..1, the two logits
+differ only if the KV written at position 0 persisted across `execute()` calls.
+The runner prints `[kv-check] PASS` and returns 0 on success, or fails if the
+two are identical (the update did not persist). It requires a CUDA device.