From 464d756af3ae69dc35fc79c6e0e0c9c64da27c51 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Wed, 29 Jul 2026 16:09:56 -0700 Subject: [PATCH 1/2] fix(dynamo): correct legacy exporter (retrace=False) submodule inlining for hybrid graphs torch_tensorrt.save(retrace=False) uses the legacy dynamo exporter, which inlines the partitioned _run_on_gpu (non-TensorRT) submodules back into the graph before building an ExportedProgram. For a hybrid graph interleaving TensorRT engines with a CUDA/pytorch delegated op, inline_torch_modules wired each submodule's inputs by MATCHING placeholder names to graph nodes (get_duplicate_nodes). Name matching binds an input to a same-named but unrelated node on a collision (e.g. a submodule input placeholder name-matching 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 (an aliased engine reports "expected N args, got N-1"); and - for a submodule mixing graph-input and computed-intermediate inputs, leaks the computed intermediates as spurious graph placeholders (misclassified USER_INPUTs). Wire submodule inputs POSITIONALLY from the call_module args (gm_node.args, which is authoritative) instead of by name: let graph_copy create a fresh placeholder for each submodule input, then rewire each to submodule_inputs[i] by position and erase it. Drop get_duplicate_nodes (now unused). Also fix two torch-version-compat gaps this path hits on recent torch: - lift(): pass an explicit persistent= flag on BUFFER InputSpecs (required since 2.3). - create_trt_exp_program(): an inlined GraphModule may carry a plain fx.CodeGen (no pytree_info); fall back to specs rebuilt from the example inputs + graph outputs. With these, retrace=False export of a hybrid TensorRT+CUDA program is bit-identical to retrace=True (validated on a 2-layer int4 MoE decode: per-step argmax + logits match). Tests: tests/py/dynamo/models/test_exporter_inlining.py -- positional input wiring under a name collision, and multi-output preservation (GPU-free fx unit tests). --- py/torch_tensorrt/dynamo/_exporter.py | 108 +++++++-------- .../dynamo/models/test_exporter_inlining.py | 127 ++++++++++++++++++ 2 files changed, 175 insertions(+), 60 deletions(-) create mode 100644 tests/py/dynamo/models/test_exporter_inlining.py diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index dacbea140a..059f32075d 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -4,6 +4,7 @@ from typing import Any, Dict, Optional, Sequence, 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,6 +20,7 @@ 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 @@ -230,6 +232,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 +253,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 +270,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) @@ -403,14 +376,29 @@ 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 = pytree.tree_flatten(tuple(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, ), ) ] 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)) From 2312f427c4e601688c5bce2645faea7411318a9a Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Fri, 31 Jul 2026 15:59:25 -0700 Subject: [PATCH 2/2] feat(executorch): caller-owned KV-cache for the TensorRT delegate Adds end-to-end caller-owned KV-cache support to the ExecuTorch TensorRT delegate: the KV buffers are owned by the caller above the delegate and threaded in as mutable-buffer delegate args, instead of being self-allocated inside a (stateless) TensorRT engine. Runtime + serialization (delegate): - serialize each engine's aliased (KV-cache / in-place) I/O into the delegate blob (serialization.py, backend.py, TensorRTBlobHeader.{h,cpp}); - at runtime bind each aliased TRT output binding to its aliased input's caller-provided pointer (in-place) and reflect the result into the delegate output EValue -- a no-op when the memory planner already aliased the two (TensorRTBackend.{h,cpp}). Export/lowering (torch_tensorrt): - expose each engine's aliased outputs as graph-level BUFFER_MUTATIONs so ExecuTorch keeps the KV buffers as caller-owned mutable buffers: at transform time for the legacy exporter (retrace=False), and via a post-export pass (_declare_aliased_kv_mutations_on_ep) for torch.export (retrace=True), which otherwise truncates the aliased outputs at the fx boundary; - keep delegate-mutated buffers above the delegate in TensorRTPartitioner (tag_constant_data would otherwise freeze them as constants). Tests cover serialization round-trip, the exposure-flag dispatch across both retrace modes, the buffer-mutation declaration, and the partitioner un-tagging. --- .../executorch/TensorRTBackend.h | 10 + .../executorch/TensorRTBlobHeader.h | 11 + .../executorch/TensorRTBackend.cpp | 152 ++++++++++- .../executorch/TensorRTBlobHeader.cpp | 72 +++++ py/torch_tensorrt/_compile.py | 7 + py/torch_tensorrt/dynamo/_exporter.py | 257 +++++++++++++++++- .../runtime/meta_ops/register_meta_ops.py | 26 +- py/torch_tensorrt/executorch/backend.py | 8 + py/torch_tensorrt/executorch/partitioner.py | 23 ++ py/torch_tensorrt/executorch/serialization.py | 18 +- .../test_executorch_blob_header.cpp | 42 +++ .../dynamo/executorch/test_kv_cache_export.py | 157 +++++++++++ .../py/dynamo/executorch/test_partitioner.py | 44 ++- .../test_partitioner_target_device.py | 1 + .../dynamo/executorch/test_serialization.py | 39 +++ 15 files changed, 848 insertions(+), 19 deletions(-) create mode 100644 tests/py/dynamo/executorch/test_kv_cache_export.py diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 31383f9e17..474914f3c4 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -59,6 +59,16 @@ 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; + 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..27edaf2151 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,64 @@ 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); + 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; + } + // AliasKind::USER aliases are not shape-enforced by TensorRT (unlike + // kv_cache_update, which IKVCacheUpdateLayer guarantees), so confirm the + // aliased output and input share a shape before binding them to the same + // storage. + if (ab.kind == "user") { + 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->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 +384,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 +462,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()); @@ -472,6 +545,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 +573,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 +632,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 +678,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 +699,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 @@ -590,7 +724,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* const bool must_sync = output_staged_to_host || input_staged_from_host || !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..ee448a6ca1 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,77 @@ 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()) { + 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/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index 32ee801886..836ede7755 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -1134,6 +1134,13 @@ 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. + exp_program = _declare_aliased_kv_mutations_on_ep(exp_program) _save_as_executorch( exp_program, file_path, diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index 059f32075d..929d6c418f 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -1,7 +1,8 @@ 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 @@ -24,6 +25,8 @@ 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 @@ -72,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: @@ -97,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. @@ -115,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) @@ -363,12 +369,29 @@ def create_trt_exp_program( assert output_nodes output_nodes = output_nodes[0].args[0] + # 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 ] @@ -389,7 +412,9 @@ def create_trt_exp_program( 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 = pytree.tree_flatten(tuple(output_nodes))[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( @@ -479,8 +504,115 @@ def create_trt_exp_program( return trt_exp_program +def _declare_aliased_kv_mutations_on_ep( + exp_program: ExportedProgram, +) -> ExportedProgram: + """retrace=True post-export pass: declare each engine's aliased KV output as a + BUFFER_MUTATION of its caller-owned buffer input. + + 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. Returns + exp_program unchanged when no engine has aliased KV outputs. + """ + 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) + + if not mutation_outputs: + return exp_program + + # BUFFER_MUTATION outputs must precede USER_OUTPUTs (ExportedProgram verifier). + out_args = list(output_node.args[0]) + output_node.args = (tuple([g for g, _ in mutation_outputs] + out_args),) + gm.graph.lint() + gm.recompile() + + new_output_specs = [ + OutputSpec(OutputKind.BUFFER_MUTATION, TensorArgument(name=g.name), fqn) + for g, fqn in mutation_outputs + ] + list(sig.output_specs) + new_signature = ExportGraphSignature( + input_specs=list(sig.input_specs), output_specs=new_output_specs + ) + 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=exp_program.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. @@ -542,12 +674,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/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..b1387c3497 --- /dev/null +++ b/tests/py/dynamo/executorch/test_kv_cache_export.py @@ -0,0 +1,157 @@ +"""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 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 == {}