feat(executorch): let a TensorRT engine write the caller's KV cache in place

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

@@ -14,16 +14,19 @@ 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 trailing arguments: one or more caller-owned KV-cache decode .pte
+#   files (see examples/torchtrt_executorch_example/export_kv_cache_decode.py).
+#   When given, kv_cache_decode_check is built and run against each of them,
+#   staged or zero-copy.
 #
-# Optional third argument: path to a coalesced TensorRT + CUDA .pte (see
+# Optional --coalesced=PATH: path to a coalesced TensorRT + CUDA .pte (see
 #   examples/torchtrt_executorch_example/export_coalesced.py). When given, the
 #   runner built from source here is run against it and its output is compared to
 #   the eager reference that export script wrote next to the model. Only that
 #   runner: the packaged binary links the TensorRT delegate alone, so it has no
-#   CUDA backend for the partition a coalesced program hands to one.
+#   CUDA backend for the partition a coalesced program hands to one. Named rather
+#   than positional because the KV-cache decode models are variadic, so a bare
+#   path after the first argument cannot be told apart from one of those.
 #
 # Optional:
 #   TensorRT_ROOT=/path/to/extracted/TensorRT
@@ -40,21 +43,35 @@ set +x
 repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
 cd "${repo_root}"
 
-if [[ $# -lt 1 || $# -gt 3 ]]; then
-  echo "Usage: $0 PATH_TO_MODEL.pte [PATH_TO_KV_CACHE_DECODE.pte [PATH_TO_COALESCED.pte]]" >&2
+if [[ $# -lt 1 ]]; then
+  echo "Usage: $0 PATH_TO_MODEL.pte [--coalesced=PATH_TO_COALESCED.pte]" \
+    "[PATH_TO_KV_CACHE_DECODE.pte ...]" >&2
   exit 1
 fi
 model_path="$1"
+shift
 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
-coalesced_model_path="${3:-}"
+coalesced_model_path=""
+kv_model_paths=()
+for arg in "$@"; do
+  case "${arg}" in
+    --coalesced=*)
+      coalesced_model_path="${arg#--coalesced=}"
+      ;;
+    *)
+      kv_model_paths+=("${arg}")
+      ;;
+  esac
+done
+for kv_model_path in "${kv_model_paths[@]:-}"; do
+  if [[ -n "${kv_model_path}" && ! -f "${kv_model_path}" ]]; then
+    echo "KV-cache decode model not found: ${kv_model_path}" >&2
+    exit 1
+  fi
+done
 if [[ -n "${coalesced_model_path}" && ! -f "${coalesced_model_path}" ]]; then
   echo "Coalesced model not found: ${coalesced_model_path}" >&2
   exit 1
@@ -318,7 +335,7 @@ fi
 cmake "${cmake_args[@]}"
 
 build_targets=(example_executorch_runner)
-if [[ -n "${kv_model_path}" ]]; then
+if [[ ${#kv_model_paths[@]} -gt 0 ]]; then
   build_targets+=(kv_cache_decode_check)
 fi
 
@@ -551,11 +568,7 @@ for _log in "${runner_log}" "${packaged_runner_log}"; do
   assert_runner_output "${_log}" "[2,3,4,4]" "2.0000" 0
 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"
+if [[ ${#kv_model_paths[@]} -gt 0 ]]; then
   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}" |
@@ -564,8 +577,24 @@ if [[ -n "${kv_model_path}" ]]; then
     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}"
+  kv_index=0
+  for kv_model_path in "${kv_model_paths[@]}"; do
+    # kv_cache_decode_check exits non-zero when a decode step does not observe the
+    # KV the previous step wrote; the greps additionally pin the assertion itself,
+    # so weakening the check inside the binary cannot quietly turn this into a
+    # no-op. Both caller-stream modes are pinned by name: the backend skips its
+    # end-of-execute synchronization only when a caller stream is set, so dropping
+    # the "own" run would leave the branch zero-copy KV relies on uncovered while
+    # the lane stayed green.
+    kv_check_log="${verify_root}/kv_cache_decode_check_${kv_index}.log"
+    "${kv_check_path}" --model_path="${kv_model_path}" 2>&1 | tee "${kv_check_log}"
+    for kv_stream_mode in none own; do
+      grep -q \
+        "PASS: decode at pos=1 observed the KV written at pos=0 across execute() calls (caller stream: ${kv_stream_mode})" \
+        "${kv_check_log}"
+    done
+    kv_index=$((kv_index + 1))
+  done
 fi
 
 if [[ -n "${coalesced_model_path}" ]]; then

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

@@ -95,11 +95,15 @@ jobs:
           --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"
+        python examples/torchtrt_executorch_example/export_kv_cache_decode.py \
+          --model_path="${RUNNER_TEMP}/torchtrt-kv-cache-decode-zero-copy.pte" \
+          --zero_copy
         python examples/torchtrt_executorch_example/export_coalesced.py \
           --model_path="${RUNNER_TEMP}/torchtrt-coalesced.pte"
         .github/scripts/verify-executorch-reference-runner.sh \
           "${RUNNER_TEMP}/torchtrt-python.pte" \
+          --coalesced="${RUNNER_TEMP}/torchtrt-coalesced.pte" \
           "${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte" \
-          "${RUNNER_TEMP}/torchtrt-coalesced.pte"
+          "${RUNNER_TEMP}/torchtrt-kv-cache-decode-zero-copy.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

@@ -61,11 +61,15 @@ struct EngineHandle {
   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.
+  // Built at init from the blob's aliased_io. Either way execute() binds the
+  // aliased TRT output binding to its aliased input's caller-provided pointer,
+  // so the engine's write lands in the caller's buffer; what differs is how the
+  // .pte carries the buffer. Threaded: the buffer is both a delegate input arg
+  // and a delegate output arg (the caller-owned mutable buffer's mutation slot),
+  // and execute() reflects the result into that output EValue for ExecuTorch's
+  // write-back copy_ to read. Elided -- zero-copy KV -- the buffer is an input
+  // arg only, the delegate has no output for it, and execute() skips the
+  // reflect: the in-place write already is the update.
   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

cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp

@@ -440,6 +440,12 @@ Result<DelegateHandle*> TensorRTBackend::init(
   // 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.
+  // The parser has already refused a blob claiming one output twice, so the count
+  // built below is one per distinct output binding, which is what execute()
+  // subtracts on. It has also refused a blob repeating a binding name, so the
+  // first-match scans below reach the only slot carrying that name -- otherwise
+  // the alias would be recorded on the first slot while execute() re-bound the
+  // same TensorRT name for the later one, replacing the caller's buffer address.
   handle->output_aliased_input_idx.assign(handle->num_outputs, -1);
   handle->input_is_alias_target.assign(handle->num_inputs, false);
   for (const auto& ab : header.aliased_io) {
@@ -551,8 +557,10 @@ Result<DelegateHandle*> TensorRTBackend::init(
 // their addresses; no separate output allocation is required.
 //
 // Args layout (mirroring the Python exporter):
-//   args[0 .. num_inputs-1]             – input EValues
-//   args[num_inputs .. num_inputs+num_outputs-1] – output EValues
+//   args[0 .. num_inputs-1]                               -- input EValues
+//   args[num_inputs .. num_inputs+num_delegate_outputs-1] -- output EValues
+// num_delegate_outputs is num_outputs, less the aliased outputs when zero-copy KV
+// elided them from the delegate; see the arity branch at the top of execute().
 // ---------------------------------------------------------------------------
 Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* handle, Span<EValue*> args) const {
   (void)context;
@@ -561,17 +569,70 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
 
   const size_t num_inputs = engine->num_inputs;
   const size_t num_outputs = engine->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;
+  // Caller-owned KV comes in two shapes. Either 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; or -- zero-copy KV -- the aliased
+  // outputs are elided from the delegate entirely, because the engine's in-place
+  // write through the aliased input already IS the buffer update. A .pte written
+  // before zero-copy existed still takes the first branch, and export will not emit
+  // the shorter arity unless zero-copy was asked for, so a short argument list
+  // cannot instead mean "the aliased outputs were never declared".
   const size_t num_delegate_inputs = num_inputs;
-  if (args.size() < num_delegate_inputs + num_delegate_outputs) {
+  const size_t num_aliased_outputs = engine->num_aliased_outputs;
+  // The blob parser refuses a second aliased_io entry for an output already
+  // claimed, and init refuses an entry whose output is not one of the recorded
+  // output bindings, so this holds for any header that reached here. It is
+  // checked anyway because the subtractions below are unsigned: on a header
+  // built some other way one of them wraps, and the length check then accepts
+  // an argument count it should not. Exactly one such count, not any: with the
+  // aliased count above the outputs, the only count that can set the elided
+  // flag is inputs plus outputs minus aliased, and once it is set
+  // num_delegate_outputs is the wrapped outputs-minus-aliased, which adding the
+  // inputs back wraps round to that same count -- so the length check passes
+  // and the call goes on to index past the end of args. Every other count below
+  // inputs plus outputs is still rejected. The two subtractions never wrap in
+  // the same call either: with the aliased count above the outputs but not
+  // above inputs plus outputs only the second wraps, and above both the first
+  // wraps to a value no argument count can match, so the flag is never set and
+  // the second never runs. A duplicate entry is the case this does not catch:
+  // it inflates the count while staying within num_outputs, passes both checks,
+  // and still indexes one past the end of args; the parser's refusal is what
+  // stops that.
+  if (num_aliased_outputs > num_outputs) {
     ET_LOG(
         Error,
-        "TensorRTBackend::execute: expected at least %zu args, got %zu",
-        num_delegate_inputs + num_delegate_outputs,
-        args.size());
+        "TensorRTBackend::execute: %zu aliased output(s) recorded for %zu output binding(s)",
+        num_aliased_outputs,
+        num_outputs);
+    return Error::InvalidProgram;
+  }
+  const bool aliased_outputs_elided =
+      num_aliased_outputs > 0 && args.size() == num_delegate_inputs + num_outputs - num_aliased_outputs;
+  const size_t num_delegate_outputs = num_outputs - (aliased_outputs_elided ? num_aliased_outputs : 0);
+  if (args.size() < num_delegate_inputs + num_delegate_outputs) {
+    // With aliased outputs there are two right answers and only one of them can
+    // ever be num_delegate_outputs here: the elided flag is set only by an
+    // argument count that already fits, so a program that arrives short is
+    // always measured against the threaded count. Reporting that alone tells a
+    // zero-copy .pte to supply the longer list, which is the shape that consumes
+    // its real outputs as mutation slots -- so both counts are named, and the
+    // aliasing that is the reason for the two.
+    if (num_aliased_outputs > 0) {
+      ET_LOG(
+          Error,
+          "TensorRTBackend::execute: expected %zu args with the engine's %zu aliased output(s) threaded, "
+          "or %zu with them elided (zero-copy KV), got %zu",
+          num_delegate_inputs + num_outputs,
+          num_aliased_outputs,
+          num_delegate_inputs + num_outputs - num_aliased_outputs,
+          args.size());
+    } else {
+      ET_LOG(
+          Error,
+          "TensorRTBackend::execute: expected at least %zu args, got %zu",
+          num_delegate_inputs + num_delegate_outputs,
+          args.size());
+    }
     return Error::InvalidArgument;
   }
 
@@ -795,9 +856,15 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
         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.
+      // Elided: setTensorAddress above pointed this output binding at the caller's
+      // buffer, so the engine writes the buffer itself; nothing to reflect into.
+      if (aliased_outputs_elided) {
+        continue;
+      }
+
+      // Otherwise 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(
@@ -945,10 +1012,25 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle*
   // next execute() and the destructor wait before reusing/freeing exec_ctx. The D2H
   // copies live in the must_sync branch: an output staged to host always sets
   // output_staged_to_host, so outputs_needing_copy is empty on the skip path.
-  // An aliased reflect enqueues the engine's in-place update into the delegate
-  // output EValue on `stream`; ExecuTorch's buffer-mutation copy_ reads that EValue
-  // after execute() returns, so the reflect must complete first. A model with
-  // aliased outputs therefore always syncs here.
+  // A non-elided aliased reflect enqueues the engine's in-place update into the
+  // delegate output EValue on `stream`; ExecuTorch's buffer-mutation copy_ reads
+  // that EValue after execute() returns, so the reflect must complete first, and a
+  // model that threads its aliased outputs as delegate output args syncs here.
+  // Under zero-copy KV skipping the sync is correct, because the aliased buffer
+  // stays device-resident and its next reader is the following engine execute() --
+  // on the same `stream`, provided the runner honours the single-shared-stream
+  // contract every coalesced .pte already depends on. A host reader that
+  // inspected it immediately after execute() returns would see stale data unless
+  // it synchronized `stream` itself; ExecuTorch's KV path never does such a read.
+  // Eliding the reflects drops aliased_reflect_pending, so an execute() with a
+  // caller stream and no host staging no longer syncs at all. What orders the
+  // engine's *other* outputs is then that same contract: the host consumer
+  // ExecuTorch inserts for a device delegate output is an et_copy::_d2h_copy,
+  // whose kernel issues its copy on getCallerStream() -- this stream -- and then
+  // synchronizes it; a device consumer of that output is another delegate,
+  // enqueued on the same stream. That covers the write-back of a copy-back buffer
+  // sitting beside the zero-copy caches. With no caller stream set that kernel
+  // falls back to a blocking cudaMemcpy, and must_sync is true here anyway.
   const bool aliased_reflect_pending = !aliased_reflects.empty();
   const bool must_sync =
       output_staged_to_host || input_staged_from_host || aliased_reflect_pending || !caller_stream_set;