diff --git a/.github/scripts/verify-executorch-reference-runner.sh b/.github/scripts/verify-executorch-reference-runner.sh index fa5881f5fc7..a4b96146927 100755 --- a/.github/scripts/verify-executorch-reference-runner.sh +++ b/.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 diff --git a/.github/workflows/executorch-test-linux.yml b/.github/workflows/executorch-test-linux.yml index 9751f2bb54a..92be5830970 100644 --- a/.github/workflows/executorch-test-linux.yml +++ b/.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 diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index b33d712d403..dc168a03367 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/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 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 diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 8408c13e881..337b7e5bd7a 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -440,6 +440,12 @@ Result 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 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 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; } @@ -615,13 +676,18 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // on the shared exec_ctx. Wait for it before reconfiguring the context below: // TensorRT forbids mutating a context while one of its enqueues is in flight, and // setInputShape/setTensorAddress run on the host, so this must be a host-side wait. + // The flag is cleared only once the wait has actually succeeded: it is the + // destructor's only reason to wait, and a failed synchronize is precisely the + // case where the enqueue may still be running. Clearing it first would hand + // the destructor a handle it believes is idle and let it free the staging + // buffers and reset exec_ctx underneath a live enqueue. if (engine->inflight_pending) { cuda_err = cudaEventSynchronize(engine->inflight_event); - engine->inflight_pending = false; if (cuda_err != cudaSuccess) { ET_LOG(Error, "TensorRTBackend::execute: cudaEventSynchronize failed: %s", cudaGetErrorString(cuda_err)); return Error::InvalidProgram; } + engine->inflight_pending = false; } const auto caller_stream = ::executorch::extension::cuda::getCallerStream(); const bool caller_stream_set = caller_stream.has_value(); @@ -795,9 +861,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 +1017,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; @@ -972,11 +1059,16 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } } cuda_err = cudaStreamSynchronize(stream); - engine->inflight_pending = false; if (cuda_err != cudaSuccess) { ET_LOG(Error, "TensorRTBackend::execute: cudaStreamSynchronize failed: %s", cudaGetErrorString(cuda_err)); return Error::InvalidProgram; } + // Same shape as the wait at the top of execute(), and for the same reason, + // though nothing here can be armed yet: this branch never records the + // event, and the wait at the top has already cleared any flag a previous + // call left. Arming happens as the last statement of the other branch, so + // the only writer that can make this clear anything is one added later. + engine->inflight_pending = false; if (copy_err != Error::Ok) { return copy_err; } diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp index 9fc5600c35a..a7959c9e805 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp @@ -3,14 +3,17 @@ #include #include #include +#include #include +#include namespace torch_tensorrt { namespace executorch_backend { namespace { -// TR02 marks a blob whose metadata carries aliased_io; TR01 is one without. This -// parser handles aliased_io, so it accepts either. +// TR02 marks a blob whose metadata carries aliased_io; TR01 is one without. +// This parser handles aliased_io, so it accepts either, and holds TR02 to its +// promise: see the cross-check at the end of parse_metadata_json. constexpr char TENSORRT_MAGIC[4] = {'T', 'R', '0', '1'}; constexpr char TENSORRT_MAGIC_ALIASED_IO[4] = {'T', 'R', '0', '2'}; constexpr uint32_t METADATA_OFFSET_FIELD_OFFSET = 4; @@ -27,7 +30,29 @@ std::size_t skip_ws(const std::string& s, std::size_t pos) { return pos; } -std::size_t parse_string(const std::string& s, std::size_t pos, std::string& out) { +// The three escapes whose JSON meaning is exactly what this parser does with +// every escape: drop the backslash, keep the character after it. json.dumps +// emits the first two whenever a name holds a quote or a backslash, and the +// third is one a hand-written blob may carry, so these are strings this parser +// and a JSON reader agree on and there is nothing to refuse. +bool escape_reads_as_written(char after_backslash) { + return after_backslash == '"' || after_backslash == '\\' || after_backslash == '/'; +} + +// A backslash is dropped and the character after it kept. For the three escapes +// above that is the correct decoding; for every other one the two ends read +// these bytes differently. "\n" comes out as the letter n where a JSON reader +// sees a newline, a \u escape as the letter u and its four digits where a JSON +// reader sees one codepoint, and "\q" comes out as the letter q where a JSON +// reader refuses the string outright. Rather than decode them, every caller +// that *compares* the text it gets -- against the engine's binding names, +// against a key, or against an alias kind -- refuses a string carrying one, and +// saw_misread_escape is how they are told. The values skip_value walks past +// need only their end, which this finds either way. +std::size_t parse_string(const std::string& s, std::size_t pos, std::string& out, bool* saw_misread_escape = nullptr) { + if (saw_misread_escape != nullptr) { + *saw_misread_escape = false; + } if (pos >= s.size() || s[pos] != '"') { return std::string::npos; } @@ -35,6 +60,9 @@ std::size_t parse_string(const std::string& s, std::size_t pos, std::string& out out.clear(); while (pos < s.size() && s[pos] != '"') { if (s[pos] == '\\' && pos + 1 < s.size()) { + if (saw_misread_escape != nullptr && !escape_reads_as_written(s[pos + 1])) { + *saw_misread_escape = true; + } ++pos; } out += s[pos++]; @@ -86,16 +114,36 @@ std::size_t skip_value(const std::string& s, std::size_t pos) { return pos; } +// The two scalar fields are found by searching the metadata text for the key, +// quotes included, so anything else quoted the same way would be matched +// instead: a binding named device_id sitting in the alias array the search runs +// over, or a string value that is exactly the key name. Requiring the match to +// be in key position -- its own colon next, whitespace aside -- is what tells +// the two apart, since a value is followed by a comma or a closing brace. A +// value that merely *contains* the key text needs no rule: json.dumps escapes +// the quotes it carries, so the closing quote of the search never lines up and +// there is no match to reject. An occurrence that is not in key position is +// passed over rather than refused: the key may still be ahead of it, and if it +// is not, the field is absent and keeps its default, which is what a blob +// written before the field existed wants. +std::size_t value_pos_after_key(const std::string& json, std::size_t search_from, const char* key) { + const std::size_t key_len = std::strlen(key); + std::size_t pos = search_from; + while ((pos = json.find(key, pos)) != std::string::npos) { + const std::size_t colon = skip_ws(json, pos + key_len); + if (colon < json.size() && json[colon] == ':') { + return skip_ws(json, colon + 1); + } + pos += key_len; + } + return std::string::npos; +} + bool parse_bool_after_key(const std::string& json, std::size_t search_from, const char* key, bool& value) { - const std::size_t key_pos = json.find(key, search_from); - if (key_pos == std::string::npos) { + const std::size_t val = value_pos_after_key(json, search_from, key); + if (val == std::string::npos) { return true; } - const std::size_t colon = json.find(':', key_pos); - if (colon == std::string::npos) { - return false; - } - const std::size_t val = skip_ws(json, colon + 1); if (json.compare(val, 4, "true") == 0) { value = true; return true; @@ -108,50 +156,91 @@ bool parse_bool_after_key(const std::string& json, std::size_t search_from, cons } bool parse_int_after_key(const std::string& json, std::size_t search_from, const char* key, int& value) { - const std::size_t key_pos = json.find(key, search_from); - if (key_pos == std::string::npos) { + std::size_t pos = value_pos_after_key(json, search_from, key); + if (pos == std::string::npos) { return true; } - const std::size_t colon = json.find(':', key_pos); - if (colon == std::string::npos) { - return false; - } - std::size_t pos = skip_ws(json, colon + 1); bool neg = false; if (pos < json.size() && json[pos] == '-') { neg = true; ++pos; } - int parsed = 0; + // Accumulated in 64 bits and bounded on every digit, because the digit count + // is the blob's to choose. Overflowing an int here is undefined behaviour -- + // a trapping build aborts on it -- and in an ordinary one it wraps, which is + // the bad case: the only field parsed this way is device_id, and a value just + // over four billion wraps onto a device that exists, so cudaSetDevice then + // succeeds and the engine deserializes on a GPU nobody asked for. Anything + // outside an int is refused instead. + int64_t parsed = 0; bool saw_digit = false; while (pos < json.size() && json[pos] >= '0' && json[pos] <= '9') { saw_digit = true; parsed = parsed * 10 + (json[pos] - '0'); + if (parsed > std::numeric_limits::max()) { + return false; + } ++pos; } if (!saw_digit) { return false; } - value = neg ? -parsed : parsed; + value = static_cast(neg ? -parsed : parsed); return true; } -bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { +// Every recorded name is resolved against the engine through c_str(), which +// stops at the first NUL, while the refusals below compare whole std::string +// values. The metadata is copied out of the blob with an explicit length, so a +// name can hold a NUL: two names differing only after one are distinct here and +// are the same tensor to TensorRT, which is exactly what those refusals exist to +// stop, and a name that is only a NUL is non-empty here and empty to TensorRT. +// A name carrying an escape this parser reads differently from a JSON reader is +// the same problem one step earlier: the name recorded here is not the one the +// writer wrote and so not one the engine has either, and two names differing +// only in such an escape collapse into one and are refused as a repeat. +// Refusing both outright makes what the parser compares be what the engine will +// compare. An escaped quote, backslash or forward slash is *not* refused: those +// three come out of parse_string exactly as the writer wrote them, and the first +// two are what json.dumps emits for a name holding a quote or a backslash. No +// writer emits a NUL -- json.dumps escapes it -- and the escapes this does +// refuse are the ones json.dumps reserves for what a binding name has no +// business holding, so it refuses only a blob assembled some other way. +// Like every refusal in this parser it is silent: parse() returns false and +// the caller reports its own generic parse failure, so a NUL is not +// distinguishable at load time from a bad magic or a wrapped extent. Refusing +// uniformly is the deliberate choice here, not a missing diagnostic. +bool usable_as_binding_name(const std::string& name, bool misread_escape) { + return !misread_escape && !name.empty() && name.find('\0') == std::string::npos; +} + +bool parse_metadata_json(const std::string& json, bool expects_aliased_io, TensorRTBlobHeader& out) { out.input_binding_names.clear(); out.output_binding_names.clear(); out.aliased_io.clear(); out.hardware_compatible = false; out.device_id = 0; - const std::size_t bindings_pos = json.find("\"io_bindings\""); - if (bindings_pos == std::string::npos) { - return false; - } - const std::size_t arr_start = json.find('[', bindings_pos); - if (arr_start == std::string::npos) { + // Found the same way as the two scalars below, so all four keys agree about + // what a key is: an occurrence followed by its own colon, and the array right + // after that colon rather than the next '[' anywhere in the text. A string + // value that spells io_bindings is then passed over instead of being taken + // for the key, and a blob whose io_bindings is not an array is refused rather + // than walked from some unrelated bracket further on. + const std::size_t arr_start = value_pos_after_key(json, 0, "\"io_bindings\""); + if (arr_start == std::string::npos || arr_start >= json.size() || json[arr_start] != '[') { return false; } + // A TensorRT engine has one name space for its tensors, so a name repeated + // across io_bindings -- in either list -- cannot be two bindings. It is + // accepted by every name lookup, which stops at the first match, and then + // contradicted by every address bind, which is keyed on the name and so + // overwrites whatever the earlier slot bound. For an aliased output the + // address overwritten is the caller's buffer, and the engine's in-place + // update lands somewhere else. Refuse the blob here, where the repeat is + // visible from the bytes alone. + std::unordered_set claimed_bindings; std::size_t pos = arr_start + 1; while (true) { pos = skip_ws(json, pos); @@ -173,7 +262,8 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { std::string name; bool is_input = false; - bool saw_name = false; + bool saw_is_input = false; + bool name_misread = false; while (true) { pos = skip_ws(json, pos); @@ -189,9 +279,19 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { continue; } + // A key is compared against name and is_input below, so it is held to the + // same rule those comparisons need: a key carrying an escape this parser + // reads differently is refused rather than matched. Without that, + // "is_i\nput" arrives here as is_input and decides which list the binding + // goes in, while a JSON reader sees a key with a newline in it, ignores + // it, and leaves is_input at its own default -- the two readers of one + // blob putting one binding in opposite lists, which is what the + // saw_is_input refusal below exists to stop for the plainly misspelled + // key. std::string key; - pos = parse_string(json, pos, key); - if (pos == std::string::npos) { + bool key_misread = false; + pos = parse_string(json, pos, key, &key_misread); + if (pos == std::string::npos || key_misread) { return false; } pos = skip_ws(json, pos); @@ -201,12 +301,12 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { pos = skip_ws(json, pos + 1); if (key == "name") { - pos = parse_string(json, pos, name); - saw_name = pos != std::string::npos; - if (!saw_name) { + pos = parse_string(json, pos, name, &name_misread); + if (pos == std::string::npos) { return false; } } else if (key == "is_input") { + saw_is_input = true; if (json.compare(pos, 4, "true") == 0) { is_input = true; pos += 4; @@ -224,12 +324,36 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { } } - if (saw_name && !name.empty()) { - if (is_input) { - out.input_binding_names.push_back(name); - } else { - out.output_binding_names.push_back(name); - } + // An entry with no usable name cannot be refused earlier because the keys + // may arrive in any order, so it is refused here, beside the repeat. + // Skipping it instead would shorten the recorded list while the delegate's + // argument list keeps its full length, and the two are only inferred from + // the engine when both are empty -- so one real name beside a blank leaves + // a short list that no longer lines up with the engine's bindings. + // An absent "name" key needs no flag of its own, unlike is_input below: + // the string is still empty here, which this predicate already refuses. + if (!usable_as_binding_name(name, name_misread)) { + return false; + } + // is_input has no safe default, so an entry without it is refused beside + // the nameless one. The initializer here reads the binding as an output + // and TensorRTIOBinding.is_input in serialization.py defaults to an input, + // so the two readers of these bytes would disagree -- and one binding + // changing list shifts every index after it, which on a static-shape + // engine leaves nothing for shape inference to object to: it runs on the + // wrong tensors. A key misspelled by one byte is the same case, since it + // falls through to skip_value and leaves the initializer standing. The + // writer always emits the key. + if (!saw_is_input) { + return false; + } + if (!claimed_bindings.insert(name).second) { + return false; + } + if (is_input) { + out.input_binding_names.push_back(name); + } else { + out.output_binding_names.push_back(name); } } @@ -239,13 +363,15 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { // // Search from pos (past the io_bindings array) so a model input literally // named "aliased_io" isn't matched as the array key. - const std::size_t alias_key = json.find("\"aliased_io\"", pos); - if (alias_key != std::string::npos) { - std::size_t apos = json.find('[', alias_key); - if (apos == std::string::npos) { + std::size_t scalars_from = pos; + std::size_t apos = value_pos_after_key(json, pos, "\"aliased_io\""); + if (apos != std::string::npos) { + if (apos >= json.size() || json[apos] != '[') { return false; } ++apos; + std::unordered_set claimed_outputs; + std::unordered_set claimed_inputs; while (true) { apos = skip_ws(json, apos); if (apos >= json.size()) { @@ -265,6 +391,9 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { ++apos; AliasedBinding ab; + bool output_misread = false; + bool input_misread = false; + bool kind_misread = false; while (true) { apos = skip_ws(json, apos); if (apos >= json.size()) { @@ -278,9 +407,13 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { ++apos; continue; } + // Held to the same rule as the io_bindings keys above, for the same + // reason: these are compared, so an escape the two ends read + // differently would match a key here that a JSON reader does not see. std::string key; - apos = parse_string(json, apos, key); - if (apos == std::string::npos) { + bool key_misread = false; + apos = parse_string(json, apos, key, &key_misread); + if (apos == std::string::npos || key_misread) { return false; } apos = skip_ws(json, apos); @@ -289,11 +422,11 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { } apos = skip_ws(json, apos + 1); if (key == "output") { - apos = parse_string(json, apos, ab.output); + apos = parse_string(json, apos, ab.output, &output_misread); } else if (key == "input") { - apos = parse_string(json, apos, ab.input); + apos = parse_string(json, apos, ab.input, &input_misread); } else if (key == "kind") { - apos = parse_string(json, apos, ab.kind); + apos = parse_string(json, apos, ab.kind, &kind_misread); } else { apos = skip_value(json, apos); } @@ -301,23 +434,80 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { return false; } } - if (!ab.output.empty() && !ab.input.empty()) { - // The current Python serializer always writes "kind" (serialization.py), - // and older blobs carry no aliased_io array at all, so this default is - // defensive: it only fires for a blob that has an aliased_io entry but - // omits "kind". 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)); + // An entry missing either name is refused rather than skipped, for the + // reason the binding walk above gives and one more: init() counts the + // entries it accepts and execute() subtracts that count from the delegate + // argument list, so a dropped entry surfaces as an argument-count error at + // every execute, which never mentions aliasing, instead of at parse, which + // the blob-header tests reach without a GPU. + if (!usable_as_binding_name(ab.output, output_misread) || !usable_as_binding_name(ab.input, input_misread)) { + return false; } + // kind is not a binding name, but it is compared -- init() reads "user" + // as the kind validated on shape alone rather than confirmed against the + // engine's own aliasing -- so an escape the two ends read differently is + // refused here too. "\user" and "use\r" both arrive as the plain string + // user, while a JSON reader refuses the first and reads the second as + // "use" and a carriage return. + if (kind_misread) { + return false; + } + // An output binding may be claimed by at most one entry. A second entry + // for the same output names the same binding, so nothing that resolves + // the names can tell the two apart -- but a reader that counts aliased + // outputs per entry, as TensorRTBackend does to size the delegate + // argument list, counts one output twice. Refuse the blob here, where + // the repeat is visible from the bytes alone. + if (!claimed_outputs.insert(ab.output).second) { + return false; + } + // An input may likewise be claimed by at most one entry. Two entries + // naming different outputs and one input record the same input index + // for both, and execute() binds each aliased output to that index's + // caller pointer -- one address with two writers, so whichever the + // engine writes second wins and the other update disappears with no + // error. TensorRT's own aliasing rules out the kv_cache_update kind + // (init cross-checks it against getAliasedInputTensor), but the user + // kind is only compared on shape, so two same-shaped outputs onto one + // input would pass. + if (!claimed_inputs.insert(ab.input).second) { + return false; + } + // The current Python serializer always writes "kind" (serialization.py), + // and older blobs carry no aliased_io array at all, so this default is + // defensive: it only fires for a blob that has an aliased_io entry but + // omits "kind". 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)); } + // Past the alias array, not merely past io_bindings: the array is written + // between the two and this walk advances a position of its own, so the + // scalar scans below would otherwise search across the alias entries and + // match an aliased binding named like one of their keys. + scalars_from = apos; + } + + // The magic says which of the two shapes this blob is, and until here nothing + // compared that against what was parsed. TR02 means the metadata carries + // aliased_io, so an empty list is a blob whose alias array this walk did not + // find -- absent, truncated away, one byte wrong in the key, or written + // before io_bindings, since the search starts past that array. The threaded + // shape does not fail afterwards: every aliased output gets storage of its + // own and the caller's cache quietly stops being updated, which is the + // outcome the aliasing code exists to make impossible. The converse is not + // refused: a TR01 blob carrying an alias array is read as aliased and run + // correctly here, and refusing it would only turn a blob this runtime handles + // into one it does not. + if (expects_aliased_io && out.aliased_io.empty()) { + return false; } - return parse_bool_after_key(json, pos, "\"hardware_compatible\"", out.hardware_compatible) && - parse_int_after_key(json, pos, "\"device_id\"", out.device_id); + return parse_bool_after_key(json, scalars_from, "\"hardware_compatible\"", out.hardware_compatible) && + parse_int_after_key(json, scalars_from, "\"device_id\"", out.device_id); } } // namespace @@ -332,8 +522,8 @@ bool TensorRTBlobHeader::parse(const void* data, std::size_t size, TensorRTBlobH } const auto* bytes = static_cast(data); - if (std::memcmp(bytes, TENSORRT_MAGIC, sizeof(TENSORRT_MAGIC)) != 0 && - std::memcmp(bytes, TENSORRT_MAGIC_ALIASED_IO, sizeof(TENSORRT_MAGIC_ALIASED_IO)) != 0) { + const bool aliased_io_magic = std::memcmp(bytes, TENSORRT_MAGIC_ALIASED_IO, sizeof(TENSORRT_MAGIC_ALIASED_IO)) == 0; + if (!aliased_io_magic && std::memcmp(bytes, TENSORRT_MAGIC, sizeof(TENSORRT_MAGIC)) != 0) { return false; } @@ -359,18 +549,32 @@ bool TensorRTBlobHeader::parse(const void* data, std::size_t size, TensorRTBlobH if (out.engine_offset % ENGINE_ALIGNMENT != 0) { return false; } - if (static_cast(out.metadata_offset) + out.metadata_size > size) { + // Every extent below is checked by subtracting the offset from the bound + // rather than by adding the size to the offset. engine_size is a 64-bit field + // read straight from the file, so the sum form wraps: a blob claiming a size + // just under 2^64 produces a small total, passes, and hands + // deserializeCudaEngine a pointer plus a length far past the end of the file. + // The two metadata extents cannot wrap on a 64-bit size_t -- both operands are + // 32-bit fields, so their sum is at most 2^33 -- but they are written the same + // way so that the form, not the width of each field, is what makes them safe. + // The metadata-against-size check below cannot be the sole reason a blob is + // refused: the other two imply it, because the metadata extent is held inside + // engine_offset and engine_offset inside size. It is kept anyway, because it + // is the metadata extent this function goes on to dereference, and bounding + // that against the file size where it is read does not depend on a chain + // through a check about the engine. + if (out.metadata_offset > size || out.metadata_size > size - out.metadata_offset) { return false; } - if (static_cast(out.engine_offset) + out.engine_size > size) { + if (out.engine_offset > size || out.engine_size > size - out.engine_offset) { return false; } - if (static_cast(out.metadata_offset) + out.metadata_size > out.engine_offset) { + if (out.metadata_offset > out.engine_offset || out.metadata_size > out.engine_offset - out.metadata_offset) { return false; } std::string json(reinterpret_cast(bytes + out.metadata_offset), out.metadata_size); - return parse_metadata_json(json, out); + return parse_metadata_json(json, aliased_io_magic, out); } } // namespace executorch_backend diff --git a/docsrc/py_api/executorch.rst b/docsrc/py_api/executorch.rst index aefb26756cb..13dfbd965a2 100644 --- a/docsrc/py_api/executorch.rst +++ b/docsrc/py_api/executorch.rst @@ -37,6 +37,8 @@ Functions .. autofunction:: export .. autofunction:: get_edge_compile_config +.. autofunction:: zero_copy_backend_config +.. autofunction:: check_zero_copy_kv Classes -------- diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 6470afd72e4..026dd7edaea 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -356,6 +356,146 @@ points but does not by itself give them shared mutable state. Neither case raises an error or a warning, so treat every shared payload as read-only. +.. _executorch_zero_copy_kv: + +**Zero-copy KV cache** + +When a TensorRT engine has aliased I/O -- a KV cache it updates through an +aliased binding -- running the engine over the cache already is the update. +ExecuTorch does not know that, so by default it pays for the update twice per +execution: it hands the delegate an ``_h2d_copy`` staging copy of the buffer +instead of the buffer itself, then copies the engine's aliased output back into +the buffer afterwards. For a KV cache both copies are cache-sized, per token. + +``zero_copy_kv=True`` removes them, so the engine writes the caller's buffer +directly. Through the two-step ``export`` + ``to_executorch`` path it takes two +calls, one at each end of the Edge boundary: + +.. code-block:: python + + from torch_tensorrt.executorch import export, zero_copy_backend_config + + edge = export( + {"prefill": prefill_program, "decode": decode_program}, + zero_copy_kv=True, + ) + + # The argument is optional. Pass your own ExecutorchBackendConfig and + # zero_copy_backend_config composes onto it; every other field (memory + # planning, passes) is preserved. + program = edge.to_executorch(zero_copy_backend_config()) + +Two fields of that config are not merely carried. The engine writes the cache +wherever memory planning put it, so ``enable_non_cpu_memory_planning=False`` -- +which plans every tensor into a single host arena whatever device its +``TensorSpec`` asks for -- cannot be combined with zero-copy KV: +``to_executorch`` raises instead of writing a ``.pte`` whose every ``execute()`` +fails on a host pointer. That refusal reads the config object +``zero_copy_backend_config`` returned, so two ways of setting the field are +outside it. Building a new config out of the returned one with +``dataclasses.replace`` copies the field by value and the pass by reference, so +the pass goes on reading the config it was built for: call +``zero_copy_backend_config`` again on the derived config. And a +``memory_planning_pass`` of your own that does not already have an attribute of +that name never receives the field at all -- ``to_executorch`` assigns it onto +the planner rather than passing it -- so where the caches land is that planner's +own business. Neither is left to the runtime to discover: ``check_zero_copy_kv`` +reads the arena planning actually chose, and the manager ``export`` returns runs +it for you. That planner does have to record the device of each arena it places, +or the check refuses the program for saying nothing about where the cache +lives -- see below. And +``propagate_device_config.skip_h2d_for_method_inputs`` is refused outright: +``PropagateDevicePass`` refuses to un-stage a method input +whose placeholder does not have exactly one user, and a zero-copy cache always +has two, so preserving that option would hand back a configuration that cannot +finalize at all. It is refused wherever it is written -- in one +``PropagateDeviceConfig`` or in a per-method dict of them -- and on every value +that pass reads as on rather than only on ``True``, because it tests the field +for truth without ever resolving it per method, so even a dict of ``False`` is +on for every method. ``False`` and the empty dict are what it reads as off, and +those are carried unchanged. + +It is opt-in rather than automatic because the resulting ``.pte`` needs a +runtime that understands a delegate whose aliased outputs are elided. Producing +one silently would break a runner built before this feature. + +.. warning:: + + **Both calls are required.** Exporting with ``zero_copy_kv=True`` removes + the copy-back; finalizing without ``zero_copy_backend_config`` leaves the + engine writing a per-call staging copy that is discarded, so the cache never + updates. For a KV cache that is wrong output, not a crash. The manager + ``export`` returns is what stops that. Its ``to_executorch`` refuses a + config that does not carry the un-staging pass *before* it finalizes + anything -- finalization rewrites the manager's edge programs in place, so a + refusal after the fact could only be acted on by exporting again -- and it + reads the finalized program back through + ``torch_tensorrt.executorch.check_zero_copy_kv``, which refuses one whose + caches ended up planned somewhere the engine cannot write them, before there + is a ``.pte`` to write. + + ``transform()`` and ``to_backend()`` return a *new* manager, which that + check does not travel to, so a program finalized off one of those needs it + by hand:: + + program = edge.transform(passes).to_executorch(zero_copy_backend_config()) + torch_tensorrt.executorch.check_zero_copy_kv(program) + + ``torch_tensorrt.save`` owns both ends and runs the same check itself. + +``torch_tensorrt.save`` finalizes the program itself, so a single +``zero_copy_kv=True`` covers both steps: + +.. code-block:: python + + torch_tensorrt.save( + trt_gm, "decode.pte", output_format="executorch", + arg_inputs=inputs, retrace=False, + zero_copy_kv=True, + ) + +It installs ``zero_copy_backend_config`` for you, so there is no need to hand it +one as ``backend_config`` as well: that installs the pass twice, which is +redundant rather than an error -- the second run finds the buffers already +un-staged. The two entry points are alternatives, not a pair. + +Three further responsibilities are the caller's. The first two nothing checks; +the third is refused where it is visible in the ``.pte``: + +* **One CUDA stream for every delegate**, if the ``.pte`` is coalesced. Getting + this wrong is a race, not a deterministic error: it is intermittent and can + surface as wrong results *or* as an illegal memory access. See + :ref:`Running a coalesced .pte `. + +* **Synchronizing that stream before reading a cache on the host**, for any + zero-copy ``.pte`` a runner drives on a caller stream, coalesced or not. A + delegate whose aliased outputs are threaded through it reflects each one into + its delegate output and so waits for the engine before returning; zero-copy + elides those outputs, so there is nothing to reflect and ``execute()`` returns + with the engine still running. A single-delegate decode loop owes the + synchronization as much as a coalesced program does. It is only the caller + stream that brings the duty: with none installed ``execute()`` synchronizes + before it returns, as it does whenever a delegate input or output is staged + through the host. + +* **Sharing one cache between methods.** Zero-copy is per method: it makes each + method's engine write that method's buffer. Giving a prefill and a decode + method *the same* cache is a memory-planning question -- their mutable buffers + have to land at the same arena offsets -- which ExecuTorch's memory planner + owns and which is deployment-specific. Supply your own + ``memory_planning_pass`` for it; ``zero_copy_backend_config`` preserves it. + Such a planner has to record the device of each arena it places, which means + running ExecuTorch's ``apply_algo`` -- the only thing that writes + ``non_const_buffer_device`` -- and running it with + ``enable_non_cpu_memory_planning=True``, since that parameter defaults to + ``False`` and with it off ``apply_algo`` plans every spec into one CPU bucket + and records nothing either. Without that record the ``.pte`` reports every + planned buffer as CPU, and a runner that honours it backs the cache with host + memory the engine cannot write. ``check_zero_copy_kv`` refuses a program in + that state, under a message of its own rather than the host-arena one: what + it can read is that nothing says where the cache lives, not that your planner + put it among the host tensors. + **Coalesced TensorRT + CUDA .pte** To run the ops TensorRT does not take on ExecuTorch's CUDA (AOTInductor) backend @@ -390,6 +530,8 @@ must be pointed at those data files to load them. ``.pte`` into the same directory overwrites the blob and the first ``.pte`` will fail to load. Save each coalesced model into its own directory. +.. _executorch_single_stream: + **Running a coalesced .pte: use a single CUDA stream** A coalesced ``.pte`` runs on more than one backend delegate (the TensorRT delegate @@ -403,21 +545,34 @@ illegal memory access. The runtime does not impose a shared stream across delegates, so it is the **runner's responsibility** to run all delegates on one CUDA stream. Create a -single stream and, for the duration of execution, direct every backend to use it -(each backend exposes a caller-stream hook). All GPU work is then enqueued in order -and every cross-boundary dependency is satisfied, while execution stays -asynchronous. +single stream and scope ``executorch::extension::cuda::CallerStreamGuard`` over it +for the duration of execution. That one guard reaches every CUDA-capable +delegate: they resolve a single shared ``libextension_cuda``, so the TensorRT +backend and the CUDA backend read the same caller-stream storage. All GPU work is +then enqueued in order and every cross-boundary dependency is satisfied, while +execution stays asynchronous. If the runner reads a delegate's outputs between calls (for example, an autoregressive decode loop), synchronize the shared stream before reading: the work may still be in flight when ``execute()`` returns, and a host-side copy on -the default stream will not wait for a non-blocking stream. +the default stream will not wait for a non-blocking stream. A model that threads +its aliased outputs through the delegate is insulated from this in practice: +reflecting each aliased output into its delegate output makes the delegate wait +for the engine before it returns. Under +:ref:`zero-copy KV ` those outputs are elided, so there +is nothing to reflect and the delegate returns with the engine still running -- +the synchronization is then the only thing making a host read see the new values. + **ExecuTorch lowering options** -When ``output_format="executorch"``, ``torch_tensorrt.save`` forwards the following -keyword arguments to ExecuTorch's ``to_edge_transform_and_lower(...)``. They are -only consulted for the ``executorch`` format; passing them with any other -``output_format`` logs a warning and is otherwise ignored. +``torch_tensorrt.save`` takes these extra keyword arguments. They are only +consulted for the ``executorch`` format; passing them with any other +``output_format`` logs a warning and is otherwise ignored. Of the six below, +``constant_methods``, ``transform_passes`` and ``compile_config`` are forwarded +to ExecuTorch's ``to_edge_transform_and_lower(...)``, and so is +``generate_etrecord``, which ``save`` also reads itself to write the record +beside the ``.pte``. ``backend_config`` goes to ``to_executorch(...)`` instead, +and ``zero_copy_kv`` is read by ``save`` on both sides of that boundary. * ``constant_methods`` — a ``dict`` of extra constant methods to embed in the ``.pte`` (e.g. ``{"get_max_seq_len": 2048}`` for an LLM runner). @@ -430,6 +585,10 @@ only consulted for the ``executorch`` format; passing them with any other your graph carries TensorRT engines, set ``_check_ir_validity=False`` explicitly. * ``backend_config`` — an ``ExecutorchBackendConfig`` forwarded to ``to_executorch(...)``. +* ``zero_copy_kv`` — a ``bool`` (default ``False``, single-method only). Lets the + TensorRT engine update an aliased KV cache in place. ``save`` owns both ends of + the Edge boundary, so this one argument covers what the two-call path spells + out; see :ref:`Zero-copy KV cache `. * ``generate_etrecord`` — a ``bool`` (default ``False``). When ``True``, an `ETRecord `_ is written next to the ``.pte`` as ``_etrecord.bin`` (e.g. ``trt.pte`` → diff --git a/examples/executorch_reference_runner/BUILD b/examples/executorch_reference_runner/BUILD index 5d85db9f5be..30025c19463 100644 --- a/examples/executorch_reference_runner/BUILD +++ b/examples/executorch_reference_runner/BUILD @@ -33,6 +33,14 @@ cc_binary( ], ) +# A convenience target for building this check out of a bazel checkout. Unlike +# the runner above, which //:bin packages into the release tar wherever this +# backend ships and so compiles on those packaging builds, nothing depends on +# this one: it ships as source in //:executorch_source_package and is compiled +# from the CMakeLists beside it, which is the path +# verify-executorch-reference-runner.sh drives. So this dependency list has to +# be kept in step with that one by hand -- an error here surfaces only when +# someone builds this target. cc_binary( name = "kv_cache_decode_check", srcs = ["kv_cache_decode_check.cpp"], @@ -43,5 +51,9 @@ cc_binary( "@cuda//:cudart", "@executorch//:executorch_core", "@executorch//:executorch_file_data_loader", + # Included and called directly (CallerStreamGuard), so declared directly + # rather than reached through a transitive header, matching the sibling + # runner above. + "@executorch//:extension_cuda", ], ) diff --git a/examples/executorch_reference_runner/README.md b/examples/executorch_reference_runner/README.md index 1b1ccba4544..3ec2f61fbb6 100644 --- a/examples/executorch_reference_runner/README.md +++ b/examples/executorch_reference_runner/README.md @@ -207,3 +207,10 @@ 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. + +That pair runs twice over, printing `caller stream: none` and then `caller +stream: own`. The second scopes a `CallerStreamGuard` over the decode loop on a +stream the runner creates, which is what lets the backend return from +`execute()` with the enqueue still in flight; the first leaves the caller stream +unset, so every `execute()` synchronizes before it returns. A zero-copy `.pte` +reaches the skip-the-sync path only in the second, so both have to pass. diff --git a/examples/executorch_reference_runner/kv_cache_decode_check.cpp b/examples/executorch_reference_runner/kv_cache_decode_check.cpp index 2be3ae1a263..bda398126e5 100644 --- a/examples/executorch_reference_runner/kv_cache_decode_check.cpp +++ b/examples/executorch_reference_runner/kv_cache_decode_check.cpp @@ -22,6 +22,14 @@ * position-0 slot is still zero). Equal logits mean the update did not persist * (cache reset per call, or the aliased output bound to scratch), so we fail. * + * Both scenarios are run twice, once with no caller stream and once with a + * CallerStreamGuard scoped over the decode loop on a stream this runner owns. + * The backend takes a different path in each: with no caller stream it + * synchronizes at the end of every execute(), and with one -- and nothing to + * stage or reflect, which is what a zero-copy .pte leaves -- it may return with + * the enqueue still in flight. Only the second exercises that skip, and only + * there does the runner owe the synchronization before it reads the logits. + * * Usage: * kv_cache_decode_check --model_path=kv_cache_decode.pte [--tol=1e-3] */ @@ -33,10 +41,12 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -72,8 +82,14 @@ static const char* get_flag(int argc, char** argv, const char* flag, const char* // Load a FRESH method (zeroed caller-owned buffers), run one decode step per // entry in `positions` (token id fixed to 1, input_pos = the position), and -// return the final step's first output as host floats. -static std::vector run_decode(Program& program, const char* method_name, const std::vector& positions) { +// return the final step's first output as host floats. With `use_caller_stream` +// the decode loop runs under a CallerStreamGuard on a stream owned here, the way +// the main runner scopes one; without it the backend sees no caller stream. +static std::vector run_decode( + Program& program, + const char* method_name, + const std::vector& positions, + bool use_caller_stream) { Result method_meta = program.method_meta(method_name); ET_CHECK_MSG(method_meta.ok(), "method_meta failed: 0x%" PRIx32, static_cast(method_meta.error())); @@ -148,6 +164,17 @@ static std::vector run_decode(Program& program, const char* method_name, exec_aten::ScalarType::Long, nd, sizes[i].data(), data[i].data(), dim_order[i].data(), strides[i].data()); } + cudaStream_t caller_stream = nullptr; + // Optional, not a guard over a null stream: an explicitly null selection is + // still a selection, so constructing one unconditionally would leave both + // modes with a caller stream set and cover the same branch twice. The guard + // scopes the whole loop rather than each step, so consecutive decodes order on + // one stream instead of on an end-of-execute sync. + std::optional caller_stream_guard; + if (use_caller_stream) { + ET_CHECK_MSG(cudaStreamCreate(&caller_stream) == cudaSuccess, "cudaStreamCreate failed"); + caller_stream_guard.emplace(caller_stream); + } for (int64_t pos : positions) { for (size_t i = 1; i < num_inputs; ++i) { std::fill(data[i].begin(), data[i].end(), pos); @@ -157,6 +184,20 @@ static std::vector run_decode(Program& program, const char* method_name, } ET_CHECK_MSG(method->execute() == Error::Ok, "execute() failed at pos %" PRId64, pos); } + caller_stream_guard.reset(); + if (use_caller_stream) { + // The last execute() may have returned with the enqueue still running, so a + // runner owning the stream owes it this before reading an output. It is what + // the contract asks for rather than what makes the read correct here: + // cudaStreamCreate returns a blocking stream, and the synchronous cudaMemcpy + // below runs on the legacy default stream, which waits for every blocking + // stream in the context. So this exercises the branch but cannot detect a + // missing synchronization -- deleting it leaves the check passing. + ET_CHECK_MSG(cudaStreamSynchronize(caller_stream) == cudaSuccess, "cudaStreamSynchronize failed"); + if (cudaStreamDestroy(caller_stream) != cudaSuccess) { + ET_LOG(Error, "cudaStreamDestroy failed"); + } + } EValue out; ET_CHECK_MSG(method->get_outputs(&out, 1) == Error::Ok, "get_outputs failed"); @@ -164,7 +205,9 @@ static std::vector run_decode(Program& program, const char* method_name, exec_aten::Tensor t = out.toTensor(); ET_CHECK_MSG(t.scalar_type() == exec_aten::ScalarType::Float, "expected float logits output"); // The output may be device-resident; cudaMemcpyDefault copies from host or - // device. execute() synchronized (no caller stream) so the result is ready. + // device. The work is finished either way: with no caller stream the backend + // synchronized at the end of execute(), and with one the stream was + // synchronized above. std::vector result(static_cast(t.numel())); ET_CHECK_MSG( cudaMemcpy(result.data(), t.const_data_ptr(), result.size() * sizeof(float), cudaMemcpyDefault) == cudaSuccess, @@ -191,37 +234,59 @@ int main(int argc, char** argv) { const char* method_name = *name; ET_LOG(Info, "Loaded '%s' method '%s'", model_path, method_name); - // A: pos=1 from a zeroed cache. B: pos=0 then pos=1 (second step sees pos 0). - std::vector a = run_decode(*program, method_name, {1}); - std::vector b = run_decode(*program, method_name, {0, 1}); - - ET_CHECK_MSG(a.size() == b.size() && !a.empty(), "output size mismatch (%zu vs %zu)", a.size(), b.size()); - double max_abs_diff = 0.0; - bool saw_nan = false; - for (size_t i = 0; i < a.size(); ++i) { - // std::max(0.0, fabs(NaN)) is 0.0 (NaN compares false), so a NaN logit would - // otherwise leave max_abs_diff at 0.0 and be misreported as "identical". - // Track NaNs explicitly and fail on them below. - if (std::isnan(a[i]) || std::isnan(b[i])) { - saw_nan = true; + // Both modes have to hold. Without a caller stream the backend synchronizes at + // the end of every execute(); with one it may skip that, which is the branch + // zero-copy KV depends on and the one a wrong result shows up in only + // sometimes. + struct Mode { + const char* label; + bool use_caller_stream; + }; + for (const Mode& mode : {Mode{"none", false}, Mode{"own", true}}) { + // A: pos=1 from a zeroed cache. B: pos=0 then pos=1 (second step sees pos 0). + std::vector a = run_decode(*program, method_name, {1}, mode.use_caller_stream); + std::vector b = run_decode(*program, method_name, {0, 1}, mode.use_caller_stream); + + ET_CHECK_MSG(a.size() == b.size() && !a.empty(), "output size mismatch (%zu vs %zu)", a.size(), b.size()); + double max_abs_diff = 0.0; + bool saw_nan = false; + for (size_t i = 0; i < a.size(); ++i) { + // std::max(0.0, fabs(NaN)) is 0.0 (NaN compares false), so a NaN logit would + // otherwise leave max_abs_diff at 0.0 and be misreported as "identical". + // Track NaNs explicitly and fail on them below. + if (std::isnan(a[i]) || std::isnan(b[i])) { + saw_nan = true; + } + max_abs_diff = std::max(max_abs_diff, std::fabs(static_cast(a[i]) - static_cast(b[i]))); } - max_abs_diff = std::max(max_abs_diff, std::fabs(static_cast(a[i]) - static_cast(b[i]))); - } - fprintf( - stderr, - "[kv-check] logits numel=%zu max|A(no-history) - B(with-history)| = %.6g (tol=%.3g)\n", - a.size(), - max_abs_diff, - tol); - if (saw_nan) { - fprintf(stderr, "[kv-check] FAIL: logits contain NaN -> the decode produced invalid output.\n"); - return 1; - } - if (max_abs_diff > tol) { - fprintf(stderr, "[kv-check] PASS: decode at pos=1 observed the KV written at pos=0 across execute() calls.\n"); - return 0; + fprintf( + stderr, + "[kv-check] caller stream: %s logits numel=%zu max|A(no-history) - B(with-history)| = %.6g (tol=%.3g)\n", + mode.label, + a.size(), + max_abs_diff, + tol); + if (saw_nan) { + fprintf( + stderr, + "[kv-check] FAIL: logits contain NaN -> the decode produced invalid output (caller stream: %s).\n", + mode.label); + return 1; + } + if (max_abs_diff <= tol) { + fprintf( + stderr, + "[kv-check] FAIL: outputs are identical -> the KV write did not persist across execute() calls " + "(caller stream: %s).\n", + mode.label); + return 1; + } + fprintf( + stderr, + "[kv-check] PASS: decode at pos=1 observed the KV written at pos=0 across execute() calls " + "(caller stream: %s).\n", + mode.label); } - fprintf(stderr, "[kv-check] FAIL: outputs are identical -> the KV write did not persist across execute() calls.\n"); - return 1; + return 0; } diff --git a/examples/torchtrt_executorch_example/export_kv_cache_decode.py b/examples/torchtrt_executorch_example/export_kv_cache_decode.py index c8590d98b03..9110efcaa84 100644 --- a/examples/torchtrt_executorch_example/export_kv_cache_decode.py +++ b/examples/torchtrt_executorch_example/export_kv_cache_decode.py @@ -14,6 +14,16 @@ ``.pte`` and asserts that a decode step observes the KV a previous step wrote (i.e. the cache is shared across ``execute()`` calls). +``--zero_copy`` exports the same model with the engine writing the cache buffer +directly, instead of ExecuTorch staging a copy for the delegate and copying the +result back. The persistence check is the same, and it is the check that matters +here: zero-copy removes the copy that was making the update stick, so if the +engine's in-place write is not reaching the caller's buffer the run fails. What +it cannot see is a ``--zero_copy`` export that degenerated into an ordinary +staged ``.pte`` -- the two are indistinguishable to it -- so the manager +``export()`` returns runs ``check_zero_copy_kv`` on the program its +``to_executorch()`` produces and raises rather than let one be written. + Prerequisites ------------- Install Torch-TensorRT with the ExecuTorch extra before running this example:: @@ -22,6 +32,7 @@ """ import argparse +import os import torch import torch_tensorrt @@ -81,11 +92,54 @@ def split_heads(proj: torch.Tensor) -> torch.Tensor: return self.lm(self.o(out)) +def _save_zero_copy(trt_gm: torch.fx.GraphModule, inputs: tuple, path: str) -> None: + """Save a .pte whose engine updates the KV cache in place. + + Zero-copy needs both ends of the Edge boundary: ``zero_copy_kv`` before + lowering, and ``zero_copy_backend_config`` on the config the program is + finalized with. Omitting the second leaves the cache staged and its updates + dropped, and the manager ``export()`` returns catches that: it refuses a + config that does not carry the un-staging pass before it finalizes anything, + and runs ``check_zero_copy_kv`` on whatever its own ``to_executorch()`` did + produce, so the wrong config raises rather than writing a silently staged + .pte. + ``torch_tensorrt.save(output_format="executorch", zero_copy_kv=True)`` owns + both steps and is the shorter way to the same .pte; this spells them out + because both halves are shown, and because a program that reaches + ``to_executorch()`` by any other route has to install the config itself. + """ + from torch_tensorrt.executorch import export, zero_copy_backend_config + + # retrace=True here, retrace=False for the plain save() below, so the two + # exporters are both covered. Which way round matters: the legacy exporter + # declares the aliased KV outputs while building the program, so on that lane + # export()'s declaration pass reads each engine's aliased_io only to find the + # mutations already declared. A retraced program arrives undeclared, so this + # is the lane where that read decides anything -- where an engine-aliased + # cache is told from an ordinary copy-back buffer. + edge = export(trt_gm, arg_inputs=inputs, retrace=True, zero_copy_kv=True) + program = edge.to_executorch(zero_copy_backend_config()) + with open(path, "wb") as output: + program.write_to_file(output) + if program._tensor_data: + # A delegate carrying external weights (the CUDA backend does) keeps them + # outside the .pte, and write_to_file does not persist them; without the + # .ptd next to it the .pte cannot load. This model has none, but a model + # built from this one may. + program.write_tensor_data_to_file(os.path.dirname(os.path.abspath(path))) + + def main() -> None: parser = argparse.ArgumentParser() parser.add_argument( "--model_path", default="kv_cache_decode.pte", help="Path to save the .pte" ) + parser.add_argument( + "--zero_copy", + action="store_true", + help="Let the engine write the KV cache in place (elides the aliased " + "delegate outputs; needs a runtime that supports them).", + ) args = parser.parse_args() with torch.no_grad(): @@ -101,13 +155,16 @@ def main() -> None: min_block_size=1, truncate_double=True, ) - torch_tensorrt.save( - trt_gm, - args.model_path, - output_format="executorch", - arg_inputs=(tokens, input_pos), - retrace=False, - ) + if args.zero_copy: + _save_zero_copy(trt_gm, (tokens, input_pos), args.model_path) + else: + torch_tensorrt.save( + trt_gm, + args.model_path, + output_format="executorch", + arg_inputs=(tokens, input_pos), + retrace=False, + ) print(f"Saved {args.model_path} successfully.") diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index fa70e13c461..d11504d4650 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -685,6 +685,23 @@ def load( ) +# The keyword arguments save() consumes only for output_format="executorch", +# each with the default it is popped with. One table, because the unexpected-keyword +# error spells the supported set out for the caller: an option added to the pops +# alone would leave that message telling someone their flag is unsupported. +_EXECUTORCH_SAVE_OPTIONS: Dict[str, Any] = { + "partitioners": None, + "compile_specs": None, + "backend_config": None, + "constant_methods": None, + "transform_passes": None, + "compile_config": None, + "generate_etrecord": False, + "weight_streaming_budget_per_engine": None, + "zero_copy_kv": False, +} + + def save( module: Any, file_path: str = "", @@ -781,10 +798,10 @@ def save( parameter takes precedence. kwargs: Additional format-specific kwargs. ``partitioners=``, ``compile_specs=``, ``backend_config=``, ``constant_methods=``, - ``transform_passes=``, ``compile_config=``, ``generate_etrecord=`` - and ``weight_streaming_budget_per_engine=`` are only used with - ``output_format="executorch"``; otherwise they are ignored with a - warning. Pass ``compile_specs=[CompileSpec("target_device", + ``transform_passes=``, ``compile_config=``, ``generate_etrecord=``, + ``weight_streaming_budget_per_engine=`` and ``zero_copy_kv=`` are + only used with ``output_format="executorch"``; otherwise they are + ignored with a warning. Pass ``compile_specs=[CompileSpec("target_device", b"cuda:")]`` to override the default target device (``cuda:0``). ``backend_config=`` takes an ``Optional[ExecutorchBackendConfig]`` and is forwarded to ``to_executorch(config=...)`` to customize @@ -815,6 +832,14 @@ def save( resident. Requires the engine to have been compiled with ``enable_weight_streaming=True``. See :func:`torch_tensorrt.executorch.export` for the full description. + ``zero_copy_kv=`` (default ``False``, single-method only) opts a + decode method's KV cache into in-place updates: the TensorRT + engine writes the aliased KV buffer directly instead of receiving + a staging copy that ExecuTorch copies back afterward. Unlike the + direct ``executorch.export()`` + ``to_executorch()`` path -- where + producing zero-copy KV takes two paired calls the caller must not + forget -- ``save()`` owns both steps and installs the finalization + config itself, so a single ``zero_copy_kv=True`` is enough. """ if isinstance(module, CudaGraphsTorchTensorRTModule): module = module.compiled_module @@ -839,16 +864,21 @@ def save( if kwarg_inputs and any(value is None for value in kwarg_inputs.values()): raise ValueError("kwargs should not include None.") - executorch_partitioners = kwargs.pop("partitioners", None) - executorch_compile_specs = kwargs.pop("compile_specs", None) - executorch_backend_config = kwargs.pop("backend_config", None) - executorch_constant_methods = kwargs.pop("constant_methods", None) - executorch_transform_passes = kwargs.pop("transform_passes", None) - executorch_compile_config = kwargs.pop("compile_config", None) - executorch_generate_etrecord = kwargs.pop("generate_etrecord", False) - executorch_weight_streaming_budget_per_engine = kwargs.pop( - "weight_streaming_budget_per_engine", None - ) + executorch_options = { + name: kwargs.pop(name, default) + for name, default in _EXECUTORCH_SAVE_OPTIONS.items() + } + executorch_partitioners = executorch_options["partitioners"] + executorch_compile_specs = executorch_options["compile_specs"] + executorch_backend_config = executorch_options["backend_config"] + executorch_constant_methods = executorch_options["constant_methods"] + executorch_transform_passes = executorch_options["transform_passes"] + executorch_compile_config = executorch_options["compile_config"] + executorch_generate_etrecord = executorch_options["generate_etrecord"] + executorch_weight_streaming_budget_per_engine = executorch_options[ + "weight_streaming_budget_per_engine" + ] + executorch_zero_copy_kv = executorch_options["zero_copy_kv"] if output_format not in accepted_formats: raise ValueError( @@ -864,11 +894,11 @@ def save( # Every executorch option is popped above, so a leftover kwarg is a typo. Fail # here rather than silently ignoring it, since nothing downstream reads kwargs. if kwargs: + supported = ", ".join(repr(name) for name in _EXECUTORCH_SAVE_OPTIONS) raise TypeError( "save() received unexpected keyword argument(s) for " f"output_format='executorch': {sorted(kwargs)}. Supported executorch " - "options are 'partitioners', 'compile_specs', 'backend_config', and " - "'weight_streaming_budget_per_engine'." + f"options are {supported}." ) # Validate the budget before the input and model-shape checks below, so a wrong # type is not reported as an unrelated failure. @@ -879,6 +909,15 @@ def save( normalize_weight_streaming_budget_per_engine( executorch_weight_streaming_budget_per_engine ) + # For the same reason, the one refusal zero-copy makes on the config + # alone. _save_as_executorch reaches it only through + # zero_copy_backend_config, which it calls after export() has partitioned + # the graph and built every engine -- a whole compile before the caller + # is told the combination is not allowed. + if executorch_zero_copy_kv and executorch_backend_config is not None: + from torch_tensorrt.executorch._zero_copy import _refuse_skip_h2d + + _refuse_skip_h2d(executorch_backend_config) def _all_are_input_objects(obj: Any) -> bool: """Recursively check if all elements in nested collections are Input objects.""" @@ -1028,6 +1067,11 @@ def _extract_tensor(obj: Any) -> Any: "output_format='executorch' and will be ignored for " f"output_format='{output_format}'." ) + if executorch_zero_copy_kv and output_format != "executorch": + logger.warning( + "zero_copy_kv= is only used with output_format='executorch' and will " + f"be ignored for output_format='{output_format}'." + ) if output_format == "aot_inductor" and platform.system() != "Linux": raise ValueError( f"The AOT Inductor format is only supported on Linux, {platform.system()} is not a supported platform for this format" @@ -1121,6 +1165,7 @@ def _extract_tensor(obj: Any) -> Any: compile_config=executorch_compile_config, generate_etrecord=executorch_generate_etrecord, weight_streaming_budget_per_engine=executorch_weight_streaming_budget_per_engine, + zero_copy_kv=executorch_zero_copy_kv, ) else: raise RuntimeError( @@ -1238,6 +1283,7 @@ def _extract_tensor(obj: Any) -> Any: compile_config=executorch_compile_config, generate_etrecord=executorch_generate_etrecord, weight_streaming_budget_per_engine=executorch_weight_streaming_budget_per_engine, + zero_copy_kv=executorch_zero_copy_kv, ) else: raise RuntimeError( @@ -1364,6 +1410,7 @@ def _extract_tensor(obj: Any) -> Any: compile_config=executorch_compile_config, generate_etrecord=executorch_generate_etrecord, weight_streaming_budget_per_engine=executorch_weight_streaming_budget_per_engine, + zero_copy_kv=executorch_zero_copy_kv, ) else: raise RuntimeError( @@ -1402,7 +1449,11 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None "(torch_tensorrt_runtime). Reinstall torch_tensorrt with the runtime extension." ) try: - from torch_tensorrt.executorch import export + from torch_tensorrt.executorch import ( + check_zero_copy_kv, + export, + zero_copy_backend_config, + ) except ImportError: raise ImportError( "ExecuTorch is not installed. Install with: pip install " @@ -1421,6 +1472,7 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None ) generate_etrecord = kwargs.get("generate_etrecord", False) + zero_copy_kv = kwargs.get("zero_copy_kv", False) # export() runs the TRT partitioner and to_edge_transform_and_lower itself; it # defaults compile_config to get_edge_compile_config() (_check_ir_validity=False, # since the TRT execute_engine placeholder graph fails edge IR validation) when a @@ -1436,8 +1488,30 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None weight_streaming_budget_per_engine=kwargs.get( "weight_streaming_budget_per_engine" ), + zero_copy_kv=zero_copy_kv, ) - executorch_program = edge_program.to_executorch(config=kwargs.get("backend_config")) + # Unlike the direct export()+to_executorch() path -- where the two steps + # belong to different owners and pairing them is the caller's job -- save() + # owns both, so it installs the finalization pass itself. Wrapping preserves + # every field of the caller's config and reads one of them: a config that + # turns non-CPU memory planning off cannot place the caches where the engine + # writes, and the pass refuses it. A backend_config that already carries + # the pass is wrapped again here; the second run finds the buffers already + # un-staged and changes nothing. + backend_config = kwargs.get("backend_config") + if zero_copy_kv: + backend_config = zero_copy_backend_config(backend_config) + executorch_program = edge_program.to_executorch(config=backend_config) + if zero_copy_kv: + # save() holds the finalized program here, which is the only place that + # shows whether the caches actually reach the engine un-staged and in an + # arena it can write. Both halves of zero-copy no-op quietly when they + # find nothing to do, so without this a save() that asked for zero-copy + # could still write an ordinary staged .pte. export() installs the same + # check on the manager it returns whenever zero-copy was asked for, so + # the two entry points refuse the same models; save() runs it here rather + # than through that manager because it finalizes the program itself. + check_zero_copy_kv(executorch_program) with open(file_path, "wb") as f: executorch_program.write_to_file(f) _write_external_tensor_data(executorch_program, file_path) diff --git a/py/torch_tensorrt/executorch/__init__.py b/py/torch_tensorrt/executorch/__init__.py index fef0943ce7b..7af29ebfb83 100644 --- a/py/torch_tensorrt/executorch/__init__.py +++ b/py/torch_tensorrt/executorch/__init__.py @@ -33,9 +33,15 @@ def __getattr__(name: str) -> NoReturn: "TensorRTPartitioner", "TensorRTBackend", "export", + "zero_copy_backend_config", + "check_zero_copy_kv", ] else: from torch_tensorrt.executorch._export import export + from torch_tensorrt.executorch._zero_copy import ( + check_zero_copy_kv, + zero_copy_backend_config, + ) from torch_tensorrt.executorch.backend import TensorRTBackend from torch_tensorrt.executorch.partitioner import TensorRTPartitioner @@ -50,4 +56,6 @@ def get_edge_compile_config() -> "EdgeCompileConfig": "TensorRTPartitioner", "TensorRTBackend", "export", + "zero_copy_backend_config", + "check_zero_copy_kv", ] diff --git a/py/torch_tensorrt/executorch/_export.py b/py/torch_tensorrt/executorch/_export.py index e793e487519..a0dace3191f 100644 --- a/py/torch_tensorrt/executorch/_export.py +++ b/py/torch_tensorrt/executorch/_export.py @@ -296,6 +296,45 @@ def _apply_weight_streaming_budget( specs.append(CompileSpec(WEIGHT_STREAMING_BUDGET_COMPILE_SPEC_KEY, spec_value)) +def _apply_zero_copy_kv( + program_map: dict[str, ExportedProgram], +) -> dict[str, list[str]]: + """Hand each method's aliased buffers to the engine to update in place. + + Runs immediately after the mutations are declared and before anything + partitions the program: the rewiring works from those declarations and has to + land before the partition boundary fixes the delegate's outputs. It operates + on the staged programs, never the caller's, so a reused ExportedProgram is + left intact. + + Returns, per method that actually lost an output, the engine output binding + names it elided -- narrower than the methods the caller asked about, and + narrower than the engine's full aliased_io. Only these names may be exempted + from the backend's output-binding check, so a method that dropped an output + for some other reason, or an aliased output export never rewired, is still + caught. + """ + from torch_tensorrt.executorch._zero_copy import ( + rewire_aliased_mutations_to_buffers, + ) + + elided = { + name: rewire_aliased_mutations_to_buffers(program) + for name, program in program_map.items() + } + if not any(elided.values()): + logger.warning( + "zero_copy_kv=True, but no aliased buffer mutation was found in %s, " + "so zero-copy KV was not applied. The returned manager still checks " + "its finalized program, so to_executorch() will refuse it; export " + "without zero_copy_kv to get an ordinary staged program.", + ", ".join(sorted(program_map)), + ) + return {} + logger.debug("zero-copy KV: elided outputs per method: %s", elided) + return {name: names for name, names in elided.items() if names} + + def _per_method_values( value: Sequence[Any] | Mapping[str, Sequence[Any] | None] | None, method_names: tuple[str, ...], @@ -328,6 +367,42 @@ def _per_method_values( return {name: list(shared) for name in method_names} +def _reject_caller_set_zero_copy_spec( + method_compile_specs: dict[str, list[Any]], +) -> None: + """Reject a caller who sets the reserved zero-copy key by hand. + + export() sets this key itself, and only for the aliased outputs it actually + elided. A hand-set value would tell the backend that outputs were taken out + of the delegate that were not, so the backend would stop rejecting a delegate + that is genuinely short an output -- dropping a real KV update silently. + + Written for this one key rather than parametrized over a reserved key: the + wording is the safety property this key carries, which no other key shares. + ``_apply_weight_streaming_budget`` rejects its own reserved key inline rather + than through this function: it also *writes* a spec, so its rejection is one + branch of a larger operation, and the two keys fail for different reasons -- + that one has a supported argument to point the caller at, this one has a + safety property that no argument can restore. + + The message names the method whose specs carry the key, as + ``_apply_weight_streaming_budget`` does. A caller who gave one shared list + gets it fanned out to every method, so the name is then whichever method came + first; that imprecision is the sibling's too. + """ + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + for name, specs in method_compile_specs.items(): + for spec in specs: + if getattr(spec, "key", None) == ZERO_COPY_KV_COMPILE_SPEC_KEY: + raise ValueError( + f"compile_specs for {name!r} may not set the reserved key " + f"'{ZERO_COPY_KV_COMPILE_SPEC_KEY}'; export() sets it " + "itself when zero_copy_kv=True elides a method's aliased " + "outputs. Setting it by hand could drop a real KV update silently." + ) + + def export( source: ExportedProgram | torch.fx.GraphModule | Mapping[str, ExportedProgram], *, @@ -345,6 +420,7 @@ def export( ) = None, compile_config: "EdgeCompileConfig | None" = None, constant_methods: Mapping[str, Any] | None = None, + zero_copy_kv: bool = False, generate_etrecord: bool = False, weight_streaming_budget_per_engine: int | None = None, ) -> "EdgeProgramManager": @@ -372,6 +448,30 @@ def export( specs name no method is not rejected here, but a backend that reads its own method name from its specs, such as the CUDA backend, then raises during lowering. + ``zero_copy_kv=True`` lets a TensorRT engine update an aliased mutable buffer + -- a KV cache -- in place, instead of ExecuTorch handing the delegate a + staging copy and copying the engine's result back afterwards. It is opt-in + rather than automatic for two reasons: the resulting ``.pte`` needs a + runtime that understands a delegate with its aliased outputs elided, so + producing one silently would break an older runner; and it is only half the + change. Finalize such a program with + ``to_executorch(torch_tensorrt.executorch.zero_copy_backend_config(config))`` + -- without it the buffer is still staged and its updates would be discarded. + The returned manager refuses that rather than handing back a program whose + caches never update: it reads the config before it finalizes anything, so a + config without the pass is refused while the manager can still be finalized + again with one, and it reads the finalized program back through + :func:`torch_tensorrt.executorch.check_zero_copy_kv` for what the config + alone cannot say. That reaches the manager returned here and no other: + ``transform()`` and ``to_backend()`` build a new one, so a program finalized + off either owes that call by hand. + + Only a buffer the engine declares aliased is affected. A method may hold both + kinds at once: a mutable buffer with no aliasing available -- a convolution + state, say -- keeps its staging copy and the copy-back that writes it, while + the aliased caches beside it go zero-copy. The two are told apart by the + engine's own ``aliased_io``, not by the graph, in which they look identical. + ``generate_etrecord=True`` is outside the payload sharing described above. It makes ExecuTorch deep copy the whole program, so peak memory grows by roughly the size of the program including engines. @@ -419,6 +519,14 @@ def export( constant_methods (Dict[str, Any]): Methods returning a constant, such as a vocab size. Keys must be valid Python identifiers and must not name a method of ``source``. + zero_copy_kv (bool): Let a TensorRT engine write its aliased mutable buffer -- + a KV cache -- in place, instead of receiving a staging copy that ExecuTorch + copies back. Requires finalizing the returned program with + :func:`torch_tensorrt.executorch.zero_copy_backend_config`; the returned + manager's ``to_executorch`` refuses a program finalized without it, since + the buffer would still be staged and every update discarded. + ``torch_tensorrt.save(output_format="executorch", + zero_copy_kv=True)`` owns both steps and needs only the one flag. generate_etrecord (bool): Ask ExecuTorch for an ETRecord for later debugging. This copies the whole program, engines included. weight_streaming_budget_per_engine (Optional[int]): Bytes of engine weights that @@ -456,6 +564,7 @@ def export( import torch_tensorrt.dynamo.runtime.meta_ops.register_meta_ops # noqa: F401 from executorch.exir import to_edge_transform_and_lower + from executorch.exir.backend.compile_spec_schema import CompileSpec as _CompileSpec from torch_tensorrt.dynamo._exporter import _declare_aliased_kv_mutations_on_ep from torch_tensorrt.executorch import TensorRTPartitioner, get_edge_compile_config from torch_tensorrt.executorch._export_utils import ( @@ -463,6 +572,14 @@ def export( stage_exported_program, validate_engine_program, ) + from torch_tensorrt.executorch._zero_copy import ( + _check_zero_copy_kv_when_finalized, + order_copyback_mutations_first, + ) + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names, + ) programs, method_names = _prepare_programs( source, @@ -501,6 +618,7 @@ def export( method_compile_specs = _per_method_values( compile_specs, method_names, "compile_specs" ) + _reject_caller_set_zero_copy_spec(method_compile_specs) _apply_weight_streaming_budget( method_compile_specs, weight_streaming_budget_per_engine ) @@ -568,6 +686,8 @@ def export( ) for name, program in program_map.items() } + zero_copy_methods = _apply_zero_copy_kv(staged_programs) if zero_copy_kv else {} + rewritten: dict[str, ExportedProgram] = {} method_partitioners: dict[str, list[Partitioner]] = {} for name, program in staged_programs.items(): @@ -589,8 +709,29 @@ def export( ) # Drop this method's engine payloads as soon as they are in the graph. rewritten[name] = replace_execute_engine(program, resolved_engines.pop(name)) + trt_compile_specs = list(method_compile_specs[name]) + if name in zero_copy_methods: + # Signal to TensorRTPartitioner that this method elided aliased + # outputs. Only the presence of the key is read: the partitioner + # drops this spec from the list it applies to every partition and + # re-derives, per engine, exactly which of a delegate's aliased + # outputs were elided, stamping only those onto only that delegate + # (see TensorRTPartitioner._partition_elided_output_names). So a + # method that lowers to several TensorRT delegates marks only the KV + # one and a plain-compute delegate beside it carries no zero-copy + # spec. Without any spec the backend rejects a delegate short of its + # bindings, which keeps an accidentally dropped output an error. + # The method-wide names go in the value so this spec reads like the + # per-partition one the partitioner writes, but nothing decodes this + # one: changing which names are listed here changes no delegate. + trt_compile_specs.append( + _CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(zero_copy_methods[name]), + ) + ) method_partitioners[name] = [ - TensorRTPartitioner(compile_specs=method_compile_specs[name]), + TensorRTPartitioner(compile_specs=trt_compile_specs), *extra_partitioners[name], ] @@ -603,7 +744,7 @@ def export( edge_programs = rewritten["forward"] partitioner_pipeline = method_partitioners["forward"] - return to_edge_transform_and_lower( + edge_manager = to_edge_transform_and_lower( edge_programs, transform_passes=transform_passes, partitioner=partitioner_pipeline, @@ -613,3 +754,31 @@ def export( ), generate_etrecord=generate_etrecord, ) + # After lowering, not beside the rewiring: to_edge re-derives the graph + # signature, so an order set earlier does not reach the finalizer. See + # order_copyback_mutations_first. + # + # Every method, not only the ones zero-copy rewired, because the crossing + # this repairs is not zero-copy's. A plain nn.Module writing one buffer from + # another buffer and a second buffer from a user input -- no TensorRT, no + # zero_copy_kv -- comes out of stock to_edge().to_executorch() with the first + # buffer's mutation spec naming the copy that writes the second. Gating this + # on zero_copy_methods would make a general repair depend on an unrelated + # opt-in, and would leave a sibling method repaired or not according to + # whether some other method asked for zero-copy. + for name in edge_manager.methods: + order_copyback_mutations_first(edge_manager.exported_program(name)) + if zero_copy_kv: + # The copy-back is gone from here on, so finalizing without the matching + # un-staging writes a .pte whose caches never update. This makes the + # manager's own to_executorch read the finalized program back and refuse + # that, which the two-call API otherwise leaves to the caller. + # + # Gated on what the caller asked for, not on what was rewired. A method + # with no aliased buffer mutation gets the warning above and a plain + # staged program, which is precisely one of the shapes the check refuses + # -- and save(..., zero_copy_kv=True) refuses that same model. Installing + # the hook only where something was rewired would leave the one entry + # point that cannot catch it the one that never looks. + _check_zero_copy_kv_when_finalized(edge_manager) + return edge_manager diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py new file mode 100644 index 00000000000..a47ff4273a2 --- /dev/null +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -0,0 +1,1639 @@ +"""Let a TensorRT engine update an aliased mutable buffer in place. + +An engine with aliased I/O (a KV cache) writes its aliased output *through* the +aliased input's pointer, so running the engine over the buffer already is the +update. Nothing in the ExecuTorch pipeline knows that, so by default the buffer +makes a full round trip on every execution: + +* ``PropagateDevicePass`` wraps every delegate input in ``et_copy._h2d_copy``, + so the delegate is handed a per-call staging copy rather than the caller's + buffer. The engine's in-place write lands in that copy. +* the aliased output is threaded back out as a delegate output, and ExecuTorch + copies it into the buffer afterwards to make the update stick. + +For a cache-sized buffer that is two copies per execution of something the +engine could have written directly. This module removes both, in the same +spirit as ``partitioner._keep_mutated_buffers_above_delegate``: let the upstream +pass run, then correct its output for the case Torch-TensorRT owns. + +The two halves are inseparable and run at different times: + +* :func:`rewire_aliased_mutations_to_buffers`, on the exported program before + partitioning, drops the copy-back by declaring that the buffer *is* the + mutation's result. The aliased output then has no user and disappears from the + partition. :func:`order_copyback_mutations_first` then repairs, on the Edge + program, the mutation-spec pairing that declaration disturbs downstream -- + a crossing this is one cause of and not the only one, which is why + ``export()`` runs that repair over every method rather than the rewired ones. +* :func:`unstage_aliased_buffers_pass`, as a ``to_out_var_pass``, drops the + staging so the engine writes the caller's buffer rather than a copy. + +Applying only the first would leave the engine writing a discarded staging copy +with nothing to copy back -- the buffer would simply never update. So neither +pass is public on its own: the rewiring is reached only through +``export(..., zero_copy_kv=True)``, and the un-staging only through +:func:`zero_copy_backend_config`. That, plus :func:`check_zero_copy_kv` -- which +reads a finalized program back and refuses one in which the engine does not end +up writing the caller's buffer -- is what this module exports. +""" + +import functools +import json +import logging +import operator +from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Set + +import torch +from executorch.exir.pass_base import PassBase +from torch.fx import Node + +if TYPE_CHECKING: + from executorch.exir import ExecutorchBackendConfig + +logger = logging.getLogger(__name__) + + +def _aliased_inputs_by_output_index( + exported_program: Any, engine_node: Node +) -> Dict[int, Node]: + """Map each aliased output index of one engine to the input it writes in place. + + Reads the engine's own ``aliased_io`` rather than inferring aliasing from the + graph. The graph cannot tell the difference: an aliased KV mutation and a + copy-back mutation are both a ``getitem`` off the engine node whose buffer is + also an engine input, and rewiring a copy-back would silently drop a real + update. An entry whose aliased *input* does not resolve -- the name is not one + of the engine's input bindings, or its index is past the delegate's argument + list -- is skipped rather than reported: + ``_declare_aliased_kv_mutations_on_ep`` warns on both of those for the same + engine, and neither leaves a mutation to rewire. + """ + from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ( + ALIASED_IO_IDX, + INPUT_BINDING_NAMES_IDX, + OUTPUT_BINDING_NAMES_IDX, + deserialize_binding_names, + ) + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + deserialize_aliased_io, + ) + from torch_tensorrt.executorch._export_utils import _resolve_engine_info + from torch_tensorrt.executorch.backend import _get_str + + # Only aliased_io and the binding names are read, never the engine itself. + engine_info = _resolve_engine_info( + exported_program, engine_node, metadata_only=True + ) + aliased_io = deserialize_aliased_io(_get_str(engine_info, ALIASED_IO_IDX)) + if not aliased_io: + return {} + input_names = deserialize_binding_names( + _get_str(engine_info, INPUT_BINDING_NAMES_IDX) + ) + output_names = deserialize_binding_names( + _get_str(engine_info, OUTPUT_BINDING_NAMES_IDX) + ) + input_nodes = list(engine_node.args[0]) + + aliased: Dict[int, Node] = {} + for output_index, output_name in enumerate(output_names): + entry = aliased_io.get(output_name) + if entry is None: + continue + input_name = entry[0] + if input_name not in input_names: + continue + input_index = input_names.index(input_name) + if input_index >= len(input_nodes): + continue + aliased[output_index] = input_nodes[input_index] + return aliased + + +def _engine_output_binding_names(exported_program: Any, engine_node: Node) -> List[str]: + """Return one engine's output binding names, in binding (index) order. + + Resolved metadata-only: reading the record without that costs a full + re-serialization of the engine through ``TRTEngine.__getstate__``, and only + the binding names are wanted here. Callers that read this repeatedly for the + same engine memoize it themselves -- ``_resolve_engine_info`` holds no cache. + """ + from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ( + OUTPUT_BINDING_NAMES_IDX, + deserialize_binding_names, + ) + from torch_tensorrt.executorch._export_utils import _resolve_engine_info + from torch_tensorrt.executorch.backend import _get_str + + engine_info = _resolve_engine_info( + exported_program, engine_node, metadata_only=True + ) + names: List[str] = deserialize_binding_names( + _get_str(engine_info, OUTPUT_BINDING_NAMES_IDX) + ) + return names + + +class _AliasedMutation(NamedTuple): + """One BUFFER_MUTATION an engine satisfies by writing the buffer in place.""" + + placeholder: Node # the buffer, as a graph input + aliased_output: Node # getitem(engine, i) currently standing in for it + engine: Node # the execute_engine call that performs the write + + +def _aliased_buffer_mutations( + exported_program: Any, +) -> Dict[int, _AliasedMutation]: + """Find the BUFFER_MUTATIONs an engine performs in place. + + Returns ``{index into graph_signature.output_specs: _AliasedMutation}``. + A mutation qualifies only when its value is ``getitem(engine_node, i)`` and + the engine declares output ``i`` as aliased onto that very buffer, so a + buffer mutated by an op outside the engine, or copied back out of one, is + left alone. + """ + from torch.export.graph_signature import OutputKind + + graph_module = exported_program.graph_module + signature = exported_program.graph_signature + execute_engine = torch.ops.tensorrt.execute_engine.default + + buffer_placeholders = { + fqn: node + for node in graph_module.graph.nodes + if node.op == "placeholder" + and (fqn := signature.inputs_to_buffers.get(node.name)) is not None + } + output_args = list(graph_module.graph.output_node().args[0]) + aliased_by_engine: Dict[Node, Dict[int, Node]] = {} + + mutations: Dict[int, _AliasedMutation] = {} + for spec_index, spec in enumerate(signature.output_specs): + if spec.kind != OutputKind.BUFFER_MUTATION or spec_index >= len(output_args): + continue + placeholder = buffer_placeholders.get(spec.target) + if placeholder is None: + continue + value = output_args[spec_index] + if ( + not isinstance(value, Node) + or value.op != "call_function" + or value.target is not operator.getitem + ): + continue + engine_node = value.args[0] + if ( + not isinstance(engine_node, Node) + or engine_node.op != "call_function" + or engine_node.target is not execute_engine + ): + continue + if engine_node not in aliased_by_engine: + aliased_by_engine[engine_node] = _aliased_inputs_by_output_index( + exported_program, engine_node + ) + if aliased_by_engine[engine_node].get(value.args[1]) is placeholder: + mutations[spec_index] = _AliasedMutation( + placeholder=placeholder, aliased_output=value, engine=engine_node + ) + return mutations + + +def rewire_aliased_mutations_to_buffers(exported_program: Any) -> List[str]: + """Declare that an aliased buffer *is* its own mutation result. + + Export declares an aliased KV mutation as a ``getitem`` off the engine node: + the engine's aliased output, surfaced as a value. ExecuTorch implements that + mutation by copying the value back into the buffer, which is the copy this + removes. Repointing the mutation at the buffer placeholder leaves nothing to + copy, and with no other user the ``getitem`` dies -- so the aliased output + also leaves the partition and the delegate never receives an argument for it. + + This must run before partitioning, because it is the partition boundary that + freezes which outputs the delegate has. It must also run after export has + declared the aliased mutations, since it works from those declarations; each + placeholder it rewires is marked for + :func:`unstage_aliased_buffers_pass`, which cannot re-derive the aliasing + once lowering has turned the engine into an opaque blob. + + On its own this is not correct: ExecuTorch still stages the buffer, so the + engine's in-place write would land in per-call scratch and, with the + copy-back gone, be lost. It is only correct paired with the un-staging pass. + + Returns the engine output binding names of the aliased outputs it elided, + one per rewired mutation. Only these names may later be exempted from the + backend's output-binding check -- every *other* aliased output (a user alias + on a plain, non-buffer input, which export never rewired) must still be a + delegate output, so a delegate that dropped one of those as well is caught + rather than silently writing that update into scratch. An engine that mixes + the two kinds does not lower at all: ``TensorRTBackend.preprocess`` refuses + it, because the runtime reads elision off a single argument count and so + cannot express eliding only part of one engine's aliased outputs. + """ + from torch.export.graph_signature import ( + ExportGraphSignature, + OutputKind, + OutputSpec, + TensorArgument, + ) + + graph_module = exported_program.graph_module + signature = exported_program.graph_signature + mutations = _aliased_buffer_mutations(exported_program) + if not mutations: + logger.debug("no aliased buffer mutations to rewire") + return [] + + engines_with_elided_outputs: Set[Node] = set() + output_names_by_engine: Dict[Node, List[str]] = {} + elided_output_names: List[str] = [] + output_node = graph_module.graph.output_node() + output_args = list(output_node.args[0]) + output_specs = list(signature.output_specs) + for spec_index, mutation in mutations.items(): + # Marked on the node rather than read back off the engine because the + # un-staging pass runs after lowering, where the engine's aliased_io is no + # longer reachable from the graph: it has become an opaque delegate blob. + mutation.placeholder.meta["_torch_tensorrt_aliased_buffer"] = True + output_args[spec_index] = mutation.placeholder + output_specs[spec_index] = OutputSpec( + OutputKind.BUFFER_MUTATION, + TensorArgument(name=mutation.placeholder.name), + output_specs[spec_index].target, + ) + engines_with_elided_outputs.add(mutation.engine) + names = output_names_by_engine.get(mutation.engine) + if names is None: + names = _engine_output_binding_names(exported_program, mutation.engine) + output_names_by_engine[mutation.engine] = names + output_index = mutation.aliased_output.args[1] + if 0 <= output_index < len(names): + elided_output_names.append(names[output_index]) + + output_node.args = (tuple(output_args),) + graph_module.graph.eliminate_dead_code() + # Leaving an engine with no output would leave its delegate with no outputs. + # Nothing downstream reports that shape: the runtime infers elision from a + # single argument count, which a zero-output delegate satisfies, and a + # delegate nothing reads is a pure node that a later graph-wide dead-code + # elimination can erase, taking the computation with it. This raise is what + # stops that. It reads the graph after the elimination above rather than + # before, so that an output kept alive only by a chain that is itself dead + # does not count: the elimination erases such a chain however long it is, + # and what survives it is what the delegate will really have. The engine + # node itself survives even with no users -- PyTorch defaults an operator + # taking a ScriptObject argument to an ORDERED effect + # (torch._library.effects), and execute_engine takes the engine as one, so + # FX reads it as impure where the delegate is not. + for engine in engines_with_elided_outputs: + if not engine.users: + raise RuntimeError( + "TensorRT zero-copy KV: eliding the aliased buffers engine node " + f"'{engine.name}' writes in place leaves it with no output any " + "node reads, so the delegate would have no outputs at all. This " + "shape is not supported; export this method without " + "zero_copy_kv." + ) + graph_module.graph.lint() + graph_module.recompile() + # The signature is replaced in place rather than by rebuilding the program: + # the graph has already been edited in place, and every other field would be + # copied across unchanged. + exported_program._graph_signature = ExportGraphSignature( + input_specs=list(signature.input_specs), output_specs=output_specs + ) + logger.debug( + "rewired %d aliased mutation(s) to their buffers, eliding outputs %s", + len(mutations), + elided_output_names, + ) + return elided_output_names + + +def _mutation_targets_a_lifted_input(signature: Any) -> Set[str]: + """The mutation targets ``insert_write_back_for_buffers_pass`` can resolve. + + Mirrors the ``lifted_inputs`` map that pass builds: a buffer, constant, + parameter or custom object contributes its ``target``, a user input its + argument name. A mutation whose target is not in here gets no copy either, + so it belongs with the copy-free ones. + """ + from torch.export.graph_signature import InputKind, TensorArgument + + lifted: Set[str] = set() + for spec in signature.input_specs: + if spec.kind in ( + InputKind.BUFFER, + InputKind.CONSTANT_TENSOR, + InputKind.PARAMETER, + InputKind.CUSTOM_OBJ, + ): + if spec.target is not None: + lifted.add(spec.target) + elif spec.kind is InputKind.USER_INPUT and isinstance(spec.arg, TensorArgument): + lifted.add(spec.arg.name) + return lifted + + +def order_copyback_mutations_first(exported_program: Any) -> int: + """Reorder one Edge method's mutations so ExecuTorch pairs them up correctly. + + ExecuTorch finalizes a mutation by inserting a ``copy_`` for it, but only + when its target is one of the lifted inputs and its value is not already + reached, through in-place ops, from a placeholder of the mutation's own kind + -- any buffer placeholder for a buffer mutation, not necessarily the one it + targets. It moves those copies to the front of the output tuple, leaves + everything else behind them in order, and then walks the mutation specs + reassigning each one's argument *by position* over the result + (``insert_write_back_for_buffers_pass``). + :func:`rewire_aliased_mutations_to_buffers` makes a mutation's value its own + placeholder, so a rewired cache gets no copy and drops out of that leading + run -- and in a method that also has a copy-back buffer, every mutation spec + from the first copy-free one on then comes out of finalization naming a + different buffer's value. The specs ahead of it are unaffected, since the + copies keep their order among themselves. The ``.pte`` is written correctly + either way, because the emitter and the memory planner read only which + buffers are mutated and not what by. What the pairing decides is the + finalized signature, which anyone inspecting the program reads, and + ExecuTorch's eager call path, which walks ``buffers_to_mutate`` writing the + graph's leading results into the state dict in that order and so updates + each buffer from another one's value. + + Zero-copy is not the only way in, which is why ``export()`` runs this over + every method rather than only the rewired ones. A plain ``nn.Module`` that + writes one buffer from another buffer -- whose mutation value is then that + other buffer's placeholder -- and a second buffer from a user input comes out + of stock ``to_edge().to_executorch()``, with no TensorRT anywhere, with the + first buffer's mutation spec naming the copy that writes the second. + + Putting the mutations that still get a copy first restores the + correspondence. Which ones those are is decided by asking upstream's own + predicate (``_inplace_lineage``, imported rather than reimplemented) rather + than by asking which mutations this module rewired, so the answer cannot + drift from the one the write-back pass will give, and any mutation the graph + already presents as in-place is covered whatever put it there. + + Only what the graph presents *here* is covered, though. ``run_reinplace_pass`` + and ``reinplace_extra_ops`` are supported ``ExecutorchBackendConfig`` fields + whose pass runs inside ``to_executorch``, after this and immediately before + the write-back: a mutation it rewrites is ordinary when this reads it and + in-place when the write-back does, so that pair comes out crossed anyway and + this reports nothing moved. Reordering cannot reach it from here. The same + crossing reproduces on a model using no zero-copy at all, so it is a pass + ordering upstream owns rather than one this creates. + + This runs on the *Edge* program rather than beside the rewiring, because + ``to_edge_transform_and_lower`` re-derives the whole graph signature -- the + order it hands back is the buffers' own order, whatever order it was given. + Nothing between here and the write-back pass re-derives it again. + + Returns the number of mutation slots whose value changed, which is zero when + the order already holds. + """ + from executorch.exir.passes.insert_write_back_for_buffers_pass import ( + _inplace_lineage, + ) + from torch.export.graph_signature import ExportGraphSignature, OutputKind + + signature = exported_program.graph_signature + specs = list(signature.output_specs) + output_node = exported_program.graph_module.graph.output_node() + args = list(output_node.args[0]) + slots = [ + index + for index, spec in enumerate(specs) + if spec.kind in (OutputKind.BUFFER_MUTATION, OutputKind.USER_INPUT_MUTATION) + and index < len(args) + ] + if not slots: + return 0 + lifted = _mutation_targets_a_lifted_input(signature) + + def gets_a_copy(index: int) -> bool: + value = args[index] + if not isinstance(value, Node): + # Upstream reads a non-Node value as needing a copy, and then raises + # walking it. Grouping it with the copies keeps this reorder from + # being what raises first. + return True + if specs[index].target not in lifted: + return False + return not _inplace_lineage(value, signature, specs[index].kind) + + copied = {index for index in slots if gets_a_copy(index)} + source = [index for index in slots if index in copied] + [ + index for index in slots if index not in copied + ] + if source == slots: + return 0 + + new_args, new_specs = list(args), list(specs) + for slot, index in zip(slots, source): + new_args[slot] = args[index] + new_specs[slot] = specs[index] + output_node.args = (tuple(new_args),) + exported_program.graph_module.recompile() + exported_program._graph_signature = ExportGraphSignature( + input_specs=list(signature.input_specs), output_specs=new_specs + ) + moved = sum(1 for slot, index in zip(slots, source) if slot != index) + logger.debug("moved %d mutation(s) so the copy-back ones come first", moved) + return moved + + +def _is_tensorrt_delegate(graph_module: torch.fx.GraphModule, node: Node) -> bool: + """True when ``node`` is a call_delegate dispatching to the TensorRT backend. + + Only a TensorRT engine promises the aliased-binding write; another backend's + delegate may legitimately need the staging copy. + """ + from executorch.exir.delegate import executorch_call_delegate + from torch_tensorrt.executorch.backend import TensorRTBackend + + if node.op != "call_function" or node.target is not executorch_call_delegate: + return False + lowered = node.args[0] if node.args else None + if not isinstance(lowered, Node) or lowered.op != "get_attr": + return False + module = getattr(graph_module, lowered.target, None) + return bool(getattr(module, "backend_id", None) == TensorRTBackend.__name__) + + +def _zero_copy_compile_spec(graph_module: torch.fx.GraphModule, node: Node) -> Any: + """One delegate's zero-copy KV compile spec, or ``None`` if it carries none.""" + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + lowered = node.args[0] if node.args else None + if not isinstance(lowered, Node) or lowered.op != "get_attr": + return None + module = getattr(graph_module, lowered.target, None) + for spec in getattr(module, "compile_specs", None) or []: + if getattr(spec, "key", None) == ZERO_COPY_KV_COMPILE_SPEC_KEY: + return spec + return None + + +def _delegate_elided_output_names( + graph_module: torch.fx.GraphModule, node: Node +) -> Set[str]: + """The aliased output binding names one delegate's zero-copy spec claims. + + ``TensorRTPartitioner`` writes one name per aliased output it elided on that + delegate's own engine, and an aliased output is elided exactly when a buffer + mutation writes it in place, so the size of this set is how many marked + buffers the delegate has to take. Empty when the delegate carries no such + spec, when it carries one whose value does not decode into a list of names, + and when that value decodes into an empty list -- the last two are shapes a + spec built by hand produces, since the partitioner writes one JSON name per + aliased output it elided. Callers read empty as "cannot tell", not as "none". + + ``backend._elided_output_names`` reads the same key on the same spec and + takes the same shapes of value. It differs in what it does with the rest: it + raises naming the key, because an undecodable spec leaves it unable to say + which outputs the delegate was allowed to drop, while here it only weakens + the two cross-checks that read it, both of which then fall back to demanding + at least one buffer. The one value the two read differently is bytes that + are not valid UTF-8: it replaces the bad units and reads a name out of them, + ``json.loads`` refuses them here, and here that is another "cannot tell". + """ + spec = _zero_copy_compile_spec(graph_module, node) + if spec is None: + return set() + value = getattr(spec, "value", None) + if not isinstance(value, (str, bytes, bytearray)): + return set() + try: + names = json.loads(value) + except ValueError: + return set() + return {str(name) for name in names} if isinstance(names, list) else set() + + +def _delegate_declares_zero_copy( + graph_module: torch.fx.GraphModule, node: Node +) -> bool: + """True when a TensorRT delegate carries the zero-copy KV compile spec. + + ``TensorRTPartitioner`` stamps this spec per partition, onto only the delegate + whose own engine had an aliased output elided (derived per engine in + ``TensorRTPartitioner._partition_elided_output_names``), so a delegate that + declares it must have had a buffer un-staged here. A method that lowers to + several TensorRT delegates therefore marks only the KV one, never the plain + compute engines beside it -- which is what keeps this cross-check from + demanding an aliased buffer from a delegate that never had one. A delegate + that declares it but ends up taking fewer marked buffers than it elided + aliased outputs has lost a KV update -- the mark that would have driven the + un-staging did not survive, or another delegate took the buffer -- and is + caught in :func:`_unstage_aliased_buffers` before planning and again in + :func:`check_zero_copy_kv` after it. + """ + return _zero_copy_compile_spec(graph_module, node) is not None + + +def _device_placement_is_safe( + graph_module: torch.fx.GraphModule, + source: Node, + h2d_copy: Any, + target_device: Any, + target_device_index: Any, + already_removed: Any = (), +) -> bool: + """True when every other consumer of ``source`` survives its device placement. + + A placeholder's device is shared by every user, so it can only be planned in + the delegate's device memory when nothing else *reads* it from somewhere + else. ExecuTorch guards the same hazard, more strictly and only under its + opt-in ``skip_h2d_for_method_inputs``: it demands the placeholder have + exactly one user. The rule here is looser because some users survive + unaffected and are allowed: the graph ``output`` node -- the buffer is its + own BUFFER_MUTATION result, which is exactly what zero-copy sets up and + which carries no device of its own -- an ``_h2d_copy`` to the same GPU that + this pass removes, because every one of its users is a TensorRT delegate + whose argument the pass rewires to the buffer, and a TensorRT delegate + already taking the buffer itself, which is the shape this pass leaves behind + and so what a second marked delegate, or a second run, finds. + + ``already_removed`` is the last of them: the staging copies this run has + detached, which are still in the graph because they are erased only once the + walk has succeeded. A detached one has no users left, which is the shape a + *foreign* dead copy has too, and that one is refused -- nothing erases it, + so the emitter keeps it and it reads device memory as a host source. Ours + are named here rather than inferred, so the two cannot be confused. + + A staging copy that outlives the pass is *not* allowed, even on the same + GPU. It would go on reading the buffer as its source once the buffer is in + device memory, and ``_h2d_copy_out`` requires a host source: the portable + kernel checks it and fails ``InvalidArgument``. + + The index is compared as well as the type, because ``spec.device`` is only + ``CUDA``/``CPU``: two engines resolved to ``cuda:0`` and ``cuda:1`` stage the + same buffer to different GPUs, and un-staging both would leave whichever ran + last owning the buffer while the other engine writes an address on the wrong + device. + + This is asked whether or not the buffer's spec already names the target + device. Where it does, the pass changes no device and the surviving copy is + already reading device memory as a host source -- a program that was broken + before this pass touched it -- but the pass is about to hand that buffer to + an engine on the strength of the same post-condition, so it is refused here + rather than left to fail on first execution. + """ + for user in source.users: + if user.op == "output": + continue + if not isinstance(user, Node) or user.op != "call_function": + return False + if user in already_removed or _is_tensorrt_delegate(graph_module, user): + continue + if user.target is not h2d_copy: + return False + spec = user.meta.get("spec") + if spec is None or spec.device != target_device: + return False + if spec.device_index != target_device_index: + return False + if not user.users or not all( + _is_tensorrt_delegate(graph_module, copy_user) for copy_user in user.users + ): + return False + return True + + +def _unstage_aliased_buffers( + graph_module: torch.fx.GraphModule, *, device_memory_planning: bool = True +) -> int: + """Route TensorRT delegate inputs from their staging copy back to the buffer. + + An input is un-staged only when it is an ``_h2d_copy`` of a placeholder + carrying the mark left by :func:`rewire_aliased_mutations_to_buffers`. Every + other input keeps its staging, including a mutable buffer the engine does + not write in place. + + The placeholder's spec takes over the staging copy's device, which is what + asks memory planning for the delegate's device arena rather than a host one + (asks, not settles -- see ``device_memory_planning`` below). That is what + makes handing the buffer straight to the engine valid at all: a host-arena + pointer is not something the engine can write. It is refused when the buffer + has a consumer this pass leaves behind that does not survive being handed + the buffer from there (see :func:`_device_placement_is_safe`) -- asked on + both routes below and whether or not the spec already names that device, + since it is the placement and not the change of device that such a consumer + does not survive. A buffer already reaching its delegate directly is in the + same position as one this pass moves there, and is the shape this pass's own + first run leaves for its second. + + What the pass has to establish is the *post-condition*: every marked buffer + is a direct argument of a TensorRT delegate *declaring zero-copy KV* -- one + whose own engine elided an aliased output -- *and* ends up planned in device + memory. Removing a staging copy is only the usual way of getting there, not + the goal, and a marked buffer that already satisfies both is left alone and + counts as satisfied -- which is what running this pass a second time over a + program it has already un-staged finds. + + The zero-copy declaration is what narrows the delegates that count, on both + routes, because the mark says an engine writes the buffer in place and only a + stamped delegate's engine did. An unstamped TensorRT delegate taking the + buffer proves nothing about the stamped one, whose write would still go to a + staging copy that is discarded. :func:`check_zero_copy_kv` narrows the same + way over the finalized program, so the two halves of this post-condition + accept and refuse the same graphs. + + Being a direct argument is not on its own enough, so it is not on its own + accepted. Two things decide where the buffer is planned, and the second is + not visible in the graph: the spec's own device, and whether memory planning + reads spec devices at all. ``enable_non_cpu_memory_planning=False`` on the + ``ExecutorchBackendConfig`` plans every tensor into the one host arena + whatever its spec says, so the engine is handed a host pointer for a buffer + it must write on the device -- and it does that to every marked buffer, the + ones this pass un-stages as much as the ones that already reach their + delegate directly. That is why it is checked once for the whole graph, up + front, rather than on one of the two branches below. + ``device_memory_planning`` carries the configuration in; the pass + :func:`unstage_aliased_buffers_pass` builds reads it off the finalization + config when it runs, resolved as :func:`_config_plans_on_devices` describes, + so a configuration whose planner never receives the flag is not refused on + it. + + A failure here is a lost KV update, so it is raised rather than logged: + export has already removed the copy-back, so a marked buffer left staged has + the engine write per-call scratch that is then discarded and the buffer never + updates. It raises when memory planning is host-only, when either end of the + move -- the staging copy or the buffer -- has no spec, when the staging copy + is not on CUDA, when a direct argument has no spec of its own + or that spec is not on CUDA, when the placement is unsafe, and -- so a + discovery miss cannot pass silently -- after the loop when the post-condition + does not hold for some marked buffer, cross-checked against each delegate's + own ``zero_copy_kv`` spec: a TensorRT delegate that takes fewer marked + buffers than the aliased outputs that spec says it elided is broken, and that + refusal names the delegate and lists what it does take. + + Returns the number of delegate inputs un-staged, which is zero for a program + that already satisfied the post-condition. + """ + from executorch.exir.schema import DeviceType + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + marked_placeholders = [ + node + for node in graph_module.graph.nodes + if node.op == "placeholder" and node.meta.get("_torch_tensorrt_aliased_buffer") + ] + if marked_placeholders and not device_memory_planning: + names = ", ".join(repr(node.name) for node in marked_placeholders) + raise RuntimeError( + "TensorRT zero-copy KV: buffer(s) " + f"{names} are marked for in-place update, but memory planning is " + "configured with enable_non_cpu_memory_planning=False, which puts " + "every tensor in the one host arena whatever its TensorSpec says. " + "The engine would be handed a host pointer it cannot write, and " + "export has already removed the copy-back, so the update would be " + "lost. Finalize over a configuration that leaves " + "enable_non_cpu_memory_planning on, so a CUDA TensorSpec is given a " + "CUDA arena." + ) + + h2d_copy = torch.ops.et_copy._h2d_copy.default + unstaged = 0 + satisfied_placeholders: Set[Node] = set() + # A dict rather than a list: the walk asks whether a staging copy has + # already been detached, and the erase below wants each one once. + orphaned_stagings: Dict[Node, None] = {} + zero_copy_delegates: List[Node] = [] + # Distinct buffers, not argument slots: one buffer occupying two of a + # delegate's slots is one cache written in place, and counting the slots + # would let it satisfy a spec naming two elided outputs. The finalized-program + # check counts the same way, so the two cannot disagree about one graph. + satisfied_per_delegate: Dict[Node, Set[Node]] = {} + + for node in list(graph_module.graph.nodes): + if not _is_tensorrt_delegate(graph_module, node): + continue + declares_zero_copy = _delegate_declares_zero_copy(graph_module, node) + if declares_zero_copy: + zero_copy_delegates.append(node) + satisfied_per_delegate[node] = set() + new_args = list(node.args) + for i, arg in enumerate(node.args[1:], start=1): + if not isinstance(arg, Node): + continue + if arg.op == "placeholder": + if arg.meta.get("_torch_tensorrt_aliased_buffer"): + direct_spec = arg.meta.get("spec") + placement = "" + remedy = "" + if direct_spec is None: + placement = "it carries no TensorSpec" + remedy = ( + "The specs exist only while this runs as the " + "ExecutorchBackendConfig to_out_var_pass, which is " + "where torch_tensorrt.executorch." + "zero_copy_backend_config() installs it." + ) + elif direct_spec.device != DeviceType.CUDA: + placement = f"its TensorSpec asks for {direct_spec.device!r}" + remedy = ( + "Give the buffer to a TensorRT delegate on CUDA, or " + "export this method without zero_copy_kv." + ) + if placement: + raise RuntimeError( + "TensorRT zero-copy KV: buffer " + f"'{arg.name}' reaches a TensorRT delegate directly, " + f"with no staging copy to remove, but {placement}, so " + "it is not planned in device memory and the engine is " + "handed a host pointer it cannot write. The engine " + "writes this buffer in place and its copy-back has " + "already been removed, so the update would be lost. " + f"{remedy}" + ) + if not _device_placement_is_safe( + graph_module, + arg, + h2d_copy, + direct_spec.device, + direct_spec.device_index, + orphaned_stagings, + ): + raise RuntimeError( + "TensorRT zero-copy KV: buffer " + f"'{arg.name}' reaches a TensorRT delegate directly " + "and is read by a consumer that does not survive it " + "being planned in this engine's device memory -- a " + "staging copy left behind reads it as a host source " + "and fails InvalidArgument. Export this method " + "without zero_copy_kv, or stop sharing the aliased " + "buffer." + ) + if declares_zero_copy: + satisfied_placeholders.add(arg) + satisfied_per_delegate[node].add(arg) + continue + if arg.target is not h2d_copy: + continue + source = arg.args[0] + if not isinstance(source, Node) or source.op != "placeholder": + continue + if not source.meta.get("_torch_tensorrt_aliased_buffer"): + continue # not written in place; it needs its staging copy + staged_spec = arg.meta.get("spec") + source_spec = source.meta.get("spec") + if staged_spec is None or source_spec is None: + missing = ( + f"the staging copy '{arg.name}'" + if staged_spec is None + else f"the buffer placeholder '{source.name}'" + ) + raise RuntimeError( + f"TensorRT zero-copy KV: no TensorSpec on {missing}, so buffer " + f"'{source.name}' cannot be moved to the delegate's device. The " + "TensorRT engine writes this buffer in place and its copy-back " + "has already been removed, so the update would be lost. This " + "pass has to run as the ExecutorchBackendConfig to_out_var_pass, " + "which is where the specs exist; " + "torch_tensorrt.executorch.zero_copy_backend_config installs it " + "there." + ) + # spec.device is an exir schema DeviceType, not a torch.device. + if staged_spec.device != DeviceType.CUDA: + raise RuntimeError( + "TensorRT zero-copy KV: the staging copy of buffer " + f"'{source.name}' targets device {staged_spec.device!r}, not " + "CUDA, so moving the buffer there would put it where the " + "TensorRT engine cannot write it. The engine writes this " + "buffer in place and its copy-back has already been removed, " + "so the update would be lost." + ) + if not _device_placement_is_safe( + graph_module, + source, + h2d_copy, + staged_spec.device, + staged_spec.device_index, + orphaned_stagings, + ): + raise RuntimeError( + "TensorRT zero-copy KV: buffer " + f"'{source.name}' is read by a consumer this pass leaves in " + "place that does not survive the buffer being planned in this " + "engine's device memory -- a staging copy left behind reads it " + "as a host source and fails InvalidArgument. Export this " + "method without zero_copy_kv, or stop sharing the aliased " + "buffer." + ) + source_spec.device = staged_spec.device + source_spec.device_index = staged_spec.device_index + new_args[i] = source + unstaged += 1 + orphaned_stagings[arg] = None + if declares_zero_copy: + satisfied_placeholders.add(source) + satisfied_per_delegate[node].add(source) + node.args = tuple(new_args) + + marked_but_unsatisfied = [ + node for node in marked_placeholders if node not in satisfied_placeholders + ] + if marked_but_unsatisfied: + names = ", ".join(repr(node.name) for node in marked_but_unsatisfied) + raise RuntimeError( + "TensorRT zero-copy KV: buffer(s) " + f"{names} were marked for in-place update but no TensorRT delegate " + f"declaring zero-copy KV (compile spec " + f"'{ZERO_COPY_KV_COMPILE_SPEC_KEY}') takes them, either directly or " + "through a staging copy this pass could remove. Export removed their " + "copy-back on the promise that a TensorRT engine writes them in " + "place, so as this program stands nothing updates them. An unstamped " + "TensorRT delegate taking the buffer does not count: the stamp is " + "what records that an engine elided its aliased output for it. " + "Export this method without zero_copy_kv, or keep the aliased buffer " + "on the delegate whose engine elided it." + ) + for delegate in zero_copy_delegates: + # One marked buffer per aliased output the spec says this delegate + # elided. Demanding only one would accept a delegate that lost all but + # one of its marks, whose remaining caches are still wired through their + # staging copies. A spec that names none -- listing none, or not + # decoding -- cannot say how many to expect, so it falls back to + # demanding at least one. + elided = _delegate_elided_output_names(graph_module, delegate) + satisfied = len(satisfied_per_delegate[delegate]) + if satisfied >= max(len(elided), 1): + continue + delegate_inputs = [ + arg.name for arg in delegate.args[1:] if isinstance(arg, Node) + ] + expectation = ( + "takes no buffer marked for in-place update" + if not elided + else ( + f"elided the aliased output(s) {sorted(elided)} but takes only " + f"{satisfied} of the {len(elided)} buffers that implies" + ) + ) + raise RuntimeError( + "TensorRT zero-copy KV: delegate " + f"'{delegate.name}' declares zero-copy KV " + f"(compile spec '{ZERO_COPY_KV_COMPILE_SPEC_KEY}') but " + f"{expectation} (inputs: {delegate_inputs}). " + "Export elided its aliased outputs, so the engine now writes " + "per-call scratch that is discarded and those caches never update." + ) + + if unstaged: + # Erase only the stagings we orphaned. A graph-wide eliminate_dead_code() + # in a to_out_var_pass could delete another backend's unused delegate. + for staging in orphaned_stagings: + if not staging.users: + graph_module.graph.erase_node(staging) + graph_module.graph.lint() + graph_module.recompile() + return unstaged + + +def _config_plans_on_devices(config: "ExecutorchBackendConfig") -> bool: + """Whether ``enable_non_cpu_memory_planning`` is anything this config decides by. + + ``to_executorch`` does not hand the flag to the memory planner it is given. + It *assigns* it, and only onto a planner that already has an attribute of + that name (``exir/program/_program.py``), which in practice means + ``MemoryPlanningPass`` or a subclass. A caller who brings their own planner + -- which the user guide tells people to do for a cache shared between + prefill and decode -- never receives the flag, and what that planner does + with the specs is its own business: ``False`` does not put the caches in the + host arena, and ``True`` does not keep them out of it. So the flag answers + nothing for such a config and ``True`` is returned, meaning only "not a + ground to refuse on"; :func:`check_zero_copy_kv` reads the arena that + planner actually chose, which is the answer this cannot give. + + ``memory_planning_pass`` may also be a per-method dict, and this pass sees + one graph module without its method name, so a dict is read as deciding + nothing unless every planner in it takes the flag. A method the dict omits + gets ExecuTorch's default planner, which does. + """ + planner = config.memory_planning_pass + planners = list(planner.values()) if isinstance(planner, dict) else [planner] + if not all(hasattr(p, "enable_non_cpu_memory_planning") for p in planners): + return True + return bool(config.enable_non_cpu_memory_planning) + + +class _UnstageThenToOutVar(PassBase): # type: ignore[misc] + """The ``to_out_var_pass`` :func:`unstage_aliased_buffers_pass` builds. + + A named type rather than a closure so that a config can be *recognised*: + :func:`_check_zero_copy_kv_when_finalized` has to answer, before it lets + finalization start, whether the config it was handed carries this pass, and + a class made afresh inside the builder gives a different type object on + every call for nothing to match against. + + ``inner`` is the ``to_out_var_pass`` that would otherwise have run. + ``device_memory_planning`` is what an unbound pass reads; once + ``finalization_config`` is set, the flag is resolved off that config on + every call instead, as :func:`_config_plans_on_devices` describes. + """ + + def __init__(self, inner: Any, device_memory_planning: bool) -> None: + self.inner = inner + self.device_memory_planning = device_memory_planning + self.finalization_config: Optional["ExecutorchBackendConfig"] = None + + def call(self, graph_module: torch.fx.GraphModule) -> Any: + config = self.finalization_config + planning = ( + self.device_memory_planning + if config is None + else _config_plans_on_devices(config) + ) + unstaged = _unstage_aliased_buffers( + graph_module, device_memory_planning=planning + ) + logger.debug("un-staged %d aliased delegate buffer(s)", unstaged) + return self.inner(graph_module) + + +def unstage_aliased_buffers_pass( + inner_pass: Optional[Any] = None, *, device_memory_planning: bool = True +) -> Any: + """Build a ``to_out_var_pass`` that un-stages aliased buffers, then delegates. + + ``to_out_var_pass`` is the last hook that runs after ``PropagateDevicePass`` + and before memory planning -- the window in which the staging copies exist + and the buffers' placement is not yet fixed. (``sym_shape_eval_pass`` is a + caller-supplied hook in that window too, but it runs first.) + + ``inner_pass`` is the ``to_out_var_pass`` that would otherwise have run; it + runs after the un-staging. Omit it for ExecuTorch's default. + + ``device_memory_planning`` is the ``enable_non_cpu_memory_planning`` the + program will be finalized with. Nothing in the graph records it, and it is + half of what decides whether a marked buffer ends up somewhere the engine can + write, so the pass has to be told: see :func:`_unstage_aliased_buffers`. + It is what a pass built here and left unbound uses. Set + ``finalization_config`` on the returned pass to the + ``ExecutorchBackendConfig`` the program will be finalized with and the pass + reads the flag off that config instead, on every call, as + :func:`_config_plans_on_devices` resolves it -- which is what memory planning + will do with it a few passes later, and is nothing at all when the config + carries a planner that does not take the flag. + ``ExecutorchBackendConfig`` is a plain mutable dataclass, so the field can be + set again on the very config being finalized after the pass was built; a pass + reading a value captured here would then accept a program the finalizer plans + into the host arena. :func:`zero_copy_backend_config` binds the attribute for + that reason, and passes no ``device_memory_planning`` at all, since the + binding would override it on every call. + """ + from executorch.exir import ExecutorchBackendConfig + + inner = ( + inner_pass + if inner_pass is not None + else ExecutorchBackendConfig().to_out_var_pass + ) + return _UnstageThenToOutVar(inner, device_memory_planning) + + +def _device_planned_arenas( + graph_module: torch.fx.GraphModule, +) -> Optional[Dict[int, Any]]: + """The finalized program's CUDA arenas, ``mem_id -> device index``, or ``None``. + + Memory planning partitions the specs by device, gives each device its own + arena, and records the non-CPU ones on the graph module as + ``non_const_buffer_device``. ``None`` means the program records no arena + devices at all, which does *not* mean the host: ``apply_algo`` is the only + thing in ExecuTorch that writes the key, and ``to_executorch`` accepts any + callable as ``memory_planning_pass``, so a caller-supplied planner -- which + the user guide tells people to bring for a cache shared between prefill and + decode -- can plan onto a device and still leave the key unwritten. + + The index is carried, not only the type, because the record is what the + runtime allocates from: an arena recorded for ``cuda:1`` holding a cache the + engine writes on ``cuda:0`` is a pointer on the wrong GPU, which the device + *type* alone cannot tell from a correct program. ``apply_algo`` writes one + entry per arena, so a ``mem_id`` names at most one device. + """ + from executorch.exir.schema import DeviceType + + entries = graph_module.meta.get("non_const_buffer_device") + if not entries: + return None + return { + entry.buffer_idx: getattr(entry, "device_index", None) + for entry in entries + if getattr(entry, "device_type", None) == DeviceType.CUDA + } + + +def _host_planned_arenas(graph_module: torch.fx.GraphModule) -> Set[int]: + """The ``mem_id``s of arenas that hold at least one host tensor. + + Planning gives each device its own arena, so an arena holding a tensor whose + spec is CPU is a host arena -- an argument from the graph rather than from + the planner's records, which is what makes it usable when those records are + absent. It is what separates the two shapes that both record no arena + devices: ``enable_non_cpu_memory_planning=False`` puts every tensor in one + bucket, so a CUDA-spec cache ends up sharing an arena with the host tensors, + while a caller-supplied device-aware planner keeps them apart. + + A method with no CPU tensor at all is the residual: nothing here can then + tell a device arena from a host one, and if the program records no arena + devices either, the buffer is accepted. + """ + from executorch.exir.schema import DeviceType + + host: Set[int] = set() + for node in graph_module.graph.nodes: + specs = node.meta.get("spec") + for spec in specs if isinstance(specs, (list, tuple)) else [specs]: + if ( + spec is not None + and getattr(spec, "device", None) == DeviceType.CPU + and getattr(spec, "mem_id", None) is not None + ): + host.add(spec.mem_id) + return host + + +def _is_host_planned( + node: Node, device_arenas: Optional[Dict[int, Any]], host_arenas: Set[int] +) -> bool: + """True when memory planning put ``node`` somewhere the engine cannot write. + + Two independent grounds, because either record may be the only one there. + Sharing an arena with a host tensor settles it whatever the program records; + otherwise, when the program does record its arena devices, an arena missing + from that record is not a CUDA one. + + Asked only for a buffer that has a ``mem_id``. One that does not was never + planned, which is not a placement this can read and is refused as its own + thing by the caller. + """ + mem_id = node.meta["spec"].mem_id + return mem_id in host_arenas or ( + device_arenas is not None and mem_id not in device_arenas + ) + + +def _planned_on_another_gpu( + node: Node, device_arenas: Optional[Dict[int, Any]] +) -> Optional[str]: + """How ``node``'s arena and its own spec disagree about which GPU, if they do. + + Being in a CUDA arena is not enough: the engine writes the cache through the + pointer the runtime allocates out of that arena, so an arena the program + records for another GPU is an address the engine cannot write, exactly as a + host one is. Only a disagreement between two recorded indices is read as + one -- either side left unrecorded says nothing, and the arena's own device + type has already been established by the caller. + """ + if device_arenas is None: + return None + spec = node.meta.get("spec") + mem_id = getattr(spec, "mem_id", None) + arena_index = None if mem_id is None else device_arenas.get(mem_id) + spec_index = getattr(spec, "device_index", None) + if arena_index is None or spec_index is None or arena_index == spec_index: + return None + return ( + f"'{node.name}' asks for cuda:{spec_index} and was planned in an arena " + f"the program records as cuda:{arena_index}" + ) + + +def _name_detail(names_by_method: Dict[str, List[str]]) -> str: + return ", ".join( + f"'{name}' in method '{method}'" + for method, names in names_by_method.items() + for name in names + ) + + +def check_zero_copy_kv(program: Any) -> None: + """Raise unless a finalized program really updates its KV buffers in place. + + ``program`` is what ``to_executorch()`` returns. Both halves of zero-copy do + nothing quietly when they find nothing to do: ``zero_copy_kv=True`` warns and + carries on when the model holds no aliased buffer mutation, and the pass + :func:`unstage_aliased_buffers_pass` builds, handed a program with nothing + marked, un-stages nothing and returns. Neither of those is wrong output -- nothing + removed a copy-back, so the ``.pte`` stages its cache and updates it like any + other -- but the optimization the caller asked for is silently not there, and + a caller who reads the successful ``save`` as proof it is gets neither an + error nor the speedup. The wrong-output case is the one below: a rewiring + that did happen and then lost its mark. + + Six shapes are refused: a marked buffer that is not a direct argument of a + TensorRT delegate carrying the zero-copy compile spec, a stamped delegate + that takes fewer marked buffers than its own spec says it elided aliased + outputs -- including one in a method with no marked buffer at all, which is + the lost-mark case -- a marked buffer that reaches such a delegate directly + but is planned in a host arena, one whose placement the program records + nowhere, one planned in a device arena the program records for another GPU, + and a program carrying neither a marked buffer nor a stamped delegate in any + of its methods. The first three are + what finalizing without :func:`zero_copy_backend_config` leaves behind -- + all but the lost-mark case folded into the second, which no finalization + choice produces. That config's pass gets to all three earlier, off the + configuration and the graph: it removes the staging copy the first two come + from, and refuses outright when there is none to remove or when the + configuration plans nothing onto a device. What it cannot see is the arena + planning then chose, which is what this reads. + + The spec is what narrows the delegates that count. The mark is put on a + buffer because one TensorRT engine writes it in place, and only a delegate + whose own engine elided an aliased output is stamped, so another backend's + delegate taking the buffer says nothing about whether the engine did, and + neither does an unrelated TensorRT engine that happens to read it. Either + would stand in for the delegate whose write was elided while that one still + reads a staging copy whose contents are discarded. Being stamped is not + enough on its own either: in a method that lowers to two of them, one + stamped delegate holding both caches leaves every marked buffer reaching + *some* stamped delegate while the other's write is still thrown away. So + each stamped delegate is also counted against its own spec, the same + cross-check :func:`_unstage_aliased_buffers` makes, against the same spec and + with the same fallback: a spec that names none -- listing none, or not + decoding -- cannot say how many to expect, so it demands at least one. What + that count cannot separate is an exact swap -- two stamped delegates each + taking one marked buffer, each the other's. The mark records only that some + engine writes the buffer in place, never which one, and the spec lists + engine output binding names rather than buffers, so there is nothing left to + match on; that pass has the same blind spot for the same reason. + + Placement is read from where memory planning actually put the buffer -- the + ``mem_id`` on its ``TensorSpec`` -- rather than from the spec's own device, + which does not settle it: ``PropagateDevicePass`` writes CUDA onto the spec + of a buffer that reaches a CUDA delegate directly, and + ``enable_non_cpu_memory_planning=False`` then plans that same buffer into the + one host arena, which is the shape whose every ``execute()`` fails on the + runtime's alias-target guard. Two independent grounds answer whether that + ``mem_id`` names a host arena. An arena that also holds one of the program's + host tensors is one whatever the program records (see + :func:`_host_planned_arenas`); failing that, an arena missing from the CUDA + ones the program *does* record in ``non_const_buffer_device`` is one too. + + A placement the program does not record at all is a refusal of its own + rather than either of those, because it is not an accusation about the + planner's choice -- a buffer with no ``mem_id`` was not planned, and a + program with no ``non_const_buffer_device`` entries says nothing about any + of its arenas. It is refused because of what the runtime makes of it: + ``MethodMeta::memory_planned_buffer_device`` answers ``CPU`` for an arena + the ``.pte`` records nothing for, so a runner that honours it -- + ``examples/executorch_reference_runner`` does -- backs that arena with host + memory and the engine then fails the alias-target guard on every call. Only + ``apply_algo`` writes that record, so a caller-supplied + ``memory_planning_pass`` that does not go through it leaves a ``.pte`` in + exactly that state whatever it planned; a planner for a zero-copy cache has + to leave the record behind, which means going through ``apply_algo`` with + ``enable_non_cpu_memory_planning=True``. That parameter defaults to + ``False``, and with it off ``apply_algo`` plans every spec into one CPU + bucket and writes no record either. + + Where the record does name a GPU it is read as one: an arena recorded for + ``cuda:1`` holding a cache whose own spec asks for ``cuda:0`` is an address + on the wrong device, which fails the same way a host one does, so the index + is compared and not only the type. An unrecorded index on either side says + nothing and is accepted, since the arena's device type is already settled by + then. + + Every method is read, not only ``forward``. ``export()`` rewires each method + on its own, so a check that stopped at ``forward`` would pass a program whose + decode had degenerated to staged -- and on the prefill/decode pair the user + guide's zero-copy example exports it would not get that far, since a + multi-method program need not have a ``forward`` at all. Within a method the + marks and the stamped delegates are enumerated independently and a method is + passed over only when it carries neither, because the disagreement between + those two records is the whole subject: starting from the marks and looking + the delegates up leaves anything recorded only on the delegate side outside + the walk. The last refusal is about the program rather than about one method, + matching the warning ``export()`` emits: a method with no aliased buffer + mutation of its own is not an error, so a model that rewires only its decode + step is accepted. + + This reads the graph and the finalized specs, so it says what the program + does rather than what the passes recorded. It still says nothing about + whether the engine's write itself is correct. + + Both entry points run it for you: ``save(..., zero_copy_kv=True)`` before it + writes the file, and the ``EdgeProgramManager`` that + ``export(..., zero_copy_kv=True)`` returns on whatever its ``to_executorch`` + produces -- that manager also refuses a config without the un-staging pass + before finalizing at all, which is the one refusal here that can be answered + without exporting again. Call this by hand for a program that reached + finalization some other way -- one derived from that manager by + ``transform()`` or ``to_backend()``, which are new managers without the + hook. + + Arguments: + program (executorch.exir.ExecutorchProgramManager): The finalized program + ``to_executorch()`` returned. Every method it holds is read. + + Returns: + None: a program that passes is left exactly as it was. + + Raises: + RuntimeError: If the program does not update its KV buffers in place, + naming the buffers or delegates and the method each is in. + """ + method_names = sorted(program.methods) + staged_by_method: Dict[str, List[str]] = {} + short_by_method: Dict[str, List[str]] = {} + unrecorded_by_method: Dict[str, List[str]] = {} + host_planned_by_method: Dict[str, List[str]] = {} + wrong_gpu_by_method: Dict[str, List[str]] = {} + marked_anywhere = False + for method_name in method_names: + graph_module = program.exported_program(method_name).graph_module + marked = [ + node + for node in graph_module.graph.nodes + if node.op == "placeholder" + and node.meta.get("_torch_tensorrt_aliased_buffer") + ] + zero_copy_delegates = [ + node + for node in graph_module.graph.nodes + if _is_tensorrt_delegate(graph_module, node) + and _delegate_declares_zero_copy(graph_module, node) + ] + # Both records have to be read before a method can be passed over. A + # method carrying neither is one zero-copy never touched; a method + # carrying a stamped delegate and no mark is the disagreement this + # exists to catch, which skipping on the marks alone would hide. + if not marked and not zero_copy_delegates: + continue + if marked: + marked_anywhere = True + zero_copy_delegate_args = { + arg for node in zero_copy_delegates for arg in node.args[1:] + } + staged = [node.name for node in marked if node not in zero_copy_delegate_args] + if staged: + staged_by_method[method_name] = staged + marked_nodes = set(marked) + short = [] + for delegate in zero_copy_delegates: + elided = _delegate_elided_output_names(graph_module, delegate) + # Distinct buffers rather than argument slots, as the un-staging + # pass counts them: one buffer in two slots is one cache. + taken = len( + { + arg + for arg in delegate.args[1:] + if isinstance(arg, Node) and arg in marked_nodes + } + ) + if taken >= max(len(elided), 1): + continue + if elided: + short.append( + f"'{delegate.name}' takes {taken} marked buffer(s), not the " + f"{len(elided)} its elided aliased output(s) {sorted(elided)} " + "imply" + ) + else: + short.append( + f"'{delegate.name}' names no aliased output this can count, so " + "it must take at least one marked buffer and takes none" + ) + if short: + short_by_method[method_name] = short + device_arenas = _device_planned_arenas(graph_module) + host_arenas = _host_planned_arenas(graph_module) + reaching = [node for node in marked if node in zero_copy_delegate_args] + unrecorded: List[str] = [] + host_planned: List[str] = [] + wrong_gpu: List[str] = [] + # One classification per buffer. The two unrecorded shapes are not + # evidence about the host arena and get a refusal of their own, but the + # positive host-arena ground is read first where it applies, since + # host-only planning both puts the cache among the host tensors and + # writes no arena record, and naming the arena it is actually in tells + # the caller more than saying nothing was recorded. + for node in reaching: + mem_id = getattr(node.meta.get("spec"), "mem_id", None) + if mem_id is None: + unrecorded.append( + f"'{node.name}' carries no mem_id, so memory planning left " + "it unplanned and nothing in the program says where it lives" + ) + elif _is_host_planned(node, device_arenas, host_arenas): + host_planned.append(node.name) + elif device_arenas is None: + unrecorded.append( + f"'{node.name}' is planned in arena {mem_id}, and the " + "program records no CUDA arena at all" + ) + else: + detail = _planned_on_another_gpu(node, device_arenas) + if detail is not None: + wrong_gpu.append(detail) + if unrecorded: + unrecorded_by_method[method_name] = unrecorded + if host_planned: + host_planned_by_method[method_name] = host_planned + if wrong_gpu: + wrong_gpu_by_method[method_name] = wrong_gpu + if staged_by_method: + raise RuntimeError( + f"TensorRT zero-copy KV: buffer(s) {_name_detail(staged_by_method)} " + "are marked for in-place update but do not reach the TensorRT " + "delegate that elided them directly, so the engine writes a staging " + "copy that is discarded and the cache never updates. Export removed " + "their copy-back, so nothing else would restore it. Finalize with " + "torch_tensorrt.executorch.zero_copy_backend_config()." + ) + if short_by_method: + raise RuntimeError( + "TensorRT zero-copy KV: delegate(s) " + + "; ".join( + f"{detail} in method '{method}'" + for method, details in short_by_method.items() + for detail in details + ) + + ". Every marked buffer does reach a delegate declaring zero-copy " + "KV, so either another one in the same method is holding this " + "engine's cache or a mark was lost; either way this engine still " + "reads a staging copy that is discarded, and export removed the " + "copy-back. Export this method without zero_copy_kv, or keep each " + "aliased buffer on the delegate whose engine elided it." + ) + # Ordered after the still-staged and short-count refusals because both of + # those fire on a disagreement between the marks and the stamped delegates, + # and this one would answer such a program with "it was probably not exported + # with zero_copy_kv=True" while its own compile specs say it was. + if not marked_anywhere: + raise RuntimeError( + "TensorRT zero-copy KV: no buffer in this program is marked for " + f"in-place update, in any of its methods ({', '.join(method_names)}), " + "so it stages its caches like any other .pte. Either it was not " + "exported with zero_copy_kv=True, or it was and no aliased buffer " + "mutation was found -- export logs a warning for that case." + ) + if host_planned_by_method: + raise RuntimeError( + f"TensorRT zero-copy KV: buffer(s) " + f"{_name_detail(host_planned_by_method)} reach their TensorRT " + "delegate directly but memory planning put them in an arena that " + "also holds the program's host tensors, or in one it does not record " + "as CUDA, so the engine is handed a host pointer it cannot write and " + "every execute() fails on the runtime's alias-target guard. Finalize " + "with torch_tensorrt.executorch.zero_copy_backend_config() over a " + "configuration that leaves enable_non_cpu_memory_planning on. If it " + "is already on, the memory_planning_pass in use is what put this " + "buffer among the host tensors, and it has to give the delegate's " + "device an arena of its own." + ) + if unrecorded_by_method: + raise RuntimeError( + "TensorRT zero-copy KV: buffer(s) " + + "; ".join( + f"{detail} in method '{method}'" + for method, details in unrecorded_by_method.items() + for detail in details + ) + + ". These buffers reach their TensorRT delegate directly, so the " + "engine writes them through the pointer the runtime allocates for " + "them, and the .pte has to say that pointer is device memory. " + "MethodMeta::memory_planned_buffer_device answers CPU for an arena " + "the program records nothing for, so a runner that honours it backs " + "the arena with host memory and every execute() fails on the " + "alias-target guard. The memory_planning_pass in use has to plan " + "these buffers and leave the record behind, which means going " + "through ExecuTorch's apply_algo with " + "enable_non_cpu_memory_planning=True -- that parameter defaults to " + "False, and with it off apply_algo plans every spec into one CPU " + "bucket and writes no record either." + ) + if wrong_gpu_by_method: + raise RuntimeError( + "TensorRT zero-copy KV: buffer(s) " + + "; ".join( + f"{detail} in method '{method}'" + for method, details in wrong_gpu_by_method.items() + for detail in details + ) + + ". The engine writes the cache through the pointer the runtime " + "allocates out of that arena, so it would write the wrong GPU. The " + "memory_planning_pass in use has to give each delegate's own device " + "an arena, and put every buffer its engine writes in place in that " + "device's one." + ) + + +def _check_zero_copy_kv_when_finalized(edge_manager: Any) -> None: + """Guard one manager's own ``to_executorch``, before and after it runs. + + Export removes the copy-back of the rewired caches before it returns, so + from that point on every way of finalizing the program that does not also + un-stage them produces a ``.pte`` whose caches never update -- silently, and + for a KV cache that is wrong output rather than a crash. ``to_executorch()`` + with ExecuTorch's defaults is one such way, and it is the call the two-step + API documents. This is the only place holding the manager that knows the + export asked for zero-copy, so the refusal lives here rather than being left + for the caller to remember. + + It asks two questions at two moments, because only the first can be answered + while there is still something to be done about it. ``to_executorch`` does + not copy the edge programs: it runs the write-back and the device passes over + the manager's own graph modules and copies the finalized graph back into + them, so once finalization has run the manager is spent -- calling it a + second time dies inside ``insert_write_back_for_buffers_pass``. A refusal + raised after the fact could therefore only be acted on by exporting again and + rebuilding every engine. So the config is read *first*: one that does not + carry the un-staging pass cannot produce a zero-copy program whatever the + graph looks like, and saying so before delegating leaves the manager + untouched and the remedy -- the same call with + :func:`zero_copy_backend_config` -- available. + + The finalized program is still read back afterwards, through + :func:`check_zero_copy_kv`, because a config can be the right one and the + program still come out wrong: memory planning chooses the arenas after this + pass has run, and a lost mark is a property of the graph. Those refusals + name what to do themselves. + + Bound to the instance rather than to a type, because ``EdgeProgramManager`` + is ExecuTorch's. That reaches the manager ``export()`` hands back and no + other: ``transform()`` and ``to_backend()`` build a *new* manager, which + this hook does not travel to, so a caller who takes either detour owes + :func:`check_zero_copy_kv` by hand. + """ + finalize = edge_manager.to_executorch + + @functools.wraps(finalize) + def to_executorch(*args: Any, **kwargs: Any) -> Any: + config = kwargs.get("config", args[0] if args else None) + if not isinstance( + getattr(config, "to_out_var_pass", None), _UnstageThenToOutVar + ): + raise RuntimeError( + "TensorRT zero-copy KV: this program was exported with " + "zero_copy_kv=True, which removed the copy-back of its aliased " + "buffers, and this configuration does not carry the pass that " + "un-stages them. Finalizing it would write a .pte whose engine " + "writes a per-call staging copy that is discarded, so the caches " + "never update -- for a KV cache, wrong output rather than a " + "crash. Finalize with " + "to_executorch(torch_tensorrt.executorch.zero_copy_backend_config" + "(config)) instead. This is raised before finalizing rather than " + "after, because to_executorch rewrites the manager's own edge " + "programs in place and a manager that has finalized once cannot " + "do it again." + ) + program = finalize(*args, **kwargs) + check_zero_copy_kv(program) + return program + + edge_manager.to_executorch = to_executorch + + +def _refuse_skip_h2d(config: "ExecutorchBackendConfig") -> None: + """Raise if the config asks ExecuTorch to un-stage method inputs as well. + + Both places the option can be written are read. ``propagate_device_config`` + is one ``PropagateDeviceConfig`` or a dict of them keyed by method, and + within either, ``skip_h2d_for_method_inputs`` is a bool or a second + per-method dict. + + What is refused is every value that pass reads as on, which is every truthy + one rather than only ``True``. ``PropagateDevicePass`` is handed the field + whole and only tests it for truth (in ``_insert_h2d_copies``), never + resolving it per method, so it reads any non-empty dict as on for every + method -- one whose entries are all ``False`` included. Refusing only the + true entries would carry such a dict through and hand back a config that + raises a layer down, which is the failure this exists to prevent. ``False`` + and the empty dict are what that pass reads as off, and both are carried. + """ + propagate = getattr(config, "propagate_device_config", None) + per_method = ( + sorted(propagate.items()) + if isinstance(propagate, dict) + else [("every method", propagate)] + ) + asked_for = [ + method + for method, entry in per_method + if getattr(entry, "skip_h2d_for_method_inputs", False) + ] + if not asked_for: + return + raise ValueError( + "TensorRT zero-copy KV: this configuration sets " + f"skip_h2d_for_method_inputs for {', '.join(asked_for)}, which cannot be " + "combined with zero-copy KV. PropagateDevicePass refuses to un-stage a " + "method input whose placeholder does not have exactly one user, and a " + "buffer zero-copy rewired has two -- the TensorRT delegate, and the graph " + "output it is its own mutation result for -- so finalization raises " + "there. Every value that pass reads as on is refused, not only True: it " + "tests skip_h2d_for_method_inputs for truth without ever resolving it " + "per method, so setting it to a dict whose entries are all False still " + "turns it on for every method. Zero-copy already un-stages the aliased " + "buffers; leave skip_h2d_for_method_inputs at False or unset -- both of " + "which are carried -- and pass the method's own inputs on the host." + ) + + +def zero_copy_backend_config( + config: Optional["ExecutorchBackendConfig"] = None, +) -> "ExecutorchBackendConfig": + """Build the ``ExecutorchBackendConfig`` a zero-copy KV program needs. + + This is the second half of ``export(..., zero_copy_kv=True)``. Export has + already removed ExecuTorch's copy-back of the aliased buffers; this installs + the pass that removes their staging, so the engine writes the caller's + buffer instead of a scratch copy that is thrown away. + + The feature is split across two calls because ``to_executorch()`` belongs to + ExecuTorch, not to Torch-TensorRT: ``export()`` hands back an + ``EdgeProgramManager`` at the Edge boundary and never sees the config the + program is finalized with. + + ``config`` is your own configuration -- every field is preserved, and a + ``to_out_var_pass`` you already set runs after the un-staging. Omit it to + start from ExecuTorch's defaults. Two fields are not merely carried: + + * ``enable_non_cpu_memory_planning`` is *read*. Zero-copy needs the caches + planned in device memory, so ``False`` -- which plans every tensor into + the one host arena -- has the pass refuse each cache it finds rather than + write a ``.pte`` whose every ``execute()`` fails. It is read off the + config returned here, at the moment the pass runs, so setting the field + *on that object* afterwards is honoured: the pass and the finalizer then + cannot disagree about it. There are two cases the field decides nothing + in, and that the pass therefore refuses nothing on. One is building a + *new* config out of this one with ``dataclasses.replace``: the field is a + bool, copied by value, while the pass is copied by reference and goes on + reading the config returned here, so call this function again on the + derived config and the pass it builds reads that one. The other is a + ``memory_planning_pass`` of your own that does not already carry an + attribute of that name -- ``to_executorch`` assigns the flag onto the + planner rather than passing it, and only onto a planner that has it, so + for any other the field reaches nothing and where the caches land is that + planner's own business. Neither is left to the runtime to discover: + :func:`check_zero_copy_kv` reads the arena memory planning actually chose + and refuses the program either mistake produces, and the manager + ``export(..., zero_copy_kv=True)`` returns runs that check itself. The one + placement it cannot settle is a method holding no host tensor to give the + shared arena away and a program that records its arena devices; a planner + that records nothing is refused rather than passed over, as + :func:`check_zero_copy_kv` describes. + * ``propagate_device_config.skip_h2d_for_method_inputs`` is *refused*, + wherever it is written -- in the single ``PropagateDeviceConfig`` or in a + per-method dict of them -- and on every value ``PropagateDevicePass`` + reads as on rather than only on ``True``: it tests the field for truth + without ever resolving it per method, so any non-empty dict is on for + every method, one of ``False`` included. It is ExecuTorch's own un-staging + of method inputs, and it requires each placeholder it un-stages to have + exactly one user. A rewired cache always has two -- the delegate, and the + graph output it is its own mutation result for -- so + ``PropagateDevicePass`` raises on every zero-copy graph. Returning the + option unchanged would hand back a config that cannot finalize at all; + this says so here instead, where the caller can act on it. + + .. warning:: + Finalizing a ``zero_copy_kv=True`` program *without* this config leaves + the engine writing a per-call staging copy that is then discarded, so + the buffer never updates -- for a KV cache, wrong output rather than a + crash. The manager ``export()`` hands back refuses that itself: it reads + the config before finalizing and refuses one that does not carry this + pass, and it reads its own finalized program through + :func:`check_zero_copy_kv` for what the config cannot say. + ``torch_tensorrt.save(..., zero_copy_kv=True)`` runs the same check + before it writes the file. A program finalized off some *other* + manager -- one ``transform()`` or ``to_backend()`` derived from that + one -- is not covered by it, so hand that program to + :func:`check_zero_copy_kv` yourself before writing the ``.pte``. + + ``save(..., zero_copy_kv=True)`` installs this pass itself, so handing + it the result of this function as ``backend_config`` applies the pass + twice. That is redundant rather than an error -- the second run finds + the buffers already wired straight to their delegates and changes + nothing -- but the two entry points are alternatives: use one or the + other. + + Arguments: + config (Optional[executorch.exir.ExecutorchBackendConfig]): The + configuration to compose onto. Omit it to start from ExecuTorch's + defaults. + + Returns: + executorch.exir.ExecutorchBackendConfig: A new config, every field of + the given one preserved, whose ``to_out_var_pass`` un-stages the aliased + buffers before running the ``to_out_var_pass`` that was there. + + Raises: + ValueError: If the configuration sets + ``propagate_device_config.skip_h2d_for_method_inputs``, which cannot + be combined with zero-copy KV. + """ + from dataclasses import replace + + from executorch.exir import ExecutorchBackendConfig + + base = config if config is not None else ExecutorchBackendConfig() + _refuse_skip_h2d(base) + # No device_memory_planning here: setting finalization_config below overrides + # it for every call, so a value passed would never be read. + unstage = unstage_aliased_buffers_pass(base.to_out_var_pass) + wrapped = replace(base, to_out_var_pass=unstage) + unstage.finalization_config = wrapped + return wrapped diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index 655905bfe74..975d934a323 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -1,7 +1,8 @@ # ExecuTorch TensorRT backend: serialize engines to a libtorch-free runtime blob. +import json import operator -from typing import Any, List, final +from typing import Any, Container, Iterable, List, Optional, Set, final import torch import torch.fx @@ -32,6 +33,19 @@ _BINDING_DELIM = "%" +# CompileSpec key naming the aliased outputs a delegate deliberately does not +# carry. Its value is the JSON list of those engine output binding names (see +# _serialize_elided_output_names), NOT a bare flag: only the aliased outputs +# backed by a registered buffer are elided, so the backend must exempt exactly +# those and still reject a delegate that dropped any other binding. export() +# appends a method-wide instance to signal the opt-in; TensorRTPartitioner strips +# that one and re-derives the per-engine value it puts on each delegate's +# DelegationSpec, the only channel that reaches preprocess (see +# TensorRTPartitioner._partition_elided_output_names). Without it a delegate +# short of its aliased outputs is a bug, not a zero-copy program, and stays an +# error. +ZERO_COPY_KV_COMPILE_SPEC_KEY = "zero_copy_kv" + def _schema_name(target: Any) -> str: """Return the qualified op schema name for an OpOverload or EdgeOpOverload.""" @@ -274,7 +288,10 @@ def _reorder_input_names_for_executorch( def _validate_output_binding_order( - edge_program: ExportedProgram, engine_node: Any, output_names: List[str] + edge_program: ExportedProgram, + engine_node: Any, + output_names: List[str], + elidable_output_names: Optional[Container[str]] = None, ) -> None: """Check the delegate's outputs are the engine's output bindings, in order. @@ -286,19 +303,70 @@ def _validate_output_binding_order( delegate -- would swap the names silently. Inputs cannot rely on position at all and recover their order by node identity in ``_reorder_input_names_for_executorch``. + + ``elidable_output_names`` names the bindings the delegate is *allowed* to + have dropped, which zero-copy KV sets to exactly the aliased outputs export + rewired to write in place: the engine's in-place write through the aliased + input already is the buffer update, so no argument is passed for them. Pass + ``None`` (the default) when elision was not asked for, and the delegate must + carry every binding. This check does not require that set to cover the + engine's whole ``aliased_io``; ``preprocess`` requires a non-empty one to, + because the runtime can only take all of an engine's aliased outputs as + elided or none of them. + + A delegate that dropped its aliased outputs because nothing declared them as + mutations looks exactly like a zero-copy one, and the runtime reads elision off + a single argument count, so it cannot tell them apart either. That is also why + a partial drop stays an error: the count cannot express which bindings went. """ + elidable_names = elidable_output_names if elidable_output_names is not None else () + all_indices = list(range(len(output_names))) + unaliased_indices = [ + i for i in all_indices if output_names[i] not in elidable_names + ] + output_node = next( node for node in edge_program.graph_module.graph.nodes if node.op == "output" ) out_args = list(output_node.args[0]) + if not out_args: + # Naming every binding elidable empties unaliased_indices too, so the + # comparison below would match. That is the zero-output delegate + # rewire_aliased_mutations_to_buffers raises about -- nothing reads it, + # so a later graph-wide dead-code elimination can erase the computation. + # The rewiring's guard runs before partitioning; this is the only one + # left after lowering. The check is unconditional because a delegate with + # no outputs is wrong however it got that way; only the remedy below is + # about zero-copy. + raise ValueError( + "TensorRT ExecuTorch backend: the delegate has no outputs at all, " + f"but the engine declares {len(output_names)} output binding(s). A " + "delegate nothing reads is a pure node a later dead-code elimination " + "can erase, taking the engine with it." + + ( + " Every one of this engine's outputs was declared elidable, so " + "nothing was left to thread out; export this method without " + "zero_copy_kv." + if elidable_output_names is not None and not unaliased_indices + else "" + ) + ) # A single-output engine is returned directly rather than through a getitem, - # and one binding has no order to get wrong. + # and one binding has no order to get wrong. The same holds under elision + # when exactly one binding is left unaliased. if len(out_args) == 1 and out_args[0] is engine_node: - if len(output_names) != 1: + if len(all_indices) != 1 and len(unaliased_indices) != 1: + remaining = ( + f", {len(unaliased_indices)} of them after eliding the in-place " + "outputs" + if unaliased_indices != all_indices + else "" + ) raise ValueError( "TensorRT ExecuTorch backend: the delegate returns the engine node " f"directly but the engine declares {len(output_names)} output " - "bindings; only a single-output engine can be returned unwrapped." + f"bindings{remaining}; only a single-output engine can be returned " + "unwrapped." ) return indices: List[Any] = [] @@ -315,15 +383,100 @@ def _validate_output_binding_order( "node; cannot establish a reliable output binding order." ) indices.append(node.args[1]) - if indices != list(range(len(output_names))): + if indices not in (all_indices, unaliased_indices): + expected = ( + f"{all_indices}, or {unaliased_indices} with the in-place outputs elided" + if unaliased_indices != all_indices + else f"{all_indices}" + ) + # Outputs are missing and nothing exempted them. That is what an export + # asking for zero_copy_kv looks like when the aliased-buffer mark did not + # reach the partitioner, so no delegate was stamped and none is exempt -- + # a failure mode with no other symptom, hence naming it here. + unexempted_drop = elidable_output_names is None and len(indices) < len( + all_indices + ) raise ValueError( "TensorRT ExecuTorch backend: delegate outputs map to engine output " - f"indices {indices}, expected {list(range(len(output_names)))} -- the " - "runtime binds output i to output_binding_names[i], so a permuted or " - "incomplete output list would bind the wrong tensors." + f"indices {indices}, expected {expected} -- the runtime binds each " + "output it is given in binding order, so a permuted, incomplete, or " + "partially elided output list would bind the wrong tensors." + + ( + " No output was declared elidable for this delegate, so if the " + "export asked for zero_copy_kv the aliased-buffer mark did not " + "survive lowering." + if unexempted_drop + else "" + ) ) +def _serialize_elided_output_names(names: Iterable[str]) -> bytes: + """Encode the elided aliased-output binding names for the compile spec. + + JSON, not the ``%`` / ``@`` delimiters that separate ``engine_info``'s + binding-name and aliased_io fields, so a binding name containing one of those + cannot corrupt the record. + """ + return json.dumps(sorted(set(names))).encode("utf-8") + + +def _elided_output_names(compile_specs: List[CompileSpec]) -> Optional[Set[str]]: + """The aliased-output binding names export declared elidable, or ``None``. + + ``None`` when no zero-copy spec is present, which keeps a missing output an + error: only a caller who asked for zero-copy may drop the aliased outputs, + and then only exactly the ones export rewired to write in place. + + A spec built by hand rather than by ``TensorRTPartitioner`` may carry + anything, so the value is type-checked and the decode is caught, and both + raise naming this key. Without that a bare ``b"1"`` comes out as an + unattributed ``TypeError``, and -- the quieter one -- a JSON *string* + decodes into a set of its own characters, exempting every one-character + binding name and not the real one. + + ``_zero_copy._delegate_elided_output_names`` reads the same key on the same + spec and takes the same shapes of value. It differs in what it does with the + rest: where this raises, it returns the empty set, because for it an + undecodable spec merely weakens a cross-check while here it leaves the + backend unable to say which outputs may be missing. The one value the two + read differently is bytes that are not valid UTF-8, replaced here and + refused there by ``json.loads``. + """ + for spec in compile_specs: + if getattr(spec, "key", None) != ZERO_COPY_KV_COMPILE_SPEC_KEY: + continue + value = spec.value + if not isinstance(value, (str, bytes, bytearray)): + raise ValueError( + "TensorRT ExecuTorch backend: compile spec " + f"'{ZERO_COPY_KV_COMPILE_SPEC_KEY}' must hold a JSON list of " + f"engine output binding names, not {type(value).__name__}. It is " + "written by TensorRTPartitioner; a hand-built spec has to match." + ) + if isinstance(value, (bytes, bytearray)): + value = bytes(value).decode("utf-8", "replace") + try: + names = json.loads(value) + except ValueError as e: + raise ValueError( + "TensorRT ExecuTorch backend: compile spec " + f"'{ZERO_COPY_KV_COMPILE_SPEC_KEY}' does not decode as JSON " + f"({e}). It holds a JSON list of engine output binding names, " + "written by TensorRTPartitioner." + ) from e + if not isinstance(names, list): + raise ValueError( + "TensorRT ExecuTorch backend: compile spec " + f"'{ZERO_COPY_KV_COMPILE_SPEC_KEY}' decoded to " + f"{type(names).__name__}, not a list of engine output binding " + "names. A JSON string would decode into its own characters and " + "exempt every one-character binding name." + ) + return {str(name) for name in names} + return None + + def _get_str(engine_info: List[Any], index: int, default: str = "") -> str: if index < 0 or index >= len(engine_info): return default @@ -375,7 +528,13 @@ def preprocess( output_names = _split_binding_names( _get_str(engine_info, OUTPUT_BINDING_NAMES_IDX) ) - _validate_output_binding_order(edge_program, engine_node, output_names) + elidable_output_names = _elided_output_names(compile_specs) + _validate_output_binding_order( + edge_program, + engine_node, + output_names, + elidable_output_names, + ) io_bindings = [ TensorRTIOBinding(name=name, is_input=True) for name in input_names ] + [TensorRTIOBinding(name=name, is_input=False) for name in output_names] @@ -384,6 +543,24 @@ def preprocess( # C++ backend binds each aliased output to its aliased input's tensor # (in-place) and reflects the update back into the delegate output. aliased_io = deserialize_aliased_io(_get_str(engine_info, ALIASED_IO_IDX)) + if elidable_output_names and set(elidable_output_names) != set(aliased_io): + raise ValueError( + "TensorRT ExecuTorch backend: engine " + f"'{engine_node.name}' aliases the outputs {sorted(aliased_io)}, " + f"but only {sorted(elidable_output_names)} of them are elided -- " + "dropped from the delegate's arguments, because the engine's " + "in-place write through the aliased input already is that " + "output. Partial elision is not expressible by the runtime: it " + "takes the aliased outputs as elided only when the argument " + "count is short by the engine's whole aliased-output count, so a " + "delegate short of only some of them reads as not elided at all " + "and every execute() fails with an argument-count error. An " + "output is elided when export rewired a buffer mutation to write " + "it in place, so this engine mixes such a buffer with an aliased " + "input that is not one -- a plain input, say. Export this method " + "without zero_copy_kv, or keep every aliased input of this engine " + "a mutated buffer." + ) metadata = TensorRTBlobMetadata( io_bindings=io_bindings, diff --git a/py/torch_tensorrt/executorch/partitioner.py b/py/torch_tensorrt/executorch/partitioner.py index 559eb271de5..f6cd1b892b6 100644 --- a/py/torch_tensorrt/executorch/partitioner.py +++ b/py/torch_tensorrt/executorch/partitioner.py @@ -1,7 +1,7 @@ # ExecuTorch partitioner: partition by execute_engine nodes. import logging -from typing import Callable, Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Set, Tuple import torch from executorch.exir.backend.compile_spec_schema import CompileSpec @@ -13,12 +13,19 @@ from executorch.exir.backend.utils import tag_constant_data from torch.export import ExportedProgram from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition +from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ( + OUTPUT_BINDING_NAMES_IDX, + deserialize_binding_names, +) from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import DEVICE_IDX from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, TensorRTBackend, _get_engine_info_for_node, _get_engine_nodes_in, + _get_str, _parse_device_id, + _serialize_elided_output_names, ) from torch_tensorrt.executorch.operator_support import TensorRTOperatorSupport @@ -127,6 +134,20 @@ def __init__( ) -> None: super().__init__() self.compile_specs = list(compile_specs) if compile_specs else [] + # The zero-copy KV spec is stamped per-partition in partition(), never + # applied to every partition like the rest of compile_specs. Its presence + # here only records that this method asked for zero-copy; the actual + # elided binding names are derived per engine at partition time, so a + # method that lowers to several TensorRT delegates marks only the one + # whose aliased outputs were elided. It has to stay out of the shared + # list: a plain-compute delegate carrying the spec would make the + # un-staging cross-check demand an aliased buffer it never had. + self._zero_copy_requested = any( + s.key == ZERO_COPY_KV_COMPILE_SPEC_KEY for s in self.compile_specs + ) + self._base_compile_specs = [ + s for s in self.compile_specs if s.key != ZERO_COPY_KV_COMPILE_SPEC_KEY + ] # Mirror CudaPartitioner: a target_device CompileSpec drives ExecuTorch's # PropagateDevicePass, which tags delegate I/O TensorSpecs with the device # and serializes it into the .pte's extra_tensor_info. When the caller pins @@ -134,8 +155,12 @@ def __init__( # its own engine node in partition() (engine nodes are not available here) # so a cuda:N engine is not mislabeled cuda:0. self._has_explicit_target_device = any( - s.key == _TARGET_DEVICE_COMPILE_SPEC_KEY for s in self.compile_specs + s.key == _TARGET_DEVICE_COMPILE_SPEC_KEY for s in self._base_compile_specs ) + # ExecuTorch partitioners conventionally hold a delegation_spec. partition() + # builds a fresh DelegationSpec per partition and never reads this one; the + # only reader is _export._declared_method_name, on a partitioner the caller + # passes to export(). self.delegation_spec = DelegationSpec( backend_id=TensorRTBackend.__name__, compile_specs=self.compile_specs, @@ -176,6 +201,81 @@ def _resolve_target_device_for_partition( ) return b"cuda:0" + def _partition_elided_output_names( + self, exported_program: ExportedProgram, partition: Partition + ) -> Set[str]: + """Engine output binding names this partition's delegate legitimately drops. + + Zero-copy KV elides an engine's aliased output when its aliased input is a + buffer export rewired to be written in place -- marked on the placeholder + with ``_torch_tensorrt_aliased_buffer``. This is derived from THIS + partition's own engine (its ``aliased_io`` paired with the marks on its own + input placeholders), never from a method-wide name list, so a second engine + that happens to share an output binding name is never told it may drop that + binding. That is what lets a real lost output on the plain delegate still + raise while the KV delegate's genuine elision is exempted. + + An extraction failure is survivable only for a method in which nothing was + rewired: the delegate then really does carry every binding, and a genuinely + missing aliased output stays an error in the backend's + ``_validate_output_binding_order``. Where a buffer *was* rewired the aliased + outputs are already gone from the graph -- that happened before partitioning + -- so an empty set stamps nothing, and the export dies downstream blaming a + lost aliased-buffer mark, which is not what went wrong. That case re-raises. + + A set that is neither empty nor the engine's whole ``aliased_io`` is the + right answer and still not a lowerable one, since the runtime reads + elision off a single argument count: ``TensorRTBackend.preprocess`` + refuses that engine. + """ + from torch_tensorrt.executorch._zero_copy import _aliased_inputs_by_output_index + + try: + engine_nodes = _get_engine_nodes_in(partition.nodes) + if len(engine_nodes) != 1: + return set() + engine = engine_nodes[0] + aliased = _aliased_inputs_by_output_index(exported_program, engine) + if not aliased: + return set() + # Only OUTPUT_BINDING_NAMES_IDX is read, never the engine itself. + engine_info = _get_engine_info_for_node( + exported_program, engine, metadata_only=True + ) + output_names = deserialize_binding_names( + _get_str(engine_info, OUTPUT_BINDING_NAMES_IDX) + ) + elided: Set[str] = set() + for output_index, input_node in aliased.items(): + if ( + isinstance(input_node, torch.fx.Node) + and input_node.meta.get("_torch_tensorrt_aliased_buffer") + and 0 <= output_index < len(output_names) + ): + elided.add(output_names[output_index]) + return elided + except Exception as e: + # Broad in the same shape as _resolve_target_device_for_partition, but + # not with the same licence: that fallback picks a default device and + # is genuinely harmless, while this one can only be harmless where no + # output was elided. The question is asked of the graph, not of the + # engine metadata that just failed to read, so it cannot fail the same + # way. + if any( + node.op == "placeholder" + and node.meta.get("_torch_tensorrt_aliased_buffer") + for node in exported_program.graph_module.graph.nodes + ): + raise + logger.warning( + "zero-copy KV: could not resolve elided outputs for partition %s " + "(%s); no buffer in this method was rewired, so the delegate " + "carries every binding.", + getattr(partition, "id", "?"), + e, + ) + return set() + def partition(self, exported_program: ExportedProgram) -> PartitionResult: capability_partitioner = CapabilityBasedPartitioner( exported_program.graph_module, @@ -189,21 +289,31 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: tag = f"tensorrt_{partition.id}" for node in partition.nodes: node.meta["delegation_tag"] = tag - if self._has_explicit_target_device: - partition_tags[tag] = self.delegation_spec - else: - partition_tags[tag] = DelegationSpec( - backend_id=TensorRTBackend.__name__, - compile_specs=self.compile_specs - + [ + specs = list(self._base_compile_specs) + if not self._has_explicit_target_device: + specs.append( + CompileSpec( + _TARGET_DEVICE_COMPILE_SPEC_KEY, + self._resolve_target_device_for_partition( + exported_program, partition + ), + ) + ) + if self._zero_copy_requested: + elided = self._partition_elided_output_names( + exported_program, partition + ) + if elided: + specs.append( CompileSpec( - _TARGET_DEVICE_COMPILE_SPEC_KEY, - self._resolve_target_device_for_partition( - exported_program, partition - ), + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(elided), ) - ], - ) + ) + partition_tags[tag] = DelegationSpec( + backend_id=TensorRTBackend.__name__, + compile_specs=specs, + ) tag_constant_data(exported_program) _keep_mutated_buffers_above_delegate(exported_program) diff --git a/py/torch_tensorrt/executorch/serialization.py b/py/torch_tensorrt/executorch/serialization.py index ba7ccac10c2..70599c4cfb2 100644 --- a/py/torch_tensorrt/executorch/serialization.py +++ b/py/torch_tensorrt/executorch/serialization.py @@ -50,8 +50,12 @@ class TensorRTBlobMetadata: target_platform: str = "" def to_json(self) -> bytes: - # Keep field order stable because the C++ parser is intentionally small - # and searches forward after io_bindings for the scalar fields. + # Keep field order stable because the C++ parser is intentionally small. + # It walks io_bindings, then aliased_io, then searches forward from the + # end of whichever of those two it last walked for the scalar fields -- + # so a scalar written before either array is not found and keeps its + # C++-side default while the parse still succeeds. Any field added here + # that the C++ side reads by key must go after both arrays. data = { "io_bindings": [ { diff --git a/tests/cpp/executorch/test_executorch_blob_header.cpp b/tests/cpp/executorch/test_executorch_blob_header.cpp index 6a7e5cd91dc..95556f6673e 100644 --- a/tests/cpp/executorch/test_executorch_blob_header.cpp +++ b/tests/cpp/executorch/test_executorch_blob_header.cpp @@ -177,6 +177,561 @@ TEST(ExecuTorchTensorRTBlobHeader, InputNamedAliasedIoWithNoAliasesStillParses) EXPECT_TRUE(header.aliased_io.empty()); } +TEST(ExecuTorchTensorRTBlobHeader, RejectsRepeatedAliasedIoOutput) { + // Two entries claiming out_k. They resolve to the same binding and the same + // engine alias, so nothing that looks the names up can tell them from one + // entry -- but a reader counting aliased outputs per entry counts out_k twice, + // which is how the backend sizes the delegate argument list. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"out_v","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"},)" + R"({"output":"out_k","input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsRepeatedAliasedIoOutputWithADifferentInput) { + // The same repeat with the second entry naming a different input. This is the + // shape the check above exists for: keying the refusal on the output/input + // pair instead of the output alone would accept it, and out_k's aliased-output + // count would be two for one output binding, which is the arity the backend + // subtracts on. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"out_v","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"},)" + R"({"output":"out_k","input":"in_v","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsRepeatedOutputBindingName) { + // One name, two output slots. Every name lookup stops at the first slot, so + // init would record the alias there, and execute() would then bind the second + // slot's ExecuTorch storage to the same TensorRT name -- replacing the address + // of the caller's buffer that the alias exists to write. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsRepeatedInputBindingName) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsOneNameUsedAsBothAnInputAndAnOutput) { + // TensorRT has one name space for its tensors, so this is the same collision + // as the two above rather than a distinct input and output that happen to + // share a spelling. + const std::string metadata = R"({"io_bindings":[{"name":"kv","is_input":true},{"name":"kv","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesAliasedIoEntriesForDistinctOutputs) { + // The minimal pair for the test above: the same blob with the second entry + // claiming its own output. A second entry is not itself the defect. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"out_v","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"},)" + R"({"output":"out_v","input":"in_v","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + ASSERT_EQ(header.aliased_io.size(), 2u); + EXPECT_EQ(header.aliased_io[0].output, "out_k"); + EXPECT_EQ(header.aliased_io[1].output, "out_v"); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsRepeatedAliasedIoInputForDifferentOutputs) { + // Two entries naming different outputs and one input. Both resolve to the same + // input index, so execute() binds both output bindings to that input's caller + // pointer -- one address with two writers, the second of which erases the + // first with no error. The kind here is "user", which init validates only by + // comparing shapes, so two same-shaped outputs get past it. + const std::string metadata = R"({"io_bindings":[{"name":"in_0","is_input":true},)" + R"({"name":"out_0","is_input":false},{"name":"out_1","is_input":false}],)" + R"("aliased_io":[{"output":"out_0","input":"in_0","kind":"user"},)" + R"({"output":"out_1","input":"in_0","kind":"user"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsEmptyBindingName) { + // Skipping a blank name instead of refusing it shortens the recorded output + // list while the delegate's argument list keeps its length: here one aliased + // output would be recorded and one real output dropped, so the aliased + // binding would consume the argument belonging to the dropped one. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +// The metadata is copied out of the blob with an explicit length, so a name can +// hold a NUL byte. The refusals above compare whole std::string values; every +// consumer hands name.c_str() to TensorRT, which stops at the first NUL. The +// three cases below that name a NUL build their metadata by std::string +// concatenation, because a raw string literal cannot carry an embedded one. +const std::string kNul(1, '\0'); + +TEST(ExecuTorchTensorRTBlobHeader, RejectsBindingNamesDifferingOnlyAfterANul) { + // Two distinct std::strings, one TensorRT tensor. The repeated-name refusal + // above compares the whole value and lets this through, and then the collision + // it exists to stop happens anyway: execute() binds an address for that one + // TensorRT name twice, the second replacing the first, so the engine's write + // lands entirely in the later output and the earlier one is never written. + // This blob declares no alias and needs none for that; with an aliased_io + // entry beside it the address replaced is the caller's cache, which is the + // case init() says it relies on this refusal for. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"out_k)" + kNul + + R"(a","is_input":false},{"name":"out_k)" + kNul + R"(b","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAnAliasedIoNameCarryingANul) { + // The binding names here are clean, so the walk above accepts them and only the + // alias walk's own check can refuse this. It is not the collision case -- init + // compares alias names to binding names whole, so this pair matches no binding + // and init would refuse it. Refusing at parse keeps the invariant that every + // name the header records is one the engine can be asked for, and puts the + // failure where the blob-header tests reach it without a GPU. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"output":"out_k)" + + kNul + R"(a","input":"in_k)" + kNul + R"(a","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsABindingNameThatIsOnlyANul) { + // Non-empty to the refusal above, empty to TensorRT -- so the emptiness check + // has to draw its line where c_str() draws it, not where size() does. + const std::string metadata = + R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":")" + kNul + R"(","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAliasedIoEntryWithABlankInput) { + // Skipping this entry rather than refusing it records one alias for the two + // the engine has, and execute() subtracts the recorded count from the delegate + // argument list, so the .pte fails its arity check on every call with a + // message that never mentions aliasing. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"out_v","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"},)" + R"({"output":"out_v","input":"","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAliasedIoEntryWithNoOutputKey) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsBindingEntryWithNoNameKey) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsWrappingEngineExtent) { + // engine_size is a 64-bit field read straight from the file, so adding it to + // engine_offset before comparing against the blob length wraps: 4096 plus + // 2^64-4086 is 10, which is comfortably inside an 8 KiB blob. The pointer and + // that length are what TensorRTBackend hands deserializeCudaEngine. + const std::string metadata = R"({"io_bindings":[{"name":"x","is_input":true}]})"; + constexpr std::size_t kBlobSize = 8192; + constexpr uint32_t kMetadataOffset = HEADER_SIZE; + constexpr uint32_t kEngineOffset = 4096; + + std::vector blob(kBlobSize, 0); + std::memcpy(blob.data(), TENSORRT_MAGIC, sizeof(TENSORRT_MAGIC)); + write_field(blob, METADATA_OFFSET_FIELD_OFFSET, kMetadataOffset); + write_field(blob, METADATA_SIZE_FIELD_OFFSET, static_cast(metadata.size())); + write_field(blob, ENGINE_OFFSET_FIELD_OFFSET, kEngineOffset); + write_field(blob, ENGINE_SIZE_FIELD_OFFSET, ~uint64_t{0} - (kEngineOffset - 11)); + std::memcpy(blob.data() + kMetadataOffset, metadata.data(), metadata.size()); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsEngineOffsetPastEndOfBlob) { + // The offset alone is out of range. The subtraction form has to refuse that + // before it evaluates size - engine_offset, which would itself wrap. + auto blob = make_blob(R"({"io_bindings":[{"name":"x","is_input":true}]})"); + const auto past_end = static_cast(align_up(blob.size() + ENGINE_ALIGNMENT, ENGINE_ALIGNMENT)); + write_field(blob, ENGINE_OFFSET_FIELD_OFFSET, past_end); + write_field(blob, ENGINE_SIZE_FIELD_OFFSET, uint64_t{0}); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsMetadataThatStartsPastTheEngine) { + // Both metadata fields are in range of the blob, so only the clause comparing + // them against engine_offset refuses this. That clause is written in the same + // subtraction form as the engine extent, and the subtraction is what needs the + // ordering test in front of it: engine_offset - metadata_offset is unsigned, + // so with the metadata past the engine it wraps to nearly 2^32 and any + // metadata_size fits under it. + const std::string metadata = R"({"io_bindings":[]})"; + constexpr std::size_t kBlobSize = 8192; + constexpr uint32_t kMetadataOffset = 4096; + constexpr uint32_t kEngineOffset = 48; + + std::vector blob(kBlobSize, 0); + std::memcpy(blob.data(), TENSORRT_MAGIC, sizeof(TENSORRT_MAGIC)); + write_field(blob, METADATA_OFFSET_FIELD_OFFSET, kMetadataOffset); + write_field(blob, METADATA_SIZE_FIELD_OFFSET, static_cast(metadata.size())); + write_field(blob, ENGINE_OFFSET_FIELD_OFFSET, kEngineOffset); + write_field(blob, ENGINE_SIZE_FIELD_OFFSET, uint64_t{4}); + std::memcpy(blob.data() + kMetadataOffset, metadata.data(), metadata.size()); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsBindingEntryWithNoIsInputKey) { + // Without the key the initializer in the parser reads the binding as an output + // while serialization.py's TensorRTIOBinding reads it as an input, so the two + // readers of these bytes disagree -- and the disagreement is not one slot: it + // moves in_v out of the input list, which shifts every index after it. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v"},)" + R"({"name":"out_k","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsBindingEntryWithAMisspelledIsInputKey) { + // One byte wrong is the same case: the key falls through to skip_value, which + // consumes the value and leaves the initializer standing. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v","is_inout":true},)" + R"({"name":"out_k","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsABindingNameThatArrivedEscaped) { + // parse_string drops the backslash and keeps what follows, so this name is + // recorded as the three characters anb and TensorRT has no such tensor. Every + // escape json.dumps reserves for a control character or a non-ASCII codepoint + // is a place the two ends read a name differently. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"a\nb","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesABindingNameEscapedTheWayThisParserReadsIt) { + // The three escapes whose JSON meaning is "the character after the + // backslash", which is what parse_string produces for every escape. A quote + // and a backslash are what json.dumps emits for a name holding one, so + // refusing these would refuse names the writer really produces and the + // merge-base parser recorded exactly as written. The forward slash is the + // third of the three; json.dumps writes it plain, but a hand-assembled blob + // may escape it and both ends read it the same way. + const std::string metadata = R"({"io_bindings":[{"name":"a\"b","is_input":true},)" + R"({"name":"c\\d","is_input":false},)" + R"({"name":"e\/f","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + ASSERT_EQ(header.input_binding_names.size(), 1u); + EXPECT_EQ(header.input_binding_names[0], "a\"b"); + ASSERT_EQ(header.output_binding_names.size(), 2u); + EXPECT_EQ(header.output_binding_names[0], "c\\d"); + EXPECT_EQ(header.output_binding_names[1], "e/f"); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAnEntryKeySpelledThroughAnEscape) { + // The key comparisons see the text parse_string produced, so a backslash + // before any letter of is_input still reads as is_input here while a JSON + // reader sees a key with a newline in it, ignores it, and leaves is_input at + // its own default -- which serialization.py's TensorRTIOBinding makes an + // input and the initializer here makes an output. That is the two readers of + // one blob putting one binding in opposite lists, which is exactly what + // RejectsBindingEntryWithAMisspelledIsInputKey stops for the plainly + // misspelled key. + const std::string metadata = R"({"io_bindings":[{"name":"x","is_i\nput":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesAnEntryKeyCarryingAnEscapeBothEndsAgreeOn) { + // The control for the refusal above: an escaped backslash in a key is decoded + // the same way here and by a JSON reader, so it names the key dty\pe for + // both, which neither matches. Nothing depends on it and the blob is good. + const std::string metadata = R"({"io_bindings":[{"name":"x","dty\\pe":"f32","is_input":true}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + ASSERT_EQ(header.input_binding_names.size(), 1u); + EXPECT_EQ(header.input_binding_names[0], "x"); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAnAliasEntryKeySpelledThroughAnEscape) { + // The alias entry's keys are compared the same way the binding entry's are. + // A backslash before any letter of output leaves this parser reading the + // output name out of it, while a JSON reader sees a key it does not know and + // an entry with no output at all -- which it would then have to refuse. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"\output":"out_k","input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAnAliasKindSpelledThroughAnEscape) { + // kind is not a binding name, but init() compares it: "user" is the kind + // validated on shape alone rather than confirmed against the engine's own + // aliasing, so reading a different string as user picks the weaker check. + // Here both entries come out as the plain string user, while a JSON reader + // refuses "\user" outright and reads "use\r" as use plus a carriage return. + for (const char* kind : {R"(\user)", R"(use\r)"}) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":")" + + std::string(kind) + R"("}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)) << kind; + } +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAnAliasedIoNameThatArrivedEscaped) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in\tk","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAliasedIoMagicWithNoAliasArrayFound) { + // TR02 says the metadata carries aliased_io. Here the key is one byte wrong, + // so the walk finds nothing and the header would come back alias-free -- and + // an alias-free header of a threaded .pte binds each aliased output to its own + // storage and stops updating the caller's cache, with nothing to fail on. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliasedXio":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAliasedIoMagicWithTheKeysInSortedOrder) { + // The alias array is searched for past the io_bindings array, so a writer that + // emitted the keys in sorted order would put it out of reach. TR02 is what + // makes that a refusal rather than a silently alias-free header. + const std::string metadata = R"({"aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}],)" + R"("io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesScalarsPastAnAliasedBindingNamedLikeAKey) { + // The two scalar scans search the metadata text, and the alias array sits + // between where io_bindings ends and where those scans used to start, so an + // aliased binding named device_id was matched as the key: the scan then walked + // to the next colon, met the kind string, and failed the whole blob. + const std::string metadata = R"({"io_bindings":[{"name":"device_id","is_input":true},)" + R"({"name":"hardware_compatible","is_input":false}],)" + R"("aliased_io":[{"output":"hardware_compatible","input":"device_id",)" + R"("kind":"kv_cache_update"}],"hardware_compatible":true,"device_id":3})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_TRUE(header.hardware_compatible); + EXPECT_EQ(header.device_id, 3); + EXPECT_EQ(header.aliased_io.size(), 1u); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesTheDeviceIdKeyOutsideAnAliasEntryCarryingOne) { + // The other half of the same window. Being in key position does not tell an + // unknown key inside an alias entry from the real one, which the alias walk + // skips and the scans would otherwise read as the field; starting them past + // the array is what does. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update",)" + R"("device_id":9}],"device_id":3})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_EQ(header.device_id, 3); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsADeviceIdPastTheIntMaximum) { + // Accumulated into an int this wrapped to 1, which is a GPU that exists on + // most machines: cudaSetDevice then succeeds and the engine deserializes on a + // device nobody asked for. + const auto blob = make_blob(R"({"io_bindings":[{"name":"x","is_input":true}],"device_id":4294967297})"); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesTheLargestDeviceIdAnIntHolds) { + // The bound is the int maximum itself, not something short of it, so the + // refusal above is about overflow and not about long-looking values. + const auto blob = make_blob(R"({"io_bindings":[{"name":"x","is_input":true}],"device_id":2147483647})"); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_EQ(header.device_id, 2147483647); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesAStringValueThatIsExactlyAScalarKeyName) { + // A string *value* that reads like the key: quoted the same way, and so a + // match for the same search. This blob carries no device_id of its own, which + // is what an older writer emits, so the value is the only match there is -- + // and reading it as the key meant walking to the next colon, meeting the + // target_platform string, and failing a blob that is perfectly good. What + // separates the two is that a key is followed by its own colon and a value is + // followed by a comma. + const std::string metadata = R"({"io_bindings":[{"name":"x","is_input":true}],)" + R"("serialized_metadata":"device_id","target_platform":"linux"})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_EQ(header.device_id, 0); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesEveryScalarFromTheWriterKeyOrder) { + // The key order TensorRTBlobMetadata.to_json emits, with every field present + // and both scalars set away from their defaults. The scalar scans start past + // whichever array they last walked, so a field moved ahead of one of them is + // not found and keeps its C++-side default while the parse still succeeds -- + // this is what fails if that order changes. The Python half of the same rule + // is test_serialization.py::test_to_json_writes_every_scalar_after_both_arrays. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","dtype":"float32","shape":[1,2],"is_input":true},)" + R"({"name":"out_k","dtype":"float32","shape":[1,2],"is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}],)" + R"("hardware_compatible":true,"device_id":6,)" + R"("serialized_metadata":"","target_platform":"linux_x86_64"})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_EQ(header.input_binding_names, std::vector{"in_k"}); + EXPECT_EQ(header.output_binding_names, std::vector{"out_k"}); + ASSERT_EQ(header.aliased_io.size(), 1u); + EXPECT_EQ(header.aliased_io[0].output, "out_k"); + EXPECT_EQ(header.aliased_io[0].input, "in_k"); + EXPECT_EQ(header.aliased_io[0].kind, "kv_cache_update"); + EXPECT_TRUE(header.hardware_compatible); + EXPECT_EQ(header.device_id, 6); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesAnArrayKeySpelledByAnEarlierStringValue) { + // The two array keys are found the way the two scalars are: an occurrence + // followed by its own colon. Taking the first occurrence anywhere and then + // the next '[' reads the value below as the key and walks the shape array + // that follows it, which refuses a blob that is perfectly good. + const std::string metadata = R"({"serialized_metadata":"io_bindings","shape":[9],)" + R"("io_bindings":[{"name":"x","is_input":true}],"device_id":6})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_EQ(header.input_binding_names, std::vector{"x"}); + EXPECT_EQ(header.device_id, 6); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesTheAliasArrayKeySpelledByAnEarlierStringValue) { + // The alias key gets the same treatment, and its window is narrower: the + // search already starts past io_bindings, so only a value between the two + // arrays can stand in for it -- which is where serialized_metadata sits. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("serialized_metadata":"aliased_io","shape":[9],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}],)" + R"("device_id":6})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + ASSERT_EQ(header.aliased_io.size(), 1u); + EXPECT_EQ(header.aliased_io[0].output, "out_k"); + EXPECT_EQ(header.device_id, 6); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsIoBindingsWhoseValueIsNotAnArray) { + // The array has to be the value of the key, not the next '[' in the text: a + // blob whose io_bindings is an object is otherwise walked from an unrelated + // bracket further on, and the entries found there are recorded as this + // engine's bindings. The array below is shaped like the real one so that + // walking it succeeds, which is what makes the wrong answer a silent one. + const std::string metadata = R"({"io_bindings":{"name":"x"},)" + R"("elsewhere":[{"name":"y","is_input":true}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + TEST(ExecuTorchTensorRTBlobHeader, RejectsUnknownFutureMagic) { constexpr char kFutureMagic[4] = {'T', 'R', '0', '3'}; const std::string metadata = R"({"io_bindings":[{"name":"x","is_input":true}]})"; diff --git a/tests/py/dynamo/executorch/test_api.py b/tests/py/dynamo/executorch/test_api.py index 7de4e563fef..191a93ea85f 100644 --- a/tests/py/dynamo/executorch/test_api.py +++ b/tests/py/dynamo/executorch/test_api.py @@ -92,18 +92,40 @@ def test_load_executorch_dispatches_to_delegate(monkeypatch): ) +_PUBLIC_API_SYMBOLS = ( + "get_edge_compile_config", + "TensorRTPartitioner", + "TensorRTBackend", + "export", + "zero_copy_backend_config", + "check_zero_copy_kv", +) + + @pytest.mark.unit def test_public_api_symbols_present(): module = importlib.import_module("torch_tensorrt.executorch") - assert "get_edge_compile_config" in module.__all__ - assert "TensorRTPartitioner" in module.__all__ - assert "TensorRTBackend" in module.__all__ - assert "export" in module.__all__ + assert set(module.__all__) == set(_PUBLIC_API_SYMBOLS) assert "Program" not in module.__all__ assert "load" not in module.__all__ assert "to_executorch" not in module.__all__ +@pytest.mark.unit +def test_public_api_symbols_are_bound_not_just_advertised(): + # __all__ is a literal written out in both branches of the + # _has_executorch_exir() guard, so reading it cannot tell whether the + # package binds what it advertises. Resolve each name instead. + module = importlib.import_module("torch_tensorrt.executorch") + if module._has_executorch_exir(): + for name in _PUBLIC_API_SYMBOLS: + assert getattr(module, name) is not None + else: + for name in _PUBLIC_API_SYMBOLS: + with pytest.raises(ImportError, match=name): + getattr(module, name) + + _REPO_ROOT = Path(__file__).resolve().parents[4] _SETUP_PY = _REPO_ROOT / "setup.py" _RUNTIME_SETUP_PY = _REPO_ROOT / "py/torch-tensorrt-executorch-runtime/setup.py" @@ -541,6 +563,10 @@ def get_etrecord(self): return _FakeETRecord() class _FakeEdge: + # export() reorders each method's mutations after lowering, over every + # method the manager holds. No method here holds a program to reorder. + methods = () + def to_executorch(self, config=None): captured["backend_config"] = config return _FakeExec() diff --git a/tests/py/dynamo/executorch/test_backend.py b/tests/py/dynamo/executorch/test_backend.py index 9c687bba6e1..6a3a04cc78e 100644 --- a/tests/py/dynamo/executorch/test_backend.py +++ b/tests/py/dynamo/executorch/test_backend.py @@ -9,6 +9,7 @@ executorch = pytest.importorskip("executorch.exir") import torch # noqa: E402 +from executorch.exir.backend.compile_spec_schema import CompileSpec # noqa: E402 from torch.export.graph_signature import InputKind # noqa: E402 from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( # noqa: E402 ALIASED_IO_IDX, @@ -437,3 +438,362 @@ def test_validate_output_binding_order_accepts_unwrapped_single_output(): g.output((engine,)) ep = SimpleNamespace(graph_module=torch.fx.GraphModule(torch.nn.Module(), g)) _validate_output_binding_order(ep, engine, ["out"]) + + +@pytest.mark.unit +def test_validate_output_binding_order_accepts_elided_aliased_outputs(): + """Zero-copy KV drops the in-place outputs from the delegate entirely. + + The engine still declares them as bindings -- the runtime binds them to + their aliased input's tensor -- but no delegate argument is passed for them. + """ + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([0]) + _validate_output_binding_order(ep, engine, ["out0", "kv0", "kv1"], {"kv0", "kv1"}) + + +@pytest.mark.unit +def test_validate_output_binding_order_rejects_elision_that_was_not_requested(): + """The same graph is a bug unless the caller asked for zero-copy. + + Aliased outputs missing because nothing declared them as buffer mutations + look identical to aliased outputs deliberately elided, here and at runtime. + So the default stays strict and only an explicit opt-in relaxes it. + + It is also the shape a zero-copy export produces when the aliased-buffer mark + never reaches the partitioner: nothing is stamped, so nothing is exempt. The + message names the feature, because that failure has no other symptom. + """ + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([0]) + with pytest.raises(ValueError, match="engine output indices") as excinfo: + _validate_output_binding_order(ep, engine, ["out0", "kv0", "kv1"]) + assert "zero_copy_kv" in str(excinfo.value) + + +@pytest.mark.unit +def test_validate_output_binding_order_still_accepts_aliased_outputs_threaded(): + """Naming a binding elidable permits the drop, it does not require it: a + delegate that still carries every aliased output is accepted too.""" + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([0, 1, 2]) + _validate_output_binding_order(ep, engine, ["out0", "kv0", "kv1"], {"kv0", "kv1"}) + + +@pytest.mark.unit +def test_validate_output_binding_order_rejects_partially_elided_outputs(): + """Half-elision is a lost buffer update, and the runtime cannot express it: + it infers elision from one argument count, so it is all-or-nothing.""" + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([0, 1]) + with pytest.raises(ValueError, match="engine output indices") as excinfo: + _validate_output_binding_order( + ep, engine, ["out0", "kv0", "kv1"], {"kv0", "kv1"} + ) + # The mark plainly did survive -- these bindings were declared elidable -- so + # the lost-mark hint would be a wrong lead here. + assert "zero_copy_kv" not in str(excinfo.value) + + +@pytest.mark.unit +def test_validate_output_binding_order_rejects_permutation_of_elided_outputs(): + """Elision removes outputs; it does not license reordering the survivors.""" + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([2, 0]) + with pytest.raises(ValueError, match="engine output indices"): + _validate_output_binding_order(ep, engine, ["out0", "kv0", "out2"], {"kv0"}) + + +def _aliased_edge_program(present_indices, out_names, aliased_io_map): + """A one-engine partition declaring aliased_io whose delegate keeps only + ``present_indices`` of its output bindings (the rest elided, as zero-copy + export produces). One input, ``tokens``, which the aliases point at.""" + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import serialize_aliased_io + + engine_info = [""] * SERIALIZATION_LEN + engine_info[ENGINE_IDX] = _engine_tensor(b"engine-bytes") + engine_info[INPUT_BINDING_NAMES_IDX] = "tokens" + engine_info[OUTPUT_BINDING_NAMES_IDX] = "%".join(out_names) + engine_info[ALIASED_IO_IDX] = serialize_aliased_io(aliased_io_map) + + graph = torch.fx.Graph() + tokens = graph.placeholder("tokens") + engine_node = graph.call_function(_ENGINE_OP, ([tokens], *engine_info)) + graph.output( + tuple( + graph.call_function(operator.getitem, (engine_node, i)) + for i in present_indices + ) + ) + return SimpleNamespace( + graph_module=SimpleNamespace(graph=graph), + graph_signature=SimpleNamespace( + input_specs=[ + SimpleNamespace( + kind=InputKind.USER_INPUT, arg=SimpleNamespace(name="tokens") + ) + ] + ), + constants={}, + ) + + +@pytest.mark.unit +def test_preprocess_rejects_elided_aliased_output_without_zero_copy_spec(): + """The gate at preprocess: a delegate that dropped its aliased output is a + bug unless a zero-copy compile spec says the drop was deliberate. Deleting + the gate (treating the drop as always allowed) would let this pass silently. + """ + from torch_tensorrt.executorch.backend import TensorRTBackend + + edge_program = _aliased_edge_program( + present_indices=[0], + out_names=["logits", "out_k"], + aliased_io_map={"out_k": ("tokens", "kv_cache_update")}, + ) + with pytest.raises(ValueError, match="engine output indices"): + TensorRTBackend.preprocess(edge_program, []) + + +@pytest.mark.unit +def test_preprocess_accepts_elided_aliased_output_with_zero_copy_spec(): + """The same delegate is accepted once the compile spec names the elided + binding, and only then.""" + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + TensorRTBackend, + _serialize_elided_output_names, + ) + + edge_program = _aliased_edge_program( + present_indices=[0], + out_names=["logits", "out_k"], + aliased_io_map={"out_k": ("tokens", "kv_cache_update")}, + ) + spec = CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, _serialize_elided_output_names(["out_k"]) + ) + # Does not raise, and produces a valid engine blob (aliased_io present bumps + # the blob format magic to TR02, so match the family, not the exact base). + result = TensorRTBackend.preprocess(edge_program, [spec]) + assert isinstance(result.processed_bytes, bytes) + assert result.processed_bytes[:2] == TENSORRT_MAGIC[:2] + + +@pytest.mark.unit +def test_preprocess_rejects_a_non_buffer_alias_elided_alongside_a_buffer_alias(): + """An engine mixing a buffer-backed aliased output (export rewired it and + named it in the spec) with a non-buffer-backed one (never rewired) must not + have the non-buffer one silently elided. The spec names only ``out_k``, so a + delegate that also dropped ``out_v`` is rejected -- exactly as it would be + without zero-copy.""" + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + TensorRTBackend, + _serialize_elided_output_names, + ) + + edge_program = _aliased_edge_program( + present_indices=[0], # both out_k (buffer) and out_v (non-buffer) dropped + out_names=["logits", "out_k", "out_v"], + aliased_io_map={ + "out_k": ("tokens", "kv_cache_update"), + "out_v": ("tokens", "kv_cache_update"), + }, + ) + spec = CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, _serialize_elided_output_names(["out_k"]) + ) + with pytest.raises(ValueError, match="engine output indices"): + TensorRTBackend.preprocess(edge_program, [spec]) + + +@pytest.mark.unit +def test_preprocess_rejects_an_engine_whose_aliased_outputs_are_partly_elided(): + """Two aliased outputs, one elided: the arity the runtime reads cannot say so. + + ``out_v`` is still a delegate output, so the binding-order check above is + satisfied and only the whole-set comparison refuses this. The runtime takes + the aliased outputs as elided only when the argument count is short by the + engine's whole aliased-output count, so a delegate short of one of two reads + as not elided at all and every ``execute()`` fails on the argument count. The + same clause is reached from the zero-copy suite through the partitioner's own + derivation; this pins it beside the other ``preprocess`` refusals. + """ + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + TensorRTBackend, + _serialize_elided_output_names, + ) + + edge_program = _aliased_edge_program( + present_indices=[0, 2], # logits and out_v kept, only out_k dropped + out_names=["logits", "out_k", "out_v"], + aliased_io_map={ + "out_k": ("tokens", "kv_cache_update"), + "out_v": ("tokens", "kv_cache_update"), + }, + ) + spec = CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, _serialize_elided_output_names(["out_k"]) + ) + with pytest.raises(ValueError, match="Partial elision is not expressible"): + TensorRTBackend.preprocess(edge_program, [spec]) + + +@pytest.mark.unit +def test_preprocess_rejects_a_delegate_with_no_outputs_at_all(): + """Naming every binding elidable leaves an empty expected index list, which + an empty delegate output list matches. That is the zero-output delegate the + rewiring pass raises about -- nothing reads it, so a later graph-wide + dead-code elimination can erase the engine with it -- and the rewiring's + guard runs before partitioning, so this is the only check left after + lowering.""" + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + TensorRTBackend, + _serialize_elided_output_names, + ) + + edge_program = _aliased_edge_program( + present_indices=[], + out_names=["out_k", "out_v"], + aliased_io_map={ + "out_k": ("tokens", "kv_cache_update"), + "out_v": ("tokens", "kv_cache_update"), + }, + ) + spec = CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(["out_k", "out_v"]), + ) + with pytest.raises(ValueError, match="no outputs at all"): + TensorRTBackend.preprocess(edge_program, [spec]) + + +# Values a hand-built zero-copy spec may carry that neither decoder can read. +# Applied to both decoders below: what each does with an unreadable value is the +# whole of the documented difference between them. +_unreadable_spec_values = pytest.mark.parametrize( + "value", + [1, True, None, b"", b"kv0", b'"kv0"', b"[1, 2", b"\xff\xfe"], + ids=[ + "int", + "bool", + "none", + "empty-bytes", + "bare-name", + "json-string", + "truncated-json", + "invalid-utf8", + ], +) + + +def _delegate_graph_with_specs(compile_specs): + """A one-node lowered graph whose delegate carries ``compile_specs``.""" + from executorch.exir.delegate import executorch_call_delegate + + graph = torch.fx.Graph() + lowered = graph.get_attr("lowered_module_0") + delegate = graph.call_function(executorch_call_delegate, (lowered,)) + graph.output((delegate,)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=compile_specs + ) + return torch.fx.GraphModule(root, graph), delegate + + +@pytest.mark.unit +@_unreadable_spec_values +def test_elided_output_names_refuses_a_value_it_cannot_read(value): + """A hand-built spec may carry anything, and every shape has to name the key. + + Two of these are the interesting ones. ``1`` is the obvious value for a key + that reads like a flag, and without a type check it surfaces as an + unattributed ``TypeError``. ``'"kv0"'`` is the quiet one: a JSON *string* + decodes into a set of its own characters, so nothing raises, the real binding + name is not exempted and every one-character binding name is. + """ + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _elided_output_names, + ) + + specs = [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, value)] + + with pytest.raises(ValueError, match=ZERO_COPY_KV_COMPILE_SPEC_KEY): + _elided_output_names(specs) + + +@pytest.mark.unit +@_unreadable_spec_values +def test_zero_copy_twin_reads_an_unreadable_value_as_cannot_tell(value): + """``_zero_copy``'s decoder on the same values the backend twin refuses. + + Where the backend raises, this returns the empty set, which its callers read + as "cannot tell" and answer by demanding at least one marked buffer. Drifting + to a raise here would abort an export these cross-checks only weaken. + """ + from torch_tensorrt.executorch._zero_copy import _delegate_elided_output_names + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + graph_module, delegate = _delegate_graph_with_specs( + [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, value)] + ) + + assert _delegate_elided_output_names(graph_module, delegate) == set() + + +@pytest.mark.unit +def test_zero_copy_twin_reads_a_missing_spec_as_cannot_tell(): + """A delegate carrying no zero-copy spec at all is the other "cannot tell". + + The one shape the two decoders do not share: for the backend a missing spec + is ``None`` and keeps a missing output an error, while here it joins the + values above. The second half is the control, so the empty sets are the + fallbacks and not a decoder that never answers. + """ + from torch_tensorrt.executorch._zero_copy import _delegate_elided_output_names + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names, + ) + + graph_module, delegate = _delegate_graph_with_specs([]) + assert _delegate_elided_output_names(graph_module, delegate) == set() + + graph_module, delegate = _delegate_graph_with_specs( + [ + CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(["out_k"]), + ) + ] + ) + assert _delegate_elided_output_names(graph_module, delegate) == {"out_k"} + + +@pytest.mark.unit +def test_elided_output_names_reads_the_list_the_partitioner_writes(): + """The control for the refusals above, so they cannot pass vacuously.""" + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _elided_output_names, + _serialize_elided_output_names, + ) + + specs = [ + CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, _serialize_elided_output_names(["kv0"]) + ) + ] + + assert _elided_output_names(specs) == {"kv0"} + assert _elided_output_names([]) is None diff --git a/tests/py/dynamo/executorch/test_edge_cases.py b/tests/py/dynamo/executorch/test_edge_cases.py index 4e5d7323d0d..a6223d67b9f 100644 --- a/tests/py/dynamo/executorch/test_edge_cases.py +++ b/tests/py/dynamo/executorch/test_edge_cases.py @@ -56,7 +56,7 @@ def test_save_as_executorch_uses_public_lowering_and_persists_data( backend_config=backend_config, ) - # The complete set of lowering options _save_as_executorch forwards. The five this + # The complete set of lowering options _save_as_executorch forwards. The six this # test does not pass are still forwarded explicitly, as None or False rather than # left out. backend_config is absent by design -- it is not a lowering option and is # routed to to_executorch() below. @@ -69,7 +69,10 @@ def test_save_as_executorch_uses_public_lowering_and_persists_data( constant_methods=None, generate_etrecord=False, weight_streaming_budget_per_engine=None, + zero_copy_kv=False, ) + # With zero_copy_kv off the caller's backend_config reaches to_executorch() + # exactly as given: save() wraps it only to install the un-staging pass. edge.to_executorch.assert_called_once_with(config=backend_config) program.write_to_file.assert_called_once() program.write_tensor_data_to_file.assert_called_once_with(str(tmp_path)) diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py index 3801d9a2cab..0bf58f40dd4 100644 --- a/tests/py/dynamo/executorch/test_export.py +++ b/tests/py/dynamo/executorch/test_export.py @@ -63,6 +63,34 @@ def __init__(self): self.graph_signature = SimpleNamespace(inputs_to_buffers={}, output_specs=[]) +class FakeEdgeProgramManager: + """The ``EdgeProgramManager`` stand-in ``to_edge_transform_and_lower`` returns. + + ``export()`` reads back the methods of the manager it is handed, to put each + zero-copy method's mutations in the order ExecuTorch finalizes them in. These + programs declare no mutation, so that call finds nothing to reorder. + + For a zero-copy method it also replaces ``to_executorch`` on the manager with + one that reads the finalized program back, so the attribute has to exist here + for that to bind to. A test that wants to watch that happen substitutes its + own on the instance; this default is here to make an accidental finalization + loud rather than to be called. + """ + + def __init__(self): + self._programs = {"forward": FakeExportedProgram()} + self.methods = set(self._programs) + + def exported_program(self, method_name="forward"): + return self._programs[method_name] + + def to_executorch(self, config=None): + raise AssertionError( + "this stub manager does not finalize; a test that needs to must " + "substitute its own to_executorch on the instance" + ) + + class FakeTensorRTPartitioner: def __init__(self, compile_specs): self.compile_specs = compile_specs @@ -241,7 +269,7 @@ def _patch_lowering(monkeypatch, engine_counts=None): ) export_module = importlib.import_module("torch_tensorrt.executorch._export") engine_counts = engine_counts or {} - lower = MagicMock(return_value=object()) + lower = MagicMock(return_value=FakeEdgeProgramManager()) monkeypatch.setattr(executorch.exir, "to_edge_transform_and_lower", lower) monkeypatch.setattr(executorch_api, "TensorRTPartitioner", FakeTensorRTPartitioner) monkeypatch.setattr(executorch_api, "get_edge_compile_config", lambda: "default") @@ -500,6 +528,313 @@ def test_export_returns_edge_and_forwards_all_options(monkeypatch): assert compile_specs == [compile_spec] +def _patch_declare(monkeypatch, log=None): + """Record the programs the declaration pass sees, and hand each one back. + + export() imports the symbol inside its own body, so the patch has to land on + the module that owns it. The stub takes ``**kw`` because export() passes + ``copyback_buffers=``. + + ``log`` is a shared, tagged call log. Two separate recorders would each be + satisfied by their own calls whichever order the two passes ran in, and the + order is the thing under test. + """ + import torch_tensorrt.dynamo._exporter as dynamo_exporter + + seen = [] + + def _declare(program, **kw): + seen.append(program) + if log is not None: + log.append(("declare", program)) + return program + + monkeypatch.setattr( + dynamo_exporter, "_declare_aliased_kv_mutations_on_ep", _declare + ) + return seen + + +def _patch_rewire(monkeypatch, elided_names=("kv",), log=None): + import torch_tensorrt.executorch._zero_copy as zero_copy + + seen = [] + + def _rewire(program): + seen.append(program) + if log is not None: + log.append(("rewire", program)) + return list(elided_names) + + monkeypatch.setattr(zero_copy, "rewire_aliased_mutations_to_buffers", _rewire) + return seen + + +@pytest.mark.unit +def test_export_zero_copy_kv_rewires_every_method(monkeypatch): + """The opt-in is what makes the aliased buffers zero-copy; nothing else does. + + It has to run per method and after the declaration, since it works from the + mutations that declaration produced. + """ + export_module, lower = _patch_lowering(monkeypatch) + calls = [] + declared = _patch_declare(monkeypatch, log=calls) + rewired = _patch_rewire(monkeypatch, log=calls) + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + + export_module.export( + {"prefill": prefill, "decode": decode}, + partitioners={"prefill": [object()], "decode": [object()]}, + zero_copy_kv=True, + ) + + assert declared == [prefill, decode] + assert rewired == [prefill, decode] + # One log, so the order between the two passes is pinned and not just the + # order within each. Rewiring works from the mutations declaration produced, + # so running it first would find nothing to rewire. + assert calls == [ + ("declare", prefill), + ("declare", decode), + ("rewire", prefill), + ("rewire", decode), + ] + # The backend rejects a delegate missing its aliased outputs unless it is + # told the omission was deliberate, and the presence of this key on the + # partitioner is the only channel that says so. The partitioner drops the + # spec itself and re-derives the names per engine, so its value is not read. + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + from torch_tensorrt.executorch.partitioner import TensorRTPartitioner + + for pipeline in lower.call_args.kwargs["partitioner"].values(): + specs = pipeline[0].compile_specs + assert any(spec.key == ZERO_COPY_KV_COMPILE_SPEC_KEY for spec in specs) + # The partitioner here is a stub, so what the real one makes of these + # specs is asserted against the real one. + real = TensorRTPartitioner(compile_specs=specs) + assert real._zero_copy_requested + assert not any( + spec.key == ZERO_COPY_KV_COMPILE_SPEC_KEY + for spec in real._base_compile_specs + ) + + +@pytest.mark.unit +def test_export_wires_the_reorder_and_the_finalize_check(monkeypatch): + """The two things export does after lowering, on the lane without a GPU. + + Both are wiring rather than computation, so the tests of the pieces + themselves say nothing about whether export calls them. The reorder has to + run after lowering because ``to_edge`` re-derives the graph signature, and + the check has to be installed on the manager export hands back because that + is the only object that knows zero-copy was asked for. + """ + import torch_tensorrt.executorch._zero_copy as zero_copy + + export_module, lower = _patch_lowering(monkeypatch) + _patch_declare(monkeypatch) + _patch_rewire(monkeypatch) + + manager = FakeEdgeProgramManager() + manager._programs = { + "prefill": FakeExportedProgram(), + "decode": FakeExportedProgram(), + } + manager.methods = set(manager._programs) + finalized = object() + manager.to_executorch = lambda config=None: finalized + lower.return_value = manager + + reordered = [] + monkeypatch.setattr( + zero_copy, + "order_copyback_mutations_first", + lambda program: (reordered.append(program), 0)[1], + ) + checked = [] + monkeypatch.setattr(zero_copy, "check_zero_copy_kv", checked.append) + + result = export_module.export( + {"prefill": FakeExportedProgram(), "decode": FakeExportedProgram()}, + partitioners={"prefill": [object()], "decode": [object()]}, + zero_copy_kv=True, + ) + + assert result is manager + assert {id(program) for program in reordered} == { + id(program) for program in manager._programs.values() + } + # Nothing is finalized yet, so the check must not have run. + assert checked == [] + + assert result.to_executorch(zero_copy.zero_copy_backend_config()) is finalized + assert checked == [finalized] + + +@pytest.mark.unit +def test_export_installs_the_finalize_check_even_when_nothing_was_rewired(monkeypatch): + """A zero_copy_kv=True that rewired nothing is the shape the check refuses. + + Export warns and carries on, so the manager it hands back would otherwise + finalize a plain staged program without a word, while + ``save(..., zero_copy_kv=True)`` refuses that same model. Gating the hook on + what was rewired would leave the case that most needs reading back as the + one case nobody reads back. + """ + import torch_tensorrt.executorch._zero_copy as zero_copy + + export_module, lower = _patch_lowering(monkeypatch) + _patch_declare(monkeypatch) + _patch_rewire(monkeypatch, elided_names=()) + + manager = FakeEdgeProgramManager() + finalized = object() + manager.to_executorch = lambda config=None: finalized + lower.return_value = manager + monkeypatch.setattr(zero_copy, "order_copyback_mutations_first", lambda program: 0) + checked = [] + monkeypatch.setattr(zero_copy, "check_zero_copy_kv", checked.append) + + result = export_module.export(FakeExportedProgram(), zero_copy_kv=True) + + assert checked == [] + assert result.to_executorch(zero_copy.zero_copy_backend_config()) is finalized + assert checked == [finalized] + + +@pytest.mark.unit +def test_export_reorders_mutations_for_a_method_that_never_asked_for_zero_copy( + monkeypatch, +): + """The reorder is not gated on zero-copy, because the crossing is not either. + + Stock ``to_edge().to_executorch()`` crosses the mutation pairing on a plain + module that writes one buffer from another buffer and a second from a user + input, with no TensorRT and no ``zero_copy_kv``. Gating the repair on + ``zero_copy_kv`` would leave that caller with the crossed program and would + repair a sibling method only because another method opted in. + """ + import torch_tensorrt.executorch._zero_copy as zero_copy + + export_module, lower = _patch_lowering(monkeypatch) + _patch_declare(monkeypatch) + + manager = FakeEdgeProgramManager() + lower.return_value = manager + + reordered = [] + monkeypatch.setattr( + zero_copy, + "order_copyback_mutations_first", + lambda program: (reordered.append(program), 0)[1], + ) + checked = [] + monkeypatch.setattr(zero_copy, "check_zero_copy_kv", checked.append) + + result = export_module.export(FakeExportedProgram(), zero_copy_kv=False) + + assert result is manager + assert [id(program) for program in reordered] == [ + id(manager.exported_program("forward")) + ] + # No zero-copy was asked for, so no finalize-time check is installed: the + # hook binds to the instance, so an unhooked manager has nothing there. + assert "to_executorch" not in vars(result) + assert checked == [] + + +@pytest.mark.unit +def test_export_does_not_exempt_a_method_that_kept_all_its_outputs(monkeypatch): + """The exemption is per method and only where an output was actually elided. + + A method that lost an output for some other reason must still be caught by + the backend's output-binding check. + """ + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + export_module, lower = _patch_lowering(monkeypatch) + import torch_tensorrt.executorch._zero_copy as zero_copy + + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + elided = {prefill: [], decode: ["k0", "k1"]} + monkeypatch.setattr( + zero_copy, "rewire_aliased_mutations_to_buffers", lambda p: elided[p] + ) + + export_module.export( + {"prefill": prefill, "decode": decode}, + partitioners={"prefill": [object()], "decode": [object()]}, + zero_copy_kv=True, + ) + + exempt = { + name: any( + spec.key == ZERO_COPY_KV_COMPILE_SPEC_KEY + for spec in pipeline[0].compile_specs + ) + for name, pipeline in lower.call_args.kwargs["partitioner"].items() + } + assert exempt == {"prefill": False, "decode": True} + + +@pytest.mark.unit +def test_export_zero_copy_kv_keeps_the_weight_streaming_spec(monkeypatch): + """Both options stamp the same partitioner, and one must not displace the other. + + The zero-copy spec is appended to a copy of the method's compile specs, so a + budget baked in earlier has to still reach the delegate. + """ + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + from torch_tensorrt.executorch.partitioner import ( + WEIGHT_STREAMING_BUDGET_COMPILE_SPEC_KEY, + ) + + export_module, lower = _patch_lowering(monkeypatch) + _patch_rewire(monkeypatch) + + export_module.export( + FakeExportedProgram(), + zero_copy_kv=True, + weight_streaming_budget_per_engine=1 << 20, + ) + + keys = {spec.key for spec in lower.call_args.kwargs["partitioner"][0].compile_specs} + assert keys == { + WEIGHT_STREAMING_BUDGET_COMPILE_SPEC_KEY, + ZERO_COPY_KV_COMPILE_SPEC_KEY, + } + + +@pytest.mark.unit +def test_export_leaves_kv_buffers_staged_by_default(monkeypatch): + """Zero-copy is opt-in: a .pte that elides its aliased outputs cannot be run + by a runtime that predates the feature, so export must not produce one + unasked.""" + export_module, lower = _patch_lowering(monkeypatch) + rewired = _patch_rewire(monkeypatch) + + export_module.export(FakeExportedProgram()) + + assert rewired == [] + + +@pytest.mark.unit +def test_export_warns_when_zero_copy_kv_has_nothing_to_do(monkeypatch, caplog): + """Asking for zero-copy on a model with no engine-aliased buffer is not an + error, but silently doing nothing would leave the caller expecting a speedup + that is not coming.""" + export_module, lower = _patch_lowering(monkeypatch) + _patch_rewire(monkeypatch, elided_names=()) + + with caplog.at_level("WARNING", logger=export_module.logger.name): + export_module.export(FakeExportedProgram(), zero_copy_kv=True) + + assert "no aliased buffer mutation was found" in caplog.text + + @pytest.mark.unit def test_export_preserves_independent_method_mapping(monkeypatch): prefill = FakeExportedProgram() @@ -1677,3 +2012,47 @@ def counting_get_engine_info_from_state(engine_obj, *, metadata_only=False): # is the tensor path fetching the bytes, which it re-resolves rather than trusting # the metadata-only record. Without the handoff the first entry would repeat. assert calls == [True, False] + + +@pytest.mark.unit +def test_export_rejects_caller_set_zero_copy_compile_spec(monkeypatch): + """The zero-copy key is reserved. export() sets it itself and only for the + outputs it elided; a hand-set value would exempt the delegate from the + output-binding check and could drop a real KV update silently. + + A bare list of compile specs is one unnamed list, but export() still fans it + out per method, so the message names the method it landed on. + """ + from executorch.exir.backend.compile_spec_schema import CompileSpec + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + export_module, lower = _patch_lowering(monkeypatch) + with pytest.raises(ValueError, match="reserved key") as excinfo: + export_module.export( + FakeExportedProgram(), + compile_specs=[CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")], + ) + assert "compile_specs for 'forward'" in str(excinfo.value) + + +@pytest.mark.unit +def test_export_rejects_caller_set_zero_copy_compile_spec_per_method(monkeypatch): + """The reserved-key rejection also covers the per-method mapping form, and + names the one method whose list carries the key rather than the method the + mapping happens to start with.""" + from executorch.exir.backend.compile_spec_schema import CompileSpec + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + export_module, lower = _patch_lowering(monkeypatch) + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + with pytest.raises(ValueError, match="reserved key") as excinfo: + export_module.export( + {"prefill": prefill, "decode": decode}, + partitioners={"prefill": [object()], "decode": [object()]}, + compile_specs={ + "prefill": [], + "decode": [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")], + }, + ) + assert "compile_specs for 'decode'" in str(excinfo.value) diff --git a/tests/py/dynamo/executorch/test_serialization.py b/tests/py/dynamo/executorch/test_serialization.py index c025b11c024..7fc45f9b24e 100644 --- a/tests/py/dynamo/executorch/test_serialization.py +++ b/tests/py/dynamo/executorch/test_serialization.py @@ -120,3 +120,47 @@ def test_deserialize_accepts_both_magics(aliased_io): engine, parsed = deserialize_engine(serialize_engine(b"engine-bytes", metadata)) assert engine == b"engine-bytes" assert parsed.aliased_io == aliased_io + + +@pytest.mark.unit +def test_to_json_writes_every_scalar_after_both_arrays(): + """The ordering rule the C++ parser's scalar scans depend on. + + ``TensorRTBlobHeader.cpp`` walks ``io_bindings``, then ``aliased_io``, then + searches forward from the end of whichever of those it last walked for the + scalar fields, so a scalar written ahead of either array is not found and + keeps its C++-side default while the parse still succeeds -- no error, and a + ``device_id`` of 0 means the engine deserializes on a GPU nobody named. The + C++ half of this rule is + ``ParsesEveryScalarFromTheWriterKeyOrder`` in + ``tests/cpp/executorch/test_executorch_blob_header.cpp``; this is the half + that fails when a field is added to ``to_json`` in the wrong place. + + Every key that side reads by key is listed, not only the two scalars: a new + scalar it learns to read is only safe in the same position. + """ + metadata = TensorRTBlobMetadata( + io_bindings=[ + TensorRTIOBinding(name="in_k", dtype="float32", shape=[1, 2]), + TensorRTIOBinding(name="out_k", dtype="float32", is_input=False), + ], + aliased_io={"out_k": ("in_k", "kv_cache_update")}, + hardware_compatible=True, + device_id=6, + target_platform="linux_x86_64", + ) + + text = metadata.to_json().decode("utf-8") + arrays_end = max( + text.index("]", text.index(key)) for key in ('"io_bindings"', '"aliased_io"') + ) + for key in ('"hardware_compatible"', '"device_id"'): + assert text.index(key) > arrays_end, f"{key} is written before an array" + + # Read back through the writer's own reader as well, so the ordering + # assertion is made about a payload that is otherwise correct. + restored = TensorRTBlobMetadata.from_json(metadata.to_json()) + assert restored.hardware_compatible is True + assert restored.device_id == 6 + assert restored.aliased_io == {"out_k": ("in_k", "kv_cache_update")} + assert [b.name for b in restored.io_bindings] == ["in_k", "out_k"] diff --git a/tests/py/dynamo/executorch/test_weight_streaming_budget.py b/tests/py/dynamo/executorch/test_weight_streaming_budget.py index cce7194f3c5..eff03a3e0c0 100644 --- a/tests/py/dynamo/executorch/test_weight_streaming_budget.py +++ b/tests/py/dynamo/executorch/test_weight_streaming_budget.py @@ -70,7 +70,9 @@ def _patch_lowering(monkeypatch, engine_counts=None): ) export_module = importlib.import_module("torch_tensorrt.executorch._export") engine_counts = engine_counts or {} - lower = MagicMock(return_value=object()) + # export() reorders each method's mutations over every method the manager + # holds after lowering, so the stand-in has to answer that much. + lower = MagicMock(return_value=SimpleNamespace(methods=())) monkeypatch.setattr(executorch.exir, "to_edge_transform_and_lower", lower) monkeypatch.setattr(executorch_api, "TensorRTPartitioner", FakeTensorRTPartitioner) monkeypatch.setattr(executorch_api, "get_edge_compile_config", lambda: "default") @@ -277,7 +279,7 @@ def test_save_rejects_negative_budget(tmp_path): @pytest.mark.unit def test_save_rejects_unknown_executorch_kwarg(tmp_path): - with pytest.raises(TypeError, match="unexpected keyword argument"): + with pytest.raises(TypeError, match="unexpected keyword argument") as excinfo: save( torch.nn.Linear(1, 1), str(tmp_path / "model.pte"), @@ -285,6 +287,15 @@ def test_save_rejects_unknown_executorch_kwarg(tmp_path): weight_streaming_budget_per_enginet=4096, ) + # The message spells the supported set out, so someone who mistyped an option + # is told what to type instead. A hand-written list drifts the moment an + # option is added, and then tells them the flag they wanted is unsupported. + from torch_tensorrt._compile import _EXECUTORCH_SAVE_OPTIONS + + assert len(_EXECUTORCH_SAVE_OPTIONS) > 1 + for name in _EXECUTORCH_SAVE_OPTIONS: + assert repr(name) in str(excinfo.value) + @pytest.mark.unit def test_save_warns_when_budget_used_with_non_executorch_format(tmp_path, caplog): diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py new file mode 100644 index 00000000000..2d88379cc75 --- /dev/null +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -0,0 +1,3501 @@ +"""Export-side coverage for zero-copy aliased KV buffers. + +Two halves have to agree for a zero-copy ``.pte`` to be correct: + + * ``rewire_aliased_mutations_to_buffers`` declares the buffer to be its own + mutation result, which removes ExecuTorch's copy-back and takes the aliased + output out of the delegate. + * ``unstage_aliased_buffers_pass`` removes the host staging copy so the + engine's in-place write lands in the caller's buffer. + +The interesting failures are all silent -- a rewired mutation whose buffer is +still staged simply never updates -- so most of what is asserted here is which +mutations are left alone, which mis-shapes raise, and that a marked buffer that +is never un-staged is caught rather than dropped. +""" + +import json +import logging +import operator +import re +from types import SimpleNamespace + +import pytest + +pytest.importorskip("executorch.exir") + +import torch # noqa: E402 +import torch_tensorrt # noqa: E402 +from executorch.exir.backend.compile_spec_schema import CompileSpec # noqa: E402 +from executorch.exir.delegate import executorch_call_delegate # noqa: E402 +from executorch.exir.schema import DeviceType # noqa: E402 +from torch.export.exported_program import ( # noqa: E402 + OutputKind, + OutputSpec, + TensorArgument, +) +from torch_tensorrt.executorch import _zero_copy as Z # noqa: E402 +from torch_tensorrt.executorch.backend import ( # noqa: E402 + ZERO_COPY_KV_COMPILE_SPEC_KEY, +) + +# The graphs below are built around torch.ops.tensorrt.execute_engine, which only +# exists once the Torch-TensorRT runtime operator library has loaded. +pytestmark = pytest.mark.skipif( + not torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime, + reason="Torch-TensorRT runtime operators are not available", +) + + +def _require_real_engine(): + """Gate the tests that build a real engine, at run time rather than collection. + + A decorator ``skipif`` resolves while pytest collects, and on a remote-GPU + runner collection happens off the GPU host, so the skip is frozen in before + any GPU is attached. These are the only tests that put this feature on a real + engine, so where that happens the lane stays green with that coverage gone. + The other CUDA gates in this directory are runtime gates for the same reason. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA + TensorRT for a real engine") + + +def _patch_engine_metadata(monkeypatch, *, aliased_io, input_names, output_names): + """Make every engine node report one fixed set of bindings and aliases.""" + import torch_tensorrt.dynamo.runtime._serialized_engine_layout as layout + import torch_tensorrt.dynamo.runtime._TorchTensorRTModule as trt_module + import torch_tensorrt.executorch._export_utils as export_utils + + info = ["x"] * (layout.ALIASED_IO_IDX + 1) + info[layout.INPUT_BINDING_NAMES_IDX] = "IN" + info[layout.OUTPUT_BINDING_NAMES_IDX] = "OUT" + + # The rewiring resolves engine info through _resolve_engine_info (the node is + # still an execute_engine at this stage), so that is what to fake. The stub + # requires metadata_only: without it the read goes through + # TRTEngine.__getstate__ and re-serializes the whole engine to recover the + # binding names and aliased_io, which are the only fields wanted here. + def _fake_resolve(ep, node, *, metadata_only=False): + assert metadata_only, "zero-copy reads binding metadata, not the engine" + return info + + monkeypatch.setattr(export_utils, "_resolve_engine_info", _fake_resolve) + monkeypatch.setattr(trt_module, "deserialize_aliased_io", lambda s: aliased_io) + monkeypatch.setattr( + layout, + "deserialize_binding_names", + lambda s: list(input_names) if s == "IN" else list(output_names), + ) + + +def _kv_program(*, mutation_value="aliased_getitem"): + """A one-engine program: engine(k_buffer, tokens) -> (logits, k_out). + + ``mutation_value`` picks what the KV buffer's BUFFER_MUTATION is bound to: + + * ``"aliased_getitem"``: the engine's aliased output (what export + declares for a caller-owned KV cache). + * ``"user_getitem"``: a non-aliased engine output, the shape a copy-back + mutation has. + * ``"external_op"``: a value produced outside the engine. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + tokens = graph.placeholder("tokens") + engine = graph.placeholder("engine") + engine_call = graph.call_function( + torch.ops.tensorrt.execute_engine.default, ([k_buffer, tokens], engine) + ) + logits = graph.call_function(operator.getitem, (engine_call, 0)) + k_out = graph.call_function(operator.getitem, (engine_call, 1)) + mutation = { + "aliased_getitem": k_out, + "user_getitem": logits, + "external_op": None, + }[mutation_value] + if mutation is None: + mutation = graph.call_function(torch.add, (k_buffer, k_buffer)) + graph.output((mutation, logits)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + signature = SimpleNamespace( + inputs_to_buffers={"b_k_0": "k_0"}, + input_specs=[], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(name=mutation.name), "k_0" + ), + OutputSpec(OutputKind.USER_OUTPUT, TensorArgument(name=logits.name), None), + ], + ) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=signature, + _graph_signature=signature, + ) + return program, k_buffer, k_out + + +@pytest.mark.unit +def test_rewire_points_the_mutation_at_its_buffer_and_marks_it(monkeypatch): + """The aliased output is replaced by the buffer itself and then dies. + + With the mutation bound to the placeholder there is nothing for ExecuTorch + to copy back, and with no other user the getitem leaves the graph -- which + is what takes the aliased output out of the delegate. The elided output's + binding name is returned so the backend can exempt exactly that one. + """ + program, k_buffer, k_out = _kv_program() + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in", "tokens"], + output_names=["logits", "out_k"], + ) + + assert Z.rewire_aliased_mutations_to_buffers(program) == ["out_k"] + + specs = program._graph_signature.output_specs + assert specs[0].kind == OutputKind.BUFFER_MUTATION + assert specs[0].target == "k_0" + assert specs[0].arg.name == k_buffer.name + output_node = program.graph_module.graph.output_node() + assert output_node.args[0][0] is k_buffer + assert k_out not in program.graph_module.graph.nodes + assert k_buffer.meta["_torch_tensorrt_aliased_buffer"] is True + + +@pytest.mark.unit +@pytest.mark.parametrize("mutation_value", ["user_getitem", "external_op"]) +def test_rewire_leaves_mutations_the_engine_does_not_alias(monkeypatch, mutation_value): + """Only a mutation the engine satisfies in place may be rewired. + + A copy-back mutation ("user_getitem") and a mutation computed outside the + engine ("external_op") both need their value copied into the buffer. Both + look exactly like an aliased mutation in the graph, so the discriminator has + to be the engine's own aliased_io -- rewiring either would delete a real + update with no error. + """ + program, k_buffer, _ = _kv_program(mutation_value=mutation_value) + original_spec = program._graph_signature.output_specs[0] + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in", "tokens"], + output_names=["logits", "out_k"], + ) + + assert Z.rewire_aliased_mutations_to_buffers(program) == [] + assert program._graph_signature.output_specs[0] is original_spec + assert "_torch_tensorrt_aliased_buffer" not in k_buffer.meta + + +@pytest.mark.unit +@pytest.mark.parametrize("unresolvable", ["unknown-input", "index-past-the-args"]) +def test_rewire_skips_an_alias_whose_input_does_not_resolve(monkeypatch, unresolvable): + """An aliased_io entry naming an input this delegate does not take is skipped. + + Two ways it can fail to resolve: the name is not one of the engine's input + bindings at all, or it is but its index is past the end of the delegate's + argument list. Neither leaves a mutation that could be rewired, and + ``_declare_aliased_kv_mutations_on_ep`` has already warned about both for the + same engine, so both are skipped rather than reported here. + """ + if unresolvable == "unknown-input": + aliased_input, input_names = "not_an_input", ["k_in", "tokens"] + else: + # A real binding name, but the third one, while the engine node takes two + # arguments -- so the index is past the end of the argument list. + aliased_input, input_names = "spare", ["k_in", "tokens", "spare"] + program, k_buffer, _ = _kv_program() + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": (aliased_input, "kv_cache_update")}, + input_names=input_names, + output_names=["logits", "out_k"], + ) + + assert Z.rewire_aliased_mutations_to_buffers(program) == [] + assert "_torch_tensorrt_aliased_buffer" not in k_buffer.meta + + +def _mixed_program(): + """One engine, one method, both kinds of mutation at once. + + ``engine(b_k_0, b_state_0, tokens) -> (logits, out_k, out_state)`` where the + engine aliases only ``out_k`` onto ``b_k_0``. ``b_state_0`` is the #4459 + shape: a mutable buffer with no aliasing available, whose new value + ``lift_mutated_buffers`` appended as a trailing output for ExecuTorch to copy + back. In the graph the two mutations are indistinguishable -- each is a + ``getitem`` off the engine node whose buffer is also an engine input. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + state_buffer = graph.placeholder("b_state_0") + tokens = graph.placeholder("tokens") + engine = graph.placeholder("engine") + engine_call = graph.call_function( + torch.ops.tensorrt.execute_engine.default, + ([k_buffer, state_buffer, tokens], engine), + ) + logits = graph.call_function(operator.getitem, (engine_call, 0)) + k_out = graph.call_function(operator.getitem, (engine_call, 1)) + state_out = graph.call_function(operator.getitem, (engine_call, 2)) + graph.output((k_out, state_out, logits)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + signature = SimpleNamespace( + inputs_to_buffers={"b_k_0": "k_0", "b_state_0": "state_0"}, + input_specs=[], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(name=k_out.name), "k_0" + ), + OutputSpec( + OutputKind.BUFFER_MUTATION, + TensorArgument(name=state_out.name), + "state_0", + ), + OutputSpec(OutputKind.USER_OUTPUT, TensorArgument(name=logits.name), None), + ], + ) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=signature, + _graph_signature=signature, + ) + return program, k_buffer, state_buffer, k_out, state_out + + +@pytest.mark.unit +def test_rewire_keeps_the_copyback_in_a_method_that_also_has_an_aliased_kv(monkeypatch): + """Zero-copy and a copy-back buffer may share one method, and must not mix. + + Rewiring the copy-back would delete a real update with no error, and refusing + the aliased one would give up the whole feature for any model carrying a + non-KV mutable buffer beside its cache. The engine's own aliased_io is what + separates them: only ``out_k`` is listed, so only ``b_k_0`` is rewired and + only its binding name is offered to the backend as elided. ``b_state_0`` + keeps its delegate output, which is the value ExecuTorch copies back. + """ + program, k_buffer, state_buffer, k_out, state_out = _mixed_program() + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in", "state_in", "tokens"], + output_names=["logits", "out_k", "out_state"], + ) + + assert Z.rewire_aliased_mutations_to_buffers(program) == ["out_k"] + + kv_spec, state_spec, _ = program._graph_signature.output_specs + assert kv_spec.arg.name == k_buffer.name + assert k_buffer.meta["_torch_tensorrt_aliased_buffer"] is True + assert k_out not in program.graph_module.graph.nodes + + assert state_spec.kind == OutputKind.BUFFER_MUTATION + assert state_spec.target == "state_0" + assert state_spec.arg.name == state_out.name + assert state_out in program.graph_module.graph.nodes + assert program.graph_module.graph.output_node().args[0][1] is state_out + # Un-staging keys on this mark, so leaving it off b_state_0 is what keeps the + # copy-back buffer's staging copy -- the engine writes that copy and + # ExecuTorch copies it back, exactly as without zero-copy. + assert "_torch_tensorrt_aliased_buffer" not in state_buffer.meta + + +@pytest.mark.unit +def test_rewire_is_a_noop_without_aliased_io(monkeypatch): + program, k_buffer, _ = _kv_program() + _patch_engine_metadata( + monkeypatch, + aliased_io={}, + input_names=["k_in", "tokens"], + output_names=["logits", "out_k"], + ) + + assert Z.rewire_aliased_mutations_to_buffers(program) == [] + assert "_torch_tensorrt_aliased_buffer" not in k_buffer.meta + + +@pytest.mark.unit +def test_rewire_rejects_an_engine_whose_every_output_is_aliased(monkeypatch): + """A delegate with no outputs at all is not a shape anything supports. + + Nothing downstream reports it: the runtime reads elision off a single + argument count, which a zero-output delegate satisfies, and the delegate + itself is a pure node a later graph-wide dead-code elimination can erase. + So the failure has to be raised here. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + engine = graph.placeholder("engine") + engine_call = graph.call_function( + torch.ops.tensorrt.execute_engine.default, ([k_buffer], engine) + ) + k_out = graph.call_function(operator.getitem, (engine_call, 0)) + graph.output((k_out,)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + signature = SimpleNamespace( + inputs_to_buffers={"b_k_0": "k_0"}, + input_specs=[], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(name=k_out.name), "k_0" + ) + ], + ) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=signature, + _graph_signature=signature, + ) + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in"], + output_names=["out_k"], + ) + + with pytest.raises(RuntimeError, match="no outputs at all"): + Z.rewire_aliased_mutations_to_buffers(program) + + +@pytest.mark.unit +@pytest.mark.parametrize("dead_chain_length", [1, 2]) +def test_rewire_rejects_an_engine_whose_only_other_output_is_dead( + monkeypatch, dead_chain_length +): + """A surviving-but-unread output does not keep the engine's delegate alive. + + The engine has two outputs: an aliased buffer this elides, and a second one + whose consumers end in nothing. Counting the second as an output would let + the check pass, and the dead-code elimination would erase the whole chain, + leaving exactly the zero-output delegate the check exists to refuse. The + two-link chain is the case a check reading only the engine's immediate users + misses: that first link does have a user, so it reads as live. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + engine = graph.placeholder("engine") + engine_call = graph.call_function( + torch.ops.tensorrt.execute_engine.default, ([k_buffer], engine) + ) + k_out = graph.call_function(operator.getitem, (engine_call, 0)) + dead = graph.call_function(operator.getitem, (engine_call, 1)) + for _ in range(dead_chain_length - 1): + dead = graph.call_function(torch.add, (dead, dead)) + graph.output((k_out,)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + signature = SimpleNamespace( + inputs_to_buffers={"b_k_0": "k_0"}, + input_specs=[], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(name=k_out.name), "k_0" + ) + ], + ) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=signature, + _graph_signature=signature, + ) + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in"], + output_names=["out_k", "out_dead"], + ) + + with pytest.raises(RuntimeError, match="no outputs at all"): + Z.rewire_aliased_mutations_to_buffers(program) + + +def _zero_copy_specs(*names): + """What ``TensorRTPartitioner`` stamps on the delegate whose engine elided. + + The value is the list of aliased output binding names, one per buffer the + engine writes in place, which is what the count checks read. + """ + return [ + CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + json.dumps(list(names or ("out_k",))).encode(), + ) + ] + + +def _staged_delegate_graph( + *, backend_id="TensorRTBackend", device=DeviceType.CUDA, compile_specs=None +): + """A lowered graph: delegate(lowered, _h2d_copy(k_buffer), _h2d_copy(tokens)).""" + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + tokens = graph.placeholder("tokens") + lowered = graph.get_attr("lowered_module_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + staged_tokens = graph.call_function(h2d, (tokens,)) + delegate = graph.call_function( + executorch_call_delegate, (lowered, staged_k, staged_tokens) + ) + graph.output((delegate,)) + + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id=backend_id, compile_specs=compile_specs + ) + graph_module = torch.fx.GraphModule(root, graph) + + for node, spec_device in ( + (k_buffer, DeviceType.CPU), + (tokens, DeviceType.CPU), + (staged_k, device), + (staged_tokens, device), + ): + node.meta["spec"] = SimpleNamespace(device=spec_device, device_index=3) + return graph_module, k_buffer, staged_k, delegate + + +@pytest.mark.unit +def test_unstage_feeds_the_buffer_straight_to_the_delegate(): + """The marked buffer replaces its staging copy and moves to the device. + + Moving the spec is not cosmetic: memory planning reads it, and a buffer left + in a host arena is somewhere the engine cannot write. + """ + graph_module, k_buffer, staged_k, delegate = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + assert Z._unstage_aliased_buffers(graph_module) == 1 + + assert delegate.args[1] is k_buffer + assert k_buffer.meta["spec"].device == DeviceType.CUDA + assert k_buffer.meta["spec"].device_index == 3 + # The other input is an ordinary one and keeps its staging copy. + assert delegate.args[2] is not None + assert delegate.args[2].target is torch.ops.et_copy._h2d_copy.default + # The orphaned staging is erased, but only it -- the other staging survives. + assert staged_k not in graph_module.graph.nodes + + +@pytest.mark.unit +def test_unstage_keeps_staging_for_an_unmarked_buffer(): + graph_module, k_buffer, staged_k, delegate = _staged_delegate_graph() + + assert Z._unstage_aliased_buffers(graph_module) == 0 + assert delegate.args[1] is staged_k + assert k_buffer.meta["spec"].device == DeviceType.CPU + + +def _direct_delegate_graph(*, compile_specs=None, device=DeviceType.CUDA): + """A lowered graph with no staging at all: delegate(lowered, k_buffer). + + The shape the un-staging pass itself leaves behind, and so the shape its own + second run is handed: the marked buffer is the delegate's argument outright + and there is no ``_h2d_copy`` to remove. + + ``device`` is what the buffer's own spec asks for, which is only half of + where memory planning ends up putting it. The other half is the + ``enable_non_cpu_memory_planning`` the program is finalized with; the graph + does not record it and the pass is told separately. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + lowered = graph.get_attr("lowered_module_0") + delegate = graph.call_function(executorch_call_delegate, (lowered, k_buffer)) + graph.output((k_buffer, delegate)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=compile_specs + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=device, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + return graph_module, k_buffer, delegate + + +@pytest.mark.unit +def test_unstage_accepts_a_buffer_that_never_had_a_staging_copy(): + """A marked buffer already handed straight to its delegate needs no work. + + What the pass has to leave behind is a marked buffer that is a delegate + argument planned in device memory; removing a staging copy is only the usual + route there. Keying success on having removed one instead rejects this + program, which is already in the shape zero-copy wants -- and it is the shape + the pass's own second run sees. + """ + graph_module, k_buffer, delegate = _direct_delegate_graph( + compile_specs=[CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")] + ) + + assert Z._unstage_aliased_buffers(graph_module) == 0 + + assert delegate.args[1] is k_buffer + assert k_buffer.meta["spec"].device == DeviceType.CUDA + + +@pytest.mark.unit +@pytest.mark.parametrize("route", ["direct", "staged"]) +def test_unstage_refuses_a_marked_buffer_whose_only_delegate_is_unstamped(route): + """The two halves of one post-condition have to answer one graph the same way. + + A TensorRT delegate carrying no zero-copy compile spec is one whose engine + elided no aliased output, so a marked buffer reaching it says nothing about + the engine that does write it in place -- that engine's write is still going + to a staging copy nothing reads back. :func:`check_zero_copy_kv` counts only + stamped delegates over the finalized program, on either route the buffer + took, so both routes are pinned here. + """ + stamped = [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")] + if route == "direct": + graph_module, _, _ = _direct_delegate_graph(compile_specs=None) + else: + graph_module, k_buffer, _, _ = _staged_delegate_graph(compile_specs=None) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="declaring zero-copy KV"): + Z._unstage_aliased_buffers(graph_module) + + # The same graph with the delegate stamped is accepted, so what the refusal + # reads is the missing stamp and not something else about these graphs. + if route == "direct": + graph_module, _, _ = _direct_delegate_graph(compile_specs=stamped) + assert Z._unstage_aliased_buffers(graph_module) == 0 + else: + graph_module, k_buffer, _, _ = _staged_delegate_graph(compile_specs=stamped) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + assert Z._unstage_aliased_buffers(graph_module) == 1 + + +@pytest.mark.unit +@pytest.mark.parametrize("spec", ["absent", "host"]) +def test_unstage_refuses_a_direct_buffer_that_is_not_on_the_device(spec): + """Reaching the delegate directly is not enough; the buffer has to be there. + + A marked buffer whose own spec asks for the host is planned in a host arena, + and the engine cannot write a host pointer in place. Accepting it on the + strength of the mark and the delegate edge alone writes a ``.pte`` whose + every ``execute()`` fails on the alias-target guard. A buffer with no spec at + all is the other half of the same refusal, and means the pass is running + somewhere the specs do not exist yet. + """ + graph_module, k_buffer, _ = _direct_delegate_graph( + device=DeviceType.CPU if spec == "host" else DeviceType.CUDA + ) + if spec == "absent": + del k_buffer.meta["spec"] + expected = ( + "carries no TensorSpec" if spec == "absent" else "its TensorSpec asks for" + ) + + with pytest.raises(RuntimeError, match=expected): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +@pytest.mark.parametrize("shape", ["direct", "staged"]) +def test_unstage_refuses_a_marked_buffer_under_host_only_memory_planning(shape): + """A CUDA spec does not mean CUDA memory when planning ignores spec devices. + + ``enable_non_cpu_memory_planning=False`` plans every tensor into the one host + arena whatever its ``TensorSpec`` says, so no marked buffer can be written in + place under it. Both graph shapes are covered because the pass reaches them + by different branches -- the buffer that already is a delegate argument, and + the one whose staging copy the pass removes -- and the refusal belongs to + neither, so it is checked once for the whole graph. + """ + if shape == "direct": + graph_module, k_buffer, _ = _direct_delegate_graph() + assert k_buffer.meta["spec"].device == DeviceType.CUDA, ( + "this shape only stands for the planning-mode hazard while the " + "buffer's spec is CUDA: what the refusal below exists for is a CUDA " + "spec that still lands in the host arena, and a host spec is refused " + "under any planning mode" + ) + else: + graph_module, k_buffer, _, _ = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): + Z._unstage_aliased_buffers(graph_module, device_memory_planning=False) + + +@pytest.mark.unit +def test_zero_copy_backend_config_carries_the_planning_mode_into_the_pass(): + """The refusal is only reachable if the config's flag actually gets there. + + ``zero_copy_backend_config`` is the only place that sees the + ``ExecutorchBackendConfig``, so building the pass without reading + ``enable_non_cpu_memory_planning`` off it leaves the check above unreachable + from any real finalization. + """ + from executorch.exir import ExecutorchBackendConfig + + config = Z.zero_copy_backend_config( + ExecutorchBackendConfig(enable_non_cpu_memory_planning=False) + ) + graph_module, _, _ = _direct_delegate_graph() + + with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): + config.to_out_var_pass(graph_module) + + +@pytest.mark.unit +def test_zero_copy_backend_config_reads_the_planning_mode_when_the_pass_runs(): + """The flag has to be read where it is in effect, not captured when it is set. + + ``ExecutorchBackendConfig`` is a plain mutable dataclass and the config + returned here is the one the caller hands to ``to_executorch``, so the value + memory planning uses is whatever the field holds by then. A pass that froze + the field when the config was built disagrees with the finalizer in both + directions, and the first of those writes a ``.pte`` whose caches are planned + in the host arena and whose every ``execute()`` fails. + """ + from executorch.exir import ExecutorchBackendConfig + + turned_off = Z.zero_copy_backend_config() + turned_off.enable_non_cpu_memory_planning = False + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): + turned_off.to_out_var_pass(graph_module) + + turned_on = Z.zero_copy_backend_config( + ExecutorchBackendConfig(enable_non_cpu_memory_planning=False) + ) + turned_on.enable_non_cpu_memory_planning = True + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + turned_on.to_out_var_pass(graph_module) + + +@pytest.mark.unit +def test_zero_copy_backend_config_does_not_refuse_a_planner_the_flag_never_reaches(): + """The flag is only a ground to refuse on where ExecuTorch delivers it. + + ``to_executorch`` does not pass ``enable_non_cpu_memory_planning`` to the + memory planner, it assigns it -- and only onto a planner that already has an + attribute of that name. A caller-supplied planner without one, which is what + the user guide tells people to bring for a cache shared between prefill and + decode, never sees the field, so where the caches land is that planner's own + business and ``False`` does not mean the host arena. Refusing on it there + turns the field into a trap that blocks a configuration that would have + worked; the ``.pte`` is still not left unchecked, since ``check_zero_copy_kv`` + reads the arena that planner actually chose. + """ + from executorch.exir import ExecutorchBackendConfig + + def a_planner_of_ones_own(graph_module): + raise AssertionError("memory planning does not run in this test") + + assert not hasattr(a_planner_of_ones_own, "enable_non_cpu_memory_planning") + config = Z.zero_copy_backend_config( + ExecutorchBackendConfig( + enable_non_cpu_memory_planning=False, + memory_planning_pass=a_planner_of_ones_own, + ) + ) + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + + config.to_out_var_pass(graph_module) + + # The same field, with a planner ExecuTorch does hand it to, is refused -- + # so what is carried above is the planner and not the flag. + stock = Z.zero_copy_backend_config( + ExecutorchBackendConfig(enable_non_cpu_memory_planning=False) + ) + assert hasattr(stock.memory_planning_pass, "enable_non_cpu_memory_planning") + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): + stock.to_out_var_pass(graph_module) + + +@pytest.mark.unit +def test_zero_copy_backend_config_rebuilt_over_a_derived_config_reads_it(): + """Deriving a config with ``dataclasses.replace`` is the case that needs it. + + The flag is a bool, copied by value; the pass is copied by reference. So a + config derived from the one this returns carries a pass still reading the + original, and the first half below is that known gap, pinned: turning + planning off on the derived config alone is not refused by the inherited + pass. It is not silent either -- ``check_zero_copy_kv`` reads the arena + memory planning chose and refuses the program it produces, for any method + holding a host tensor to give the shared arena away, and the manager + ``export`` returns runs that check on whatever its ``to_executorch`` + finalizes -- and the remedy the docstring gives is the second half: call the + function again on the derived config and the pass it builds is bound to + that one. + """ + import dataclasses + + derived = dataclasses.replace( + Z.zero_copy_backend_config(), enable_non_cpu_memory_planning=False + ) + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + derived.to_out_var_pass(graph_module) + + rebuilt = Z.zero_copy_backend_config(derived) + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): + rebuilt.to_out_var_pass(graph_module) + + +@pytest.mark.unit +def test_unstage_runs_a_second_time_without_raising(): + """Installing the pass twice is redundant rather than an error. + + ``save(zero_copy_kv=True)`` installs it, so a caller who also passes + ``zero_copy_backend_config()`` as ``backend_config`` gets two of them. The + second run finds the buffer already wired to the delegate and returns + without un-staging anything. + """ + graph_module, k_buffer, _, delegate = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + assert Z._unstage_aliased_buffers(graph_module) == 1 + + assert Z._unstage_aliased_buffers(graph_module) == 0 + assert delegate.args[1] is k_buffer + + +@pytest.mark.unit +def test_unstage_raises_for_a_marked_buffer_on_another_backends_delegate(): + """Only a TensorRT engine promises the in-place write, so a marked buffer + routed to another backend's delegate is never un-staged -- and because + export has already dropped its copy-back, that is a broken program, not a + silent no-op.""" + graph_module, k_buffer, staged_k, delegate = _staged_delegate_graph( + backend_id="CudaBackend" + ) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises( + RuntimeError, match="no TensorRT delegate declaring zero-copy KV" + ): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_raises_when_a_marked_buffer_is_never_unstaged(): + """A marked buffer that reaches no TensorRT delegate at all must raise, not + return 0. Its copy-back is already gone, so leaving it staged would silently + discard every update.""" + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + graph.output((k_buffer,)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises( + RuntimeError, match="no TensorRT delegate declaring zero-copy KV" + ): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_raises_when_a_zero_copy_delegate_unstaged_nothing(): + """A TensorRT delegate that declares zero-copy but had no buffer un-staged + (its mark did not survive to this pass) is unambiguously broken and must + raise, naming the delegate.""" + graph_module, k_buffer, staged_k, delegate = _staged_delegate_graph( + compile_specs=[CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")] + ) + # k_buffer deliberately left unmarked: nothing gets un-staged for the delegate. + + with pytest.raises(RuntimeError, match="declares zero-copy KV"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_raises_when_a_zero_copy_delegate_unstaged_only_some(): + """One surviving mark must not stand in for the ones that were lost. + + The delegate's spec names both aliased outputs it elided, so it has to take + two buffers written in place. Only one still carries the mark, and demanding + merely one would un-stage that one, raise nothing, and leave the second cache + wired through the staging copy whose contents are discarded -- with its + copy-back already gone, a silently frozen cache. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + v_buffer = graph.placeholder("b_v_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + staged_v = graph.call_function(h2d, (v_buffer,)) + lowered = graph.get_attr("lowered_module_0") + delegate = graph.call_function( + executorch_call_delegate, (lowered, staged_k, staged_v) + ) + graph.output((k_buffer, v_buffer, delegate)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs("out_k", "out_v") + ) + graph_module = torch.fx.GraphModule(root, graph) + for node, dev in ( + (k_buffer, DeviceType.CPU), + (v_buffer, DeviceType.CPU), + (staged_k, DeviceType.CUDA), + (staged_v, DeviceType.CUDA), + ): + node.meta["spec"] = SimpleNamespace(device=dev, device_index=0) + # v_buffer's mark did not survive lowering; k_buffer's did. + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="takes only 1 of the 2 buffers"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_refuses_a_buffer_another_backend_stages_on_the_same_gpu(): + """A second backend's staging copy on the *same* GPU is refused, not kept. + + The other backend's ``_h2d_copy`` outlives this pass and goes on reading the + buffer as its source, but the move has just put the buffer in device memory. + ``_h2d_copy_out`` requires a host source and fails ``InvalidArgument`` on a + device one, so leaving the copy in place produces a program that does not + run. Same GPU is what makes this shape distinct: the device and index both + match, so neither of ``_device_placement_is_safe``'s spec comparisons rejects it. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_trt = graph.call_function(h2d, (k_buffer,)) + staged_other = graph.call_function(h2d, (k_buffer,)) + lowered_trt = graph.get_attr("lowered_module_0") + lowered_other = graph.get_attr("lowered_module_1") + delegate_trt = graph.call_function( + executorch_call_delegate, (lowered_trt, staged_trt) + ) + delegate_other = graph.call_function( + executorch_call_delegate, (lowered_other, staged_other) + ) + graph.output((k_buffer, delegate_trt, delegate_other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="CudaBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + staged_trt.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_other.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): + Z._unstage_aliased_buffers(graph_module) + + # Refused before anything moved: the buffer is where the other backend's + # staging copy expects to read it. + assert k_buffer.meta["spec"].device == DeviceType.CPU + assert delegate_trt.args[1] is staged_trt + assert delegate_other.args[1] is staged_other + + +@pytest.mark.unit +def test_unstage_refuses_a_left_behind_staging_of_a_buffer_already_on_the_device(): + """The same refusal when the buffer's spec already names the engine's GPU. + + Nothing moves in this shape, so a check asked only about the move skips it + entirely -- and then the pass rewires the delegate anyway and hands the + engine a buffer whose other consumer, another backend's ``_h2d_copy``, reads + device memory as a host source and fails ``InvalidArgument`` on every call. + What the pass has to establish is where the buffer ends up, which is the same + place either way, so the question is asked either way. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_trt = graph.call_function(h2d, (k_buffer,)) + staged_other = graph.call_function(h2d, (k_buffer,)) + lowered_trt = graph.get_attr("lowered_module_0") + lowered_other = graph.get_attr("lowered_module_1") + delegate_trt = graph.call_function( + executorch_call_delegate, (lowered_trt, staged_trt) + ) + delegate_other = graph.call_function( + executorch_call_delegate, (lowered_other, staged_other) + ) + graph.output((k_buffer, delegate_trt, delegate_other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="CudaBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + # The buffer is already where the staging copy would have put it, which is + # the one thing that separates this from the test above. + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_trt.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_other.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): + Z._unstage_aliased_buffers(graph_module) + + assert delegate_trt.args[1] is staged_trt + + +@pytest.mark.unit +def test_unstage_moves_a_buffer_two_tensorrt_delegates_stage_from(): + """Two TensorRT delegates staging one cache is not a surviving consumer. + + Both staging copies go, so what is left reading the buffer is two delegates + taking it directly, which is the shape this pass exists to produce. The walk + reaches the second one after the first has already been rewired, so a + surviving-consumer test that did not allow a rewired delegate would refuse + the program the pass had just built. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_first = graph.call_function(h2d, (k_buffer,)) + staged_second = graph.call_function(h2d, (k_buffer,)) + lowered_first = graph.get_attr("lowered_module_0") + lowered_second = graph.get_attr("lowered_module_1") + first = graph.call_function(executorch_call_delegate, (lowered_first, staged_first)) + second = graph.call_function( + executorch_call_delegate, (lowered_second, staged_second) + ) + graph.output((k_buffer, first, second)) + root = torch.nn.Module() + for name in ("lowered_module_0", "lowered_module_1"): + setattr( + root, + name, + SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs() + ), + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + for staged in (staged_first, staged_second): + staged.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + assert Z._unstage_aliased_buffers(graph_module) == 2 + assert first.args[1] is k_buffer + assert second.args[1] is k_buffer + + +@pytest.mark.unit +def test_unstage_raises_when_the_staging_copy_is_not_on_cuda(): + """Following the staging to the CPU would put the buffer out of the engine's + reach, and the copy-back that would have saved it is already gone.""" + graph_module, k_buffer, _, _ = _staged_delegate_graph(device=DeviceType.CPU) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="not.*CUDA"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +@pytest.mark.parametrize("missing", ["staging-copy", "buffer"]) +def test_unstage_raises_when_either_side_of_the_move_has_no_spec(missing): + """The move reads a spec on both nodes, and the message names the bare one. + + One condition covers both, so a message that always blamed the staging copy + would send someone whose copy has a spec to look at the wrong node. + """ + graph_module, k_buffer, staged_k, _ = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + bare, kind = ( + (staged_k, "staging copy") + if missing == "staging-copy" + else (k_buffer, "buffer placeholder") + ) + del bare.meta["spec"] + + # Anchored on the node kind as well as the name: the buffer's name appears + # again later in the message, so a loose pattern would match either wording. + with pytest.raises( + RuntimeError, match=re.escape(f"no TensorSpec on the {kind} '{bare.name}'") + ): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_allows_a_buffer_that_is_also_its_mutation_output(): + """The zero-copy shape itself: the marked buffer is both the delegate's + staged input and its own BUFFER_MUTATION graph output. The output-node + reference carries no device of its own, so the device move must be allowed -- + the real lowered KV graph looks exactly like this.""" + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + lowered = graph.get_attr("lowered_module_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + delegate = graph.call_function(executorch_call_delegate, (lowered, staged_k)) + graph.output((k_buffer, delegate)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs() + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=3) + staged_k.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=3) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + assert Z._unstage_aliased_buffers(graph_module) == 1 + assert delegate.args[1] is k_buffer + assert k_buffer.meta["spec"].device == DeviceType.CUDA + + +@pytest.mark.unit +@pytest.mark.parametrize("route", ["staged", "direct"]) +def test_unstage_refuses_a_buffer_a_surviving_consumer_reads(route): + """The surviving-consumer refusal belongs to the placement, not to the move. + + A buffer already reaching its delegate directly is in the same position as + one this pass un-stages: it is planned in the engine's device memory, and a + consumer that reads it as a host source fails ``InvalidArgument`` on the + first call. Asking only on the branch that removes a staging copy would let + the identical graph through by the other door -- including the graph this + pass's own first run produces, which its second run reads directly. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + lowered = graph.get_attr("lowered_module_0") + h2d = torch.ops.et_copy._h2d_copy.default + delegate_arg = ( + graph.call_function(h2d, (k_buffer,)) if route == "staged" else k_buffer + ) + # Another backend's host copy of the same buffer, which nothing removes. + foreign = graph.call_function(h2d, (k_buffer,)) + lowered_other = graph.get_attr("lowered_module_1") + other = graph.call_function(executorch_call_delegate, (lowered_other, foreign)) + delegate = graph.call_function(executorch_call_delegate, (lowered, delegate_arg)) + graph.output((k_buffer, delegate, other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs() + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="CudaBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace( + device=DeviceType.CPU if route == "staged" else DeviceType.CUDA, + device_index=0, + ) + if route == "staged": + delegate_arg.meta["spec"] = SimpleNamespace( + device=DeviceType.CUDA, device_index=0 + ) + # On the GPU the engine is not on, so it is the placement and not merely the + # presence of a second reader that this refuses. + foreign.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=1) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="does not survive|leaves in place"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_refuses_to_move_a_shared_buffer(): + """A buffer read by a consumer other than its TensorRT delegate staging + cannot have its device moved -- that would silently retarget the other + consumer too, exactly what ExecuTorch's PropagateDevicePass rejects. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + lowered = graph.get_attr("lowered_module_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + other = graph.call_function(torch.add, (k_buffer, k_buffer)) + delegate = graph.call_function(executorch_call_delegate, (lowered, staged_k)) + graph.output((delegate, other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=3) + staged_k.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=3) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus(): + """One buffer staged to two TensorRT delegates on *different* GPUs cannot be + un-staged for either: a spec carries one device index, so whichever engine + lost the race would be handed an address on the other's GPU. ``spec.device`` + is only CUDA/CPU, so it is the device-index comparison in + ``_device_placement_is_safe`` that refuses the first delegate here -- both + stagings feed a TensorRT delegate, which is what separates this from the + two-backends shapes and leaves the index the only comparison that can catch + it. + + Which is why the refusal alone would not pin it. Were the index comparison + gone, delegate 0 would be un-staged and the direct-consumer branch would then + refuse delegate 1, and this test would still see a RuntimeError. What it + checks is that nothing moved: the buffer is still on the host and delegate 0 + still reads its staging copy. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_0 = graph.call_function(h2d, (k_buffer,)) + staged_1 = graph.call_function(h2d, (k_buffer,)) + lowered_0 = graph.get_attr("lowered_module_0") + lowered_1 = graph.get_attr("lowered_module_1") + delegate_0 = graph.call_function(executorch_call_delegate, (lowered_0, staged_0)) + delegate_1 = graph.call_function(executorch_call_delegate, (lowered_1, staged_1)) + graph.output((k_buffer, delegate_0, delegate_1)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + staged_0.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_1.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=1) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): + Z._unstage_aliased_buffers(graph_module) + + assert k_buffer.meta["spec"].device == DeviceType.CPU + assert delegate_0.args[1] is staged_0 + + +@pytest.mark.unit +def test_unstage_refuses_to_move_a_buffer_a_second_consumer_stages_to_the_host(): + """The device-*type* half of ``_device_placement_is_safe``'s spec comparison. + + Its sibling, the device index, is pinned by + ``test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus``. Here the second + consumer is another ``_h2d_copy`` that stays on the host, so the indices + agree and only the type comparison separates the two: deleting it accepts + this graph, leaving a copy that reads a buffer the move has put on the GPU. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_trt = graph.call_function(h2d, (k_buffer,)) + staged_host = graph.call_function(h2d, (k_buffer,)) + lowered_trt = graph.get_attr("lowered_module_0") + lowered_other = graph.get_attr("lowered_module_1") + delegate_trt = graph.call_function( + executorch_call_delegate, (lowered_trt, staged_trt) + ) + delegate_other = graph.call_function( + executorch_call_delegate, (lowered_other, staged_host) + ) + graph.output((k_buffer, delegate_trt, delegate_other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + staged_trt.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_host.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): + Z._unstage_aliased_buffers(graph_module) + + assert k_buffer.meta["spec"].device == DeviceType.CPU + assert delegate_trt.args[1] is staged_trt + + +@pytest.mark.unit +def test_unstage_refuses_a_buffer_another_backend_stages_to_a_different_gpu(): + """A marked buffer staged to a TensorRT delegate on cuda:0 and to a + *non*-TensorRT delegate on cuda:1 cannot be moved either. + + Un-staging skips the other backend's delegate, so its staging copy keeps + reading the buffer while staging it to cuda:1, and re-homing the buffer onto + the TensorRT engine's cuda:0 would move the source of that read to the wrong + GPU. Two of ``_device_placement_is_safe``'s comparisons refuse this shape -- the + device index, which runs first, and the surviving copy's non-TensorRT user + -- so this test does not discriminate between them. The index comparison on + its own is pinned by + ``test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus``, where every + staging does feed a TensorRT delegate. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_trt = graph.call_function(h2d, (k_buffer,)) + staged_other = graph.call_function(h2d, (k_buffer,)) + lowered_trt = graph.get_attr("lowered_module_0") + lowered_other = graph.get_attr("lowered_module_1") + delegate_trt = graph.call_function( + executorch_call_delegate, (lowered_trt, staged_trt) + ) + delegate_other = graph.call_function( + executorch_call_delegate, (lowered_other, staged_other) + ) + graph.output((k_buffer, delegate_trt, delegate_other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="CudaBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + staged_trt.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_other.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=1) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_refuses_to_rehome_a_buffer_already_on_another_gpu(): + """A buffer already resident on cuda:0 is not re-homed to a second TensorRT + delegate's cuda:1. + + Whether the move needs checking at all is decided by comparing the buffer's + device *and index* against the staging copy's. Comparing the device alone + would call this buffer already placed -- both ends are CUDA -- skip the + check, and overwrite the index with the second engine's, leaving the first + engine holding an address on the other GPU. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_0 = graph.call_function(h2d, (k_buffer,)) + staged_1 = graph.call_function(h2d, (k_buffer,)) + lowered_0 = graph.get_attr("lowered_module_0") + lowered_1 = graph.get_attr("lowered_module_1") + delegate_0 = graph.call_function(executorch_call_delegate, (lowered_0, staged_0)) + delegate_1 = graph.call_function(executorch_call_delegate, (lowered_1, staged_1)) + graph.output((k_buffer, delegate_0, delegate_1)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_0.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_1.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=1) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): + Z._unstage_aliased_buffers(graph_module) + assert k_buffer.meta["spec"].device_index == 0 + + +@pytest.mark.unit +def test_zero_copy_backend_config_keeps_the_callers_config(): + """It composes onto a config rather than replacing one: a caller finalizing + a zero-copy program still needs their own memory planning and passes.""" + from executorch.exir import ExecutorchBackendConfig + + inner = object() + base = ExecutorchBackendConfig(to_out_var_pass=inner, emit_stacktrace=True) + + config = torch_tensorrt.executorch.zero_copy_backend_config(base) + + assert config.emit_stacktrace is True + assert config.memory_planning_pass is base.memory_planning_pass + assert config.to_out_var_pass is not inner + # The caller's to_out_var_pass is not dropped, it is run after the un-staging. + graph_module, k_buffer, _, delegate = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + seen = [] + base = ExecutorchBackendConfig(to_out_var_pass=lambda gm: seen.append(gm)) + torch_tensorrt.executorch.zero_copy_backend_config(base).to_out_var_pass.call( + graph_module + ) + assert seen == [graph_module] + assert delegate.args[1] is k_buffer + + +def _finalized_program(forward=None, **methods): + """The shape ``check_zero_copy_kv`` reads: to_executorch()'s return value. + + One positional graph module makes a single-method ``forward`` program; the + keywords name a method each. ``exported_program`` defaults to ``forward`` + and raises ``KeyError`` on a method the program does not have, like + ``ExecutorchProgramManager``'s -- which is what a program with no ``forward`` + does to a caller that never asked for one. + """ + if forward is not None: + methods = {"forward": forward, **methods} + return SimpleNamespace( + methods=set(methods), + exported_program=lambda method_name="forward": SimpleNamespace( + graph_module=methods[method_name] + ), + ) + + +CUDA_ARENA = 2 + + +HOST_ARENA = 1 + + +def _planned( + graph_module, + *, + arena=CUDA_ARENA, + on_device=True, + host_tensor_arena=None, + device_type=DeviceType.CUDA, + device_index=None, +): + """Add what memory planning leaves behind, which the checker reads. + + The graphs above are built for the passes that run before planning, so they + carry no ``mem_id`` and the module records no arena devices. Planning assigns + both. ``on_device=False`` drops the arena-device record, which is what + ``enable_non_cpu_memory_planning=False`` leaves -- and also what a + caller-supplied planner that does not go through ``apply_algo`` leaves, that + being the only thing that writes it. What separates those two is where the + program's *host* tensors ended up: host-only planning puts every tensor in + one bucket, so they share the buffer's arena, while a device-aware planner + keeps them apart. + ``host_tensor_arena`` is where the CPU-spec tensors go, and it is also put on + the output node's spec, which is where a real finalized program carries one. + Leaving it unset gives the two shapes above: a separate arena when the + program records devices, and the buffer's own when it does not. + ``device_index`` is the GPU the arena is recorded for; unset, it is the one + the graph's own CUDA specs ask for, which is what a planner that honoured + them would record. Passing a different one is the multi-GPU mistake. + """ + from executorch.exir.schema import NonConstBufferDevice + + if host_tensor_arena is None: + host_tensor_arena = HOST_ARENA if on_device else arena + asked_for = [] + for node in graph_module.graph.nodes: + spec = node.meta.get("spec") + if node.op == "placeholder" and spec is not None: + spec.mem_id = arena if spec.device == DeviceType.CUDA else host_tensor_arena + if spec.device == DeviceType.CUDA: + asked_for.append(spec.device_index) + if node.op == "output": + node.meta["spec"] = [ + SimpleNamespace( + device=DeviceType.CPU, device_index=0, mem_id=host_tensor_arena + ) + ] + if on_device: + graph_module.meta["non_const_buffer_device"] = [ + NonConstBufferDevice( + buffer_idx=arena, + device_type=device_type, + device_index=( + next(iter(asked_for), 0) if device_index is None else device_index + ), + ) + ] + return graph_module + + +def _unstaged_graph(): + """A graph whose marked buffer already reaches its delegate directly.""" + graph_module, k_buffer, _, _ = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + Z._unstage_aliased_buffers(graph_module) + return _planned(graph_module) + + +@pytest.mark.unit +def test_check_zero_copy_kv_accepts_an_unstaged_buffer(): + Z.check_zero_copy_kv(_finalized_program(_unstaged_graph())) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_still_staged_buffer(): + """The shape a program finalized without zero_copy_backend_config has: the + buffer is marked, so export dropped its copy-back, but it still reaches the + delegate through a staging copy the engine's write is thrown away with.""" + graph_module, k_buffer, _, _ = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="do not reach the TensorRT delegate"): + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module))) + + +@pytest.mark.unit +def test_check_zero_copy_kv_accepts_a_buffer_that_never_had_a_staging_copy(): + """The checker and the un-staging pass read the same post-condition. + + A program the pass has already un-staged hands the buffer straight to the + delegate with no staging copy left. The pass accepts that shape + (``test_unstage_accepts_a_buffer_that_never_had_a_staging_copy``) and so must + this: the two disagreeing is what would let one path refuse a program the + other calls correct. + """ + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module))) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_direct_buffer_planned_on_the_host(): + """Reaching the delegate directly is only half of it; placement is the rest. + + Finalizing with ``enable_non_cpu_memory_planning=False`` and no zero-copy + config gives exactly this: no staging copy is inserted, so the wiring is what + zero-copy wants, and every tensor still lands in the one host arena. Nothing + else refuses it -- the un-staging pass never ran -- and the engine is handed a + host pointer it cannot write, so the ``.pte`` fails its first ``execute()``. + The graph is the accepted one above with only the arena changed, so the + refusal can come from nothing else. The program records no arena devices -- + host-only planning writes none -- so what identifies the arena as the host's + is that the program's host tensors are in it too. + """ + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + + with pytest.raises(RuntimeError, match="also holds the program's host tensors"): + Z.check_zero_copy_kv( + _finalized_program( + _planned(graph_module, arena=HOST_ARENA, on_device=False) + ) + ) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_an_arena_no_planner_recorded(): + """An unrecorded arena is what the runtime reads as the host, whoever planned it. + + ``apply_algo`` is the only thing in ExecuTorch that writes + ``non_const_buffer_device``, and ``to_executorch`` takes any callable as + ``memory_planning_pass``, so a caller-supplied planner can put the cache in + device memory and leave the ``.pte`` saying nothing. It is refused all the + same: ``MethodMeta::memory_planned_buffer_device`` answers ``CPU`` for an + arena with no entry, so the runner backs it with host memory and the engine + fails the alias-target guard on the first call. + + This is the host-arena test's graph with the host tensors moved out of the + buffer's arena, so the two refusals are told apart by which of them fires: + the arena here is not one the program's host tensors are in, and the message + says the record is absent rather than accusing the planner of putting the + cache among the host tensors. + """ + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + + with pytest.raises(RuntimeError, match="records no CUDA arena at all"): + Z.check_zero_copy_kv( + _finalized_program( + _planned( + graph_module, + arena=CUDA_ARENA, + on_device=False, + host_tensor_arena=HOST_ARENA, + ) + ) + ) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_an_unplanned_buffer_without_blaming_the_planner(): + """A buffer with no ``mem_id`` was not planned, which is not the host arena. + + A planner that excludes mutable buffers, or one built with graph-input + allocation off, leaves the cache with no ``mem_id`` at all. Nothing then + says where it lives, so it is refused -- but reporting it among the host + tensors would tell the caller their planner made a placement it never made, + and the two are separated here by the message. The control is the same + program with the cache in the recorded CUDA arena, which is accepted, so the + refusal is the missing ``mem_id`` and nothing else about this graph. + """ + graph_module, k_buffer, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + program = _finalized_program(_planned(graph_module, host_tensor_arena=HOST_ARENA)) + Z.check_zero_copy_kv(program) + + del k_buffer.meta["spec"].mem_id + with pytest.raises(RuntimeError, match="carries no mem_id") as raised: + Z.check_zero_copy_kv(program) + assert "host tensors" not in str(raised.value) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_an_arena_the_program_records_as_non_cuda(): + """When the record *is* present it is read, and only its CUDA entries count. + + Nothing in ExecuTorch emits a CPU entry today -- the builder filters them -- + so this pins the filter against a program that carries one, hand-built or + from a future planner, rather than against the stock one. + """ + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + + with pytest.raises(RuntimeError, match="does not record as CUDA"): + Z.check_zero_copy_kv( + _finalized_program( + _planned( + graph_module, + device_type=DeviceType.CPU, + host_tensor_arena=HOST_ARENA, + ) + ) + ) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_an_arena_recorded_for_another_gpu(): + """A CUDA arena is not enough; it has to be the GPU the cache asks for. + + The runtime allocates the cache out of the arena the program records, so an + arena recorded for another device hands the engine an address on a GPU it is + not running on -- which fails exactly as a host pointer does, and which the + device *type* on its own cannot tell from a correct program. + """ + graph_module, k_buffer, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + assert k_buffer.meta["spec"].device_index == 0 + + with pytest.raises(RuntimeError, match="asks for cuda:0 and was planned"): + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module, device_index=1))) + + +@pytest.mark.unit +def test_check_zero_copy_kv_counts_one_buffer_in_two_slots_once(): + """Two argument slots holding one buffer are one cache, not two. + + The delegate's spec names two elided aliased outputs, so it owes two caches + written in place, and it has one. Counting slots satisfies that count with + the same buffer twice -- which is the arrangement the count exists to catch, + a delegate whose second mark did not survive. The un-staging pass counts the + same way, so the two are pinned together: they must not start disagreeing + about one graph. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + lowered = graph.get_attr("lowered_module_0") + delegate = graph.call_function( + executorch_call_delegate, (lowered, k_buffer, k_buffer) + ) + graph.output((k_buffer, delegate)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs("out_k", "out_v") + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="takes only 1 of the 2 buffers"): + Z._unstage_aliased_buffers(graph_module) + with pytest.raises(RuntimeError, match="takes 1 marked buffer"): + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module))) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_buffer_staged_at_its_own_delegate(): + """An unstamped TensorRT delegate must not stand in for the one that elided. + + The engine whose aliased output was elided -- the one carrying the zero-copy + spec -- still reads a staging copy, so its write is discarded and the cache + never updates. Another TensorRT engine happens to read the same buffer + directly, which says nothing about that write. Taking the union over every + TensorRT delegate accepts this program; narrowing to the stamped ones is what + refuses it here, and ``..._a_second_zero_copy_delegate_standing_in`` is the + case where both are stamped and that narrowing is not enough. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + lowered_kv = graph.get_attr("lowered_module_0") + lowered_plain = graph.get_attr("lowered_module_1") + delegate_kv = graph.call_function(executorch_call_delegate, (lowered_kv, staged_k)) + delegate_plain = graph.call_function( + executorch_call_delegate, (lowered_plain, k_buffer) + ) + graph.output((k_buffer, delegate_kv, delegate_plain)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs() + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="do not reach the TensorRT delegate"): + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module))) + + +def _two_zero_copy_delegate_graph(*, crossed, first_specs=None): + """Two stamped TensorRT delegates in one method, one elided output each. + + ``crossed`` gives the second delegate both caches and leaves the first + reading a staging copy of the one it elided, which is the shape a check that + reads the stamped delegates as one set cannot see: every marked buffer does + reach a stamped delegate, just not the one whose write was removed. + ``first_specs`` overrides what the first delegate's spec claims it elided. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + v_buffer = graph.placeholder("b_v_0") + lowered_k = graph.get_attr("lowered_module_0") + lowered_v = graph.get_attr("lowered_module_1") + if crossed: + staged_k = graph.call_function(torch.ops.et_copy._h2d_copy.default, (k_buffer,)) + k_args, v_args = (lowered_k, staged_k), (lowered_v, k_buffer, v_buffer) + else: + k_args, v_args = (lowered_k, k_buffer), (lowered_v, v_buffer) + delegate_k = graph.call_function(executorch_call_delegate, k_args) + delegate_v = graph.call_function(executorch_call_delegate, v_args) + graph.output((k_buffer, v_buffer, delegate_k, delegate_v)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", + compile_specs=( + _zero_copy_specs("out_k") if first_specs is None else first_specs + ), + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs("out_v") + ) + graph_module = torch.fx.GraphModule(root, graph) + for buffer in (k_buffer, v_buffer): + buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + buffer.meta["_torch_tensorrt_aliased_buffer"] = True + return _planned(graph_module) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_second_zero_copy_delegate_standing_in(): + """A stamped delegate is counted against its own spec, not pooled with the rest. + + Both delegates here declare zero-copy KV, so narrowing by the spec does not + separate them, and both marked buffers reach one of the two directly -- so + the still-staged refusal has nothing to say. What is wrong is which delegate + took which: ``lowered_module_0`` elided ``out_k`` and reads a staging copy of + the cache that write was removed for, so that cache never updates. The + matched half is the same graph with each delegate holding the cache it + elided, and it is accepted, so the refusal can come only from the crossing. + """ + Z.check_zero_copy_kv( + _finalized_program(_two_zero_copy_delegate_graph(crossed=False)) + ) + + with pytest.raises(RuntimeError, match="takes 0 marked buffer"): + Z.check_zero_copy_kv( + _finalized_program(_two_zero_copy_delegate_graph(crossed=True)) + ) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_stamped_delegate_with_no_names_to_count(): + """A spec listing no name still says its engine elided an aliased output. + + Only a delegate whose own engine had one is stamped, so the count it cannot + read off the spec falls back to at least one -- the same fallback + ``_unstage_aliased_buffers`` makes. Reading "no names" as "no buffers owed" + would accept the crossing above whenever the partitioner's list is empty or + unreadable. + """ + graph_module = _two_zero_copy_delegate_graph( + crossed=True, + first_specs=[CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")], + ) + + with pytest.raises(RuntimeError, match="must take at least one marked buffer"): + Z.check_zero_copy_kv(_finalized_program(graph_module)) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_buffer_only_another_backend_takes(): + """Another backend's delegate taking the buffer directly is not zero-copy. + + The mark is on this buffer because a TensorRT engine writes it in place, and + that engine here is still reading a staging copy whose contents are thrown + away. Both delegates carry the zero-copy spec, so the backend is the only + thing separating them: counting any backend's delegate passes this program. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + lowered_trt = graph.get_attr("lowered_module_0") + lowered_other = graph.get_attr("lowered_module_1") + delegate_trt = graph.call_function( + executorch_call_delegate, (lowered_trt, staged_k) + ) + delegate_other = graph.call_function( + executorch_call_delegate, (lowered_other, k_buffer) + ) + graph.output((k_buffer, delegate_trt, delegate_other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs() + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="CudaBackend", compile_specs=_zero_copy_specs() + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="do not reach the TensorRT delegate"): + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module))) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_program_with_nothing_marked(): + """zero_copy_kv=True on a model with no engine-aliased buffer only warns, so + the .pte that comes out is an ordinary staged one. Refuse it rather than let + a caller who asked for zero-copy ship a program that never got it.""" + graph_module, _, _, _ = _staged_delegate_graph() + + with pytest.raises(RuntimeError, match="marked for in-place update"): + Z.check_zero_copy_kv(_finalized_program(graph_module)) + + +@pytest.mark.unit +def test_check_zero_copy_kv_accepts_a_program_with_no_forward_method(): + """The shape the user guide's zero-copy example exports: prefill and decode, + no ``forward``. Reading the default method would raise KeyError naming a + method the caller never asked for. A method that rewired nothing of its own + is not an error either, so only ``decode`` here carries a marked buffer.""" + unmarked, _, _, _ = _staged_delegate_graph() + + Z.check_zero_copy_kv(_finalized_program(prefill=unmarked, decode=_unstaged_graph())) + + +@pytest.mark.unit +def test_check_zero_copy_kv_catches_a_method_other_than_forward(): + """The silent case: ``forward`` got zero-copy and ``decode`` degenerated to + staged. Stopping at ``forward`` would write a .pte whose decode cache never + updates, so the failure has to name the method that lost it.""" + staged, k_buffer, _, _ = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="in method 'decode'"): + Z.check_zero_copy_kv(_finalized_program(_unstaged_graph(), decode=staged)) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_multi_method_program_with_nothing_marked(): + """Nothing marked anywhere is about the program, not about one method: a + method with no aliased buffer mutation is an error only when no other method + has one, and the failure lists every method it looked in.""" + first, _, _, _ = _staged_delegate_graph() + second, _, _, _ = _staged_delegate_graph() + + with pytest.raises(RuntimeError, match=r"\(decode, prefill\)"): + Z.check_zero_copy_kv(_finalized_program(prefill=first, decode=second)) + + +def _stamped_delegate_with_no_mark(): + """A stamped zero-copy delegate whose buffer lost its mark. + + Nothing but the compile spec is left recording that this engine's aliased + output was elided, the mark being the only other witness to it. + """ + graph_module, k_buffer, _ = _direct_delegate_graph( + compile_specs=_zero_copy_specs("out_k") + ) + del k_buffer.meta["_torch_tensorrt_aliased_buffer"] + return _planned(graph_module) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_stamped_delegate_whose_method_lost_its_mark(): + """A method is read for both records, not passed over on the marks alone. + + ``_unstage_aliased_buffers`` refuses this graph -- the spec says an aliased + output was elided and no marked buffer arrives to hold it -- so a check that + passed it would call correct a method the pass calls broken. Both halves + matter: the two-method program pins that a marked method does not vouch for + an unmarked one, and the single-method program pins the message, which would + otherwise be the program-wide "probably not exported with zero_copy_kv=True" + on a program whose own compile spec says it was. + """ + with pytest.raises(RuntimeError, match=r"takes 0 marked buffer"): + Z.check_zero_copy_kv( + _finalized_program( + prefill=_unstaged_graph(), decode=_stamped_delegate_with_no_mark() + ) + ) + + with pytest.raises(RuntimeError, match=r"takes 0 marked buffer"): + Z.check_zero_copy_kv( + _finalized_program(decode=_stamped_delegate_with_no_mark()) + ) + + +@pytest.mark.unit +def test_export_manager_checks_the_program_its_to_executorch_returns(): + """The two-call API cannot hand back a silently non-updating program. + + Export removed the copy-back before it returned, so every way of finalizing + that does not also un-stage the caches writes a ``.pte`` whose caches never + update -- and ``to_executorch()`` with ExecuTorch's defaults is one of those + and is the call the documentation spells out. The manager reads its own + finalized program back through the same check ``save`` runs, so the mistake + is an error rather than a file. A program that is right is handed back + untouched, which is the second half here: the check must not be a toll on + the working path. + """ + zero_copy_config = torch_tensorrt.executorch.zero_copy_backend_config() + graph_module, k_buffer, _, _ = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + staged = _finalized_program(_planned(graph_module)) + edge = SimpleNamespace(to_executorch=lambda config=None: staged) + Z._check_zero_copy_kv_when_finalized(edge) + + with pytest.raises(RuntimeError, match="do not reach the TensorRT delegate"): + edge.to_executorch(config=zero_copy_config) + + unstaged = _finalized_program(_unstaged_graph()) + good = SimpleNamespace(to_executorch=lambda config=None: unstaged) + Z._check_zero_copy_kv_when_finalized(good) + assert good.to_executorch(config=zero_copy_config) is unstaged + + +@pytest.mark.unit +@pytest.mark.parametrize("how", ["omitted", "positional", "keyword"]) +def test_export_manager_refuses_a_config_without_the_pass_before_finalizing(how): + """The refusal has to come before the manager is spent, not after. + + ``to_executorch`` rewrites the manager's own edge programs in place, so a + manager that has finalized once cannot finalize again -- the second call + dies inside ``insert_write_back_for_buffers_pass``. A refusal raised on the + finalized program could therefore only be acted on by exporting again and + rebuilding every engine, while the config alone already settles it. The + config reaches ``to_executorch`` either way round, so both are read. + """ + from executorch.exir import ExecutorchBackendConfig + + program = _finalized_program(_unstaged_graph()) + finalized = [] + edge = SimpleNamespace( + to_executorch=lambda config=None: (finalized.append(config), program)[1] + ) + Z._check_zero_copy_kv_when_finalized(edge) + + with pytest.raises(RuntimeError, match="does not carry the pass"): + if how == "omitted": + edge.to_executorch() + elif how == "positional": + edge.to_executorch(ExecutorchBackendConfig()) + else: + edge.to_executorch(config=ExecutorchBackendConfig()) + + # Nothing was finalized, so the caller can follow the remedy on this manager. + assert finalized == [] + zero_copy_config = torch_tensorrt.executorch.zero_copy_backend_config() + assert edge.to_executorch(zero_copy_config) is program + + +@pytest.mark.unit +def test_zero_copy_backend_config_defaults_to_executorch_defaults(): + """Called with no config it starts from ExecuTorch's defaults, and the one + field it replaces is to_out_var_pass, wrapped in the un-staging pass.""" + from executorch.exir import ExecutorchBackendConfig + + config = torch_tensorrt.executorch.zero_copy_backend_config() + + defaults = ExecutorchBackendConfig() + # The un-staging pass specifically, not merely "some object that is not the + # default" -- which is all any wrapper would have to be. Compared by type + # rather than by name, which is also what the finalization hook does when it + # asks a config whether it is one of these. + assert isinstance(config.to_out_var_pass, Z._UnstageThenToOutVar) + assert type(config.memory_planning_pass) is type(defaults.memory_planning_pass) + assert type(config.sym_shape_eval_pass) is type(defaults.sym_shape_eval_pass) + assert config.emit_stacktrace == defaults.emit_stacktrace + + +@pytest.mark.unit +@pytest.mark.parametrize( + "skip", + [True, {"decode": True}, {"prefill": False, "decode": True}, {"decode": False}], + ids=["bool", "dict-one-true", "dict-one-of-two-true", "dict-all-false"], +) +def test_zero_copy_backend_config_refuses_skip_h2d_for_method_inputs(skip): + """The one option that cannot be carried through, refused wherever it is on. + + ``skip_h2d_for_method_inputs`` is ExecuTorch's own un-staging of method + inputs and it demands each placeholder it un-stages have exactly one user. A + rewired cache has two -- the delegate, and the graph output it is its own + mutation result for -- so ``PropagateDevicePass`` raises on every zero-copy + graph. Preserving the option hands back a config that cannot finalize at + all, which is a failure a long way from the line that caused it. + + ``dict-all-false`` is the case that makes the refusal key on truthiness + rather than on ``True``: that pass is handed the field whole and only tests + it for truth, so it reads a dict of ``False`` as on for every method, and + finalizing such a config raises with ``placeholder 'b_k_cache' to have + exactly one user``. ``False`` and ``{}`` are read as off there and are + carried, in ``..._carries_skip_h2d_left_falsy``. + + The ids name what each value *is*, not who it applies to: that pass never + resolves this field per method, so every dict here is on for every method + whatever key it carries. + """ + from executorch.exir import ExecutorchBackendConfig + from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig + + base = ExecutorchBackendConfig( + propagate_device_config=PropagateDeviceConfig(skip_h2d_for_method_inputs=skip) + ) + + with pytest.raises(ValueError, match="skip_h2d_for_method_inputs"): + Z.zero_copy_backend_config(base) + + +@pytest.mark.unit +def test_zero_copy_backend_config_refuses_skip_h2d_in_a_per_method_config(): + """``propagate_device_config`` is itself one config or a dict of them. + + ExecuTorch resolves a dict of ``PropagateDeviceConfig`` by method name + (``_program.py``, ``edge_to_executorch_passes``) and hands the chosen one's + ``skip_h2d_for_method_inputs`` to the pass, so the option reaches + ``PropagateDevicePass`` from here exactly as it does from the single-config + form. Reading only the single form leaves that route open: measured, such a + config finalizes into the same ``exactly one user`` failure. + """ + from executorch.exir import ExecutorchBackendConfig + from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig + + base = ExecutorchBackendConfig( + propagate_device_config={ + "prefill": PropagateDeviceConfig(), + "decode": PropagateDeviceConfig(skip_h2d_for_method_inputs=True), + } + ) + + with pytest.raises(ValueError, match="skip_h2d_for_method_inputs for decode"): + Z.zero_copy_backend_config(base) + + +@pytest.mark.unit +@pytest.mark.parametrize("skip", [False, {}], ids=["bool", "empty-dict"]) +def test_zero_copy_backend_config_carries_skip_h2d_left_falsy(skip): + """A field left falsy *to PropagateDevicePass* is carried through. + + That pass only tests the field for truth, so what it reads as off is exactly + ``False`` and the empty dict -- and those two are what this carries. A + non-empty dict of ``False`` is not one of them; it is refused, in + ``..._refuses_skip_h2d_for_method_inputs[dict-all-false]``. + """ + from executorch.exir import ExecutorchBackendConfig + from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig + + base = ExecutorchBackendConfig( + propagate_device_config=PropagateDeviceConfig(skip_h2d_for_method_inputs=skip) + ) + + config = Z.zero_copy_backend_config(base) + + assert config.propagate_device_config.skip_h2d_for_method_inputs == skip + + +@pytest.mark.unit +def test_unstage_pass_runs_the_inner_pass_after_unstaging(): + """A caller's own to_out_var_pass has to survive being composed with.""" + graph_module, k_buffer, _, delegate = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + seen = [] + + def inner(gm): + # The un-staging is already done by the time the inner pass sees the graph. + seen.append(delegate.args[1] is k_buffer) + return "inner-result" + + result = Z.unstage_aliased_buffers_pass(inner).call(graph_module) + + assert seen == [True] + assert result == "inner-result" + + +# -------------------------------------------------------------------------- +# Multi-delegate: a method that lowers to two TensorRT engines -- one with an +# aliased+elided KV buffer, one plain-compute engine with none. The zero-copy +# CompileSpec is appended once, to a single TensorRTPartitioner, and the +# partitioner must stamp it onto ONLY the delegate whose own engine had an +# aliased output elided. Stamped partition-wide instead, the plain delegate +# declares zero-copy while un-staging nothing, and the un-staging cross-check +# then rejects an otherwise-correct program. +# -------------------------------------------------------------------------- + + +def _no_op_engine_node( + graph, input_nodes, *, aliased_io, input_names, output_names, engine="" +): + """A no_op_placeholder_for_execute_engine node with inlined engine info. + + Mirrors what replace_execute_engine() produces before partitioning: args are + ``(input_list, *engine_info)`` with the binding names and aliased_io in their + serialized wire form, so the partitioner's real per-engine resolution + (_resolve_engine_info / _aliased_inputs_by_output_index) runs unmocked. + + ``engine`` is the serialized-plan slot, which the partitioner never reads. + Give it bytes to hand the same node to ``TensorRTBackend.preprocess``, which + does. + """ + from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ( + ALIASED_IO_IDX, + DEVICE_IDX, + ENGINE_IDX, + INPUT_BINDING_NAMES_IDX, + OUTPUT_BINDING_NAMES_IDX, + SERIALIZATION_LEN, + SERIALIZED_ENGINE_BINDING_DELIM, + ) + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import serialize_aliased_io + + info = [""] * SERIALIZATION_LEN + info[ENGINE_IDX] = engine + info[DEVICE_IDX] = "0" + info[INPUT_BINDING_NAMES_IDX] = SERIALIZED_ENGINE_BINDING_DELIM.join(input_names) + info[OUTPUT_BINDING_NAMES_IDX] = SERIALIZED_ENGINE_BINDING_DELIM.join(output_names) + info[ALIASED_IO_IDX] = serialize_aliased_io(aliased_io) + return graph.call_function( + torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default, + (list(input_nodes), *info), + ) + + +def _two_engine_program(): + """engine_a(k_buffer, tokens) -> (logits, out_k[aliased]); engine_b(x) -> (y). + + ``k_buffer`` carries ``_torch_tensorrt_aliased_buffer`` (rewiring already + ran); engine_a aliases its ``out_k`` output onto it, engine_b aliases nothing. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + tokens = graph.placeholder("tokens") + x = graph.placeholder("x") + engine_a = _no_op_engine_node( + graph, + [k_buffer, tokens], + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in", "tokens"], + output_names=["logits", "out_k"], + ) + engine_b = _no_op_engine_node( + graph, + [x], + aliased_io={}, + input_names=["x_in"], + output_names=["y"], + ) + graph.output((engine_a, engine_b)) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=SimpleNamespace(buffers_to_mutate={}, inputs_to_buffers={}), + constants={}, + ) + return program, engine_a, engine_b + + +def _partition_two_engines(program, engine_a, engine_b, monkeypatch): + """Run the real TensorRTPartitioner, one partition per engine node.""" + from torch_tensorrt.executorch.backend import _serialize_elided_output_names + from torch_tensorrt.executorch.partitioner import TensorRTPartitioner + + class _FakeCap: + def __init__(self, graph_module, *args, **kwargs): + self._engines = [engine_a, engine_b] + + def propose_partitions(self): + return [ + SimpleNamespace(id=i, nodes=[node]) + for i, node in enumerate(self._engines) + ] + + monkeypatch.setattr( + "torch_tensorrt.executorch.partitioner.CapabilityBasedPartitioner", _FakeCap + ) + monkeypatch.setattr( + "torch_tensorrt.executorch.partitioner.tag_constant_data", + lambda exported_program: None, + ) + # Appended once, method-wide -- exactly how export() builds the partitioner. + partitioner = TensorRTPartitioner( + compile_specs=[ + CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(["out_k"]), + ) + ] + ) + return partitioner.partition(program) + + +def _zero_copy_names(compile_specs): + from torch_tensorrt.executorch.backend import _elided_output_names + + return _elided_output_names(compile_specs) + + +@pytest.mark.unit +def test_partition_stamps_zero_copy_only_on_the_kv_delegate(monkeypatch): + """The KV delegate carries the zero-copy spec naming its own elided binding, + and the plain-compute delegate carries no zero-copy spec at all. + + The method-wide spec the partitioner is constructed with must not reach every + partition: the names it holds are the method's, and only this engine's own + aliased_io says which of them are its. + """ + program, engine_a, engine_b = _two_engine_program() + result = _partition_two_engines(program, engine_a, engine_b, monkeypatch) + + kv_specs = result.partition_tags["tensorrt_0"].compile_specs + plain_specs = result.partition_tags["tensorrt_1"].compile_specs + assert _zero_copy_names(kv_specs) == {"out_k"} + assert _zero_copy_names(plain_specs) is None + + +@pytest.mark.unit +def test_multi_delegate_zero_copy_lowers_without_false_raise(monkeypatch): + """A correct two-delegate zero-copy program survives the whole pipeline: run + the real partitioner, build the lowered two-delegate graph from the specs it + produced, and un-stage. + + The KV buffer is un-staged and the plain delegate is left alone. A plain + delegate stamped zero-copy would instead make _unstage_aliased_buffers raise + "declares zero-copy KV ... but takes no buffer marked for in-place update" + over a program that is correct. + """ + program, engine_a, engine_b = _two_engine_program() + result = _partition_two_engines(program, engine_a, engine_b, monkeypatch) + kv_specs = result.partition_tags["tensorrt_0"].compile_specs + plain_specs = result.partition_tags["tensorrt_1"].compile_specs + + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + x = graph.placeholder("x") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + staged_x = graph.call_function(h2d, (x,)) + kv_lowered = graph.get_attr("lowered_module_0") + plain_lowered = graph.get_attr("lowered_module_1") + kv_delegate = graph.call_function(executorch_call_delegate, (kv_lowered, staged_k)) + plain_delegate = graph.call_function( + executorch_call_delegate, (plain_lowered, staged_x) + ) + graph.output((k_buffer, kv_delegate, plain_delegate)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=kv_specs + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=plain_specs + ) + graph_module = torch.fx.GraphModule(root, graph) + for node, dev in ( + (k_buffer, DeviceType.CPU), + (x, DeviceType.CPU), + (staged_k, DeviceType.CUDA), + (staged_x, DeviceType.CUDA), + ): + node.meta["spec"] = SimpleNamespace(device=dev, device_index=3) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + assert Z._unstage_aliased_buffers(graph_module) == 1 + assert kv_delegate.args[1] is k_buffer + # The plain delegate keeps its staging and is never demanded to un-stage. + assert plain_delegate.args[1] is staged_x + + +@pytest.mark.unit +def test_unstage_raises_when_the_plain_delegate_is_wrongly_stamped(): + """The other side of the per-partition stamping, in isolation: a delegate + that carries the zero-copy spec and un-stages nothing must raise, whatever + put the spec there. Narrowing which delegates get stamped must not weaken + this -- it is the lost-update guard for the KV delegate too. + """ + zero_copy_spec = [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b'["out_k"]')] + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + x = graph.placeholder("x") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + staged_x = graph.call_function(h2d, (x,)) + kv_lowered = graph.get_attr("lowered_module_0") + plain_lowered = graph.get_attr("lowered_module_1") + kv_delegate = graph.call_function(executorch_call_delegate, (kv_lowered, staged_k)) + plain_delegate = graph.call_function( + executorch_call_delegate, (plain_lowered, staged_x) + ) + graph.output((k_buffer, kv_delegate, plain_delegate)) + root = torch.nn.Module() + # Both delegates wrongly carry the spec -- the shape that per-engine stamping + # in TensorRTPartitioner exists to prevent. + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=list(zero_copy_spec) + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=list(zero_copy_spec) + ) + graph_module = torch.fx.GraphModule(root, graph) + for node, dev in ( + (k_buffer, DeviceType.CPU), + (x, DeviceType.CPU), + (staged_k, DeviceType.CUDA), + (staged_x, DeviceType.CUDA), + ): + node.meta["spec"] = SimpleNamespace(device=dev, device_index=3) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="declares zero-copy KV"): + Z._unstage_aliased_buffers(graph_module) + + +# -------------------------------------------------------------------------- +# Single-engine, on the same engine-node helper: the same per-engine derivation +# also has to narrow *within* one engine, from every aliased output down to the +# ones whose aliased input is a buffer export rewired. +# -------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_a_mixed_alias_engine_derives_the_narrower_set_and_is_then_refused(): + """One engine, two aliased outputs, one marked buffer: the narrower elidable + set is right, and an engine that needs it cannot be lowered. + + An aliased output whose input is not a buffer export rewired -- a user alias, + which nothing rewires and whose placeholder therefore carries no + ``_torch_tensorrt_aliased_buffer`` -- is still a delegate output. Deriving the + elidable set from the engine's aliased_io alone would exempt it too, and the + backend would then accept a delegate that dropped a mutation nothing writes + back. + + The derivation is right and the engine is still unusable, because the runtime + reads elision off one argument count and subtracts the engine's *whole* + aliased-output count. A .pte written for this shape loads and then fails + every ``execute()`` with an argument-count error, so ``preprocess`` refuses + it where the export can still be re-run. + """ + from torch_tensorrt.executorch._zero_copy import _aliased_inputs_by_output_index + from torch_tensorrt.executorch.backend import ( + TensorRTBackend, + _serialize_elided_output_names, + ) + from torch_tensorrt.executorch.partitioner import TensorRTPartitioner + + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + user_alias = graph.placeholder("u") + engine = _no_op_engine_node( + graph, + [k_buffer, user_alias], + aliased_io={ + "out_k": ("k_in", "kv_cache_update"), + "out_u": ("u_in", "user"), + }, + input_names=["k_in", "u_in"], + output_names=["out_k", "out_u"], + engine=b"engine-bytes", + ) + # out_k has already left the partition with the mutation that was rewired + # onto the buffer, so the delegate returns the one surviving binding. + graph.output((engine,)) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=SimpleNamespace(buffers_to_mutate={}, inputs_to_buffers={}), + constants={}, + ) + partition = SimpleNamespace(id=0, nodes=[engine]) + + # Both outputs are aliased, or the narrowing below would have nothing to do. + assert set(_aliased_inputs_by_output_index(program, engine)) == {0, 1} + + # The spec deliberately names the output the derivation must NOT pick, so a + # result of {"out_k"} can only have come from the engine's aliased_io and the + # marks on its inputs. + partitioner = TensorRTPartitioner( + compile_specs=[CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b'["out_u"]')] + ) + elided = partitioner._partition_elided_output_names(program, partition) + assert elided == {"out_k"} + + with pytest.raises(ValueError, match="Partial elision is not expressible"): + TensorRTBackend.preprocess( + program, + [ + CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(elided), + ) + ], + ) + + +def _one_aliased_engine_partition(marked): + """One engine with one aliased output, its buffer input marked or not.""" + from torch_tensorrt.executorch.partitioner import TensorRTPartitioner + + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + engine = _no_op_engine_node( + graph, + [k_buffer], + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in"], + output_names=["out_k"], + ) + graph.output((engine,)) + if marked: + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + program = SimpleNamespace( + graph_module=torch.fx.GraphModule(torch.nn.Module(), graph), + graph_signature=SimpleNamespace(buffers_to_mutate={}, inputs_to_buffers={}), + constants={}, + ) + return TensorRTPartitioner(), program, SimpleNamespace(id=0, nodes=[engine]) + + +def _make_engine_info_unreadable(monkeypatch): + def _boom(*args, **kwargs): + raise RuntimeError("engine record unreadable") + + monkeypatch.setattr( + "torch_tensorrt.executorch.partitioner._get_engine_info_for_node", _boom + ) + + +@pytest.mark.unit +def test_unreadable_engine_propagates_when_a_buffer_was_rewired(monkeypatch): + """A method with a rewired buffer must not fall back to eliding nothing. + + The aliased outputs of a rewired buffer left the graph before partitioning, + so an empty set stamps no delegate and the export dies further down blaming a + lost aliased-buffer mark -- a failure that names neither this partition nor + the record that would not read. + """ + partitioner, program, partition = _one_aliased_engine_partition(marked=True) + _make_engine_info_unreadable(monkeypatch) + + with pytest.raises(RuntimeError, match="engine record unreadable"): + partitioner._partition_elided_output_names(program, partition) + + +@pytest.mark.unit +def test_unreadable_engine_elides_nothing_when_no_buffer_was_rewired( + monkeypatch, caplog +): + """With nothing rewired the delegate really does carry every binding. + + The pair with the test above is the whole point of the handler: the same + failure is survivable here and not there, and only the graph says which. + """ + partitioner, program, partition = _one_aliased_engine_partition(marked=False) + _make_engine_info_unreadable(monkeypatch) + + with caplog.at_level( + logging.WARNING, logger="torch_tensorrt.executorch.partitioner" + ): + assert partitioner._partition_elided_output_names(program, partition) == set() + assert "could not resolve elided outputs" in caplog.text + + +@pytest.mark.unit +def test_a_partition_holding_two_engines_elides_nothing(): + """Elision is derived from one engine's aliased_io, so two is not answerable. + + ``TensorRTBackend.preprocess`` refuses a multi-engine partition outright, but + this runs first, and guessing here would stamp one engine's binding names + onto a delegate that also carries another's. + """ + from torch_tensorrt.executorch.partitioner import TensorRTPartitioner + + program, engine_a, engine_b = _two_engine_program() + both = SimpleNamespace(id=0, nodes=[engine_a, engine_b]) + + partitioner = TensorRTPartitioner() + assert partitioner._partition_elided_output_names(program, both) == set() + # The same partitioner does answer for engine_a alone, so the empty set above + # is the two-engine shape and not a graph that had nothing to elide. + single = SimpleNamespace(id=0, nodes=[engine_a]) + assert partitioner._partition_elided_output_names(program, single) == {"out_k"} + + +# -------------------------------------------------------------------------- +# GPU integration: the mark set during rewiring must survive real lowering, or +# the un-staging pass has nothing to act on and every KV update is lost. Only a +# real export exercises that -- the stub graphs above set the mark by hand. +# -------------------------------------------------------------------------- +VOCAB = 64 +DIM = 32 +HEADS = 2 +HEAD_DIM = 16 +MAX_LEN = 16 + + +class _KVDecodeStep(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.embed = torch.nn.Embedding(VOCAB, DIM) + self.pos_embed = torch.nn.Embedding(MAX_LEN, DIM) + self.q = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.k = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.v = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.o = torch.nn.Linear(HEADS * HEAD_DIM, DIM, bias=False) + self.lm = torch.nn.Linear(DIM, VOCAB, bias=False) + self.register_buffer("k_cache", torch.zeros(1, HEADS, MAX_LEN, HEAD_DIM)) + self.register_buffer("v_cache", torch.zeros(1, HEADS, MAX_LEN, HEAD_DIM)) + + def forward(self, tokens: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor: + pos_idx = input_pos.reshape(-1) + pos = input_pos.reshape(()) + x = self.embed(tokens) + self.pos_embed(input_pos.reshape(1, 1)) + + def split_heads(proj: torch.Tensor) -> torch.Tensor: + return proj.view(1, 1, HEADS, HEAD_DIM).transpose(1, 2) + + q = split_heads(self.q(x)) + k = split_heads(self.k(x)) + v = split_heads(self.v(x)) + self.k_cache.index_copy_(2, pos_idx, k) + self.v_cache.index_copy_(2, pos_idx, v) + scores = (q @ self.k_cache.transpose(-1, -2)) / (HEAD_DIM**0.5) + allowed = torch.arange(MAX_LEN, device=x.device) <= pos + bias = torch.where( + allowed, + torch.zeros((), dtype=x.dtype, device=x.device), + torch.full((), torch.finfo(x.dtype).min, dtype=x.dtype, device=x.device), + ) + attn = torch.softmax(scores + bias.view(1, 1, 1, MAX_LEN), dim=-1) + out = (attn @ self.v_cache).transpose(1, 2).reshape(1, 1, HEADS * HEAD_DIM) + return self.lm(self.o(out)) + + +@pytest.mark.parametrize("generate_etrecord", [False, True], ids=["plain", "etrecord"]) +def test_aliased_buffer_mark_survives_real_lowering(generate_etrecord): + """After a real export(..., zero_copy_kv=True), the KV buffer placeholder in + the lowered edge program still carries ``_torch_tensorrt_aliased_buffer`` -- + the token the to_out_var_pass keys the un-staging on. + + ``generate_etrecord=True`` is covered because it makes ExecuTorch deep copy + the whole program, and the mark rides on node meta. Losing it there would not + raise here: it surfaces later as the un-staging pass finding a marked buffer + it never un-staged, by which point the connection to this option is gone. + """ + _require_real_engine() + with torch.no_grad(): + torch.manual_seed(0) + model = _KVDecodeStep().eval().cuda() + tokens = torch.zeros(1, 1, dtype=torch.long).cuda() + input_pos = torch.tensor([0], dtype=torch.long).cuda() + + exported_program = torch.export.export(model, (tokens, input_pos)) + trt_gm = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=(tokens, input_pos), + min_block_size=1, + truncate_double=True, + ) + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=(tokens, input_pos), + retrace=False, + zero_copy_kv=True, + generate_etrecord=generate_etrecord, + ) + + ep = edge.exported_program() + marked = { + ep.graph_signature.inputs_to_buffers[node.name] + for node in ep.graph_module.graph.nodes + if node.op == "placeholder" and node.meta.get("_torch_tensorrt_aliased_buffer") + } + # Both, by name. The model registers two caches, so asserting the list is + # non-empty would pass on partial marker loss -- which is the very shape the + # per-delegate count check elsewhere in this feature exists for. + assert marked == {"k_cache", "v_cache"} + + +def test_finalizing_a_real_export_with_executorch_defaults_raises(): + """The documented two-call path, with the second call left as the default. + + ``to_executorch()`` with ExecuTorch's own defaults runs no un-staging, and + export has already removed the copy-back, so what it produces is a whole + ``.pte`` whose caches never update -- wrong output for a KV cache, and + nothing about the call says so. This is the real-engine end of + ``test_export_manager_checks_the_program_its_to_executorch_returns``: the + manager export returns refuses it, on a program that really was lowered and + really was finalized. + """ + _require_real_engine() + with torch.no_grad(): + torch.manual_seed(0) + model = _KVDecodeStep().eval().cuda() + tokens = torch.zeros(1, 1, dtype=torch.long).cuda() + input_pos = torch.tensor([0], dtype=torch.long).cuda() + + exported_program = torch.export.export(model, (tokens, input_pos)) + trt_gm = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=(tokens, input_pos), + min_block_size=1, + truncate_double=True, + ) + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=(tokens, input_pos), + retrace=False, + zero_copy_kv=True, + ) + + with pytest.raises(RuntimeError, match="do not reach the TensorRT delegate"): + edge.to_executorch() + + +class _MixedDecodeStep(torch.nn.Module): + """A decode step with an engine-aliased KV cache and a copy-back buffer. + + ``k_cache``/``v_cache`` are written by ``index_copy_`` on the sequence axis, + which the converter turns into an aliased engine binding. ``conv_state`` is a + ring shift -- a whole-buffer rewrite with no position to alias on -- so + ``lift_mutated_buffers`` records it in ``_copyback_mutation_buffers`` and its + new value comes back as a trailing delegate output instead. + """ + + def __init__(self) -> None: + super().__init__() + self.embed = torch.nn.Embedding(VOCAB, DIM) + self.q = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.k = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.v = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.o = torch.nn.Linear(HEADS * HEAD_DIM, DIM, bias=False) + self.lm = torch.nn.Linear(DIM, VOCAB, bias=False) + self.register_buffer("k_cache", torch.zeros(1, HEADS, MAX_LEN, HEAD_DIM)) + self.register_buffer("v_cache", torch.zeros(1, HEADS, MAX_LEN, HEAD_DIM)) + self.register_buffer("conv_state", torch.zeros(1, DIM, 4)) + + def forward(self, tokens: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor: + pos_idx = input_pos.reshape(-1) + pos = input_pos.reshape(()) + x = self.embed(tokens) + + shifted = torch.cat([self.conv_state[:, :, 1:], x.reshape(1, DIM, 1)], dim=2) + self.conv_state.copy_(shifted) + x = x + self.conv_state.sum(dim=2).reshape(1, 1, DIM) + + def split_heads(proj: torch.Tensor) -> torch.Tensor: + return proj.view(1, 1, HEADS, HEAD_DIM).transpose(1, 2) + + q = split_heads(self.q(x)) + k = split_heads(self.k(x)) + v = split_heads(self.v(x)) + self.k_cache.index_copy_(2, pos_idx, k) + self.v_cache.index_copy_(2, pos_idx, v) + scores = (q @ self.k_cache.transpose(-1, -2)) / (HEAD_DIM**0.5) + allowed = torch.arange(MAX_LEN, device=x.device) <= pos + bias = torch.where( + allowed, + torch.zeros((), dtype=x.dtype, device=x.device), + torch.full((), torch.finfo(x.dtype).min, dtype=x.dtype, device=x.device), + ) + attn = torch.softmax(scores + bias.view(1, 1, 1, MAX_LEN), dim=-1) + out = (attn @ self.v_cache).transpose(1, 2).reshape(1, 1, HEADS * HEAD_DIM) + return self.lm(self.o(out)) + + +def _real_delegates(graph_module): + return [ + node + for node in graph_module.graph.nodes + if node.op == "call_function" and node.target is executorch_call_delegate + ] + + +def _lowered_module(graph_module, delegate): + return getattr(graph_module, delegate.args[0].target) + + +def _assert_marked_buffers_reach_the_engine_unstaged(program): + """Every marked buffer is a direct argument of a TensorRT delegate. + + Finalizing a zero-copy program without raising is a weak signal, because both + of the completeness raises at the end of ``_unstage_aliased_buffers`` fire off + its own bookkeeping: a pass that records each un-staging and then leaves the + argument pointing at the staging copy still satisfies them. Only the graph says + whether the rewiring happened, and getting it wrong is silent -- the engine + writes per-call scratch that is discarded and the cache never updates. + + The library's own check runs first, on the whole program. It is the stronger + of the two -- it also requires the delegate to carry the zero-copy compile + spec and the buffer to be planned somewhere the engine can write it, neither + of which the assertions below read -- and it is the only place the container + API it reads, ``methods`` and ``exported_program(name)``, meets a real + ``ExecutorchProgramManager``: every other test of it builds the program + itself, so an upstream rename would leave those green and break + ``save(zero_copy_kv=True)`` for every caller. What follows is kept because it + reads the graph rather than the marks, which is what says the rewiring + actually happened. + """ + torch_tensorrt.executorch.check_zero_copy_kv(program) + graph_module = program.exported_program().graph_module + marked = [ + node + for node in graph_module.graph.nodes + if node.op == "placeholder" and node.meta.get("_torch_tensorrt_aliased_buffer") + ] + assert marked, "no buffer was marked for in-place update" + reached = { + arg + for node in _real_delegates(graph_module) + if _lowered_module(graph_module, node).backend_id == "TensorRTBackend" + for arg in node.args[1:] + if isinstance(arg, torch.fx.Node) + } + for node in marked: + assert node in reached, ( + f"buffer '{node.name}' is marked for in-place update but is not a " + "direct argument of any TensorRT delegate -- it either still reaches " + "one through a staging copy, or reaches none at all. Either way " + "nothing writes the caller's buffer and the cache never updates" + ) + + +class _StubProgram: + """The three attributes the reorder and upstream's write-back pass read. + + ``graph_signature`` is a property over ``_graph_signature`` because that is + how ``ExportedProgram`` exposes it, and the reorder replaces the signature by + assigning the private name. + """ + + def __init__(self, graph_module, signature): + self.graph_module = graph_module + self._graph_signature = signature + + @property + def graph_signature(self): + return self._graph_signature + + @property + def graph(self): + return self.graph_module.graph + + +def _two_mutation_program(first_value_is_inplace): + """A two-mutation program in the shape the write-back pass sees. + + Slot 0 mutates ``first``; its value is either an in-place op on its own + buffer -- which upstream reads as needing no copy -- or an ordinary + functional result, which does. Slot 1 is always an ordinary copy-back on + ``cb``. Nothing here carries the zero-copy mark, so a reorder keyed on that + mark cannot see slot 0 at all. + """ + from torch.export.exported_program import ExportGraphSignature + from torch.export.graph_signature import InputKind, InputSpec + + graph = torch.fx.Graph() + b_first = graph.placeholder("b_first") + b_cb = graph.placeholder("b_cb") + x = graph.placeholder("x") + first_value = ( + graph.call_function(torch.ops.aten.add_.Tensor, (b_first, x)) + if first_value_is_inplace + else graph.call_function(torch.ops.aten.add.Tensor, (b_first, x)) + ) + cb_value = graph.call_function(torch.ops.aten.add.Tensor, (b_cb, x)) + user = graph.call_function(torch.ops.aten.mul.Tensor, (x, x)) + graph.output((first_value, cb_value, user)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + signature = ExportGraphSignature( + input_specs=[ + InputSpec(InputKind.BUFFER, TensorArgument("b_first"), "first", False), + InputSpec(InputKind.BUFFER, TensorArgument("b_cb"), "cb", False), + InputSpec(InputKind.USER_INPUT, TensorArgument("x"), None), + ], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(first_value.name), "first" + ), + OutputSpec(OutputKind.BUFFER_MUTATION, TensorArgument(cb_value.name), "cb"), + OutputSpec(OutputKind.USER_OUTPUT, TensorArgument(user.name), None), + ], + ) + return _StubProgram(graph_module, signature) + + +@pytest.mark.unit +@pytest.mark.parametrize("reorder", [False, True], ids=["without", "with"]) +def test_reorder_moves_an_inplace_mutation_this_feature_did_not_create(reorder): + """The reorder keys on upstream's predicate, not on this feature's own mark. + + Slot 0 is in-place in the graph and carries no zero-copy mark, so upstream + inserts no copy for it exactly as for a rewired cache, and a reorder that + asked "did zero-copy rewire this?" instead of "will upstream copy this?" + would leave the pair crossed while reporting that it moved nothing. A + mutation ``reinplace_pass`` rewrites is *not* this case: that pass runs after + the reorder, so such a mutation is still ordinary when the predicate is asked + and no reorder here can pre-empt it -- see ``order_copyback_mutations_first``. + """ + from executorch.exir.passes.insert_write_back_for_buffers_pass import ( + insert_write_back_for_buffers_pass, + ) + + program = _two_mutation_program(first_value_is_inplace=True) + moved = Z.order_copyback_mutations_first(program) if reorder else 0 + _, signature = insert_write_back_for_buffers_pass(program) + value_of = {buffer: value for value, buffer in signature.buffers_to_mutate.items()} + + if not reorder: + # Pin the defect too, so the assertions below cannot pass vacuously. + assert value_of["first"].startswith("copy_") + return + assert moved == 2 + assert not value_of["first"].startswith("copy_"), ( + "'first' is mutated in place, so upstream inserts no copy for it and its " + f"finalized value must not be one; got {value_of['first']!r}" + ) + assert value_of["cb"].startswith("copy_"), ( + "'cb' is copied back, so its finalized value is the copy upstream " + f"inserted; got {value_of['cb']!r}" + ) + + +@pytest.mark.unit +def test_reorder_leaves_an_already_correct_order_alone(): + """Two copy-back mutations need no move, and the function says so.""" + program = _two_mutation_program(first_value_is_inplace=False) + before = [spec.arg.name for spec in program.graph_signature.output_specs] + + assert Z.order_copyback_mutations_first(program) == 0 + assert [spec.arg.name for spec in program.graph_signature.output_specs] == before + + +@pytest.mark.unit +def test_reorder_treats_an_unlifted_mutation_target_as_needing_no_copy(): + """A mutation upstream cannot resolve to a lifted input gets no copy. + + ``insert_write_back_for_buffers_pass`` only copies mutations whose target is + in the map it builds from the input specs, so one that is not belongs with + the copy-free mutations however its value was produced -- and its value here + is an ordinary functional result, which is what the lineage test reads as + needing a copy. Asking only the lineage test would leave this slot ahead of + the real copy-back and cross the finalized pairing. + """ + from torch.export.exported_program import ExportGraphSignature + from torch.export.graph_signature import InputKind, InputSpec + + graph = torch.fx.Graph() + b_cb = graph.placeholder("b_cb") + x = graph.placeholder("x") + unlifted_value = graph.call_function(torch.ops.aten.add.Tensor, (x, x)) + cb_value = graph.call_function(torch.ops.aten.add.Tensor, (b_cb, x)) + graph.output((unlifted_value, cb_value)) + program = _StubProgram( + torch.fx.GraphModule(torch.nn.Module(), graph), + ExportGraphSignature( + input_specs=[ + InputSpec(InputKind.BUFFER, TensorArgument("b_cb"), "cb", False), + InputSpec(InputKind.USER_INPUT, TensorArgument("x"), None), + ], + output_specs=[ + # No input spec targets "ghost", so it is not in the lifted map. + OutputSpec( + OutputKind.BUFFER_MUTATION, + TensorArgument(unlifted_value.name), + "ghost", + ), + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(cb_value.name), "cb" + ), + ], + ), + ) + + assert Z.order_copyback_mutations_first(program) == 2 + assert [spec.target for spec in program.graph_signature.output_specs] == [ + "cb", + "ghost", + ] + + +@pytest.mark.unit +def test_reorder_groups_a_non_node_mutation_value_with_the_copies(): + """A mutation slot holding a literal is ordered as though it were copied. + + Upstream reads a non-Node value as needing a copy and then raises walking it, + so putting it anywhere else would make this reorder the thing that raises and + hide the program upstream is actually complaining about. + """ + from torch.export.exported_program import ExportGraphSignature + from torch.export.graph_signature import ConstantArgument, InputKind, InputSpec + + graph = torch.fx.Graph() + b_first = graph.placeholder("b_first") + b_cb = graph.placeholder("b_cb") + x = graph.placeholder("x") + inplace_value = graph.call_function(torch.ops.aten.add_.Tensor, (b_first, x)) + graph.output((inplace_value, 7)) + program = _StubProgram( + torch.fx.GraphModule(torch.nn.Module(), graph), + ExportGraphSignature( + input_specs=[ + InputSpec(InputKind.BUFFER, TensorArgument("b_first"), "first", False), + InputSpec(InputKind.BUFFER, TensorArgument("b_cb"), "cb", False), + InputSpec(InputKind.USER_INPUT, TensorArgument("x"), None), + ], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, + TensorArgument(inplace_value.name), + "first", + ), + OutputSpec( + OutputKind.BUFFER_MUTATION, + ConstantArgument(name="literal", value=7), + "cb", + ), + ], + ), + ) + + assert Z.order_copyback_mutations_first(program) == 2 + assert [spec.target for spec in program.graph_signature.output_specs] == [ + "cb", + "first", + ] + assert program.graph_module.graph.output_node().args[0][0] == 7 + + +def _assert_each_mutation_names_its_own_value(program): + """The finalized signature pairs every mutated buffer with its own new value. + + ExecuTorch finalizes the mutations by inserting a copy for each one whose + value is not already reached from a buffer placeholder through in-place ops, + moving those copies to the front of the output tuple, and then reassigning + the mutation specs' arguments by position. Rewiring a cache to its own + placeholder takes it out of that leading run, and so does any other + mutation upstream reads as in-place, so a method that mixes the two kinds + comes out of finalization with each buffer named against another buffer's + value unless the mutations were declared in the order the pass assumes. + Nothing inside the finalizer reads the pairing, so the ``.pte`` is written + either way -- what reads it is + anyone inspecting the program, and the eager call path that copies mutated + values back into the state dict in this order. + """ + signature = program.exported_program().graph_signature + placeholder_of = {fqn: name for name, fqn in signature.inputs_to_buffers.items()} + mutated = { + spec.target: spec.arg.name + for spec in signature.output_specs + if spec.kind == OutputKind.BUFFER_MUTATION + } + assert set(mutated) == {"k_cache", "v_cache", "conv_state"} + for cache in ("k_cache", "v_cache"): + assert mutated[cache] == placeholder_of[cache], ( + f"the finalized signature gives {cache} the value " + f"{mutated[cache]!r}, but zero-copy left that cache as its own " + "mutation result, so its value is its own placeholder " + f"{placeholder_of[cache]!r}" + ) + assert mutated["conv_state"] not in placeholder_of.values(), ( + "the finalized signature gives conv_state a buffer placeholder as its " + f"value ({mutated['conv_state']!r}); it is copied back, so its value is " + "the copy ExecuTorch inserted" + ) + + +@pytest.mark.parametrize("retrace", [False, True], ids=["legacy", "retrace"]) +def test_zero_copy_kv_keeps_a_copyback_buffer_in_the_same_method(retrace): + """A real method holding both kinds of mutable buffer exports and keeps both. + + The KV caches end up bound to their own placeholders -- no value for + ExecuTorch to copy, which is the zero copy -- while ``conv_state`` stays bound + to a delegate output, which is the value ExecuTorch copies back into it. + Losing that distinction in either direction is silent wrong output, so it is + pinned on a real engine rather than a stub: the aliased_io the discriminator + reads is produced by the converter, not by this test. + + Both exporters are covered because they reach that distinction by different + routes, and only one of them is the ``save()`` default. The legacy exporter + declares all three mutations while it inlines the engines, so + ``_declare_aliased_kv_mutations_on_ep`` finds nothing left to do and the + discriminator never runs. Under ``retrace=True`` the retraced program arrives + with no mutations declared at all -- torch.export drops the aliased outputs at + the fx boundary and leaves the copy-back value as a plain return -- so that + pass is what separates the two kinds, by reading each engine's ``aliased_io``. + """ + _require_real_engine() + with torch.no_grad(): + torch.manual_seed(0) + model = _MixedDecodeStep().eval().cuda() + tokens = torch.zeros(1, 1, dtype=torch.long).cuda() + input_pos = torch.tensor([0], dtype=torch.long).cuda() + + exported_program = torch.export.export(model, (tokens, input_pos)) + trt_gm = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=(tokens, input_pos), + min_block_size=1, + truncate_double=True, + ) + assert trt_gm.meta.get("_copyback_mutation_buffers") == ["conv_state"], ( + "the model no longer produces a copy-back buffer, so this test would " + "pass without exercising the combination it exists for" + ) + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=(tokens, input_pos), + retrace=retrace, + zero_copy_kv=True, + ) + + ep = edge.exported_program() + output_args = list(ep.graph_module.graph.output_node().args[0]) + bound = { + spec.target: value + for spec, value in zip(ep.graph_signature.output_specs, output_args) + if spec.kind == OutputKind.BUFFER_MUTATION + } + assert set(bound) == {"k_cache", "v_cache", "conv_state"} + for name in ("k_cache", "v_cache"): + assert bound[name].op == "placeholder", ( + f"{name} is still satisfied by a delegate output, so ExecuTorch will " + "copy it back and zero-copy bought nothing" + ) + assert bound[name].meta.get("_torch_tensorrt_aliased_buffer") is True + assert bound["conv_state"].op == "call_function", ( + "conv_state was rewired to its own placeholder, which deletes the " + "copy-back of a buffer no engine writes in place -- a lost update" + ) + assert "_torch_tensorrt_aliased_buffer" not in bound["conv_state"].meta + + # Everything above is the export half. The staging the other half removes does + # not exist until PropagateDevicePass runs inside to_executorch, so this is the + # earliest point at which the caches can be seen reaching the engine directly. + # Composed onto a caller's own config, the optional form the user guide + # describes. It is the only shape under which a preserved field can make the + # returned config unfinalizable; a no-argument zero_copy_backend_config() + # starts from the defaults and so has nothing to preserve. + from executorch.exir import ExecutorchBackendConfig + + program = edge.to_executorch( + config=torch_tensorrt.executorch.zero_copy_backend_config( + ExecutorchBackendConfig(extract_delegate_segments=False) + ) + ) + _assert_marked_buffers_reach_the_engine_unstaged(program) + _assert_each_mutation_names_its_own_value(program) + + +class _SplitRolesDecodeStep(_MixedDecodeStep): + """The same two buffer kinds, but on two different TensorRT engines. + + ``torch.sinh`` is pinned out of TensorRT by the test, so the attention half + -- which holds the engine-aliased caches -- and the ``conv_state`` half end up + in separate partitions. Only the first engine has aliased outputs, and the + copy-back output rides on the second. + """ + + def forward(self, tokens: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor: + pos_idx = input_pos.reshape(-1) + pos = input_pos.reshape(()) + x = self.embed(tokens) + + def split_heads(proj: torch.Tensor) -> torch.Tensor: + return proj.view(1, 1, HEADS, HEAD_DIM).transpose(1, 2) + + q = split_heads(self.q(x)) + k = split_heads(self.k(x)) + v = split_heads(self.v(x)) + self.k_cache.index_copy_(2, pos_idx, k) + self.v_cache.index_copy_(2, pos_idx, v) + scores = (q @ self.k_cache.transpose(-1, -2)) / (HEAD_DIM**0.5) + allowed = torch.arange(MAX_LEN, device=x.device) <= pos + bias = torch.where( + allowed, + torch.zeros((), dtype=x.dtype, device=x.device), + torch.full((), torch.finfo(x.dtype).min, dtype=x.dtype, device=x.device), + ) + attn = torch.softmax(scores + bias.view(1, 1, 1, MAX_LEN), dim=-1) + out = (attn @ self.v_cache).transpose(1, 2).reshape(1, 1, HEADS * HEAD_DIM) + + h = torch.sinh(self.o(out) + x) + shifted = torch.cat([self.conv_state[:, :, 1:], h.reshape(1, DIM, 1)], dim=2) + self.conv_state.copy_(shifted) + return self.lm(h + self.conv_state.sum(dim=2).reshape(1, 1, DIM)) + + +@pytest.mark.parametrize("retrace", [False, True], ids=["legacy", "retrace"]) +def test_zero_copy_kv_with_the_copyback_on_a_second_delegate(retrace): + """Two TensorRT delegates, one with the aliased caches and one with the copy-back. + + This is the shape ``_delegate_declares_zero_copy`` reasons about: the + partitioner must stamp ``zero_copy_kv`` on the KV delegate only, or + ``_unstage_aliased_buffers``'s cross-check demands an aliased buffer from the + plain compute delegate and the export dies. Finalizing here is the assertion: + a wrongly stamped delegate raises inside ``to_executorch``. + + Under ``retrace=True`` this is also the only shape where + ``_declare_aliased_kv_mutations_on_ep`` has to pick the aliased engine out of + several: it scans every ``execute_engine`` node and skips the ones whose + ``aliased_io`` is empty, and the copy-back value it detaches comes off a + different engine than the caches it declares. The legacy exporter declares all + of that while inlining, so that scan runs only on this parameter. + """ + _require_real_engine() + with torch.no_grad(): + torch.manual_seed(0) + model = _SplitRolesDecodeStep().eval().cuda() + tokens = torch.zeros(1, 1, dtype=torch.long).cuda() + input_pos = torch.tensor([0], dtype=torch.long).cuda() + + exported_program = torch.export.export(model, (tokens, input_pos)) + trt_gm = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=(tokens, input_pos), + min_block_size=1, + truncate_double=True, + torch_executed_ops={"torch.ops.aten.sinh.default"}, + ) + aliased_per_engine = [ + bool(getattr(sub, "aliased_io", None)) for _, sub in trt_gm.named_children() + ] + assert len(aliased_per_engine) > 1 and sum(aliased_per_engine) == 1, ( + "the model no longer lowers to several engines with the aliasing on " + f"exactly one of them ({aliased_per_engine}), so it does not exercise " + "the multi-delegate split this test exists for" + ) + assert trt_gm.meta.get("_copyback_mutation_buffers") == ["conv_state"] + + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=(tokens, input_pos), + retrace=retrace, + zero_copy_kv=True, + ) + + ep = edge.exported_program() + graph_module = ep.graph_module + output_args = list(graph_module.graph.output_node().args[0]) + bound = { + spec.target: value + for spec, value in zip(ep.graph_signature.output_specs, output_args) + if spec.kind == OutputKind.BUFFER_MUTATION + } + assert set(bound) == {"k_cache", "v_cache", "conv_state"} + for name in ("k_cache", "v_cache"): + assert bound[name].op == "placeholder" + assert bound[name].meta.get("_torch_tensorrt_aliased_buffer") is True + assert bound["conv_state"].target is operator.getitem + + delegates = _real_delegates(graph_module) + assert len(delegates) > 1 + kv_delegate = next( + node + for node in delegates + if any( + isinstance(arg, torch.fx.Node) + and arg.meta.get("_torch_tensorrt_aliased_buffer") + for arg in node.args[1:] + ) + ) + copyback_delegate = bound["conv_state"].args[0] + assert copyback_delegate in delegates + assert copyback_delegate is not kv_delegate, ( + "the copy-back landed on the same delegate as the aliased caches, so this " + "test is running the single-delegate shape again" + ) + + stamped = [ + node + for node in delegates + if any( + spec.key == ZERO_COPY_KV_COMPILE_SPEC_KEY + for spec in _lowered_module(graph_module, node).compile_specs + ) + ] + assert stamped == [kv_delegate], ( + "the zero-copy spec must sit on the delegate whose engine lost an output " + "and on no other; a plain compute delegate carrying it is asked for an " + "aliased buffer it never had" + ) + + # The un-staging cross-check runs here, not above -- and so does the + # un-staging itself, which only the finalized graph shows. + program = edge.to_executorch( + config=torch_tensorrt.executorch.zero_copy_backend_config() + ) + _assert_marked_buffers_reach_the_engine_unstaged(program) + + +@pytest.mark.parametrize("retrace", [False, True], ids=["legacy", "retrace"]) +def test_zero_copy_kv_beside_an_executorch_cuda_delegate(retrace): + """An aliased KV cache in a method that also holds an ExecuTorch CUDA delegate. + + ``erfinv`` has no TensorRT converter, so with a ``CudaPartitioner`` catch-all + the method lowers to TensorRT, CudaBackend and TensorRT delegates in sequence. + The un-staging must reach into the TensorRT delegate only. What is asserted + here is the reachable half -- that only the KV TensorRT delegate is stamped, + and that the caches reach it un-staged with a CUDA delegate in the middle. + That the gate itself refuses a marked buffer on another backend is pinned by + ``test_unstage_raises_for_a_marked_buffer_on_another_backends_delegate``. + """ + _require_real_engine() + cuda_backend = pytest.importorskip("executorch.backends.cuda.cuda_backend") + cuda_partitioner = pytest.importorskip("executorch.backends.cuda.cuda_partitioner") + + class _CudaNeighbourDecodeStep(_KVDecodeStep): + def __init__(self): + super().__init__() + # A TensorRT-supported op AFTER erfinv, so the CUDA delegate is + # sandwiched between two TensorRT ones. Ending on erfinv would leave + # the method with a single TensorRT delegate, and the assertion below + # that only the KV delegate carries the zero-copy spec would then hold + # whatever the partitioner did. + self.tail = torch.nn.Linear(VOCAB, VOCAB, bias=False) + + def forward(self, tokens, input_pos): + h = super().forward(tokens, input_pos) + return self.tail(torch.erfinv(torch.tanh(h))) + + with torch.no_grad(): + torch.manual_seed(0) + model = _CudaNeighbourDecodeStep().eval().cuda() + tokens = torch.zeros(1, 1, dtype=torch.long).cuda() + input_pos = torch.tensor([0], dtype=torch.long).cuda() + + exported_program = torch.export.export(model, (tokens, input_pos)) + trt_gm = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=(tokens, input_pos), + min_block_size=1, + truncate_double=True, + ) + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=(tokens, input_pos), + retrace=retrace, + zero_copy_kv=True, + partitioners=[ + cuda_partitioner.CudaPartitioner( + [ + cuda_backend.CudaBackend.generate_method_name_compile_spec( + "forward" + ) + ] + ) + ], + ) + + graph_module = edge.exported_program().graph_module + delegates = _real_delegates(graph_module) + backends = { + node: _lowered_module(graph_module, node).backend_id for node in delegates + } + assert sorted(backends.values()) == [ + "CudaBackend", + "TensorRTBackend", + "TensorRTBackend", + ], ( + f"the method no longer lowers to TensorRT/CudaBackend/TensorRT " + f"({sorted(backends.values())}), so it does not cover the sandwiched " + "CUDA delegate this test exists for" + ) + + def _marked_args(node): + return [ + arg.name + for arg in node.args[1:] + if isinstance(arg, torch.fx.Node) + and arg.meta.get("_torch_tensorrt_aliased_buffer") + ] + + kv_delegates = [node for node in delegates if _marked_args(node)] + assert len(kv_delegates) == 1 and backends[kv_delegates[0]] == "TensorRTBackend", ( + "the aliased buffers must reach exactly one TensorRT delegate; " + f"got {[(backends[n], _marked_args(n)) for n in kv_delegates]}" + ) + + stamped = [ + node + for node in delegates + if any( + spec.key == ZERO_COPY_KV_COMPILE_SPEC_KEY + for spec in _lowered_module(graph_module, node).compile_specs + ) + ] + assert stamped == kv_delegates, ( + "only the delegate whose engine lost an aliased output may carry the " + "zero-copy spec; the CUDA delegate and the trailing TensorRT one had no " + f"aliased buffer, yet {[backends[n] for n in stamped]} are stamped" + ) + + program = edge.to_executorch( + config=torch_tensorrt.executorch.zero_copy_backend_config() + ) + _assert_marked_buffers_reach_the_engine_unstaged(program) + + +# -------------------------------------------------------------------------- +# save() path: unlike the direct export()+to_executorch() contract -- two paired +# calls the caller must not forget -- torch_tensorrt.save() owns both steps, so a +# single zero_copy_kv=True must both hand export() the opt-in and install the +# finalization config before to_executorch(). These are CPU-only: the TensorRT +# lowering and the ExecuTorch finalization are stubbed so the wiring is checked +# without a GPU. The passes they invoke have their own coverage above; a real +# end-to-end run is exercised by kv_cache_decode_check on GPU. +# -------------------------------------------------------------------------- +def _trivial_exported_program(): + """A tiny CPU ExportedProgram -- enough for save() to reach _save_as_executorch. + + It carries no execute_engine node, so the retrace=True KV-declaration pass is + a no-op on it and the stubs below stand in for the real lowering. + """ + + class _Add(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + 1 + + return torch.export.export(_Add(), (torch.randn(3),)) + + +def _install_save_stubs(monkeypatch, *, wrap_config=True): + """Stub the executorch lowering that save() drives and record how it is called. + + Returns a namespace capturing the kwargs export() received, the arguments + zero_copy_backend_config() was wrapped with, the config finally handed to + to_executorch(), and the programs check_zero_copy_kv() was given. When + ``wrap_config`` is False the real zero_copy_backend_config runs, so the + recorded config is the genuine one. + """ + import torch_tensorrt._compile as compile_module + import torch_tensorrt.executorch as executorch_api + + monkeypatch.setattr( + compile_module, + "ENABLED_FEATURES", + compile_module.ENABLED_FEATURES._replace(torch_tensorrt_runtime=True), + ) + + calls = SimpleNamespace( + export_kwargs=None, + wrap_args=[], + to_executorch_config="unset", + program=None, + checked=[], + ) + + def _to_executorch(config=None): + calls.to_executorch_config = config + calls.program = SimpleNamespace( + _tensor_data=None, write_to_file=lambda f: f.write(b"stub-pte") + ) + return calls.program + + edge = SimpleNamespace(to_executorch=_to_executorch) + + # The stub program has no graph, so the real check cannot read it; what these + # tests pin is that save() runs it, on the finalized program, before writing. + monkeypatch.setattr( + executorch_api, + "check_zero_copy_kv", + lambda program: calls.checked.append(program), + ) + + def _export(exp_program, **kwargs): + calls.export_kwargs = kwargs + return edge + + monkeypatch.setattr(executorch_api, "export", _export) + + if wrap_config: + wrapped = object() + + def _wrap(config=None): + calls.wrap_args.append(config) + return wrapped + + monkeypatch.setattr(executorch_api, "zero_copy_backend_config", _wrap) + calls.wrapped_sentinel = wrapped + else: + real_wrap = executorch_api.zero_copy_backend_config + + def _wrap(config=None): + calls.wrap_args.append(config) + return real_wrap(config) + + monkeypatch.setattr(executorch_api, "zero_copy_backend_config", _wrap) + + return calls + + +@pytest.mark.unit +def test_save_zero_copy_kv_true_threads_flag_and_installs_config(monkeypatch, tmp_path): + """save(zero_copy_kv=True, backend_config=cfg) opts export() in and wraps the + caller's config exactly once, forwarding the wrapped one to to_executorch().""" + calls = _install_save_stubs(monkeypatch) + user_cfg = object() + + torch_tensorrt.save( + _trivial_exported_program(), + str(tmp_path / "model.pte"), + output_format="executorch", + zero_copy_kv=True, + backend_config=user_cfg, + ) + + assert calls.export_kwargs["zero_copy_kv"] is True + # The user's config is wrapped once (preserving their fields), not double-wrapped. + assert calls.wrap_args == [user_cfg] + assert calls.to_executorch_config is calls.wrapped_sentinel + assert calls.checked == [calls.program] + + +@pytest.mark.unit +def test_save_zero_copy_kv_true_wraps_defaults_without_a_config(monkeypatch, tmp_path): + """With no backend_config, zero_copy_backend_config(None) starts from ET + defaults; the finalization config is still installed.""" + calls = _install_save_stubs(monkeypatch) + + torch_tensorrt.save( + _trivial_exported_program(), + str(tmp_path / "model.pte"), + output_format="executorch", + zero_copy_kv=True, + ) + + assert calls.export_kwargs["zero_copy_kv"] is True + assert calls.wrap_args == [None] + assert calls.to_executorch_config is calls.wrapped_sentinel + + +@pytest.mark.unit +def test_save_zero_copy_kv_true_installs_the_real_unstaging_pass(monkeypatch, tmp_path): + """End of the wiring with the real config builder: the config reaching + to_executorch() carries the un-staging to_out_var_pass, not ET's default.""" + calls = _install_save_stubs(monkeypatch, wrap_config=False) + + torch_tensorrt.save( + _trivial_exported_program(), + str(tmp_path / "model.pte"), + output_format="executorch", + zero_copy_kv=True, + ) + + assert calls.wrap_args == [None] + assert ( + type(calls.to_executorch_config.to_out_var_pass).__name__ + == "_UnstageThenToOutVar" + ) + + +@pytest.mark.unit +def test_save_refuses_skip_h2d_before_it_compiles_anything(monkeypatch, tmp_path): + """The one refusal that reads nothing but the config fires before the compile. + + Reached only through ``zero_copy_backend_config``, it lands after export has + partitioned the graph and built every engine, so a caller who set a field + that was never going to be allowed pays the whole compile to find out. + ``save`` already validates the weight-streaming budget up front for that + reason; this is the same rule applied to the same kind of field. + """ + from executorch.exir import ExecutorchBackendConfig + from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig + + calls = _install_save_stubs(monkeypatch) + + with pytest.raises(ValueError, match="skip_h2d_for_method_inputs"): + torch_tensorrt.save( + _trivial_exported_program(), + str(tmp_path / "model.pte"), + output_format="executorch", + zero_copy_kv=True, + backend_config=ExecutorchBackendConfig( + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True + ) + ), + ) + + # Refused before the lowering ran, which is the whole of this: the same + # ValueError comes out either way. + assert calls.export_kwargs is None + + +@pytest.mark.unit +def test_save_defaults_leave_kv_staged(monkeypatch, tmp_path): + """Default save() (zero_copy_kv omitted) never wraps the config, so the KV + buffer keeps its staging and its copy-back: the caller's config reaches + to_executorch() untouched and export() is told zero_copy_kv=False.""" + calls = _install_save_stubs(monkeypatch) + user_cfg = object() + + torch_tensorrt.save( + _trivial_exported_program(), + str(tmp_path / "model.pte"), + output_format="executorch", + backend_config=user_cfg, + ) + + assert calls.export_kwargs["zero_copy_kv"] is False + assert calls.wrap_args == [] + assert calls.to_executorch_config is user_cfg + assert calls.checked == [] + + +@pytest.mark.unit +@pytest.mark.parametrize( + "retrace", [False, True], ids=["retrace-false", "retrace-true"] +) +def test_save_forwards_zero_copy_kv_from_a_graph_module(monkeypatch, tmp_path, retrace): + """A compiled module is a GraphModule, and save reaches ExecuTorch by a + different branch for each value of ``retrace``. + + The tests above hand save an ``ExportedProgram``, which is the third branch. + Dropping the option from either of these two leaves a caller who asked for + zero-copy with an ordinary staged ``.pte`` and no error, since export is + never told and so nothing is rewired for the checker to miss. + """ + calls = _install_save_stubs(monkeypatch) + + class _Add(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + 1 + + torch_tensorrt.save( + torch.fx.symbolic_trace(_Add()), + str(tmp_path / "model.pte"), + output_format="executorch", + arg_inputs=[torch.randn(3)], + retrace=retrace, + # Neither branch's default exporter reads a plain traced GraphModule: the + # legacy one wants the engine-node shape a real compile produces. + use_legacy_exporter=False, + zero_copy_kv=True, + ) + + assert calls.export_kwargs["zero_copy_kv"] is True + assert calls.to_executorch_config is calls.wrapped_sentinel + assert calls.checked == [calls.program]