diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 31383f9e17..7bdc38a1e0 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -59,6 +59,21 @@ struct EngineHandle { std::vector cached_output_sizes; size_t num_inputs = 0; size_t num_outputs = 0; + // Per output binding [0..num_outputs): index into input_binding_names of the + // input it aliases (in-place KV-cache / user alias), or -1 for a normal output. + // Built at init from the blob's aliased_io. The KV buffers are threaded by + // ExecuTorch as caller-owned mutable-buffer delegate args (input AND aliased + // output): execute() binds each aliased TRT output binding to its aliased + // input's caller-provided pointer (in-place) and reflects the result into the + // delegate output EValue (a no-op when the memory planner already aliased the + // two -> zero-copy). + std::vector output_aliased_input_idx; + // Per input binding [0..num_inputs): true if any output aliases this input, so + // its in-place (KV/user) update must land in the caller-owned storage. Built at + // init from aliased_io; execute() uses it to reject a non-device-resident + // aliased input instead of silently staging its update into delegate scratch. + std::vector input_is_alias_target; + size_t num_aliased_outputs = 0; int device_id = 0; bool unified_memory = false; std::mutex mu; diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h b/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h index b3e22755d0..ce1dfaa9b9 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBlobHeader.h @@ -8,6 +8,16 @@ namespace torch_tensorrt { namespace executorch_backend { +// One aliased output->input binding pair (KV-cache in-place update, or a +// user-declared alias). The engine's output binding shares device memory with +// the named input binding; the runtime binds the output to the input's tensor +// so the update lands in-place in the caller-owned buffer. +struct AliasedBinding { + std::string output; // output binding name + std::string input; // input binding name it aliases + std::string kind; // "kv_cache_update" (TRT-enforced) or "user" +}; + struct TensorRTBlobHeader { uint32_t metadata_offset = 0; uint32_t metadata_size = 0; @@ -15,6 +25,7 @@ struct TensorRTBlobHeader { uint64_t engine_size = 0; std::vector input_binding_names; std::vector output_binding_names; + std::vector aliased_io; bool hardware_compatible = false; int device_id = 0; diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index b2e3b08232..fa86903666 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -289,6 +290,93 @@ Result TensorRTBackend::init( return err; } + // Map each aliased output binding to the index of the input it aliases so + // execute() can bind it to that input's device pointer (in-place). + // Non-aliased models have an empty header.aliased_io -> all -1, unchanged path. + handle->output_aliased_input_idx.assign(handle->num_outputs, -1); + handle->input_is_alias_target.assign(handle->num_inputs, false); + for (const auto& ab : header.aliased_io) { + int oi = -1; + for (size_t k = 0; k < handle->output_binding_names.size(); ++k) { + if (handle->output_binding_names[k] == ab.output) { + oi = static_cast(k); + break; + } + } + int ii = -1; + for (size_t k = 0; k < handle->input_binding_names.size(); ++k) { + if (handle->input_binding_names[k] == ab.input) { + ii = static_cast(k); + break; + } + } + if (oi < 0 || ii < 0) { + ET_LOG( + Error, + "TensorRTBackend::init: aliased_io names not found (output='%s', input='%s')", + ab.output.c_str(), + ab.input.c_str()); + return Error::InvalidProgram; + } + // Validate the alias kind against the two we understand. The blob parser + // defaults a missing "kind" to "kv_cache_update"; any other value is a + // corrupt or newer-than-us wire format we can't safely bind, so fail loudly + // rather than fall through and treat it as a KV alias (which would bind two + // tensors to the same storage). Mirrors the Python _reconcile_aliased_io. + if (ab.kind != "kv_cache_update" && ab.kind != "user") { + ET_LOG( + Error, + "TensorRTBackend::init: aliased_io entry (output='%s') has unknown kind '%s'", + ab.output.c_str(), + ab.kind.c_str()); + return Error::InvalidProgram; + } + if (ab.kind == "kv_cache_update") { + // TensorRT's IKVCacheUpdateLayer aliasing is the source of truth for + // kv_cache_update; the persisted map must agree with what the engine + // reports, else the blob is inconsistent with its own engine. + const char* trt_alias = handle->engine->getAliasedInputTensor(ab.output.c_str()); + if (trt_alias == nullptr || ab.input != trt_alias) { + ET_LOG( + Error, + "TensorRTBackend::init: kv_cache_update alias for output '%s' disagrees with the " + "engine (persisted input='%s', engine input='%s')", + ab.output.c_str(), + ab.input.c_str(), + trt_alias == nullptr ? "" : trt_alias); + return Error::InvalidProgram; + } + } else { + // AliasKind::USER aliases are declared by Torch-TensorRT and not tracked + // by TensorRT, so it can't validate them; confirm the aliased output and + // input share a shape before binding them to the same storage. + const nvinfer1::Dims od = handle->engine->getTensorShape(ab.output.c_str()); + const nvinfer1::Dims id = handle->engine->getTensorShape(ab.input.c_str()); + bool compatible = od.nbDims == id.nbDims; + for (int d = 0; compatible && d < od.nbDims; ++d) { + compatible = od.d[d] == id.d[d]; + } + if (!compatible) { + ET_LOG( + Error, + "TensorRTBackend::init: user alias output '%s' shape is incompatible with input '%s'", + ab.output.c_str(), + ab.input.c_str()); + return Error::InvalidProgram; + } + } + handle->output_aliased_input_idx[static_cast(oi)] = ii; + handle->input_is_alias_target[static_cast(ii)] = true; + ++handle->num_aliased_outputs; + } + + if (handle->num_aliased_outputs > 0) { + ET_LOG( + Info, + "TensorRTBackend::init: %zu aliased output(s) bound in-place to caller-owned inputs", + handle->num_aliased_outputs); + } + err = initialize_input_profiles(*handle); if (err != Error::Ok) { return err; @@ -325,9 +413,17 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* const size_t num_inputs = engine->num_inputs; const size_t num_outputs = engine->num_outputs; - if (args.size() < num_inputs + num_outputs) { + // Caller-owned KV: every input is a delegate arg, and each aliased output is + // threaded as a delegate output arg (the caller-owned mutable buffer's mutation + // slot), so all engine bindings map 1:1 to delegate args. + const size_t num_delegate_outputs = num_outputs; + const size_t num_delegate_inputs = num_inputs; + if (args.size() < num_delegate_inputs + num_delegate_outputs) { ET_LOG( - Error, "TensorRTBackend::execute: expected at least %zu args, got %zu", num_inputs + num_outputs, args.size()); + Error, + "TensorRTBackend::execute: expected at least %zu args, got %zu", + num_delegate_inputs + num_delegate_outputs, + args.size()); return Error::InvalidArgument; } @@ -395,16 +491,22 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // ------------------------------------------------------------------ // 1. Bind input shapes and addresses // ------------------------------------------------------------------ + // Device pointer each input binding was bound to; aliased outputs reuse the + // pointer of the input they alias so their update lands in-place. + std::vector input_bind_ptrs(num_inputs, nullptr); + size_t arg_idx = 0; // running index into delegate args for (size_t i = 0; i < num_inputs; ++i) { - EValue* arg = args[i]; - TORCHTRT_ET_CHECK_NOT_NULL(arg, Error::InvalidArgument, "TensorRTBackend::execute: input %zu is not a tensor", i); + const std::string& name = engine->input_binding_names[i]; + + EValue* arg = args[arg_idx++]; + TORCHTRT_ET_CHECK_NOT_NULL( + arg, Error::InvalidArgument, "TensorRTBackend::execute: input arg %zu is not a tensor", i); if (!arg->isTensor()) { ET_LOG(Error, "TensorRTBackend::execute: input %zu is not a tensor", i); return Error::InvalidArgument; } exec_aten::Tensor et_in = arg->toTensor(); - const std::string& name = engine->input_binding_names[i]; nvinfer1::Dims dims = to_trt_dims(et_in); if (dims.nbDims > nvinfer1::Dims::MAX_DIMS) { ET_LOG(Error, "TensorRTBackend::execute: input '%s' rank exceeds TensorRT limit", name.c_str()); @@ -433,6 +535,26 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return Error::InvalidState; } + // Caller-owned aliased input: an aliased output binds in-place to this + // input's device pointer, so its update must land in the caller's storage. + // If it isn't device-resident the branches below would stage it through + // delegate scratch, and the in-place update (bound to that scratch) would be + // silently lost on the next execute() when the staging copy re-reads the + // caller's unchanged buffer. Fail loudly instead. + if (engine->input_is_alias_target[i]) { + const bool device_resident = + et_in.nbytes() > 0 && (engine->unified_memory || is_cuda_accessible_ptr(et_in.const_data_ptr())); + if (!device_resident) { + ET_LOG( + Error, + "TensorRTBackend::execute: aliased input '%s' must be device-resident (non-empty and " + "CUDA-accessible or unified memory); its caller-owned in-place update cannot be staged " + "through host scratch", + name.c_str()); + return Error::InvalidArgument; + } + } + void* bind_ptr = nullptr; if (et_in.nbytes() == 0) { if (engine->cached_input_sizes[i] == 0) { @@ -472,6 +594,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } } + input_bind_ptrs[i] = bind_ptr; if (!ctx->setTensorAddress(name.c_str(), bind_ptr)) { ET_LOG(Error, "TensorRTBackend::execute: setTensorAddress failed for input '%s'", name.c_str()); return Error::InvalidState; @@ -499,9 +622,58 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // nbytes() and before the Python binding reads back the shape. // If the buffer is CPU, stage through a temporary CUDA allocation. // ------------------------------------------------------------------ + // (arg index, device_src ptr) for outputs staged through a device buffer. std::vector> outputs_needing_copy; + // Caller-owned KV: (dst = delegate output EValue ptr, src = aliased input ptr, + // nbytes). The engine updates the aliased input in place; reflect that into the + // delegate output EValue after enqueue so ExecuTorch's write-back copy_ sees the + // updated cache. Skipped when dst == src (memory planner aliased them: zero-copy). + std::vector> aliased_reflects; for (size_t o = 0; o < num_outputs; ++o) { - EValue* arg = args[num_inputs + o]; + const std::string& name = engine->output_binding_names[o]; + + // Aliased output (KV-cache / user): the engine updates the aliased input in + // place, so bind this output binding to the aliased input's device pointer. + const int alias_in = engine->output_aliased_input_idx[o]; + if (alias_in >= 0) { + void* bind_ptr = input_bind_ptrs[static_cast(alias_in)]; + if (bind_ptr == nullptr) { + ET_LOG(Error, "TensorRTBackend::execute: aliased output '%s' has no bound input pointer", name.c_str()); + return Error::InvalidState; + } + if (!ctx->setTensorAddress(name.c_str(), bind_ptr)) { + ET_LOG(Error, "TensorRTBackend::execute: setTensorAddress failed for aliased output '%s'", name.c_str()); + return Error::InvalidState; + } + // The aliased output IS a delegate output arg (the caller-owned mutable + // buffer's mutation slot). Consume it and record a reflect so ExecuTorch's + // write-back copy_ sees the engine's in-place update. + const size_t arg_i = arg_idx++; + EValue* out_arg = args[arg_i]; + TORCHTRT_ET_CHECK_NOT_NULL( + out_arg, Error::InvalidArgument, "TensorRTBackend::execute: aliased output %zu is not a tensor", o); + if (!out_arg->isTensor()) { + ET_LOG(Error, "TensorRTBackend::execute: aliased output %zu is not a tensor", o); + return Error::InvalidArgument; + } + exec_aten::Tensor et_alias_out = out_arg->toTensor(); + nvinfer1::Dims a_dims = ctx->getTensorShape(name.c_str()); + if (a_dims.nbDims >= 0 && a_dims.nbDims <= nvinfer1::Dims::MAX_DIMS) { + SizesType a_sizes[nvinfer1::Dims::MAX_DIMS]; + for (int d = 0; d < a_dims.nbDims; ++d) { + a_sizes[d] = static_cast(a_dims.d[d]); + } + (void)executorch::runtime::resize_tensor(et_alias_out, {a_sizes, static_cast(a_dims.nbDims)}); + } + void* dst = et_alias_out.nbytes() > 0 ? et_alias_out.mutable_data_ptr() : nullptr; + if (dst != nullptr && dst != bind_ptr) { + aliased_reflects.emplace_back(dst, bind_ptr, et_alias_out.nbytes()); + } + continue; + } + + const size_t arg_i = arg_idx++; // continue the shared running arg index after the inputs + EValue* arg = args[arg_i]; TORCHTRT_ET_CHECK_NOT_NULL(arg, Error::InvalidArgument, "TensorRTBackend::execute: output %zu is not a tensor", o); if (!arg->isTensor()) { ET_LOG(Error, "TensorRTBackend::execute: output %zu is not a tensor", o); @@ -509,7 +681,6 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } exec_aten::Tensor et_out = arg->toTensor(); - const std::string& name = engine->output_binding_names[o]; // Update the ExecuTorch tensor shape to the actual TRT output shape. // getTensorShape() is valid after inferShapes() has been called. @@ -556,7 +727,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } bind_ptr = engine->cached_output_ptrs[o]; output_staged_to_host = true; - outputs_needing_copy.push_back({o, bind_ptr}); + outputs_needing_copy.push_back({arg_i, bind_ptr}); } if (!ctx->setTensorAddress(name.c_str(), bind_ptr)) { @@ -577,6 +748,18 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return Error::InvalidState; } + // Caller-owned KV: reflect each engine in-place update into its delegate output + // EValue (D2D on the same stream, after the engine work). No-op list under + // zero-copy (dst == src filtered out at bind time). + for (const auto& r : aliased_reflects) { + cuda_err = cudaMemcpyAsync(std::get<0>(r), std::get<1>(r), std::get<2>(r), cudaMemcpyDeviceToDevice, stream); + if (cuda_err != cudaSuccess) { + ET_LOG( + Error, "TensorRTBackend::execute: aliased-output reflect D2D copy failed: %s", cudaGetErrorString(cuda_err)); + return Error::InvalidProgram; + } + } + // The engine work is now in flight on `stream`. Decide whether to wait for it: // must_sync = an output is staged to host (the caller reads the D2H result on // return), an input was staged from host (its async H2D read the caller's host @@ -587,10 +770,17 @@ 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. - const bool must_sync = output_staged_to_host || input_staged_from_host || !g_user_stream_set; + // A non-zero-copy 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. Zero-copy aliases (dst == src) record no reflect, so the common + // caller-owned KV fast path is untouched. + const bool aliased_reflect_pending = !aliased_reflects.empty(); + const bool must_sync = + output_staged_to_host || input_staged_from_host || aliased_reflect_pending || !g_user_stream_set; if (must_sync) { for (auto& output : outputs_needing_copy) { - exec_aten::Tensor et_out = args[num_inputs + output.first]->toTensor(); + exec_aten::Tensor et_out = args[output.first]->toTensor(); cuda_err = cudaMemcpyAsync(et_out.mutable_data_ptr(), output.second, et_out.nbytes(), cudaMemcpyDeviceToHost, stream); if (cuda_err != cudaSuccess) { diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp index 64c60ddf79..c5c49e9a26 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp @@ -136,6 +136,7 @@ bool parse_int_after_key(const std::string& json, std::size_t search_from, const bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { out.input_binding_names.clear(); out.output_binding_names.clear(); + out.aliased_io.clear(); out.hardware_compatible = false; out.device_id = 0; @@ -229,6 +230,84 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { } } + // Optional aliased_io array: [{"output":..,"input":..,"kind":..}, ...]. + // Absent in older blobs -> leave empty (backward compatible). Mirrors the + // io_bindings walk above using the same string helpers. + const std::size_t alias_key = json.find("\"aliased_io\""); + if (alias_key != std::string::npos) { + std::size_t apos = json.find('[', alias_key); + if (apos == std::string::npos) { + return false; + } + ++apos; + while (true) { + apos = skip_ws(json, apos); + if (apos >= json.size()) { + return false; + } + if (json[apos] == ']') { + ++apos; + break; + } + if (json[apos] == ',') { + ++apos; + continue; + } + if (json[apos] != '{') { + return false; + } + ++apos; + + AliasedBinding ab; + while (true) { + apos = skip_ws(json, apos); + if (apos >= json.size()) { + return false; + } + if (json[apos] == '}') { + ++apos; + break; + } + if (json[apos] == ',') { + ++apos; + continue; + } + std::string key; + apos = parse_string(json, apos, key); + if (apos == std::string::npos) { + return false; + } + apos = skip_ws(json, apos); + if (apos >= json.size() || json[apos] != ':') { + return false; + } + apos = skip_ws(json, apos + 1); + if (key == "output") { + apos = parse_string(json, apos, ab.output); + } else if (key == "input") { + apos = parse_string(json, apos, ab.input); + } else if (key == "kind") { + apos = parse_string(json, apos, ab.kind); + } else { + apos = skip_value(json, apos); + } + if (apos == std::string::npos) { + return false; + } + } + if (!ab.output.empty() && !ab.input.empty()) { + // A missing "kind" key means an older blob (the Python serializer omits + // it for KV aliases); default to the TRT-enforced kind so init()'s kind + // validation treats an absent key the same as the Python runtime rather + // than rejecting it as unknown. + if (ab.kind.empty()) { + ab.kind = "kv_cache_update"; + } + out.aliased_io.push_back(std::move(ab)); + } + } + } + return parse_bool_after_key(json, pos, "\"hardware_compatible\"", out.hardware_compatible) && parse_int_after_key(json, pos, "\"device_id\"", out.device_id); } diff --git a/examples/executorch_reference_runner/BUILD b/examples/executorch_reference_runner/BUILD index 4f53bc6b91..84d87986d5 100644 --- a/examples/executorch_reference_runner/BUILD +++ b/examples/executorch_reference_runner/BUILD @@ -7,6 +7,7 @@ filegroup( srcs = [ "CMakeLists.txt", "README.md", + "kv_cache_decode_check.cpp", "main.cpp", ], ) @@ -20,3 +21,14 @@ cc_binary( "@executorch//:executorch_file_data_loader", ], ) + +cc_binary( + name = "kv_cache_decode_check", + srcs = ["kv_cache_decode_check.cpp"], + deps = [ + "//cpp:tensorrt_executorch_backend", + "@cuda//:cudart", + "@executorch//:executorch_core", + "@executorch//:executorch_file_data_loader", + ], +) diff --git a/examples/executorch_reference_runner/CMakeLists.txt b/examples/executorch_reference_runner/CMakeLists.txt index 2bd3544d67..aa877164a9 100644 --- a/examples/executorch_reference_runner/CMakeLists.txt +++ b/examples/executorch_reference_runner/CMakeLists.txt @@ -61,3 +61,18 @@ target_link_libraries( executorch::extensions executorch::kernels torchtrt::executorch_backend) + +# Caller-owned KV-cache persistence check (see kv_cache_decode_check.cpp). It +# cudaMalloc's the device-tagged planned arenas that hold the KV buffers and +# copies the logits back to host, so it links the CUDA runtime directly. +find_package(CUDAToolkit REQUIRED) +add_executable(kv_cache_decode_check kv_cache_decode_check.cpp) +target_link_libraries( + kv_cache_decode_check + PRIVATE + executorch + executorch::backends + executorch::extensions + executorch::kernels + torchtrt::executorch_backend + CUDA::cudart) diff --git a/examples/executorch_reference_runner/README.md b/examples/executorch_reference_runner/README.md index a518353bb6..6d412585e1 100644 --- a/examples/executorch_reference_runner/README.md +++ b/examples/executorch_reference_runner/README.md @@ -103,3 +103,30 @@ Loading the method initializes the TensorRT ExecuTorch backend for any Torch-TensorRT delegate subgraphs embedded in the `.pte`. The Python `torch_tensorrt` package is needed when exporting the `.pte`; it is not needed by this native runner at inference time. + +## Caller-Owned KV-Cache Persistence Check + +`kv_cache_decode_check` is a small self-asserting runner for a caller-owned +KV-cache decode `.pte` (its aliased KV output is bound in place to the caller's +mutable buffer, which persists across `execute()` calls). + +Export a minimal single-layer decode model: + +```bash +python examples/torchtrt_executorch_example/export_kv_cache_decode.py \ + --model_path=kv_cache_decode.pte +``` + +The same CMake build produces the check runner (`kv_cache_decode_check` +target). Run it: + +```bash +./build-executorch-reference-runner/kv_cache_decode_check --model_path=kv_cache_decode.pte +``` + +It loads the method twice (each starting from a zeroed cache) and runs a decode +at `input_pos=1` once with no prior step and once after a step at `input_pos=0`. +Because the causal attention at position 1 covers positions 0..1, the two logits +differ only if the KV written at position 0 persisted across `execute()` calls. +The runner prints `[kv-check] PASS` and returns 0 on success, or fails if the +two are identical (the update did not persist). It requires a CUDA device. diff --git a/examples/executorch_reference_runner/kv_cache_decode_check.cpp b/examples/executorch_reference_runner/kv_cache_decode_check.cpp new file mode 100644 index 0000000000..fcf04144eb --- /dev/null +++ b/examples/executorch_reference_runner/kv_cache_decode_check.cpp @@ -0,0 +1,209 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + * + * Caller-owned KV-cache persistence check for a Torch-TensorRT ExecuTorch .pte. + * + * Exercises the caller-owned KV contract: the engine's aliased KV output is + * bound in place to the caller's mutable buffer, which persists across + * execute() calls. Given a single-layer decode .pte (see + * examples/torchtrt_executorch_example/export_kv_cache_decode.py) with signature + * forward(tokens[1,1], input_pos[1]) -> logits, this runs two scenarios on + * FRESH method loads (each starts from a zeroed cache): + * + * A) one decode at input_pos=1 (no prior write at pos 0) + * B) a decode at input_pos=0, then at input_pos=1 + * + * At input_pos=1 the causal attention covers positions 0..1. If the cache is + * shared across execute() calls, scenario B's second step sees the key/value + * step 0 wrote at position 0, so its logits differ from scenario A (whose + * 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. + * + * Usage: + * kv_cache_decode_check --model_path=kv_cache_decode.pte [--tol=1e-3] + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using executorch::extension::FileDataLoader; +using executorch::runtime::Error; +using executorch::runtime::EValue; +using executorch::runtime::HierarchicalAllocator; +using executorch::runtime::MemoryAllocator; +using executorch::runtime::MemoryManager; +using executorch::runtime::Method; +using executorch::runtime::MethodMeta; +using executorch::runtime::Program; +using executorch::runtime::Result; +using executorch::runtime::Span; +using executorch::runtime::TensorInfo; + +static const char* get_flag(int argc, char** argv, const char* flag, const char* def) { + const size_t n = strlen(flag); + for (int i = 1; i < argc; ++i) { + if (strncmp(argv[i], flag, n) == 0 && argv[i][n] == '=') { + return argv[i] + n + 1; + } + } + return def; +} + +// 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) { + Result method_meta = program.method_meta(method_name); + ET_CHECK_MSG(method_meta.ok(), "method_meta failed: 0x%" PRIx32, static_cast(method_meta.error())); + + auto method_pool = std::make_unique(4 * 1024U * 1024U); + auto temp_pool = std::make_unique(1 * 1024U * 1024U); + MemoryAllocator method_allocator{4 * 1024U * 1024U, method_pool.get()}; + MemoryAllocator temp_allocator{1 * 1024U * 1024U, temp_pool.get()}; + + // Caller-owned KV buffers live in the memory-planned arenas. Arenas tagged + // CUDA (by PropagateDevicePass, for device tensors a delegate reads/writes) + // must be backed by real device memory -- otherwise the aliased KV input is + // not device-resident and the backend rejects it. + std::vector> host_arenas; + std::vector cuda_arenas; + std::vector> planned_spans; + const size_t num_planned = method_meta->num_memory_planned_buffers(); + for (size_t i = 0; i < num_planned; ++i) { + const size_t sz = static_cast(method_meta->memory_planned_buffer_size(i).get()); + auto dev = method_meta->memory_planned_buffer_device(i); + if (dev.ok() && dev.get().type() == executorch::runtime::etensor::DeviceType::CUDA) { + void* p = nullptr; + ET_CHECK_MSG(cudaMalloc(&p, sz) == cudaSuccess, "cudaMalloc planned buffer %zu failed", i); + cuda_arenas.push_back(p); + planned_spans.push_back({reinterpret_cast(p), sz}); + } else { + host_arenas.push_back(std::make_unique(sz)); + planned_spans.push_back({host_arenas.back().get(), sz}); + } + } + HierarchicalAllocator planned_memory{{planned_spans.data(), planned_spans.size()}}; + MemoryManager memory_manager{&method_allocator, &planned_memory, &temp_allocator}; + + Result method = program.load_method(method_name, &memory_manager, nullptr); + ET_CHECK_MSG(method.ok(), "load_method failed: 0x%" PRIx32, static_cast(method.error())); + + // One int64 tensor per declared input (numel==1 for a decode step): input 0 is + // the token id, the rest carry input_pos. Rank is read from method_meta so this + // works whether input_pos is rank-1 ([1]) or rank-2 ([1,1]). + const size_t num_inputs = method_meta->num_inputs(); + std::vector> data(num_inputs); + std::vector> sizes(num_inputs); + std::vector> dim_order(num_inputs); + std::vector> strides(num_inputs); + std::vector impls; + impls.reserve(num_inputs); + for (size_t i = 0; i < num_inputs; ++i) { + Result ti = method_meta->input_tensor_meta(i); + ET_CHECK_MSG(ti.ok(), "input_tensor_meta(%zu) failed", i); + const auto& s = ti->sizes(); + const ssize_t nd = static_cast(s.size()); + sizes[i].assign(s.begin(), s.end()); + dim_order[i].resize(nd); + strides[i].resize(nd); + exec_aten::StridesType stride = 1; + for (ssize_t d = nd - 1; d >= 0; --d) { + dim_order[i][d] = static_cast(d); + strides[i][d] = stride; + stride *= static_cast(sizes[i][d]); + } + size_t numel = 1; + for (auto x : sizes[i]) + numel *= static_cast(x); + data[i].assign(numel, i == 0 ? 1 : 0); + impls.emplace_back( + exec_aten::ScalarType::Long, nd, sizes[i].data(), data[i].data(), dim_order[i].data(), strides[i].data()); + } + + for (int64_t pos : positions) { + for (size_t i = 1; i < num_inputs; ++i) { + std::fill(data[i].begin(), data[i].end(), pos); + } + for (size_t i = 0; i < num_inputs; ++i) { + ET_CHECK(method->set_input(EValue(exec_aten::Tensor(&impls[i])), i) == Error::Ok); + } + ET_CHECK_MSG(method->execute() == Error::Ok, "execute() failed at pos %" PRId64, pos); + } + + EValue out; + ET_CHECK_MSG(method->get_outputs(&out, 1) == Error::Ok, "get_outputs failed"); + ET_CHECK_MSG(out.isTensor(), "output 0 is not a tensor"); + 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. + std::vector result(static_cast(t.numel())); + ET_CHECK_MSG( + cudaMemcpy(result.data(), t.const_data_ptr(), result.size() * sizeof(float), cudaMemcpyDefault) == cudaSuccess, + "cudaMemcpy of logits to host failed"); + for (void* p : cuda_arenas) { + cudaFree(p); + } + return result; +} + +int main(int argc, char** argv) { + executorch::runtime::runtime_init(); + const char* model_path = get_flag(argc, argv, "--model_path", "kv_cache_decode.pte"); + const double tol = atof(get_flag(argc, argv, "--tol", "1e-3")); + + Result loader = FileDataLoader::from(model_path); + ET_CHECK_MSG(loader.ok(), "FileDataLoader::from('%s') failed", model_path); + auto loader_ptr = std::make_unique(std::move(loader.get())); + Result program = Program::load(loader_ptr.get()); + ET_CHECK_MSG(program.ok(), "Failed to parse model '%s'", model_path); + + auto name = program->get_method_name(0); + ET_CHECK_MSG(name.ok(), "Program has no methods"); + 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; + for (size_t i = 0; i < a.size(); ++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 (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] FAIL: outputs are identical -> the KV write did not persist across execute() calls.\n"); + return 1; +} diff --git a/examples/torchtrt_executorch_example/export_kv_cache_decode.py b/examples/torchtrt_executorch_example/export_kv_cache_decode.py new file mode 100644 index 0000000000..c8590d98b0 --- /dev/null +++ b/examples/torchtrt_executorch_example/export_kv_cache_decode.py @@ -0,0 +1,115 @@ +""" +.. _executorch_export_kv_cache: + +Exporting a Caller-Owned KV-Cache Decode Model to ExecuTorch (.pte) +================================================================== + +This example exports a minimal single-layer attention decode step whose KV cache +is a registered buffer updated in place with ``index_copy_``. Torch-TensorRT +carries the cache as a *caller-owned* mutable buffer through the ExecuTorch +delegate, so the engine's aliased KV output is bound in place to the caller's +buffer and persists across ``execute()`` calls. + +The companion ``kv_cache_decode_check`` reference runner loads the resulting +``.pte`` and asserts that a decode step observes the KV a previous step wrote +(i.e. the cache is shared across ``execute()`` calls). + +Prerequisites +------------- +Install Torch-TensorRT with the ExecuTorch extra before running this example:: + + pip install -e ".[executorch]" +""" + +import argparse + +import torch +import torch_tensorrt + +VOCAB = 64 +DIM = 32 +HEADS = 2 +HEAD_DIM = 16 +MAX_LEN = 16 + + +class KVDecodeStep(torch.nn.Module): + """One attention layer with an in-place (index_copy_) KV cache. + + ``forward(tokens[1,1], input_pos[1]) -> logits[1,1,VOCAB]``. The ``k_cache`` / + ``v_cache`` buffers are written at ``input_pos`` and attended over up to + ``input_pos`` (causal), so a later step's output depends on earlier steps' + writes -- which only holds if the cache persists across ``execute()`` calls. + """ + + 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)) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--model_path", default="kv_cache_decode.pte", help="Path to save the .pte" + ) + args = parser.parse_args() + + 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, + ) + 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.") + + +if __name__ == "__main__": + main() diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index 32ee801886..9dbdde48aa 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -1134,6 +1134,16 @@ def _extract_tensor(obj: Any) -> Any: package_path=file_path, ) elif output_format == "executorch": + from torch_tensorrt.dynamo._exporter import ( + _declare_aliased_kv_mutations_on_ep, + ) + + # retrace=True: torch.export truncates the engines' aliased KV + # outputs, so declare them as buffer mutations before lowering. + _copyback_bufs = module.meta.get("_copyback_mutation_buffers", []) + exp_program = _declare_aliased_kv_mutations_on_ep( + exp_program, copyback_buffers=_copyback_bufs + ) _save_as_executorch( exp_program, file_path, diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index 3ebac2a21f..cf45ccd4f4 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -5,7 +5,6 @@ import os import platform import warnings - from typing import Any, Collection, Dict, List, Optional, Sequence, Tuple, Union import sympy @@ -45,6 +44,7 @@ pre_export_lowering, ) from torch_tensorrt.dynamo.lowering._buffer_lifting import ( + assert_predicted_kv_aliased, inline_lifted_buffers_into_gm, lift_mutated_buffers, ) @@ -792,6 +792,8 @@ def compile( # module-held cache). Returns a fresh GraphModule whose forward signature # reflects the new placeholders. gm, lifted_buffers = lift_mutated_buffers(gm) + _copyback_mutation_buffers = gm.meta.get("_copyback_mutation_buffers", []) + _predicted_kv_bindings = gm.meta.get("_predicted_kv_bindings", []) if lifted_buffers: # Append each lifted buffer as an engine input AFTER the user inputs. # Buffer tensors live on the gm's state; prepare an Input spec for @@ -832,6 +834,12 @@ def compile( engine_cache, graph_signature=exported_program.graph_signature, ) + if _copyback_mutation_buffers: + trt_gm.meta["_copyback_mutation_buffers"] = _copyback_mutation_buffers + # Ground-truth check: every write lift classified as KV (engine-aliased, so its + # copy_ was dropped) must actually appear in a compiled engine's aliased_io, + # else its write-back would be silently lost -- fail loudly instead. + assert_predicted_kv_aliased(trt_gm, _predicted_kv_bindings) if lifted_buffers: # Inline buffers into the compiled gm as get_attr nodes + registered # buffers. The resulting gm's forward takes only user inputs; buffers diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index dacbea140a..113fde44b8 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -1,9 +1,11 @@ import base64 import copy +import logging import operator -from typing import Any, Dict, Optional, Sequence, Tuple +from typing import Any, Dict, List, Optional, Sequence, Set, Tuple import torch +import torch.utils._pytree as pytree from torch._export.non_strict_utils import make_constraints from torch._guards import detect_fake_mode from torch._library.fake_class_registry import FakeScriptObject @@ -19,9 +21,12 @@ OutputSpec, TensorArgument, ) +from torch.fx.graph import _PyTreeCodeGen from torch_tensorrt._features import ENABLED_FEATURES from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ENGINE_IDX, NAME_IDX +logger = logging.getLogger(__name__) + def _resolve_lifted_custom_obj( exp_program: ExportedProgram, node: torch.fx.Node @@ -70,7 +75,9 @@ def export( inputs (torch.Tensor): Torch input tensors cross_compile_module (bool): Flag to indicated whether it is cross_compilation enabled or not """ - patched_module = transform(gm, cross_compile_module) + patched_module = transform( + gm, cross_compile_module, expose_aliased_mutations=bool(use_legacy_exporter) + ) if not use_legacy_exporter: args = () if arg_inputs is not None: @@ -95,6 +102,7 @@ def export( def transform( gm: torch.fx.GraphModule, cross_compile_module: Optional[bool] = False, + expose_aliased_mutations: bool = True, ) -> torch.fx.GraphModule: """ Transforms the graphmodule by inlining Pytorch and TensorRT submodules. @@ -113,7 +121,7 @@ def transform( gm = copy.deepcopy(gm) # Inline TensorRT submodules - inline_trt_modules(gm, cross_compile_module) + inline_trt_modules(gm, cross_compile_module, expose_aliased_mutations) # Inline pytorch submodules inline_torch_modules(gm) @@ -230,6 +238,12 @@ def lift( kind=input_kind, arg=input_spec_arg, target=node.target, + # torch>=2.3 requires an explicit persistent flag on BUFFER + # specs. state_dict() excludes non-persistent buffers by + # construction, so any buffer reaching this in-state_dict + # branch is persistent (non-persistent buffers take the + # not-in-state_dict path above and are lifted as constants). + persistent=(True if input_kind == InputKind.BUFFER else None), ), ) non_user_input_idx += 1 @@ -245,29 +259,6 @@ def lift( return gm, graph_signature, state_dict, constants -def get_duplicate_nodes( - gm: torch.fx.GraphModule, submodule: torch.fx.GraphModule -) -> Tuple[Sequence[Any], Sequence[Any]]: - """ - We check if there are duplicate nodes when we copy submodule graph into gm. - Handle the case where the subgraph input placeholders are same as - gm placeholders. This happens when the first submodule in the graph is - a pytorch submodule - """ - submodule_placeholder_inputs = [ - node for node in submodule.graph.nodes if node.op == "placeholder" - ] - submodule_input_node_names = [node.name for node in submodule_placeholder_inputs] - gm_node_names = [node.name for node in gm.graph.nodes] - submodule_duplicate_inputs = [ - node for node in submodule_placeholder_inputs if node.name in gm_node_names - ] - gm_duplicate_inputs = [ - node for node in gm.graph.nodes if node.name in submodule_input_node_names - ] - return submodule_duplicate_inputs, gm_duplicate_inputs - - def inline_torch_modules(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: """ Inline a submodule within the parent graph (gm). All `call_module` nodes @@ -285,43 +276,31 @@ def inline_torch_modules(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: # or a placeholder of the main graph submodule_inputs = gm_node.args - submodule_duplicate_inputs, gm_duplicate_inputs = get_duplicate_nodes( - gm, submodule - ) - assert len(submodule_duplicate_inputs) == len(gm_duplicate_inputs) - # Avoid creating new copies of duplicate inputs by creating a mapping - val_map = {} - for i in range(len(submodule_duplicate_inputs)): - val_map[submodule_duplicate_inputs[i]] = gm_duplicate_inputs[i] - - # Copy all nodes in the submodule into gm and - # store the output node of this submodule which is now present in gm + # Copy the submodule's nodes into gm, then wire its inputs POSITIONALLY. + # + # We deliberately do NOT pre-seed val_map by matching submodule input + # placeholders to gm nodes by NAME. Name matching silently binds an + # input to the WRONG node when names collide (e.g. a _run_on_gpu input + # placeholder whose name matches a different engine's getitem), which + # rewires a consumer to the wrong producer and orphans the real one -- + # the orphan is then pruned by dead-code elimination, leaving a delegate + # short an output at runtime. It also leaks a mixed submodule's + # computed-intermediate inputs as spurious graph inputs. gm_node.args is + # the authoritative, ordered list of the real inputs, so we let + # graph_copy create a fresh (auto-renamed on collision) placeholder for + # every submodule input and rewire each to submodule_inputs[i] by + # position, then erase it. + val_map: Dict[Any, Any] = {} submodule_output = gm.graph.graph_copy(submodule.graph, val_map) - # Get their references (since we copied) in the parent graph (gm) - if len(submodule_duplicate_inputs) == 0: - submodule_placeholder_input_names = [ - node.name - for node in submodule.graph.nodes - if node.op == "placeholder" - ] - gm_added_placeholder_inputs = [ - node - for node in gm.graph.nodes - if node.name in submodule_placeholder_input_names - ] - - assert len(submodule_inputs) == len(gm_added_placeholder_inputs) - - # Replace the added placeholder inputs with original inputs to this submodule node - for idx in range(len(gm_added_placeholder_inputs)): - gm_added_placeholder_inputs[idx].replace_all_uses_with( - submodule_inputs[idx] - ) - - # Erase the placeholder input nodes in the gm - for idx in range(len(gm_added_placeholder_inputs)): - gm.graph.erase_node(gm_added_placeholder_inputs[idx]) + submodule_placeholders = [ + node for node in submodule.graph.nodes if node.op == "placeholder" + ] + assert len(submodule_placeholders) == len(submodule_inputs) + for idx, submodule_placeholder in enumerate(submodule_placeholders): + copied_placeholder = val_map[submodule_placeholder] + copied_placeholder.replace_all_uses_with(submodule_inputs[idx]) + gm.graph.erase_node(copied_placeholder) # Replace the pytorch submodule node (call_module) with the inlined subgraph output # Special handling when submodule returns multiple outputs (tuple) @@ -388,14 +367,54 @@ def create_trt_exp_program( input_nodes = [node for node in gm.graph.nodes if node.op == "placeholder"] output_nodes = [node for node in gm.graph.nodes if node.op == "output"] assert output_nodes - output_nodes = output_nodes[0].args[0] + _output_node = output_nodes[0] + output_nodes = _output_node.args[0] + + # Copy-back mutable buffers (non-KV, e.g. a convolution-state ring-buffer): + # lift_mutated_buffers appended each buffer's new-value as a trailing graph output. Tag it + # BUFFER_MUTATION (reusing the KV `_kv_mutation_target` path) so the runtime + # threads it and ExecuTorch copies it back to the caller-owned buffer, then move + # all mutation outputs before the user outputs (ExportedProgram verifier requires + # BUFFER_MUTATIONs to be contiguous at the front). + _copyback = gm.meta.get("_copyback_mutation_buffers", []) + if _copyback: + _outs = list(output_nodes) + for _n, _buf in zip(_outs[-len(_copyback) :], _copyback): + if isinstance(_n, torch.fx.Node): + _n.meta["_kv_mutation_target"] = _buf + + def _is_mut(_n: Any) -> bool: + return isinstance(_n, torch.fx.Node) and "_kv_mutation_target" in _n.meta + + output_nodes = tuple( + [_n for _n in _outs if _is_mut(_n)] + + [_n for _n in _outs if not _is_mut(_n)] + ) + _output_node.args = (output_nodes,) + + # Outputs tagged by `_expose_aliased_buffer_mutations` become BUFFER_MUTATION + # specs (their `_kv_mutation_target` meta names the backing buffer); the rest + # are ordinary user outputs, used below to rebuild the user-facing out_spec. + user_output_nodes = [ + node for node in output_nodes if "_kv_mutation_target" not in node.meta + ] input_specs = [ InputSpec(InputKind.USER_INPUT, TensorArgument(name=node.name), node.target) for node in input_nodes ] output_specs = [ - OutputSpec(OutputKind.USER_OUTPUT, TensorArgument(name=node.name), node.target) + ( + OutputSpec( + OutputKind.BUFFER_MUTATION, + TensorArgument(name=node.name), + node.meta["_kv_mutation_target"], + ) + if "_kv_mutation_target" in node.meta + else OutputSpec( + OutputKind.USER_OUTPUT, TensorArgument(name=node.name), node.target + ) + ) for node in output_nodes ] @@ -403,14 +422,31 @@ def create_trt_exp_program( input_specs=input_specs, output_specs=output_specs ) + # A hybrid TRT+CUDA GraphModule from dynamo.compile carries a plain fx.CodeGen + # (no pytree_info): the module already returns a flat tuple, so rebuild + # in_spec/out_spec from the example inputs and the flat graph outputs. This + # matches retrace=True, which traces the same flat module and likewise cannot + # re-nest -- so the fallback never diverges from it. + codegen = gm.graph._codegen + if isinstance(codegen, _PyTreeCodeGen): + in_spec = codegen.pytree_info.in_spec + out_spec = codegen.pytree_info.out_spec + else: + example_args = tuple(arg_inputs) if arg_inputs is not None else () + example_kwargs = kwarg_inputs or {} + in_spec = pytree.tree_flatten((example_args, example_kwargs))[1] + # out_spec describes the user-visible return structure only; buffer + # mutations are stripped before unflatten. + out_spec = pytree.tree_flatten(tuple(user_output_nodes))[1] + module_call_graph = [ ModuleCallEntry( "", ModuleCallSignature( inputs=[], outputs=[], - in_spec=gm.graph._codegen.pytree_info.in_spec, - out_spec=gm.graph._codegen.pytree_info.out_spec, + in_spec=in_spec, + out_spec=out_spec, ), ) ] @@ -491,8 +527,179 @@ def create_trt_exp_program( return trt_exp_program +def _declare_aliased_kv_mutations_on_ep( + exp_program: ExportedProgram, + copyback_buffers: Optional[List[str]] = None, +) -> ExportedProgram: + """retrace=True post-export pass: declare each engine's aliased KV output as a + BUFFER_MUTATION of its caller-owned buffer input, and reclassify any copy-back + outputs as BUFFER_MUTATIONs of their buffers. + + torch.export produces execute_engine nodes whose meta['val'] covers only the + user outputs (the aliased KV outputs are network bindings excluded at the fx + boundary), so the KV buffers -- though BUFFER inputs -- are never recorded as + mutated and get frozen downstream. This surfaces each aliased output as a + getitem and declares it a BUFFER_MUTATION of the aliased input's buffer, + mirroring create_trt_exp_program's handling on the retrace=False path. + + Non-KV mutable buffers (e.g. a convolution-state ring-buffer) have no engine + aliasing: lift_mutated_buffers appended their new values as trailing user outputs, which + torch.export kept as the last ``len(copyback_buffers)`` outputs. Those are + reclassified from USER_OUTPUT to BUFFER_MUTATION here so ExecuTorch copies them + back to the caller-owned buffers after the delegate runs. + + Args: + exp_program: the retraced ExportedProgram to rewrite. + copyback_buffers: buffer FQNs, in output order, for the trailing copy-back + outputs to reclassify as BUFFER_MUTATION. Threaded from + ``gm.meta['_copyback_mutation_buffers']`` by ``save()``; empty/None means + KV-only (no copy-back). + + Returns exp_program unchanged when no engine has aliased KV outputs and there are + no copy-back buffers. + """ + 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.backend import _get_engine_info_for_node + + def _estr(engine_info: List[Any], idx: int) -> str: + if idx < 0 or idx >= len(engine_info) or engine_info[idx] is None: + return "" + v = engine_info[idx] + return v.decode("utf-8", "replace") if isinstance(v, bytes) else str(v) + + gm = exp_program.graph_module + sig = exp_program.graph_signature + inputs_to_buffers = sig.inputs_to_buffers + output_node = next(n for n in gm.graph.nodes if n.op == "output") + exec_target = torch.ops.tensorrt.execute_engine.default + + already_exposed: Set[str] = set() + mutation_outputs: List[Tuple[torch.fx.Node, str]] = [] + for node in gm.graph.nodes: + if node.op != "call_function" or node.target is not exec_target: + continue + engine_info = _get_engine_info_for_node(exp_program, node) + aliased_io = deserialize_aliased_io(_estr(engine_info, ALIASED_IO_IDX)) + if not aliased_io: + continue + in_names = deserialize_binding_names( + _estr(engine_info, INPUT_BINDING_NAMES_IDX) + ) + out_names = deserialize_binding_names( + _estr(engine_info, OUTPUT_BINDING_NAMES_IDX) + ) + input_nodes = list(node.args[0]) + val_list = list(node.meta["val"]) + for out_name in out_names: + if out_name not in aliased_io: + continue + in_name = aliased_io[out_name][0] + if in_name not in in_names: + continue + ii = in_names.index(in_name) + if ii >= len(input_nodes): + continue + buf_node = input_nodes[ii] + buf_fqn = inputs_to_buffers.get(getattr(buf_node, "name", None)) + if buf_fqn is None or buf_fqn in already_exposed: + continue + oi = out_names.index(out_name) + while len(val_list) <= oi: + val_list.append(buf_node.meta["val"]) + val_list[oi] = buf_node.meta["val"] + with gm.graph.inserting_after(node): + getitem_node = gm.graph.call_function(operator.getitem, (node, oi)) + getitem_node.meta["val"] = buf_node.meta["val"] + already_exposed.add(buf_fqn) + mutation_outputs.append((getitem_node, buf_fqn)) + node.meta["val"] = tuple(val_list) + + copyback_buffers = copyback_buffers or [] + num_copyback = len(copyback_buffers) + if not mutation_outputs and not num_copyback: + return exp_program + + out_args = list(output_node.args[0]) + orig_specs = list(sig.output_specs) + + kv_getitems = [g for g, _ in mutation_outputs] + kv_specs = [ + OutputSpec(OutputKind.BUFFER_MUTATION, TensorArgument(name=g.name), fqn) + for g, fqn in mutation_outputs + ] + + # Copy-back mutable buffers (non-KV, e.g. a convolution-state ring-buffer): lift + # appended each buffer's new-value as a trailing user output; torch.export kept them as the + # last num_copyback outputs. Reclassify those from USER_OUTPUT to BUFFER_MUTATION + # so ET copies them back to the caller-owned buffers. + copyback_getitems: List[torch.fx.Node] = [] + copyback_specs: List[OutputSpec] = [] + remaining_specs = orig_specs + if num_copyback: + copyback_getitems = list(out_args[-num_copyback:]) + remaining_args = out_args[:-num_copyback] + remaining_specs = orig_specs[:-num_copyback] + copyback_specs = [ + OutputSpec(OutputKind.BUFFER_MUTATION, TensorArgument(name=g.name), buf) + for g, buf in zip(copyback_getitems, copyback_buffers) + if isinstance(g, torch.fx.Node) + ] + out_args = remaining_args + + # BUFFER_MUTATIONs (KV + copy-back) must precede USER_OUTPUTs (verifier). + output_node.args = (tuple(kv_getitems + copyback_getitems + out_args),) + gm.graph.lint() + gm.recompile() + + new_output_specs = kv_specs + copyback_specs + list(remaining_specs) + new_signature = ExportGraphSignature( + input_specs=list(sig.input_specs), output_specs=new_output_specs + ) + + # Copy-back outputs were user returns in the retraced program, so torch.export's + # out_spec counts them; after reclassifying them as BUFFER_MUTATION only the real + # user outputs remain. Rebuild the top-level out_spec (mirroring + # create_trt_exp_program) so to_edge's unflatten sees the right leaf count. + module_call_graph = exp_program.module_call_graph + if num_copyback and module_call_graph: + new_out_spec = pytree.tree_flatten(tuple(out_args))[1] + _e0 = module_call_graph[0] + _sig0 = _e0.signature + module_call_graph = [ + ModuleCallEntry( + _e0.fqn, + ModuleCallSignature( + inputs=_sig0.inputs if _sig0 is not None else [], + outputs=[], + in_spec=_sig0.in_spec if _sig0 is not None else None, + out_spec=new_out_spec, + ), + ) + ] + list(module_call_graph[1:]) + + return ExportedProgram( + root=gm, + graph=gm.graph, + graph_signature=new_signature, + state_dict=exp_program.state_dict, + range_constraints=exp_program.range_constraints, + module_call_graph=module_call_graph, + constants=exp_program.constants, + ) + + def inline_trt_modules( - gm: torch.fx.GraphModule, cross_compile_module: Optional[bool] = False + gm: torch.fx.GraphModule, + cross_compile_module: Optional[bool] = False, + expose_aliased_mutations: bool = True, ) -> torch.fx.GraphModule: """ Replace TRT submodules with trt engine nodes. @@ -554,12 +761,125 @@ def inline_trt_modules( for idx, getitem_node in enumerate(getitem_nodes): getitem_node.meta["val"] = trt_node.meta["val"][idx] + # Expose the engine's aliased (KV-cache) outputs as graph-level buffer + # mutations so the ExecuTorch path sees a real mutable buffer instead of + # a frozen constant. Only on the legacy (create_trt_exp_program) path, + # which declares the BUFFER_MUTATION specs; on the torch.export path the + # extra outputs would just perturb the user outputs (see save()'s + # post-export declaration for retrace=True). Non-cross-compile only. + if not cross_compile_module and expose_aliased_mutations: + _expose_aliased_buffer_mutations(gm, trt_node, trt_module, num_outputs) + # Erase the TRT submodule (call_module) node. gm.graph.erase_node(trt_module_node) return gm +def _expose_aliased_buffer_mutations( + gm: torch.fx.GraphModule, + trt_node: torch.fx.Node, + trt_module: Any, + num_user_outputs: int, +) -> None: + """Surface an engine's aliased KV-cache outputs as graph buffer mutations. + + The interpreter appends aliased layer outputs (e.g. ``IKVCacheUpdateLayer``) + to the engine's network bindings *after* the fx output boundary, so + ``trt_node.meta["val"]`` (and the partitioner-emitted getitems) only cover + the user outputs. Here we add a ``getitem`` for each aliased output binding + and route it to the graph output tagged as a buffer mutation of the aliased + input's backing buffer. ``create_trt_exp_program`` turns the tag into a + ``BUFFER_MUTATION`` OutputSpec, so ``torch.export``/``to_edge`` record the + cache in ``buffers_to_mutate`` -- without a graph ``copy_`` that + functionalization would fold away (the aliased output shares the buffer's + storage, so a ``copy_`` from it is a no-op self-copy). + """ + aliased_io = getattr(trt_module, "aliased_io", None) + if not aliased_io: + return + + in_names = list(getattr(trt_module, "input_binding_names", [])) + out_names = list(getattr(trt_module, "output_binding_names", [])) + input_arg_nodes = list(trt_node.args[0]) + + # Only get_attr nodes backed by a *registered buffer* can be declared + # BUFFER_MUTATION targets; a get_attr that lift() would classify as a + # constant (not in named_buffers) is not a valid mutation target. + registered_buffers = {name for name, _ in gm.named_buffers()} + + # A buffer can be declared mutated at most once in the graph signature. + # Multiple engines can alias the same backing buffer (they share its + # storage), so dedup exposures across engines. + already_exposed = gm.meta.setdefault("_kv_exposed_mutation_targets", set()) + + output_node = next(node for node in gm.graph.nodes if node.op == "output") + + val_list = list(trt_node.meta["val"]) + new_mutation_outputs: List[torch.fx.Node] = [] + for oi, out_name in enumerate(out_names): + if out_name not in aliased_io: + continue + in_name = aliased_io[out_name][0] + if in_name not in in_names: + continue + ii = in_names.index(in_name) + if ii >= len(input_arg_nodes): + continue + buffer_node = input_arg_nodes[ii] + buf_target = getattr(buffer_node, "target", None) + if buffer_node.op != "get_attr" or not isinstance(buf_target, str): + logger.warning( + "Aliased input %s for engine output %s is not a buffer get_attr " + "(op=%s); skipping buffer-mutation exposure.", + in_name, + out_name, + buffer_node.op, + ) + continue + if buf_target not in registered_buffers: + logger.warning( + "Aliased input %s for engine output %s resolves to get_attr %s " + "which is not a registered buffer; skipping buffer-mutation exposure.", + in_name, + out_name, + buf_target, + ) + continue + if buf_target in already_exposed: + logger.warning( + "Buffer %s (engine output %s / input %s) already exposed as a " + "mutation by another engine; skipping duplicate.", + buf_target, + out_name, + in_name, + ) + continue + already_exposed.add(buf_target) + + # Ensure the engine node advertises at least oi+1 outputs so getitem(oi) + # is in range; the aliased output has the shape/dtype of its input buffer. + while len(val_list) <= oi: + val_list.append(buffer_node.meta["val"]) + val_list[oi] = buffer_node.meta["val"] + + with gm.graph.inserting_after(trt_node): + getitem_node = gm.graph.call_function(operator.getitem, (trt_node, oi)) + getitem_node.meta["val"] = buffer_node.meta["val"] + getitem_node.meta["_kv_mutation_target"] = buf_target + new_mutation_outputs.append(getitem_node) + + if not new_mutation_outputs: + return + + trt_node.meta["val"] = tuple(val_list) + # BUFFER_MUTATION outputs must precede USER_OUTPUTs (the ExportedProgram + # verifier treats output_nodes[num_tokens:num_tokens+num_mutations] as the + # mutations), so prepend. + out_args = list(output_node.args[0]) + output_node.args = (tuple(new_mutation_outputs + out_args),) + + def replace_execute_engine_no_op_node( exp_program: ExportedProgram, ) -> ExportedProgram: diff --git a/py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py b/py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py index c99387cdba..4c10c98995 100644 --- a/py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py +++ b/py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py @@ -35,6 +35,92 @@ logger = logging.getLogger(__name__) +def _kv_write_will_alias(value_node: object, cache_shape: Tuple[int, ...]) -> bool: + """Whether the converter will emit an ``IKVCacheUpdateLayer`` (with in-place + aliased I/O) for this mutated buffer's new-value node. + + Eligibility depends on the write dim, static shapes, and the cache being a + direct network input, so this reuses the converters' own eligibility + predicates. A ``slice_scatter`` / ``index_copy`` that is not eligible is + lowered to a non-aliasing scatter, so its write-back is kept as a + BUFFER_MUTATION output (copy-back) rather than dropped. Imports are local to + avoid a lowering<->conversion import cycle. + """ + if not (isinstance(value_node, torch.fx.Node) and value_node.op == "call_function"): + return False + args = value_node.args + # The KV layer aliases the cache only if it is a direct network input; after + # lifting, the mutated buffer is a placeholder the write op reads from. + if not ( + args and isinstance(args[0], torch.fx.Node) and args[0].op == "placeholder" + ): + return False + + if value_node.target is torch.ops.aten.index_copy.default: + from torch_tensorrt.dynamo.conversion.aten_ops_converters import ( + _index_copy_kv_eligible, + ) + + return _index_copy_kv_eligible(value_node) + + if value_node.target is torch.ops.aten.slice_scatter.default: + from torch_tensorrt.dynamo.conversion.impl.slice_scatter import _kv_eligible + + dim = args[2] if len(args) > 2 and isinstance(args[2], int) else 0 + start = args[3] if len(args) > 3 and isinstance(args[3], int) else 0 + # Prefer the source's extent along `dim` for the write length; fall back + # to end-start from the slice bounds. + update_len = None + src = args[1] if len(args) > 1 else None + if isinstance(src, torch.fx.Node): + src_val = src.meta.get("val") + if src_val is not None and dim < len(src_val.shape): + s = src_val.shape[dim] + update_len = s if isinstance(s, int) else None + if update_len is None: + end = ( + args[4] + if len(args) > 4 and isinstance(args[4], int) + else (cache_shape[dim] if dim < len(cache_shape) else 0) + ) + update_len = max(end - start, 0) + eligible, _reason = _kv_eligible(tuple(cache_shape), dim, start, update_len) + return eligible + + return False + + +def assert_predicted_kv_aliased( + gm: torch.fx.GraphModule, predicted_kv_bindings: List[str] +) -> None: + """Ground-truth check for the KV predictions :func:`_kv_write_will_alias` made. + + Each write ``lift_mutated_buffers`` classified as KV-aliased had its ``copy_`` + dropped in the expectation that the engine would alias it in place. If the + converter did not actually emit an ``IKVCacheUpdateLayer`` for it, that + write-back is silently lost. So assert every predicted-KV input binding + appears in a compiled engine's ``aliased_io``, and raise loudly otherwise. + Keyed on the ``buf_*`` binding name, which is stable across the later buffer + rename and is what ``aliased_io`` records on the input side. + """ + if not predicted_kv_bindings: + return + aliased_in: set = set() + for _name, sub in gm.named_children(): + amap = getattr(sub, "aliased_io", None) + if amap: + for v in amap.values(): + aliased_in.add(v[0] if isinstance(v, (tuple, list)) else v) + missing = [b for b in predicted_kv_bindings if b not in aliased_in] + if missing: + raise RuntimeError( + "lift_mutated_buffers classified these buffer writes as KV-cache " + "(engine-aliased) and dropped their copy_, but the compiled engine " + f"did not alias them (absent from aliased_io): {missing}. Their " + "write-back would be silently dropped." + ) + + def lift_mutated_buffers( gm: torch.fx.GraphModule, ) -> Tuple[torch.fx.GraphModule, List[Tuple[str, str, torch.Tensor]]]: @@ -55,6 +141,19 @@ def lift_mutated_buffers( tuples, in the order placeholders were appended (which matches the order they appear in the new gm's forward signature, after the original user inputs). + + Side effects: the trailing ``copy_`` of each mutated buffer is erased. A write + the converter will lower to an ``IKVCacheUpdateLayer`` with in-place aliased + I/O (an eligible ``slice_scatter`` / ``index_copy``, per + :func:`_kv_write_will_alias`) relies on that engine aliasing for its + write-back, so nothing further is added. Every other mutation -- a non-KV + buffer, or a ``slice_scatter`` / ``index_copy`` that fails eligibility and is + lowered to a non-aliasing scatter -- has no engine aliasing, so its new value + is re-appended as a graph output ("copy-back") and its buffer name recorded, + in output order, in ``gm.meta['_copyback_mutation_buffers']`` -- the + downstream exporters (``create_trt_exp_program`` / + ``_declare_aliased_kv_mutations_on_ep``) read that list to reclassify those + outputs as BUFFER_MUTATIONs. """ # Find all aten.copy_(get_attr_X, _) calls. The first arg's target is # the buffer name. Some EPs emit copy_.default, others copy_. @@ -82,6 +181,27 @@ def lift_mutated_buffers( lifted: List[Tuple[str, str, torch.Tensor]] = [] seen_buffers: Dict[str, torch.fx.Node] = {} # buffer name -> new placeholder node + # Each mutated buffer's write-back is handled one of two ways downstream: + # - Engine-level aliasing (zero-copy): the slice_scatter / index_copy + # converters emit an IKVCacheUpdateLayer whose output is aliased in-place + # to the cache input. We drop the copy_ and rely on that aliasing. + # - Copy-back: any other mutation (a non-KV buffer such as a convolution- + # state ring-buffer, OR a slice_scatter / index_copy the converter cannot + # turn into an IKVCacheUpdateLayer) has no aliasing, so its new value is + # re-attached as an ordinary BUFFER_MUTATION output that ExecuTorch copies + # back after the delegate runs. + # `_kv_write_will_alias` decides between them by reusing the converters' own + # eligibility predicates (not the op target alone), so a non-aliasable + # slice_scatter / index_copy falls to copy-back rather than being dropped. + # index_put has no aliasing converter, so it always falls to copy-back too. + copyback: List[Tuple[torch.fx.Node, str]] = [] # (new_value_node, buffer_name) + # Input-binding names of writes predicted to alias (KV). compile() asserts each + # actually appears in the engine's aliased_io, turning a mis-prediction into a + # loud error instead of a silently dropped write-back. The binding name (buf_*) + # is the stable key: it survives the buffer renaming inline does later, and it + # is exactly what aliased_io records on the input side. + predicted_kv_bindings: set[str] = set() + for copy_node, get_attr_node in mutation_pairs: buffer_name = get_attr_node.target if not hasattr(gm, buffer_name): @@ -133,17 +253,39 @@ def lift_mutated_buffers( # to the new placeholder. get_attr_node.replace_all_uses_with(replacement) - # Drop the trailing copy_ (it's now redundant — the mutation lands on the - # placeholder's storage via engine-level aliasing). + # KV writes rely on engine-level aliasing, so the trailing copy_ is + # redundant and dropped. Other (non-KV) mutations have no aliasing: + # record their new value so we can re-attach it as a copy-back + # BUFFER_MUTATION output below, then drop the (now input-target) copy_. + new_value = copy_node.args[1] if len(copy_node.args) > 1 else None + if isinstance(new_value, torch.fx.Node): + if _kv_write_will_alias(new_value, tuple(buffer_tensor.shape)): + predicted_kv_bindings.add(replacement.name) + else: + copyback.append((new_value, buffer_name)) gm.graph.erase_node(copy_node) # Erase the now-unused get_attr. if not get_attr_node.users: gm.graph.erase_node(get_attr_node) + # Re-attach non-KV mutation new-values as graph outputs so they survive + # compilation (otherwise, with the copy_ gone, they are dead and eliminated). + # Appended in order; recorded so create_trt_exp_program / _declare can tag the + # corresponding engine outputs as BUFFER_MUTATION (copy-back) downstream. + copyback_buffers: List[str] = [] + if copyback: + output_node = next(n for n in gm.graph.nodes if n.op == "output") + out_args = list(output_node.args[0]) + out_args.extend(nv for nv, _ in copyback) + output_node.args = (tuple(out_args),) + copyback_buffers = [buf for _, buf in copyback] + gm.graph.lint() if not lifted: + gm.meta["_copyback_mutation_buffers"] = copyback_buffers + gm.meta["_predicted_kv_bindings"] = sorted(predicted_kv_bindings) return gm, [] # ExportedProgram.module() produces a GraphModule whose forward is @@ -174,11 +316,14 @@ def lift_mutated_buffers( except AttributeError: pass new_gm.recompile() + new_gm.meta["_copyback_mutation_buffers"] = copyback_buffers + new_gm.meta["_predicted_kv_bindings"] = sorted(predicted_kv_bindings) logger.debug( - "Lifted %d mutated buffer(s) to placeholders: %s", + "Lifted %d mutated buffer(s) to placeholders: %s; copy-back buffers: %s", len(lifted), [(p, b) for p, b, _ in lifted], + copyback_buffers, ) return new_gm, lifted diff --git a/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py b/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py index 916c5dcf3f..15ae608e52 100644 --- a/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py +++ b/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py @@ -352,12 +352,34 @@ def fake_no_op_placeholder_for_execute_engine( C++ schema validator. Output shapes are inferred from the serialized metadata embedded in the op's string args, same as fake_tensorrt_execute_engine. """ - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule + from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ( + deserialize_binding_names, + ) + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + TorchTensorRTModule, + deserialize_aliased_io, + ) metadata = TorchTensorRTModule.decode_metadata(serialized_metadata) shape_info = metadata.get("inout_symexprs") if metadata else None if shape_info: - return _apply_symbolic_shape_expressions(inputs, shape_info) + outputs = _apply_symbolic_shape_expressions(inputs, shape_info) + # Append the engine's aliased (KV-cache) outputs so the getitem indices + # produced when to_edge re-traces this op stay in range: the aliased + # outputs are network bindings appended after the fx output boundary, so + # their shape/dtype come from the aliased input binding. + aliased_io = deserialize_aliased_io(serialized_aliased_io) + if aliased_io: + in_names = deserialize_binding_names(serialized_in_binding_names) + out_names = deserialize_binding_names(serialized_out_binding_names) + for out_name in out_names: + if out_name in aliased_io: + in_name = aliased_io[out_name][0] + if in_name in in_names: + outputs.append( + torch.empty_like(inputs[in_names.index(in_name)]) + ) + return outputs else: raise RuntimeError( "No symbolic shape expressions found in TensorRT engine metadata. " diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index b73d50eea2..fbc973ae7d 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -12,6 +12,7 @@ from torch.export.exported_program import ExportedProgram from torch_tensorrt.dynamo._exporter import _resolve_lifted_custom_obj from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + ALIASED_IO_IDX, DEVICE_IDX, ENGINE_IDX, HW_COMPATIBLE_IDX, @@ -20,6 +21,7 @@ REQUIRES_OUTPUT_ALLOCATOR_IDX, SERIALIZED_METADATA_IDX, TARGET_PLATFORM_IDX, + deserialize_aliased_io, ) from torch_tensorrt.executorch.serialization import ( TensorRTBlobMetadata, @@ -306,8 +308,14 @@ def preprocess( TensorRTIOBinding(name=name, is_input=True) for name in input_names ] + [TensorRTIOBinding(name=name, is_input=False) for name in output_names] + # Carry the KV-cache / user aliasing (out->in, kind) into the blob so the + # 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)) + metadata = TensorRTBlobMetadata( io_bindings=io_bindings, + aliased_io=aliased_io, hardware_compatible=_get_str(engine_info, HW_COMPATIBLE_IDX) == "1", device_id=_parse_device_id(engine_info[DEVICE_IDX]), serialized_metadata=_get_str(engine_info, SERIALIZED_METADATA_IDX), diff --git a/py/torch_tensorrt/executorch/partitioner.py b/py/torch_tensorrt/executorch/partitioner.py index 11e0e6aa1b..8e03ed48d6 100644 --- a/py/torch_tensorrt/executorch/partitioner.py +++ b/py/torch_tensorrt/executorch/partitioner.py @@ -39,6 +39,28 @@ logger = logging.getLogger(__name__) +def _keep_mutated_buffers_above_delegate(exported_program: ExportedProgram) -> None: + """Undo tag_constant_data freezing a delegate-mutated buffer as constant. + + tag_constant_data detects a mutated buffer only via its *direct* users, so a + buffer whose mutation is produced inside the delegate (the mutation is a + getitem off the call_delegate, not a direct user of the buffer placeholder) + is misclassified as constant data and tagged into the delegate. A TensorRT + engine is stateless across executions, so an absorbed mutable buffer would be + a frozen constant (the KV-cache update would be lost). Strip the delegation + tag from any buffer that is a mutation target so it stays a caller-owned + mutable buffer owned above the delegate. + """ + sig = exported_program.graph_signature + mutated_buffer_targets = set(sig.buffers_to_mutate.values()) + for node in exported_program.graph_module.graph.nodes: + if ( + node.op == "placeholder" + and sig.inputs_to_buffers.get(node.name) in mutated_buffer_targets + ): + node.meta.pop("delegation_tag", None) + + class TensorRTPartitioner(Partitioner): # type: ignore[misc] """Partitions the graph for TensorRT delegation. @@ -140,6 +162,7 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: ) tag_constant_data(exported_program) + _keep_mutated_buffers_above_delegate(exported_program) return PartitionResult( tagged_exported_program=exported_program, diff --git a/py/torch_tensorrt/executorch/serialization.py b/py/torch_tensorrt/executorch/serialization.py index bd9e0e6619..47d8fe9f4b 100644 --- a/py/torch_tensorrt/executorch/serialization.py +++ b/py/torch_tensorrt/executorch/serialization.py @@ -10,7 +10,7 @@ import json import struct from dataclasses import dataclass, field -from typing import List, Tuple +from typing import Dict, List, Tuple TENSORRT_MAGIC = b"TR01" HEADER_FORMAT = "<4sIIIQ8s" @@ -32,6 +32,11 @@ class TensorRTIOBinding: @dataclass class TensorRTBlobMetadata: io_bindings: List[TensorRTIOBinding] = field(default_factory=list) + # Aliased output->input bindings: out_name -> (in_name, kind). "kind" is an + # AliasKind value ("kv_cache_update" or "user"); the C++ backend binds each + # aliased engine output to its aliased input's tensor (in-place) so the + # update lands in the caller-owned buffer. + aliased_io: Dict[str, Tuple[str, str]] = field(default_factory=dict) hardware_compatible: bool = False device_id: int = 0 serialized_metadata: str = "" @@ -50,6 +55,12 @@ def to_json(self) -> bytes: } for binding in self.io_bindings ], + # List form (not a dict) so the small C++ parser can walk it like + # io_bindings. Emitted right after io_bindings, before the scalars. + "aliased_io": [ + {"output": out, "input": inp, "kind": kind} + for out, (inp, kind) in self.aliased_io.items() + ], "hardware_compatible": self.hardware_compatible, "device_id": self.device_id, "serialized_metadata": self.serialized_metadata, @@ -67,8 +78,13 @@ def from_json(cls, data: bytes) -> "TensorRTBlobMetadata": ) for binding in parsed.get("io_bindings", []) ] + aliased_io = { + b["output"]: (b["input"], b.get("kind", "kv_cache_update")) + for b in parsed.get("aliased_io", []) + } return cls( io_bindings=io_bindings, + aliased_io=aliased_io, hardware_compatible=parsed.get("hardware_compatible", False), device_id=parsed.get("device_id", 0), serialized_metadata=parsed.get("serialized_metadata", ""), diff --git a/tests/cpp/executorch/test_executorch_blob_header.cpp b/tests/cpp/executorch/test_executorch_blob_header.cpp index 4146f5c907..e527e990e1 100644 --- a/tests/cpp/executorch/test_executorch_blob_header.cpp +++ b/tests/cpp/executorch/test_executorch_blob_header.cpp @@ -107,6 +107,48 @@ TEST(ExecuTorchTensorRTBlobHeader, RejectsMissingIoBindingsMetadata) { EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); } +TEST(ExecuTorchTensorRTBlobHeader, ParsesAliasedIo) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"out_k","is_input":false},)" + R"({"name":"in_u","is_input":true},{"name":"out_u","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"},)" + R"({"output":"out_u","input":"in_u","kind":"user"}],)" + R"("hardware_compatible":false,"device_id":0})"; + const auto blob = make_blob(metadata); + + 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[0].input, "in_k"); + EXPECT_EQ(header.aliased_io[0].kind, "kv_cache_update"); + EXPECT_EQ(header.aliased_io[1].output, "out_u"); + EXPECT_EQ(header.aliased_io[1].input, "in_u"); + EXPECT_EQ(header.aliased_io[1].kind, "user"); +} + +TEST(ExecuTorchTensorRTBlobHeader, DefaultsMissingAliasedIo) { + // Blobs written before aliased_io existed omit the key; parsing must still + // succeed and leave aliased_io empty (backward compatible). + const auto blob = + make_blob(R"({"io_bindings":[{"name":"input_0","is_input":true},{"name":"output_0","is_input":false}],)" + R"("hardware_compatible":false,"device_id":0})"); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_TRUE(header.aliased_io.empty()); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesEmptyAliasedIo) { + const auto blob = + make_blob(R"({"io_bindings":[{"name":"input_0","is_input":true},{"name":"output_0","is_input":false}],)" + R"("aliased_io":[],"hardware_compatible":false,"device_id":0})"); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_TRUE(header.aliased_io.empty()); +} + } // namespace } // namespace executorch_backend } // namespace torch_tensorrt diff --git a/tests/py/dynamo/executorch/test_kv_cache_export.py b/tests/py/dynamo/executorch/test_kv_cache_export.py new file mode 100644 index 0000000000..8f596b6d42 --- /dev/null +++ b/tests/py/dynamo/executorch/test_kv_cache_export.py @@ -0,0 +1,265 @@ +"""Export-side coverage for caller-owned KV-cache buffer mutations. + +Both retrace modes must surface an engine's aliased KV outputs as graph-level +BUFFER_MUTATIONs so the ExecuTorch delegate keeps them as caller-owned mutable +buffers instead of freezing them: + + * retrace=False (legacy exporter): ``inline_trt_modules`` exposes them at + transform time (guarded by ``expose_aliased_mutations``), then + ``create_trt_exp_program`` declares the specs. + * retrace=True (torch.export): ``torch.export`` truncates the aliased outputs + at the fx boundary, so ``_declare_aliased_kv_mutations_on_ep`` re-declares + them on the exported program before lowering. +""" + +import operator +from types import SimpleNamespace + +import pytest +import torch +from torch.export.exported_program import ( + InputKind, + InputSpec, + OutputKind, + OutputSpec, + TensorArgument, +) +from torch_tensorrt.dynamo import _exporter as E + + +@pytest.mark.unit +@pytest.mark.parametrize("use_legacy, expected_expose", [(True, True), (False, False)]) +def test_export_exposes_aliased_mutations_only_for_legacy_exporter( + monkeypatch, use_legacy, expected_expose +): + """The transform-time KV exposure runs on the retrace=False (legacy) path + only; retrace=True defers to the post-export declaration pass, so it must not + perturb the user outputs at transform time. + """ + captured = {} + + def fake_transform(gm, cross_compile_module=False, expose_aliased_mutations=True): + captured["expose"] = expose_aliased_mutations + return gm + + monkeypatch.setattr(E, "transform", fake_transform) + monkeypatch.setattr(E, "create_trt_exp_program", lambda *a, **k: "legacy-ep") + monkeypatch.setattr(torch.export, "export", lambda *a, **k: "retrace-ep") + + result = E.export(torch.nn.Module(), use_legacy_exporter=use_legacy) + + assert captured["expose"] is expected_expose + assert result == ("legacy-ep" if use_legacy else "retrace-ep") + + +@pytest.mark.unit +def test_declare_aliased_kv_mutations_is_noop_without_engines(): + """With no execute_engine node carrying aliased I/O, the pass returns the + exported program unchanged (same object).""" + pytest.importorskip("executorch.exir") + g = torch.fx.Graph() + x = g.placeholder("x") + g.output((x,)) + gm = torch.fx.GraphModule(torch.nn.Module(), g) + + ep = SimpleNamespace( + graph_module=gm, + graph_signature=SimpleNamespace(inputs_to_buffers={}), + ) + assert E._declare_aliased_kv_mutations_on_ep(ep) is ep + + +@pytest.mark.unit +def test_declare_aliased_kv_mutations_declares_buffer_mutation(monkeypatch): + """An engine whose aliased KV output is dropped from meta['val'] gets that + output surfaced as a getitem and declared a BUFFER_MUTATION of the aliased + input's buffer, ordered before the user outputs (verifier requirement).""" + pytest.importorskip("executorch.exir") + import torch_tensorrt.dynamo.runtime._serialized_engine_layout as L + import torch_tensorrt.dynamo.runtime._TorchTensorRTModule as M + import torch_tensorrt.executorch.backend as B + + exec_target = torch.ops.tensorrt.execute_engine.default + + # b_k_0 (KV buffer) + tokens feed the engine; meta['val'] covers only the one + # user output -- the aliased KV output ("out_k") is truncated at the boundary. + g = torch.fx.Graph() + b_k_0 = g.placeholder("b_k_0") + tokens = g.placeholder("tokens") + engine = g.placeholder("engine") + eng = g.call_function(exec_target, ([b_k_0, tokens], engine)) + user_out = g.call_function(operator.getitem, (eng, 0)) + g.output((user_out,)) + + buf_val = torch.zeros(2, 2) + out_val = torch.zeros(1) + b_k_0.meta["val"] = buf_val + tokens.meta["val"] = torch.zeros(1) + engine.meta["val"] = None + eng.meta["val"] = [out_val] + user_out.meta["val"] = out_val + gm = torch.fx.GraphModule(torch.nn.Module(), g) + + sig = SimpleNamespace( + inputs_to_buffers={"b_k_0": "k_0"}, + input_specs=[ + InputSpec(InputKind.BUFFER, TensorArgument(name="b_k_0"), "k_0", True), + ], + output_specs=[ + OutputSpec(OutputKind.USER_OUTPUT, TensorArgument(name="user_out"), None), + ], + ) + ep = SimpleNamespace( + graph_module=gm, + graph_signature=sig, + state_dict={}, + range_constraints={}, + module_call_graph=[], + constants={}, + ) + + info = ["x"] * (L.ALIASED_IO_IDX + 1) + info[L.INPUT_BINDING_NAMES_IDX] = "IN" + info[L.OUTPUT_BINDING_NAMES_IDX] = "OUT" + monkeypatch.setattr(B, "_get_engine_info_for_node", lambda ep_, n: info) + monkeypatch.setattr( + M, "deserialize_aliased_io", lambda s: {"out_k": ("k_in", "kv_cache_update")} + ) + monkeypatch.setattr( + L, + "deserialize_binding_names", + lambda s: ["k_in", "tokens"] if s == "IN" else ["user_out", "out_k"], + ) + + captured = {} + + class _CapturingEP: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(E, "ExportedProgram", _CapturingEP) + + E._declare_aliased_kv_mutations_on_ep(ep) + + new_specs = captured["graph_signature"].output_specs + # BUFFER_MUTATION for k_0 is declared first, ahead of the user output. + assert len(new_specs) == 2 + assert new_specs[0].kind == OutputKind.BUFFER_MUTATION + assert new_specs[0].target == "k_0" + assert new_specs[1].kind == OutputKind.USER_OUTPUT + + # The mutation getitem is prepended to the graph output (mutations first). + out_node = next(n for n in gm.graph.nodes if n.op == "output") + assert out_node.args[0][0].name == new_specs[0].arg.name + assert out_node.args[0][1] is user_out + + # The engine's meta['val'] is extended to cover the previously-dropped output. + assert len(eng.meta["val"]) == 2 + + +@pytest.mark.unit +def test_declare_aliased_kv_mutations_declares_copyback(monkeypatch): + """retrace=True: a trailing copy-back output (a non-KV mutable buffer with no + engine aliasing) is reclassified from USER_OUTPUT to BUFFER_MUTATION of its + buffer and ordered ahead of the user outputs.""" + pytest.importorskip("executorch.exir") + + # No execute_engine node -> the pure copy-back case (num_copyback drives the + # pass). ``user_out`` is a real return; ``state_new`` is the copy-back new value + # that lift_mutated_buffers appended as the trailing output. + g = torch.fx.Graph() + x = g.placeholder("x") + state_in = g.placeholder("state_in") + user_out = g.call_function(torch.add, (x, x)) + state_new = g.call_function(torch.add, (state_in, x)) + g.output((user_out, state_new)) + gm = torch.fx.GraphModule(torch.nn.Module(), g) + + sig = SimpleNamespace( + inputs_to_buffers={"state_in": "state_0"}, + input_specs=[ + InputSpec( + InputKind.BUFFER, TensorArgument(name="state_in"), "state_0", True + ), + InputSpec(InputKind.USER_INPUT, TensorArgument(name="x"), None), + ], + output_specs=[ + OutputSpec(OutputKind.USER_OUTPUT, TensorArgument(name="user_out"), None), + OutputSpec(OutputKind.USER_OUTPUT, TensorArgument(name="state_new"), None), + ], + ) + ep = SimpleNamespace( + graph_module=gm, + graph_signature=sig, + state_dict={}, + range_constraints={}, + module_call_graph=[], + constants={}, + ) + + captured = {} + + class _CapturingEP: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(E, "ExportedProgram", _CapturingEP) + + E._declare_aliased_kv_mutations_on_ep(ep, copyback_buffers=["state_0"]) + + new_specs = captured["graph_signature"].output_specs + # The trailing copy-back output becomes a BUFFER_MUTATION of state_0, first. + assert len(new_specs) == 2 + assert new_specs[0].kind == OutputKind.BUFFER_MUTATION + assert new_specs[0].target == "state_0" + assert new_specs[0].arg.name == state_new.name + assert new_specs[1].kind == OutputKind.USER_OUTPUT + + # Mutation is prepended to the graph output; the user output follows. + out_node = next(n for n in gm.graph.nodes if n.op == "output") + assert out_node.args[0][0] is state_new + assert out_node.args[0][1] is user_out + + +@pytest.mark.unit +def test_create_trt_exp_program_declares_copyback(monkeypatch): + """retrace=False: create_trt_exp_program reads + gm.meta['_copyback_mutation_buffers'], tags the trailing outputs with their + buffer target, and emits them as BUFFER_MUTATION specs ahead of the user + outputs.""" + pytest.importorskip("executorch.exir") + + g = torch.fx.Graph() + x = g.placeholder("x") + state_in = g.placeholder("state_in") + user_out = g.call_function(torch.add, (x, x)) + state_new = g.call_function(torch.add, (state_in, x)) + g.output((user_out, state_new)) + gm = torch.fx.GraphModule(torch.nn.Module(), g) + gm.recompile() + gm.meta["_copyback_mutation_buffers"] = ["state_0"] + + captured = {} + + class _CapturingEP: + def __init__(self, **kwargs): + captured.update(kwargs) + + # Stub the heavy tail: lift() (param/buffer lifting) and the real EP ctor. + monkeypatch.setattr(E, "ExportedProgram", _CapturingEP) + monkeypatch.setattr(E, "lift", lambda gm_, sig_: (gm_, sig_, {}, {})) + + E.create_trt_exp_program(gm) + + specs = captured["graph_signature"].output_specs + assert len(specs) == 2 + assert specs[0].kind == OutputKind.BUFFER_MUTATION + assert specs[0].target == "state_0" + assert specs[0].arg.name == state_new.name + assert specs[1].kind == OutputKind.USER_OUTPUT + + # The trailing output is tagged with its buffer and reordered mutation-first. + assert state_new.meta["_kv_mutation_target"] == "state_0" + out_node = next(n for n in gm.graph.nodes if n.op == "output") + assert out_node.args[0][0] is state_new + assert out_node.args[0][1] is user_out diff --git a/tests/py/dynamo/executorch/test_partitioner.py b/tests/py/dynamo/executorch/test_partitioner.py index f7f2ab6b6a..48d030a1d7 100644 --- a/tests/py/dynamo/executorch/test_partitioner.py +++ b/tests/py/dynamo/executorch/test_partitioner.py @@ -40,9 +40,51 @@ def fake_tag_constant_data(exported_program): ) graph_module = SimpleNamespace(graph=SimpleNamespace(nodes=[])) - exported_program = SimpleNamespace(graph_module=graph_module) + exported_program = SimpleNamespace( + graph_module=graph_module, + graph_signature=SimpleNamespace(buffers_to_mutate={}, inputs_to_buffers={}), + ) result = TensorRTPartitioner().partition(exported_program) assert tagged["called"] assert sorted(result.partition_tags.keys()) == ["tensorrt_1", "tensorrt_2"] + + +@pytest.mark.unit +def test_keep_mutated_buffers_above_delegate_untags_only_mutation_targets(): + """The un-tag post-pass keeps a delegate-mutated buffer above the delegate + (strips its delegation_tag) while leaving non-mutated constants tagged into + the delegate and non-placeholder nodes untouched. + """ + from torch_tensorrt.executorch.partitioner import ( + _keep_mutated_buffers_above_delegate, + ) + + mutated_buf = SimpleNamespace( + op="placeholder", name="b_k_0", meta={"delegation_tag": "tensorrt_0"} + ) + const_buf = SimpleNamespace( + op="placeholder", name="b_w", meta={"delegation_tag": "tensorrt_0"} + ) + engine_node = SimpleNamespace( + op="call_function", name="tensorrt_0", meta={"delegation_tag": "tensorrt_0"} + ) + exported_program = SimpleNamespace( + graph_module=SimpleNamespace( + graph=SimpleNamespace(nodes=[mutated_buf, const_buf, engine_node]) + ), + graph_signature=SimpleNamespace( + buffers_to_mutate={"getitem_5": "k_0"}, + inputs_to_buffers={"b_k_0": "k_0", "b_w": "w"}, + ), + ) + + _keep_mutated_buffers_above_delegate(exported_program) + + # k_0 is a mutation target -> its buffer placeholder is kept above the delegate + assert "delegation_tag" not in mutated_buf.meta + # w is not mutated -> still frozen into the delegate + assert const_buf.meta["delegation_tag"] == "tensorrt_0" + # non-placeholder nodes are untouched + assert engine_node.meta["delegation_tag"] == "tensorrt_0" diff --git a/tests/py/dynamo/executorch/test_partitioner_target_device.py b/tests/py/dynamo/executorch/test_partitioner_target_device.py index f7d6b3e481..106f73f82f 100644 --- a/tests/py/dynamo/executorch/test_partitioner_target_device.py +++ b/tests/py/dynamo/executorch/test_partitioner_target_device.py @@ -45,6 +45,7 @@ def _engine_node(device_id): def _edge_program(*nodes): return SimpleNamespace( graph_module=SimpleNamespace(graph=SimpleNamespace(nodes=list(nodes))), + graph_signature=SimpleNamespace(buffers_to_mutate={}, inputs_to_buffers={}), constants={}, ) diff --git a/tests/py/dynamo/executorch/test_serialization.py b/tests/py/dynamo/executorch/test_serialization.py index 65c9e46ad4..9eddee2cd8 100644 --- a/tests/py/dynamo/executorch/test_serialization.py +++ b/tests/py/dynamo/executorch/test_serialization.py @@ -1,3 +1,5 @@ +import json + import pytest from torch_tensorrt.executorch.serialization import ( HEADER_SIZE, @@ -37,3 +39,40 @@ def test_serialize_engine_writes_tr01_blob(): def test_deserialize_engine_rejects_bad_magic(): with pytest.raises(ValueError, match="Invalid magic"): deserialize_engine(b"NOPE" + b"\x00" * (HEADER_SIZE - 4)) + + +@pytest.mark.unit +def test_serialize_engine_round_trips_aliased_io(): + metadata = TensorRTBlobMetadata( + io_bindings=[ + TensorRTIOBinding(name="in_k", is_input=True), + TensorRTIOBinding(name="out_k", is_input=False), + TensorRTIOBinding(name="in_u", is_input=True), + TensorRTIOBinding(name="out_u", is_input=False), + ], + aliased_io={ + "out_k": ("in_k", "kv_cache_update"), + "out_u": ("in_u", "user"), + }, + ) + + engine, parsed = deserialize_engine(serialize_engine(b"eng", metadata)) + assert engine == b"eng" + assert parsed.aliased_io == { + "out_k": ("in_k", "kv_cache_update"), + "out_u": ("in_u", "user"), + } + + +@pytest.mark.unit +def test_metadata_from_json_without_aliased_io_defaults_empty(): + # Blobs written before aliased_io existed omit the key entirely; parsing must + # default to an empty mapping rather than raising (backward compatibility). + metadata = TensorRTBlobMetadata( + io_bindings=[TensorRTIOBinding(name="x", is_input=True)] + ) + data = json.loads(metadata.to_json().decode("utf-8")) + del data["aliased_io"] + + restored = TensorRTBlobMetadata.from_json(json.dumps(data).encode("utf-8")) + assert restored.aliased_io == {} diff --git a/tests/py/dynamo/lowering/test_buffer_lifting.py b/tests/py/dynamo/lowering/test_buffer_lifting.py index 4dda3bfd6e..dbeb7fe089 100644 --- a/tests/py/dynamo/lowering/test_buffer_lifting.py +++ b/tests/py/dynamo/lowering/test_buffer_lifting.py @@ -30,6 +30,7 @@ from torch.export import export from torch.testing._internal.common_utils import TestCase, run_tests from torch_tensorrt.dynamo.lowering._buffer_lifting import ( + assert_predicted_kv_aliased, inline_lifted_buffers_into_gm, lift_mutated_buffers, ) @@ -245,5 +246,245 @@ def test_inline_preserves_user_input_order(self): self.assertEqual(out.item(), 33.0) +class TestCopyBackClassification(TestCase): + """``lift_mutated_buffers`` splits mutated buffers into two kinds. + + An *eligible* ``slice_scatter`` / ``index_copy`` KV write (one the converter + lowers to an ``IKVCacheUpdateLayer`` with in-place aliased I/O) relies on that + engine aliasing and is NOT recorded for copy-back. Every other mutation -- + including a ``slice_scatter`` / ``index_copy`` that fails the converter's + eligibility (wrong rank/dim/shape) and is lowered to a non-aliasing scatter -- + has no engine aliasing, so its new value is re-appended as a trailing graph + output and its buffer name recorded in + ``gm.meta['_copyback_mutation_buffers']`` for the exporters to reclassify as + a BUFFER_MUTATION. + """ + + def test_kv_slice_scatter_write_no_copyback(self): + """A slice-assignment KV write is aliased downstream, not copied back.""" + + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("cache", torch.zeros(2, 4, 16, 8)) + + def forward(self, x): + self.cache[:, :, 3:4, :] = x + return self.cache.sum() + + gm = _ep_module_decomposed(M(), (torch.ones(2, 4, 1, 8),)) + new_gm, lifted = lift_mutated_buffers(gm) + self.assertEqual(len(lifted), 1) + self.assertEqual(new_gm.meta["_copyback_mutation_buffers"], []) + # Predicted-KV binding recorded so compile() can assert it actually aliases. + self.assertEqual(new_gm.meta["_predicted_kv_bindings"], ["buf_cache"]) + + def test_kv_index_copy_write_no_copyback(self): + """An eligible ``index_copy`` KV write (4-D static cache, dim=2, batch 1, + single-position source) is aliased by the KV converter, so it is not + copied back.""" + + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("cache", torch.zeros(1, 4, 16, 8)) + + def forward(self, x): + self.cache.index_copy_(2, torch.tensor([3]), x) + return self.cache.sum() + + gm = _ep_module_decomposed(M(), (torch.ones(1, 4, 1, 8),)) + new_gm, lifted = lift_mutated_buffers(gm) + self.assertEqual(len(lifted), 1) + self.assertEqual(new_gm.meta["_copyback_mutation_buffers"], []) + + def test_non_kv_mutation_recorded_for_copyback(self): + """A non-KV in-place mutation is recorded for copy-back by buffer name.""" + + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("state", torch.zeros(4)) + + def forward(self, x): + self.state.add_(x) + return self.state.sum() + + gm = _ep_module_decomposed(M(), (torch.ones(4),)) + new_gm, lifted = lift_mutated_buffers(gm) + self.assertEqual(len(lifted), 1) + self.assertEqual(new_gm.meta["_copyback_mutation_buffers"], ["state"]) + # A non-KV mutation is not predicted to alias, so nothing to assert later. + self.assertEqual(new_gm.meta["_predicted_kv_bindings"], []) + + def test_copyback_value_appended_as_last_output(self): + """The non-KV new value is re-attached as a trailing graph output so it + survives DCE; at the lift stage it is the LAST output and equals the + updated buffer.""" + + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("state", torch.zeros(4)) + + def forward(self, x): + self.state.add_(x) + return self.state.sum() + + x = torch.arange(4, dtype=torch.float32) + gm = _ep_module_decomposed(M(), (x.clone(),)) + new_gm, lifted = lift_mutated_buffers(gm) + _, buf_name, buf_tensor = lifted[0] + self.assertEqual(buf_name, "state") + + out = new_gm(x.clone(), buf_tensor.clone()) + self.assertIsInstance(out, tuple) + # Last output is the copy-back value == state + x. + self.assertTrue(torch.allclose(out[-1], buf_tensor + x)) + + def test_index_put_is_copyback_not_kv(self): + """Regression: ``index_put`` has no aliasing converter, so it must fall + into copy-back rather than being dropped in expectation of aliasing.""" + + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("state", torch.zeros(4, 8)) + + def forward(self, x): + self.state[torch.tensor([1, 3])] = x + return self.state.sum() + + gm = _ep_module_decomposed(M(), (torch.ones(2, 8),)) + new_gm, lifted = lift_mutated_buffers(gm) + self.assertEqual(len(lifted), 1) + self.assertEqual(new_gm.meta["_copyback_mutation_buffers"], ["state"]) + + def test_mixed_kv_and_copyback(self): + """One KV buffer + one non-KV buffer: only the non-KV one is recorded.""" + + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("cache", torch.zeros(2, 4, 16, 8)) + self.register_buffer("state", torch.zeros(4)) + + def forward(self, x_kv, x_state): + self.cache[:, :, 3:4, :] = x_kv + self.state.add_(x_state) + return self.cache.sum() + self.state.sum() + + gm = _ep_module_decomposed(M(), (torch.ones(2, 4, 1, 8), torch.ones(4))) + new_gm, lifted = lift_mutated_buffers(gm) + self.assertEqual(len(lifted), 2) + self.assertEqual(new_gm.meta["_copyback_mutation_buffers"], ["state"]) + + def test_ineligible_index_copy_is_copyback(self): + """Regression: an ``index_copy`` the KV converter cannot alias (here a 2-D + cache / dim 0, not the 4-D dim=2 layout ``IKVCacheUpdateLayer`` requires) + is lowered to a non-aliasing scatter, so its write-back must be preserved + as copy-back rather than dropped.""" + + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("cache", torch.zeros(4, 8)) + + def forward(self, x): + self.cache.index_copy_(0, torch.tensor([2]), x) + return self.cache.sum() + + gm = _ep_module_decomposed(M(), (torch.ones(1, 8),)) + new_gm, lifted = lift_mutated_buffers(gm) + self.assertEqual(len(lifted), 1) + self.assertEqual(new_gm.meta["_copyback_mutation_buffers"], ["cache"]) + + def test_ineligible_slice_scatter_is_copyback(self): + """Regression: a ``slice_scatter`` on a KV-shaped 4-D cache but the wrong + axis (dim 1, not dim 2) is not IKVCacheUpdateLayer-eligible, so it is + lowered to a non-aliasing scatter and must fall to copy-back rather than + being dropped.""" + + class M(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("cache", torch.zeros(2, 16, 4, 8)) + + def forward(self, x): + self.cache[:, 3:4, :, :] = x # write on dim 1, not the seq dim 2 + return self.cache.sum() + + gm = _ep_module_decomposed(M(), (torch.ones(2, 1, 4, 8),)) + new_gm, lifted = lift_mutated_buffers(gm) + self.assertEqual(len(lifted), 1) + self.assertEqual(new_gm.meta["_copyback_mutation_buffers"], ["cache"]) + + +class _FakeEngine: + """Stand-in for a compiled TRT submodule exposing an ``aliased_io`` map.""" + + def __init__(self, aliased_io): + self.aliased_io = aliased_io + + +class _FakeGM: + """Stand-in for a compiled GraphModule whose children are TRT engines.""" + + def __init__(self, children): + self._children = children + + def named_children(self): + return list(self._children.items()) + + +class TestPredictedKvAssertion(TestCase): + """`assert_predicted_kv_aliased` is the ground-truth backstop for the + pre-conversion KV prediction: every write predicted to alias must actually + appear in a compiled engine's `aliased_io`, else its write-back would be + silently dropped.""" + + def test_passes_when_predicted_kv_is_aliased(self): + gm = _FakeGM( + { + "_run_on_acc_0": _FakeEngine( + {"out_k": ("buf_k_cache", "kv_cache_update")} + ) + } + ) + # buf_k_cache is aliased -> no error. + assert_predicted_kv_aliased(gm, ["buf_k_cache"]) + + def test_raises_when_predicted_kv_not_aliased(self): + # Predicted KV for buf_conv_state, but the engine aliased only buf_k_cache + # (the converter emitted no IKVCacheUpdateLayer for conv_state) -> must + # raise rather than silently drop the write-back. + gm = _FakeGM( + { + "_run_on_acc_0": _FakeEngine( + {"out_k": ("buf_k_cache", "kv_cache_update")} + ) + } + ) + with self.assertRaises(RuntimeError): + assert_predicted_kv_aliased(gm, ["buf_conv_state"]) + + def test_aggregates_aliased_io_across_engines(self): + gm = _FakeGM( + { + "_run_on_acc_0": _FakeEngine( + {"out_k": ("buf_k_cache", "kv_cache_update")} + ), + "_run_on_acc_1": _FakeEngine( + {"out_v": ("buf_v_cache", "kv_cache_update")} + ), + } + ) + # Both predicted-KV bindings are aliased across the two engines -> no error. + assert_predicted_kv_aliased(gm, ["buf_k_cache", "buf_v_cache"]) + + def test_noop_when_no_prediction(self): + assert_predicted_kv_aliased(_FakeGM({}), []) + + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/models/test_exporter_inlining.py b/tests/py/dynamo/models/test_exporter_inlining.py new file mode 100644 index 0000000000..0793c24213 --- /dev/null +++ b/tests/py/dynamo/models/test_exporter_inlining.py @@ -0,0 +1,127 @@ +"""Unit tests for the legacy dynamo exporter's submodule inlining +(torch_tensorrt.dynamo._exporter). These run on plain fx graphs and need neither a +GPU nor a TensorRT build.""" + +import operator + +import pytest +import torch +from torch_tensorrt.dynamo._exporter import inline_torch_modules + + +@pytest.mark.unit +def test_inline_torch_modules_wires_inputs_by_position(): + """inline_torch_modules must wire a _run_on_gpu submodule's inputs from the + call_module args by POSITION, not by matching placeholder names to graph nodes. + + Regression: the old name-matching path bound a submodule input to a same-named + but unrelated graph node, rewiring a consumer to the wrong producer (and, for a + submodule mixing graph-input and computed-intermediate inputs, leaking the + latter as spurious graph placeholders). Here the submodule's first input + placeholder is named "y", colliding with the parent's second input "y" even + though the first *argument* is the parent's "x"; positional wiring must ignore + the collision. Subtraction makes the input order observable. + """ + # Submodule: out = first - second. First placeholder deliberately named "y". + sub_graph = torch.fx.Graph() + first = sub_graph.placeholder("y") + second = sub_graph.placeholder("z") + sub_graph.output(sub_graph.call_function(torch.sub, (first, second))) + submodule = torch.fx.GraphModule(torch.nn.Module(), sub_graph) + + # Parent: inputs (x, y); call _run_on_gpu_0(x, y) -> expected x - y. + parent_graph = torch.fx.Graph() + x = parent_graph.placeholder("x") + y = parent_graph.placeholder("y") + root = torch.nn.Module() + root.add_module("_run_on_gpu_0", submodule) + call = parent_graph.call_module("_run_on_gpu_0", (x, y)) + parent_graph.output(call) + parent = torch.fx.GraphModule(root, parent_graph) + + n_placeholders_before = sum(1 for n in parent.graph.nodes if n.op == "placeholder") + + inline_torch_modules(parent) + parent.recompile() + + # No spurious placeholders leaked by the inlining. + assert ( + sum(1 for n in parent.graph.nodes if n.op == "placeholder") + == n_placeholders_before + ) + # No call_module node survives (the submodule was inlined). + assert not any(n.op == "call_module" for n in parent.graph.nodes) + # Positional wiring: first input <- x, second input <- y, so out == x - y. + out = parent(torch.tensor(5.0), torch.tensor(3.0)) + assert torch.allclose(out, torch.tensor(2.0)) + + +@pytest.mark.unit +def test_inline_torch_modules_preserves_all_submodule_outputs(): + """A multi-output _run_on_gpu submodule must keep every output wired to its + consumer after inlining. Regression: a mis-wired input orphaned one submodule + output, which dead-code elimination then pruned, leaving a downstream consumer + (or, in the hybrid case, a TensorRT engine) short an output at runtime. + """ + # Submodule returns (a + b, a - b); both outputs are consumed downstream. + sub_graph = torch.fx.Graph() + a = sub_graph.placeholder("a") + b = sub_graph.placeholder("b") + add = sub_graph.call_function(torch.add, (a, b)) + sub = sub_graph.call_function(torch.sub, (a, b)) + sub_graph.output((add, sub)) + submodule = torch.fx.GraphModule(torch.nn.Module(), sub_graph) + + parent_graph = torch.fx.Graph() + x = parent_graph.placeholder("x") + y = parent_graph.placeholder("y") + root = torch.nn.Module() + root.add_module("_run_on_gpu_0", submodule) + call = parent_graph.call_module("_run_on_gpu_0", (x, y)) + o0 = parent_graph.call_function(operator.getitem, (call, 0)) + o1 = parent_graph.call_function(operator.getitem, (call, 1)) + # Consume both outputs: (a+b) * (a-b). + parent_graph.output(parent_graph.call_function(torch.mul, (o0, o1))) + parent = torch.fx.GraphModule(root, parent_graph) + + inline_torch_modules(parent) + parent.recompile() + + # (x+y)*(x-y) == x^2 - y^2 ; with x=5, y=3 -> 25 - 9 = 16. + out = parent(torch.tensor(5.0), torch.tensor(3.0)) + assert torch.allclose(out, torch.tensor(16.0)) + + +@pytest.mark.unit +def test_inline_torch_modules_computed_intermediate_inputs(): + """A _run_on_gpu submodule whose inputs are computed intermediates (not top-level + graph placeholders, and not name-matching any graph node) must inline correctly. + This is the case the old zero-duplicate path handled; positional wiring preserves + it (and it is the shape that leaked spurious placeholders in the mixed case). + """ + # Submodule: out = m + n. Names don't collide with anything in the parent. + sub_graph = torch.fx.Graph() + m = sub_graph.placeholder("m") + n = sub_graph.placeholder("n") + sub_graph.output(sub_graph.call_function(torch.add, (m, n))) + submodule = torch.fx.GraphModule(torch.nn.Module(), sub_graph) + + # Parent: x -> c0 = x*2, c1 = x+1; call _run_on_gpu_0(c0, c1) -> c0 + c1. + parent_graph = torch.fx.Graph() + x = parent_graph.placeholder("x") + c0 = parent_graph.call_function(torch.mul, (x, 2)) + c1 = parent_graph.call_function(torch.add, (x, 1)) + root = torch.nn.Module() + root.add_module("_run_on_gpu_0", submodule) + call = parent_graph.call_module("_run_on_gpu_0", (c0, c1)) + parent_graph.output(call) + parent = torch.fx.GraphModule(root, parent_graph) + + inline_torch_modules(parent) + parent.recompile() + + # No spurious placeholders leaked; the computed intermediates stay in-graph. + assert sum(1 for node in parent.graph.nodes if node.op == "placeholder") == 1 + # out = (x*2) + (x+1); x=5 -> 10 + 6 = 16. + out = parent(torch.tensor(5.0)) + assert torch.allclose(out, torch.tensor(16.0))