From 98cc621c3e1c73b4b421153b9a0da7c1307455ce Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Wed, 29 Jul 2026 16:09:56 -0700 Subject: [PATCH 01/11] 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 | 119 ++++++------ .../dynamo/models/test_exporter_inlining.py | 175 ++++++++++++++++++ 2 files changed, 234 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..f3514c68c3 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,40 @@ 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 (or, when none are supplied, + # positionally from the graph placeholders) and the flat graph outputs. With + # inputs supplied this matches retrace=True, which traces the same flat module + # and likewise cannot re-nest. + 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 {} + if not example_args and not example_kwargs and input_nodes: + # Reachable with no example inputs (save(retrace=False) passes + # arg_inputs=()); flattening them gives a 0-leaf spec that mismatches + # the placeholders and fails only later, so size it from them instead. + in_spec = pytree.tree_flatten((tuple(range(len(input_nodes))), {}))[1] + else: + in_spec = pytree.tree_flatten((example_args, example_kwargs))[1] + out_spec = pytree.tree_flatten(tuple(output_nodes))[1] + assert in_spec.num_leaves == len(input_nodes), ( + f"create_trt_exp_program: in_spec has {in_spec.num_leaves} leaves but " + f"the graph has {len(input_nodes)} input placeholder(s)" + ) + 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..351bd84ae4 --- /dev/null +++ b/tests/py/dynamo/models/test_exporter_inlining.py @@ -0,0 +1,175 @@ +"""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 ( + create_trt_exp_program, + 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)) + + +@pytest.mark.unit +def test_create_trt_exp_program_rebuilds_in_spec_without_inputs(): + """create_trt_exp_program must rebuild a correct in_spec on the plain-CodeGen + fallback even when no example inputs are supplied. + + Regression: torch_tensorrt.save(retrace=False) passes arg_inputs=() by + contract. Flattening an empty () produced a 0-leaf in_spec while the graph kept + its placeholders, so the ExportedProgram built and saved but failed later -- + ep.module()(x) raised "Trying to flatten user inputs ..." and + output_format="executorch" failed inside to_edge. The no-input path must + instead rebuild the spec positionally from the placeholders. Reachable in the + normal flow because lift_mutated_buffers sets a plain CodeGen for any + mutated-buffer (KV-cache) model. + """ + + class M(torch.nn.Module): + def forward(self, x): + return x.relu() + + # ExportedProgram.module() carries a _PyTreeCodeGen plus a _guards_fn node; + # lift_mutated_buffers strips both (replacing the codegen with a plain one) for + # mutated-buffer models. Replicate that so the fallback branch is exercised. + gm = torch.export.export(M().eval(), (torch.randn(3, 4),)).module() + for node in list(gm.graph.nodes): + if node.op == "call_module" and node.target == "_guards_fn": + gm.graph.erase_node(node) + break + gm.graph.set_codegen(torch.fx.graph.CodeGen()) + gm.graph.lint() + gm.recompile() + + # retrace=False passes arg_inputs=() -- the no-input path. + ep = create_trt_exp_program(gm, arg_inputs=()) + + n_user_inputs = sum( + 1 for s in ep.graph_signature.input_specs if s.kind.name == "USER_INPUT" + ) + assert n_user_inputs == 1 + + x = torch.randn(3, 4) + out = ep.module()(x) + out = out[0] if isinstance(out, (tuple, list)) else out + assert torch.allclose(out, x.relu()) From 753a845b481d10eca068d6ece9aaf7a54d7f0b7a Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Thu, 13 Aug 2026 16:51:25 -0700 Subject: [PATCH 02/11] feat(executorch): caller-owned KV-cache for the TensorRT delegate Adds end-to-end caller-owned KV-cache support to the ExecuTorch TensorRT delegate: the KV buffers are owned by the caller above the delegate and threaded in as mutable-buffer delegate args, instead of being self-allocated inside a (stateless) TensorRT engine. Runtime + serialization (delegate): - serialize each engine's aliased (KV-cache / in-place) I/O into the delegate blob (serialization.py, backend.py, TensorRTBlobHeader.{h,cpp}); - at runtime bind each aliased TRT output binding to its aliased input's caller-provided pointer (in-place) and reflect the result into the delegate output EValue -- a no-op when the memory planner already aliased the two (TensorRTBackend.{h,cpp}). Export/lowering (torch_tensorrt): - expose each engine's aliased outputs as graph-level BUFFER_MUTATIONs so ExecuTorch keeps the KV buffers as caller-owned mutable buffers: at transform time for the legacy exporter (retrace=False), and via a post-export pass (_declare_aliased_kv_mutations_on_ep) for torch.export (retrace=True), which otherwise truncates the aliased outputs at the fx boundary; - keep delegate-mutated buffers above the delegate in TensorRTPartitioner (tag_constant_data would otherwise freeze them as constants). The retrace=True pass runs for exported_program as well as executorch. The truncation happens at the fx boundary for every output format, so declaring only on the executorch path left an exported_program saved with the mutation absent from its signature while the engine still updated the cache in place. It is declared before _normalize_engine_constants_to_python, which rewrites the engine constants the pass reads aliased_io from. retrace=False was already correct for every format via create_trt_exp_program. aot_inductor stays undeclared and warns: whether an aliased in-place mutation survives functionalization under inductor is unverified. Tests cover serialization round-trip, the exposure-flag dispatch across both retrace modes, the buffer-mutation declaration, and the partitioner un-tagging. --- .../executorch/TensorRTBackend.h | 15 + .../executorch/TensorRTBlobHeader.h | 11 + .../executorch/TensorRTBackend.cpp | 225 ++++++++++++++- .../executorch/TensorRTBlobHeader.cpp | 79 +++++ examples/executorch_reference_runner/BUILD | 12 + .../CMakeLists.txt | 15 + .../executorch_reference_runner/README.md | 27 ++ .../kv_cache_decode_check.cpp | 209 ++++++++++++++ .../export_kv_cache_decode.py | 115 ++++++++ py/torch_tensorrt/_compile.py | 26 ++ py/torch_tensorrt/dynamo/_exporter.py | 269 +++++++++++++++++- .../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 +++ 20 files changed, 1341 insertions(+), 20 deletions(-) create mode 100644 examples/executorch_reference_runner/kv_cache_decode_check.cpp create mode 100644 examples/torchtrt_executorch_example/export_kv_cache_decode.py 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..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 e335748b33..afd3fe5d8e 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 @@ -307,6 +308,98 @@ 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. + // ICudaEngine::getAliasedInputTensor is TensorRT 10.15+; on older TRT (e.g. + // Jetson's 10.13) skip the cross-check and trust the persisted aliased_io + // from the blob (recorded by Torch-TensorRT at export). +#if NV_TENSORRT_MAJOR > 10 || (NV_TENSORRT_MAJOR == 10 && NV_TENSORRT_MINOR >= 15) + 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; + } +#endif + } 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; @@ -343,9 +436,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; } @@ -413,16 +514,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()); @@ -451,6 +558,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) { @@ -490,6 +617,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; @@ -517,9 +645,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); @@ -527,7 +704,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. @@ -574,7 +750,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)) { @@ -595,6 +771,24 @@ 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)); + // enqueueV3 already submitted engine work to `stream`, and inflight_pending + // is not armed until the end of the happy path -- drain now so a later + // execute() or the destructor never reconfigures/frees exec_ctx while this + // enqueue is still running. + (void)cudaStreamSynchronize(stream); + engine->inflight_pending = false; + 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 @@ -605,10 +799,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) { @@ -617,6 +818,10 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* "TensorRTBackend::execute: D2H copy failed for output %zu: %s", output.first, cudaGetErrorString(cuda_err)); + // Same drain as the reflect path: the enqueue is in flight and the + // stream sync below hasn't run yet, so drain before the early return. + (void)cudaStreamSynchronize(stream); + engine->inflight_pending = false; return Error::InvalidProgram; } } 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 8335092725..861f0dc108 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", "load_model.py", "main.cpp", ], @@ -21,3 +22,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 ddf635f297..3cdadaa992 100644 --- a/examples/executorch_reference_runner/README.md +++ b/examples/executorch_reference_runner/README.md @@ -125,3 +125,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 64b7412200..908c6265e0 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -1147,7 +1147,30 @@ def _extract_tensor(obj: Any) -> Any: strict=False, ) + from torch_tensorrt.dynamo._exporter import ( + _declare_aliased_kv_mutations_on_ep, + ) + + # aot_inductor is left undeclared: whether an aliased in-place mutation + # survives functionalization under inductor is unverified. + if output_format == "aot_inductor" and any( + getattr(sub, "aliased_io", None) + for _sub_name, sub in module.named_modules() + ): + logger.warning( + "Module has TensorRT engine(s) with aliased I/O (e.g. KV-cache), " + "but output_format='aot_inductor' does not declare those aliased " + "outputs as buffer mutations. The saved program's signature will " + "not reflect the in-place update." + ) + if output_format == "exported_program": + # torch.export truncates the engines' aliased KV outputs at the fx + # boundary for every format, so the mutation has to be re-declared or + # the signature omits an update the engine performs. Must precede + # normalization, which rewrites the engine constants this pass reads + # aliased_io from. + exp_program = _declare_aliased_kv_mutations_on_ep(exp_program) _normalize_engine_constants_to_python(exp_program) function_overload_with_kwargs( torch.export.save, @@ -1168,6 +1191,9 @@ def _extract_tensor(obj: Any) -> Any: package_path=file_path, ) elif output_format == "executorch": + # 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 f3514c68c3..3b105286b3 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 ] @@ -396,7 +419,9 @@ def create_trt_exp_program( in_spec = pytree.tree_flatten((tuple(range(len(input_nodes))), {}))[1] else: 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] assert in_spec.num_leaves == len(input_nodes), ( f"create_trt_exp_program: in_spec has {in_spec.num_leaves} leaves but " f"the graph has {len(input_nodes)} input placeholder(s)" @@ -490,8 +515,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. @@ -553,12 +685,137 @@ 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] + # These two are internal invariant violations: aliased_io is built from the + # engine's own bindings, so an aliased output must map to a real input arg. + # If it doesn't, the mutation can't be wired and its in-place update would be + # silently dropped (a corrupted cache) -- fail loudly rather than degrade. + if in_name not in in_names: + raise RuntimeError( + f"Aliased output {out_name!r} references input {in_name!r} which is " + "not an engine input binding -- the engine's aliased_io map is " + "inconsistent with its input bindings." + ) + ii = in_names.index(in_name) + if ii >= len(input_arg_nodes): + raise RuntimeError( + f"Aliased output {out_name!r} -> input {in_name!r} maps to arg index " + f"{ii}, out of range for {len(input_arg_nodes)} delegate args -- the " + "engine's aliased_io map is inconsistent with the delegate args." + ) + 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 == {} From 3b5b0c52c4eeac6ce01e014b5205d31122673b5a Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Thu, 13 Aug 2026 17:31:20 -0700 Subject: [PATCH 03/11] fix(executorch): keep example_inputs and verifiers when re-declaring mutations _declare_aliased_kv_mutations_on_ep rebuilds the ExportedProgram to attach the new output specs, but reconstructed only root/graph/signature/state_dict/ range_constraints/module_call_graph/constants. example_inputs and verifiers are not recoverable from the graph and reset to their defaults when omitted, so the pass was not a drop-in replacement for the program it rewrites. That was invisible while the pass ran only on the executorch path, since to_edge does not read either. Declaring mutations for exported_program as well makes it reachable: torch.export.save then persists a program whose example inputs are silently gone. AOTI refuses such a program outright ("exported_program. example_inputs is required to be set in order for AOTInductor compilation"), so this also has to be fixed before that format can ever declare mutations. Carry both through. The stub programs in test_kv_cache_export.py now model them, and the capturing test asserts they reach the constructor. Reported by cehongwang in review. --- py/torch_tensorrt/dynamo/_exporter.py | 2 ++ .../dynamo/executorch/test_kv_cache_export.py | 19 +++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index 3b105286b3..dda2940186 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -616,7 +616,9 @@ def _estr(engine_info: List[Any], idx: int) -> str: state_dict=exp_program.state_dict, range_constraints=exp_program.range_constraints, module_call_graph=exp_program.module_call_graph, + example_inputs=exp_program.example_inputs, constants=exp_program.constants, + verifiers=exp_program.verifiers, ) diff --git a/tests/py/dynamo/executorch/test_kv_cache_export.py b/tests/py/dynamo/executorch/test_kv_cache_export.py index b1387c3497..623197767f 100644 --- a/tests/py/dynamo/executorch/test_kv_cache_export.py +++ b/tests/py/dynamo/executorch/test_kv_cache_export.py @@ -17,6 +17,7 @@ import pytest import torch +from torch._export.verifier import Verifier from torch.export.exported_program import ( InputKind, InputSpec, @@ -26,6 +27,10 @@ ) from torch_tensorrt.dynamo import _exporter as E +# A real ExportedProgram carries these, so the stubs below do too -- the pass has to +# hand them on rather than let them reset to their defaults. +_EXAMPLE_INPUTS = ((torch.randn(2),), {}) + @pytest.mark.unit @pytest.mark.parametrize("use_legacy, expected_expose", [(True, True), (False, False)]) @@ -116,6 +121,8 @@ def test_declare_aliased_kv_mutations_declares_buffer_mutation(monkeypatch): range_constraints={}, module_call_graph=[], constants={}, + example_inputs=_EXAMPLE_INPUTS, + verifiers=[Verifier], ) info = ["x"] * (L.ALIASED_IO_IDX + 1) @@ -155,3 +162,15 @@ def __init__(self, **kwargs): # The engine's meta['val'] is extended to cover the previously-dropped output. assert len(eng.meta["val"]) == 2 + + # A rewrite replaces only the graph and signature; every other field of the + # source program has to come through untouched. + for field in ( + "state_dict", + "range_constraints", + "module_call_graph", + "constants", + "example_inputs", + "verifiers", + ): + assert captured[field] is getattr(ep, field) From 340b42df3f46da3898bef032e8dc5fdecbe45e4c Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Fri, 14 Aug 2026 11:55:30 -0700 Subject: [PATCH 04/11] fix(executorch): make the aliased-mutation declaration idempotent Whether an aliased KV output gets declared a BUFFER_MUTATION is decided in two places: the exporter (the legacy one declares at transform time, via create_trt_exp_program) and save()'s per-output-format branch. Nothing reconciles them, so a program that arrives already declared is declared a second time. The duplicate spec then fails the ExportedProgram verifier's output ordering check. Seed already_exposed from the incoming signature's BUFFER_MUTATION targets rather than an empty set, so the pass skips buffers that are already declared and returns the program untouched when nothing new remains. That covers every combination that reaches it, including exported_program with use_legacy_exporter=True, which the preceding commit made reachable. The added test drives the pass on an already-declared program with ExportedProgram monkeypatched to raise, so a regression fails on "rebuilt the program" rather than on some later verifier complaint. The no-op fixture grows an output_specs field, which a real ExportGraphSignature always has. Reported by shoumikhin in review. --- py/torch_tensorrt/dynamo/_exporter.py | 11 ++- .../dynamo/executorch/test_kv_cache_export.py | 88 ++++++++++++++++++- 2 files changed, 97 insertions(+), 2 deletions(-) diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index dda2940186..3863de6394 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -552,7 +552,16 @@ def _estr(engine_info: List[Any], idx: int) -> str: 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() + # Seeded from the incoming signature, not empty: exposure is decided both by + # the exporter (the legacy one declares at transform time) and by save()'s + # per-format branch, so this pass can run on a program whose mutations are + # already declared. Re-declaring appends a second spec for the same buffer and + # the ExportedProgram verifier then rejects the output ordering. + already_exposed: Set[str] = { + spec.target + for spec in sig.output_specs + if spec.kind == OutputKind.BUFFER_MUTATION and spec.target + } 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: diff --git a/tests/py/dynamo/executorch/test_kv_cache_export.py b/tests/py/dynamo/executorch/test_kv_cache_export.py index 623197767f..9255f6345a 100644 --- a/tests/py/dynamo/executorch/test_kv_cache_export.py +++ b/tests/py/dynamo/executorch/test_kv_cache_export.py @@ -69,11 +69,97 @@ def test_declare_aliased_kv_mutations_is_noop_without_engines(): ep = SimpleNamespace( graph_module=gm, - graph_signature=SimpleNamespace(inputs_to_buffers={}), + graph_signature=SimpleNamespace(inputs_to_buffers={}, output_specs=[]), ) assert E._declare_aliased_kv_mutations_on_ep(ep) is ep +@pytest.mark.unit +def test_declare_aliased_kv_mutations_is_idempotent(monkeypatch): + """Running on a program whose mutation is already declared must be a no-op. + + Exposure is decided both by the exporter (the legacy one declares at transform + time) and by save()'s per-format branch, so this pass can receive a program that + already carries the spec. Declaring again appends a second BUFFER_MUTATION for + the same buffer, which fails the ExportedProgram verifier's output ordering. + """ + 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 + + 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)) + kv_out = g.call_function(operator.getitem, (eng, 1)) + g.output((kv_out, 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, buf_val] + user_out.meta["val"] = out_val + kv_out.meta["val"] = buf_val + gm = torch.fx.GraphModule(torch.nn.Module(), g) + + # k_0 already declared -- what the legacy exporter leaves behind. + 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.BUFFER_MUTATION, TensorArgument(name=kv_out.name), "k_0" + ), + 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={}, + example_inputs=_EXAMPLE_INPUTS, + verifiers=[Verifier], + ) + + 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"], + ) + + def _must_not_rebuild(**kwargs): + raise AssertionError( + "the pass rebuilt the program even though k_0 was already declared" + ) + + monkeypatch.setattr(E, "ExportedProgram", _must_not_rebuild) + + assert E._declare_aliased_kv_mutations_on_ep(ep) is ep + assert [spec.kind for spec in sig.output_specs] == [ + OutputKind.BUFFER_MUTATION, + OutputKind.USER_OUTPUT, + ] + + @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 From 8155c315cf187362cbc393b5b503e0664132b408 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sat, 15 Aug 2026 01:01:07 -0700 Subject: [PATCH 05/11] fix(executorch): correct the aliased-output reflect comments Four comments describe a zero-copy path where the memory planner places the delegate's output slot on the aliased input, so the reflect is skipped. That never happens: the delegate input and its aliased output are live at the same time, so the planner cannot co-locate them. Instrumenting the branch over a 30-layer export confirms it -- dst == bind_ptr was false for all 120 aliased outputs across a prefill and a decode step. Two of them go further and call the skipped case the common fast path, which inverts what the code does: every aliased output reflects, and a model with aliased outputs therefore always syncs before returning. Say what actually happens in all four (TensorRTBackend.h, and the reflect list, the reflect loop and the must_sync rationale in TensorRTBackend.cpp). The dst != bind_ptr check itself stays -- it is unreachable today, but it is what stops a self-copy if planning ever changes -- and its comment now says so instead of advertising a fast path. Also two comment fixes found in the same pass: "These two" in _exporter.py referred forward to checks the reader had not reached yet, and a test comment said "previously-dropped" where it meant the output torch.export truncates. No behaviour change. Reported by shoumikhin in review. --- .../torch_tensorrt/executorch/TensorRTBackend.h | 3 +-- .../executorch/TensorRTBackend.cpp | 17 +++++++++-------- py/torch_tensorrt/dynamo/_exporter.py | 5 +++-- .../dynamo/executorch/test_kv_cache_export.py | 2 +- 4 files changed, 14 insertions(+), 13 deletions(-) diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index 7bdc38a1e0..a32ea8c967 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -65,8 +65,7 @@ struct EngineHandle { // 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). + // delegate output EValue, which ExecuTorch's write-back copy_ then reads. std::vector output_aliased_input_idx; // Per input binding [0..num_inputs): true if any output aliases this input, so // its in-place (KV/user) update must land in the caller-owned storage. Built at diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index afd3fe5d8e..fb60bf5aa9 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -650,7 +650,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // 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). + // updated cache. std::vector> aliased_reflects; for (size_t o = 0; o < num_outputs; ++o) { const std::string& name = engine->output_binding_names[o]; @@ -689,6 +689,9 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* (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; + // dst != bind_ptr guards against issuing a self-copy. The memory planner does + // not currently place the delegate's output slot on the aliased input -- the + // two are live at the same time -- so this holds for every aliased output. if (dst != nullptr && dst != bind_ptr) { aliased_reflects.emplace_back(dst, bind_ptr, et_alias_out.nbytes()); } @@ -772,8 +775,7 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } // 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). + // EValue (D2D on the same stream, after the engine work). 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) { @@ -799,11 +801,10 @@ 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. - // 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. + // An aliased reflect enqueues the engine's in-place update into the delegate + // output EValue on `stream`; ExecuTorch's buffer-mutation copy_ reads that EValue + // after execute() returns, so the reflect must complete first. A model with + // aliased outputs therefore always syncs here. 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; diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index 3863de6394..95578dbb6e 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -756,8 +756,9 @@ def _expose_aliased_buffer_mutations( if out_name not in aliased_io: continue in_name = aliased_io[out_name][0] - # These two are internal invariant violations: aliased_io is built from the - # engine's own bindings, so an aliased output must map to a real input arg. + # The two checks below are internal invariant violations: aliased_io is + # built from the engine's own bindings, so an aliased output must map to a + # real input arg. # If it doesn't, the mutation can't be wired and its in-place update would be # silently dropped (a corrupted cache) -- fail loudly rather than degrade. if in_name not in in_names: diff --git a/tests/py/dynamo/executorch/test_kv_cache_export.py b/tests/py/dynamo/executorch/test_kv_cache_export.py index 9255f6345a..5d8888494d 100644 --- a/tests/py/dynamo/executorch/test_kv_cache_export.py +++ b/tests/py/dynamo/executorch/test_kv_cache_export.py @@ -246,7 +246,7 @@ def __init__(self, **kwargs): 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. + # The engine's meta['val'] is extended to cover the truncated aliased output. assert len(eng.meta["val"]) == 2 # A rewrite replaces only the graph and signature; every other field of the From bc7142c7f83837c9c240657b29d8a9c1afd04440 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sat, 15 Aug 2026 01:12:19 -0700 Subject: [PATCH 06/11] fix(executorch): treat aliased-output resize failures as fatal The aliased-output branch discarded the resize_tensor result and skipped the resize altogether when the rank was out of range, while the sibling non-aliased branch treats both as fatal. et_alias_out.nbytes() is read on the next line and sizes both the reflect D2D copy and, through the delegate output EValue, ExecuTorch's write-back copy_. A dynamic aliased output that outgrows its planned size would therefore move the stale planned byte count in both, silently truncating the cache update rather than failing. Mirror the sibling branch: reject an out-of-range rank and propagate a resize_tensor error. Reported by shoumikhin in review. --- .../executorch/TensorRTBackend.cpp | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index fb60bf5aa9..b576b4b746 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -680,13 +680,22 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return Error::InvalidArgument; } exec_aten::Tensor et_alias_out = out_arg->toTensor(); + // nbytes() below sizes both the reflect copy and ExecuTorch's write-back, so + // the tensor has to carry the shape TRT inferred before either of them reads it. 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)}); + if (a_dims.nbDims < 0 || a_dims.nbDims > nvinfer1::Dims::MAX_DIMS) { + ET_LOG(Error, "TensorRTBackend::execute: invalid rank for aliased output '%s'", name.c_str()); + return Error::InvalidState; + } + 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]); + } + Error a_resize_err = + executorch::runtime::resize_tensor(et_alias_out, {a_sizes, static_cast(a_dims.nbDims)}); + if (a_resize_err != Error::Ok) { + ET_LOG(Error, "TensorRTBackend::execute: resize_tensor failed for aliased output '%s'", name.c_str()); + return a_resize_err; } void* dst = et_alias_out.nbytes() > 0 ? et_alias_out.mutable_data_ptr() : nullptr; // dst != bind_ptr guards against issuing a self-copy. The memory planner does From 17a814a30537b39d453c3eb7dae71e082e94870b Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sat, 15 Aug 2026 01:24:24 -0700 Subject: [PATCH 07/11] fix(executorch): declare aliased KV mutations on the retrace=False path too Whether an aliased KV output is declared a BUFFER_MUTATION depends on two independent switches. save() picks the exporter (retrace, plus an optional use_legacy_exporter override) and only the legacy exporter exposes the mutations, at transform time; save()'s per-format branch declares them for everything else. The retrace=False branch did neither, so retrace=False with use_legacy_exporter=False produced a program that silently omits an update the engine performs -- no declaration and no diagnostic. Run the declaration pass on the exported_program and executorch branches there as well. The preceding commit made the pass skip buffers that already carry a spec, so this is correct for either exporter: the legacy one keeps declaring at transform time and the pass returns its program untouched. aot_inductor stays undeclared on both paths, as before. The added test drives save() over both formats and both exporters and asserts the pass runs exactly once; against the previous commit all four parametrizations fail. Reported by shoumikhin in review. --- py/torch_tensorrt/_compile.py | 28 +++++++---- .../dynamo/executorch/test_kv_cache_export.py | 46 +++++++++++++++++++ 2 files changed, 65 insertions(+), 9 deletions(-) diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index 908c6265e0..fcb01c2201 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -1040,6 +1040,14 @@ def _extract_tensor(obj: Any) -> Any: **kwargs, ) else: + # torch.export truncates the engines' aliased KV outputs at the fx + # boundary, so a retraced program's signature omits an update the engine + # performs unless the mutations are re-declared; the legacy exporter + # instead exposes them at transform time. The declaration pass skips + # buffers that already carry a spec, so every branch below can run it + # whichever exporter produced the program. aot_inductor is left + # undeclared: whether an aliased in-place mutation survives + # functionalization under inductor is unverified. if not retrace: from torch_tensorrt.dynamo._exporter import export @@ -1060,7 +1068,15 @@ def _extract_tensor(obj: Any) -> Any: dynamic_shapes=dynamic_shapes, use_legacy_exporter=_use_legacy, ) + + from torch_tensorrt.dynamo._exporter import ( + _declare_aliased_kv_mutations_on_ep, + ) + if output_format == "exported_program": + # Must precede normalization, which rewrites the engine constants + # this pass reads aliased_io from. + exp_program = _declare_aliased_kv_mutations_on_ep(exp_program) _normalize_engine_constants_to_python(exp_program) function_overload_with_kwargs( torch.export.save, @@ -1081,6 +1097,7 @@ def _extract_tensor(obj: Any) -> Any: package_path=file_path, ) elif output_format == "executorch": + exp_program = _declare_aliased_kv_mutations_on_ep(exp_program) _save_as_executorch( exp_program, file_path, @@ -1151,8 +1168,6 @@ def _extract_tensor(obj: Any) -> Any: _declare_aliased_kv_mutations_on_ep, ) - # aot_inductor is left undeclared: whether an aliased in-place mutation - # survives functionalization under inductor is unverified. if output_format == "aot_inductor" and any( getattr(sub, "aliased_io", None) for _sub_name, sub in module.named_modules() @@ -1165,11 +1180,8 @@ def _extract_tensor(obj: Any) -> Any: ) if output_format == "exported_program": - # torch.export truncates the engines' aliased KV outputs at the fx - # boundary for every format, so the mutation has to be re-declared or - # the signature omits an update the engine performs. Must precede - # normalization, which rewrites the engine constants this pass reads - # aliased_io from. + # Must precede normalization, which rewrites the engine constants + # this pass reads aliased_io from. exp_program = _declare_aliased_kv_mutations_on_ep(exp_program) _normalize_engine_constants_to_python(exp_program) function_overload_with_kwargs( @@ -1191,8 +1203,6 @@ def _extract_tensor(obj: Any) -> Any: package_path=file_path, ) elif output_format == "executorch": - # 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, diff --git a/tests/py/dynamo/executorch/test_kv_cache_export.py b/tests/py/dynamo/executorch/test_kv_cache_export.py index 5d8888494d..d2d49aecc6 100644 --- a/tests/py/dynamo/executorch/test_kv_cache_export.py +++ b/tests/py/dynamo/executorch/test_kv_cache_export.py @@ -260,3 +260,49 @@ def __init__(self, **kwargs): "verifiers", ): assert captured[field] is getattr(ep, field) + + +@pytest.mark.unit +@pytest.mark.parametrize("output_format", ["exported_program", "executorch"]) +@pytest.mark.parametrize("use_legacy", [True, False]) +def test_save_declares_aliased_mutations_without_retrace( + monkeypatch, tmp_path, output_format, use_legacy +): + """retrace=False must declare the aliased KV mutations as well. + + Only the legacy exporter exposes them at transform time, so with + use_legacy_exporter=False nothing declares them and the saved program omits an + update the engine performs. The pass skips buffers already declared, so save() + can run it for either exporter. + """ + pytest.importorskip("executorch.exir") + import torch_tensorrt + from torch_tensorrt import _compile as C + from torch_tensorrt.dynamo import _exporter as E + + sentinel = object() + declared = [] + + def _declare(ep, **kwargs): + declared.append(ep) + return ep + + monkeypatch.setattr(E, "export", lambda *a, **k: sentinel) + monkeypatch.setattr(E, "_declare_aliased_kv_mutations_on_ep", _declare) + monkeypatch.setattr(C, "_normalize_engine_constants_to_python", lambda ep: None) + monkeypatch.setattr(C, "_save_as_executorch", lambda *a, **k: None) + monkeypatch.setattr(torch.export, "save", lambda *a, **k: None) + + g = torch.fx.Graph() + g.output((g.placeholder("x"),)) + gm = torch.fx.GraphModule(torch.nn.Module(), g) + + torch_tensorrt.save( + gm, + str(tmp_path / "out.pte"), + output_format=output_format, + retrace=False, + use_legacy_exporter=use_legacy, + ) + + assert declared == [sentinel] From 09b8701ab095a0f8ed3df0e7e612e0f6df38342d Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sat, 15 Aug 2026 01:50:04 -0700 Subject: [PATCH 08/11] feat(executorch): check delegate outputs are in engine binding order The runtime binds delegate output i to output_binding_names[i]. That holds because the partition's outputs are getitem(engine_node, i) in index order, and nothing verifies it: inputs are checked in _reorder_input_names_for_executorch, outputs were not. arrange_graph_outputs moves buffer mutations ahead of user outputs and is a no-op here only while the mutated buffers stay above the delegate, so a regression there would swap the serialized names silently. Validate the correspondence in preprocess. A single-output engine returned unwrapped is accepted -- one binding has no order to get wrong -- and anything else must be one getitem per binding, in index order. _build_edge_program only ever emitted `output((engine_node,))`, including for its three-output-binding case, which is not a shape that can occur: a three-tuple cannot be consumed as one value. It now emits one getitem per output binding, so the fixtures model what the backend actually receives. Reported by shoumikhin in review. --- py/torch_tensorrt/executorch/backend.py | 53 +++++++++++++++ tests/py/dynamo/executorch/test_backend.py | 78 +++++++++++++++++++++- 2 files changed, 128 insertions(+), 3 deletions(-) diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index fbc973ae7d..5f9ab00809 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -1,5 +1,6 @@ # ExecuTorch TensorRT backend: serialize engines to a libtorch-free runtime blob. +import operator from typing import Any, List, final import torch @@ -253,6 +254,57 @@ def _reorder_input_names_for_executorch( return [input_names[i] for i in order] +def _validate_output_binding_order( + edge_program: ExportedProgram, engine_node: Any, output_names: List[str] +) -> None: + """Check the delegate's outputs are the engine's output bindings, in order. + + The runtime binds output ``i`` to ``output_binding_names[i]``, and nothing + downstream re-derives that correspondence: it holds because the partition's + outputs are ``getitem(engine_node, i)`` in index order. A pass that reordered + them -- ``arrange_graph_outputs`` moves buffer mutations ahead of user outputs, + which only stays a no-op here while the mutated buffers are kept above the + delegate -- would swap the names silently. Inputs cannot rely on position at + all and recover their order by node identity in + ``_reorder_input_names_for_executorch``. + """ + output_node = next( + node for node in edge_program.graph_module.graph.nodes if node.op == "output" + ) + out_args = list(output_node.args[0]) + # A single-output engine is returned directly rather than through a getitem, + # and one binding has no order to get wrong. + if len(out_args) == 1 and out_args[0] is engine_node: + if len(output_names) != 1: + raise ValueError( + "TensorRT ExecuTorch backend: the delegate returns the engine node " + f"directly but the engine declares {len(output_names)} output " + "bindings; only a single-output engine can be returned unwrapped." + ) + return + indices: List[Any] = [] + for node in out_args: + if ( + not isinstance(node, torch.fx.Node) + or node.op != "call_function" + or node.target is not operator.getitem + or node.args[0] is not engine_node + ): + raise ValueError( + "TensorRT ExecuTorch backend: delegate output " + f"{getattr(node, 'name', node)!r} is not a getitem of the engine " + "node; cannot establish a reliable output binding order." + ) + indices.append(node.args[1]) + if indices != list(range(len(output_names))): + raise ValueError( + "TensorRT ExecuTorch backend: delegate outputs map to engine output " + f"indices {indices}, expected {list(range(len(output_names)))} -- the " + "runtime binds output i to output_binding_names[i], so a permuted or " + "incomplete output list would bind the wrong tensors." + ) + + def _get_str(engine_info: List[Any], index: int, default: str = "") -> str: if index < 0 or index >= len(engine_info): return default @@ -304,6 +356,7 @@ def preprocess( output_names = _split_binding_names( _get_str(engine_info, OUTPUT_BINDING_NAMES_IDX) ) + _validate_output_binding_order(edge_program, engine_node, output_names) io_bindings = [ TensorRTIOBinding(name=name, is_input=True) for name in input_names ] + [TensorRTIOBinding(name=name, is_input=False) for name in output_names] diff --git a/tests/py/dynamo/executorch/test_backend.py b/tests/py/dynamo/executorch/test_backend.py index 1c070f318e..c48408c982 100644 --- a/tests/py/dynamo/executorch/test_backend.py +++ b/tests/py/dynamo/executorch/test_backend.py @@ -1,4 +1,5 @@ import ast +import operator from pathlib import Path from types import SimpleNamespace @@ -18,12 +19,12 @@ SERIALIZATION_LEN, ) from torch_tensorrt.executorch.backend import ( # noqa: E402 - _get_engine_info_from_edge_program, TensorRTBackend, + _get_engine_info_from_edge_program, ) from torch_tensorrt.executorch.serialization import ( # noqa: E402 - deserialize_engine, TENSORRT_MAGIC, + deserialize_engine, ) @@ -73,7 +74,18 @@ def _build_edge_program( _ENGINE_OP, (engine_inputs, *engine_info), ) - graph.output((engine_node,)) + # A multi-output engine is consumed through one getitem per output binding, in + # index order; a single-output engine is returned unwrapped. + out_names = [n for n in str(engine_info[OUTPUT_BINDING_NAMES_IDX]).split("%") if n] + if len(out_names) > 1: + graph.output( + tuple( + graph.call_function(operator.getitem, (engine_node, i)) + for i in range(len(out_names)) + ) + ) + else: + graph.output((engine_node,)) graph_signature = SimpleNamespace( input_specs=[ @@ -338,3 +350,63 @@ def test_preprocess_preserves_output_binding_order(): False, False, ] + + +def _engine_partition(output_indices): + """A one-engine partition whose graph outputs are getitem(engine, i) for each i + in `output_indices`, in that order.""" + import operator + + g = torch.fx.Graph() + x = g.placeholder("x") + engine = g.call_function( + torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default, ([x],) + ) + g.output( + tuple(g.call_function(operator.getitem, (engine, i)) for i in output_indices) + ) + gm = torch.fx.GraphModule(torch.nn.Module(), g) + return SimpleNamespace(graph_module=gm), engine + + +@pytest.mark.unit +def test_validate_output_binding_order_accepts_index_order(): + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([0, 1, 2]) + _validate_output_binding_order(ep, engine, ["out0", "out1", "out2"]) + + +@pytest.mark.unit +def test_validate_output_binding_order_rejects_permuted_outputs(): + """The runtime binds output i to output_binding_names[i]. A pass that moved a + mutation output ahead of the user outputs would rename them silently.""" + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([2, 0, 1]) + with pytest.raises(ValueError, match="engine output indices"): + _validate_output_binding_order(ep, engine, ["out0", "out1", "out2"]) + + +@pytest.mark.unit +def test_validate_output_binding_order_rejects_dropped_output(): + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([0, 1]) + with pytest.raises(ValueError, match="engine output indices"): + _validate_output_binding_order(ep, engine, ["out0", "out1", "out2"]) + + +@pytest.mark.unit +def test_validate_output_binding_order_accepts_unwrapped_single_output(): + """A single-output engine is returned directly, not through a getitem.""" + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + g = torch.fx.Graph() + x = g.placeholder("x") + engine = g.call_function( + torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default, ([x],) + ) + g.output((engine,)) + ep = SimpleNamespace(graph_module=torch.fx.GraphModule(torch.nn.Module(), g)) + _validate_output_binding_order(ep, engine, ["out"]) From 2384ff6de1d1fbcc8c3b8928138aa2fac083d04b Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sat, 15 Aug 2026 02:05:55 -0700 Subject: [PATCH 09/11] fix(executorch): bump the blob magic when the metadata carries aliased_io aliased_io changes what a blob means. A parser that predates it binds each aliased output to its own allocation instead of the input it aliases, so it does not fail -- it returns wrong results. The magic is the only field that parser validates, so it is the only thing that can make the skew fail closed. Emit TR02 when metadata.aliased_io is non-empty and keep TR01 otherwise, rather than bumping unconditionally: a blob with no alias map means exactly what it meant before, so it stays loadable by an older runtime. Both magics are accepted on read, so new runtimes still load existing artifacts. Verified against a real older build rather than a simulated one: a TR02 blob loads and produces the expected KV-persistence result on a runtime built from this branch, and a runtime built before the change rejects the same blob with "failed to parse TensorRT blob". Reported by shoumikhin in review. --- .../executorch/TensorRTBackend.h | 2 +- .../executorch/TensorRTBackend.cpp | 2 +- .../executorch/TensorRTBlobHeader.cpp | 6 ++- py/torch_tensorrt/executorch/backend.py | 2 +- py/torch_tensorrt/executorch/serialization.py | 11 ++++- .../test_executorch_blob_header.cpp | 29 +++++++++++- .../dynamo/executorch/test_serialization.py | 44 +++++++++++++++++++ 7 files changed, 88 insertions(+), 8 deletions(-) diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index a32ea8c967..de84014edf 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -5,7 +5,7 @@ * LICENSE file in the root directory of this source tree. * * ExecuTorch backend delegate that runs TensorRT engines serialized by - * torch_tensorrt. The processed blob uses the standalone TR01 wire format from + * torch_tensorrt. The processed blob uses the standalone wire format from * py/torch_tensorrt/executorch/serialization.py and is parsed directly here. * This runtime path intentionally does not depend on the legacy * Torch-TensorRT C++ runtime or libtorch. diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index b576b4b746..e790fa5bdf 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -251,7 +251,7 @@ Result TensorRTBackend::init( TensorRTBlobHeader header; if (!TensorRTBlobHeader::parse(processed->data(), processed->size(), header)) { - ET_LOG(Error, "TensorRTBackend::init: failed to parse TR01 TensorRT blob"); + ET_LOG(Error, "TensorRTBackend::init: failed to parse TensorRT blob"); return Error::InvalidProgram; } diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp index c5c49e9a26..1dede1776e 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp @@ -9,7 +9,10 @@ namespace torch_tensorrt { namespace executorch_backend { namespace { +// TR02 marks a blob whose metadata carries aliased_io; TR01 is one without. This +// parser handles aliased_io, so it accepts either. constexpr char TENSORRT_MAGIC[4] = {'T', 'R', '0', '1'}; +constexpr char TENSORRT_MAGIC_ALIASED_IO[4] = {'T', 'R', '0', '2'}; constexpr uint32_t METADATA_OFFSET_FIELD_OFFSET = 4; constexpr uint32_t METADATA_SIZE_FIELD_OFFSET = 8; constexpr uint32_t ENGINE_OFFSET_FIELD_OFFSET = 12; @@ -324,7 +327,8 @@ bool TensorRTBlobHeader::parse(const void* data, std::size_t size, TensorRTBlobH } const auto* bytes = static_cast(data); - if (std::memcmp(bytes, TENSORRT_MAGIC, sizeof(TENSORRT_MAGIC)) != 0) { + if (std::memcmp(bytes, TENSORRT_MAGIC, sizeof(TENSORRT_MAGIC)) != 0 && + std::memcmp(bytes, TENSORRT_MAGIC_ALIASED_IO, sizeof(TENSORRT_MAGIC_ALIASED_IO)) != 0) { return false; } diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index 5f9ab00809..f533eb2adc 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -321,7 +321,7 @@ class TensorRTBackend(BackendDetails): # type: ignore[misc] """Backend that serializes TensorRT engines for the native ExecuTorch runtime. The partition contains a single execute_engine node; we extract the engine - and metadata and encode them as a standalone TR01 blob. The C++ runtime + and metadata and encode them as a standalone blob. The C++ runtime backend parses that blob directly without the legacy Torch-TensorRT C++ runtime. """ diff --git a/py/torch_tensorrt/executorch/serialization.py b/py/torch_tensorrt/executorch/serialization.py index 47d8fe9f4b..ba7ccac10c 100644 --- a/py/torch_tensorrt/executorch/serialization.py +++ b/py/torch_tensorrt/executorch/serialization.py @@ -12,7 +12,14 @@ from dataclasses import dataclass, field from typing import Dict, List, Tuple +# A blob carrying aliased_io means something different to a parser that ignores it: +# the runtime would bind each aliased output to its own allocation instead of the +# input it aliases, and return wrong results rather than fail. The only field a +# pre-aliasing parser validates is the magic, so aliased blobs carry TR02 and it +# rejects them. Blobs without aliased_io keep TR01 and stay readable everywhere. TENSORRT_MAGIC = b"TR01" +TENSORRT_MAGIC_ALIASED_IO = b"TR02" +SUPPORTED_MAGICS = (TENSORRT_MAGIC, TENSORRT_MAGIC_ALIASED_IO) HEADER_FORMAT = "<4sIIIQ8s" HEADER_SIZE = struct.calcsize(HEADER_FORMAT) @@ -99,7 +106,7 @@ def serialize_engine(engine_bytes: bytes, metadata: TensorRTBlobMetadata) -> byt reserved = b"\x01" + b"\x00" * 7 header = struct.pack( HEADER_FORMAT, - TENSORRT_MAGIC, + TENSORRT_MAGIC_ALIASED_IO if metadata.aliased_io else TENSORRT_MAGIC, metadata_offset, len(metadata_json), engine_offset, @@ -116,7 +123,7 @@ def deserialize_engine(blob: bytes) -> Tuple[bytes, TensorRTBlobMetadata]: magic, metadata_offset, metadata_size, engine_offset, engine_size, _ = ( struct.unpack(HEADER_FORMAT, blob[:HEADER_SIZE]) ) - if magic != TENSORRT_MAGIC: + if magic not in SUPPORTED_MAGICS: raise ValueError(f"Invalid magic: {magic!r}") if engine_offset % 16 != 0: raise ValueError(f"Engine offset is not 16-byte aligned: {engine_offset}") diff --git a/tests/cpp/executorch/test_executorch_blob_header.cpp b/tests/cpp/executorch/test_executorch_blob_header.cpp index e527e990e1..74543e88b4 100644 --- a/tests/cpp/executorch/test_executorch_blob_header.cpp +++ b/tests/cpp/executorch/test_executorch_blob_header.cpp @@ -13,6 +13,7 @@ namespace executorch_backend { namespace { constexpr char TENSORRT_MAGIC[4] = {'T', 'R', '0', '1'}; +constexpr char TENSORRT_MAGIC_ALIASED_IO[4] = {'T', 'R', '0', '2'}; constexpr uint32_t METADATA_OFFSET_FIELD_OFFSET = 4; constexpr uint32_t METADATA_SIZE_FIELD_OFFSET = 8; constexpr uint32_t ENGINE_OFFSET_FIELD_OFFSET = 12; @@ -29,13 +30,16 @@ std::size_t align_up(std::size_t value, std::size_t alignment) { return ((value + alignment - 1) / alignment) * alignment; } -std::vector make_blob(const std::string& metadata, std::size_t engine_size = 4) { +std::vector make_blob( + const std::string& metadata, + std::size_t engine_size = 4, + const char* magic = TENSORRT_MAGIC) { const auto metadata_offset = static_cast(HEADER_SIZE); const auto metadata_size = static_cast(metadata.size()); const auto engine_offset = static_cast(align_up(metadata_offset + metadata_size, ENGINE_ALIGNMENT)); std::vector blob(static_cast(engine_offset) + engine_size, 0); - std::memcpy(blob.data(), TENSORRT_MAGIC, sizeof(TENSORRT_MAGIC)); + std::memcpy(blob.data(), magic, sizeof(TENSORRT_MAGIC)); write_field(blob, METADATA_OFFSET_FIELD_OFFSET, metadata_offset); write_field(blob, METADATA_SIZE_FIELD_OFFSET, metadata_size); write_field(blob, ENGINE_OFFSET_FIELD_OFFSET, engine_offset); @@ -149,6 +153,27 @@ TEST(ExecuTorchTensorRTBlobHeader, ParsesEmptyAliasedIo) { EXPECT_TRUE(header.aliased_io.empty()); } +TEST(ExecuTorchTensorRTBlobHeader, ParsesAliasedIoMagic) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + ASSERT_EQ(header.aliased_io.size(), 1u); + EXPECT_EQ(header.aliased_io[0].output_name, "out_k"); + EXPECT_EQ(header.aliased_io[0].input_name, "in_k"); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsUnknownFutureMagic) { + constexpr char kFutureMagic[4] = {'T', 'R', '0', '3'}; + const std::string metadata = R"({"io_bindings":[{"name":"x","is_input":true}]})"; + const auto blob = make_blob(metadata, 4, kFutureMagic); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + } // namespace } // namespace executorch_backend } // namespace torch_tensorrt diff --git a/tests/py/dynamo/executorch/test_serialization.py b/tests/py/dynamo/executorch/test_serialization.py index 9eddee2cd8..c025b11c02 100644 --- a/tests/py/dynamo/executorch/test_serialization.py +++ b/tests/py/dynamo/executorch/test_serialization.py @@ -4,6 +4,7 @@ from torch_tensorrt.executorch.serialization import ( HEADER_SIZE, TENSORRT_MAGIC, + TENSORRT_MAGIC_ALIASED_IO, TensorRTBlobMetadata, TensorRTIOBinding, deserialize_engine, @@ -76,3 +77,46 @@ def test_metadata_from_json_without_aliased_io_defaults_empty(): restored = TensorRTBlobMetadata.from_json(json.dumps(data).encode("utf-8")) assert restored.aliased_io == {} + + +@pytest.mark.unit +def test_aliased_io_blob_uses_the_bumped_magic(): + """aliased_io changes what a blob means: a parser that ignores it binds each + aliased output to its own allocation instead of the input it aliases, and + returns wrong results. The magic is the only field such a parser validates, so + aliased blobs must not present as TR01.""" + metadata = TensorRTBlobMetadata( + io_bindings=[ + TensorRTIOBinding(name="x", is_input=True), + TensorRTIOBinding(name="out_k", is_input=False), + ], + aliased_io={"out_k": ("x", "kv_cache_update")}, + ) + blob = serialize_engine(b"engine-bytes", metadata) + assert blob[:4] == TENSORRT_MAGIC_ALIASED_IO + assert blob[:4] != TENSORRT_MAGIC + + +@pytest.mark.unit +def test_blob_without_aliased_io_keeps_the_original_magic(): + """Nothing about a non-aliased blob is new, so it stays loadable by a parser + that predates aliased_io.""" + metadata = TensorRTBlobMetadata( + io_bindings=[TensorRTIOBinding(name="x", is_input=True)] + ) + assert serialize_engine(b"engine-bytes", metadata)[:4] == TENSORRT_MAGIC + + +@pytest.mark.unit +@pytest.mark.parametrize("aliased_io", [{}, {"out_k": ("x", "kv_cache_update")}]) +def test_deserialize_accepts_both_magics(aliased_io): + metadata = TensorRTBlobMetadata( + io_bindings=[ + TensorRTIOBinding(name="x", is_input=True), + TensorRTIOBinding(name="out_k", is_input=False), + ], + aliased_io=aliased_io, + ) + engine, parsed = deserialize_engine(serialize_engine(b"engine-bytes", metadata)) + assert engine == b"engine-bytes" + assert parsed.aliased_io == aliased_io From 6b6374688824e21019c72a2facf87e4030fafc96 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sat, 15 Aug 2026 02:07:34 -0700 Subject: [PATCH 10/11] ci(executorch): exercise kv_cache_decode_check in the reference-runner verify The caller-owned KV path had no CI coverage: kv_cache_decode_check ships in the release tarball and is defined in the packaged CMake project, but nothing built or ran it, so a regression in the aliased binding would only surface downstream. Export a decode .pte alongside the static-shape one and pass it to the verify script as an optional second argument. When present the script builds kv_cache_decode_check from the unpacked tarball, runs it, and requires the persistence assertion to pass; the same no-libtorch link check the example runner gets is applied to it. Without the argument the script behaves as before. Also assert the tarball ships kv_cache_decode_check.cpp, next to the existing entries, so the packaging contract is checked rather than assumed. Reported by shoumikhin in review. --- .../verify-executorch-reference-runner.sh | 38 +++++++++++++++++-- .github/workflows/executorch-test-linux.yml | 5 ++- 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/.github/scripts/verify-executorch-reference-runner.sh b/.github/scripts/verify-executorch-reference-runner.sh index b59ed905ec..c195b5dfa8 100755 --- a/.github/scripts/verify-executorch-reference-runner.sh +++ b/.github/scripts/verify-executorch-reference-runner.sh @@ -14,6 +14,10 @@ set +x # First argument: path to an existing .pte model. # EXECUTORCH_SOURCE_DIR=/path/to/executorch # +# Optional second argument: path to a caller-owned KV-cache decode .pte (see +# examples/torchtrt_executorch_example/export_kv_cache_decode.py). When given, +# kv_cache_decode_check is built and run against it as well. +# # Optional: # TensorRT_ROOT=/path/to/extracted/TensorRT # If unset, the script reuses Bazel's fetched TensorRT SDK when available @@ -29,8 +33,8 @@ set +x repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" cd "${repo_root}" -if [[ $# -ne 1 ]]; then - echo "Usage: $0 PATH_TO_MODEL.pte" >&2 +if [[ $# -lt 1 || $# -gt 2 ]]; then + echo "Usage: $0 PATH_TO_MODEL.pte [PATH_TO_KV_CACHE_DECODE.pte]" >&2 exit 1 fi model_path="$1" @@ -38,6 +42,11 @@ if [[ ! -f "${model_path}" ]]; then echo "ExecuTorch model not found: ${model_path}" >&2 exit 1 fi +kv_model_path="${2:-}" +if [[ -n "${kv_model_path}" && ! -f "${kv_model_path}" ]]; then + echo "KV-cache decode model not found: ${kv_model_path}" >&2 + exit 1 +fi python_executable="${PYTHON_EXECUTABLE:-}" if [[ -z "${python_executable}" ]]; then @@ -261,6 +270,7 @@ require_tar_entry() { require_tar_entry "torch_tensorrt/src/torch_tensorrt/executorch/CMakeLists.txt" require_tar_entry "torch_tensorrt/examples/executorch_reference_runner/CMakeLists.txt" +require_tar_entry "torch_tensorrt/examples/executorch_reference_runner/kv_cache_decode_check.cpp" require_tar_entry "torch_tensorrt/BUILD" export TORCH_TENSORRT_ROOT="${verify_root}/torch_tensorrt" @@ -282,8 +292,13 @@ fi cmake "${cmake_args[@]}" +build_targets=(example_executorch_runner) +if [[ -n "${kv_model_path}" ]]; then + build_targets+=(kv_cache_decode_check) +fi + cmake --build "${verify_root}/build-executorch-reference-runner" \ - --target example_executorch_runner \ + --target "${build_targets[@]}" \ -j"${MAX_JOBS:-$(nproc)}" runner_log="${verify_root}/my_runner.log" @@ -304,3 +319,20 @@ fi grep -q "Inference completed" "${runner_log}" grep -q "output\\[0\\] shape=" "${runner_log}" grep -Eq "first [0-9]+ values:.* 2\\.0000" "${runner_log}" + +if [[ -n "${kv_model_path}" ]]; then + # kv_cache_decode_check exits non-zero when a decode step does not observe the KV + # the previous step wrote; the grep additionally pins the assertion itself, so + # weakening the check inside the binary cannot quietly turn this into a no-op. + kv_check_log="${verify_root}/kv_cache_decode_check.log" + kv_check_path="${verify_root}/build-executorch-reference-runner/kv_cache_decode_check" + if command -v ldd >/dev/null 2>&1 && + ldd "${kv_check_path}" | + grep -E "libtorch|libtorch_cpu|libtorch_cuda|libc10" >&2; then + echo "kv_cache_decode_check links PyTorch/libtorch shared libraries" >&2 + exit 1 + fi + + "${kv_check_path}" --model_path="${kv_model_path}" 2>&1 | tee "${kv_check_log}" + grep -q "PASS: decode at pos=1 observed the KV written at pos=0" "${kv_check_log}" +fi diff --git a/.github/workflows/executorch-test-linux.yml b/.github/workflows/executorch-test-linux.yml index 38f37d663c..a480fee4c3 100644 --- a/.github/workflows/executorch-test-linux.yml +++ b/.github/workflows/executorch-test-linux.yml @@ -82,7 +82,10 @@ jobs: export EXECUTORCH_ROOT="${EXECUTORCH_SOURCE_DIR}" python examples/torchtrt_executorch_example/export_static_shape.py \ --model_path="${RUNNER_TEMP}/torchtrt-python.pte" + python examples/torchtrt_executorch_example/export_kv_cache_decode.py \ + --model_path="${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte" .github/scripts/verify-executorch-reference-runner.sh \ - "${RUNNER_TEMP}/torchtrt-python.pte" + "${RUNNER_TEMP}/torchtrt-python.pte" \ + "${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte" python examples/executorch_reference_runner/load_model.py \ --model_path="${RUNNER_TEMP}/torchtrt-python.pte" --num_runs=1 From 28a3ec8cb9919326ab8a5922ea97c5535b3cff17 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sat, 15 Aug 2026 03:41:04 -0700 Subject: [PATCH 11/11] docs(executorch): log the skipped aliased outputs and correct the ordering claim Four `continue`s in _declare_aliased_kv_mutations_on_ep left an aliased output undeclared without saying so, and the mismatch only surfaced later as a delegate arity error at execute. Log at each, at the level the case warrants: warn when the persisted alias map disagrees with the engine's bindings (unknown input name, or an index past the delegate args) and when the aliased input is not a registered buffer, since all three leave the engine with an output binding the delegate cannot satisfy; debug when the buffer already carries a spec, which is the expected idempotent skip. The non-aliased output path stays silent -- it is the common case, not a fault. _reorder_input_names_for_executorch's docstring also justified skipping the output reordering by claiming a TensorRT partition has no mutation outputs. That has not been true since aliased I/O landed. The order does survive lowering, but for a different reason: _keep_mutated_buffers_above_delegate keeps mutated buffers out of the delegate, so ExecuTorch records the mutation as a USER_OUTPUT and arrange_graph_outputs computes the identity permutation. Say that, and note the guarantee is conditional -- _validate_output_binding_order is what enforces it. Reported by cehongwang in review. --- py/torch_tensorrt/dynamo/_exporter.py | 38 ++++++++++++++++++++++++- py/torch_tensorrt/executorch/backend.py | 28 ++++++++++++------ 2 files changed, 57 insertions(+), 9 deletions(-) diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index 95578dbb6e..17f3a0cc0e 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -582,14 +582,50 @@ def _estr(engine_info: List[Any], idx: int) -> str: if out_name not in aliased_io: continue in_name = aliased_io[out_name][0] + # The alias map comes from the engine's own bindings, so an entry that does + # not resolve to a delegate input arg means the map and the engine disagree. + # The mutation is then left undeclared and only surfaces as a delegate arity + # error at execute, so log which output was dropped. if in_name not in in_names: + logger.warning( + "Aliased output %s references input %s, which is not an engine " + "input binding; leaving the mutation undeclared.", + out_name, + in_name, + ) continue ii = in_names.index(in_name) if ii >= len(input_nodes): + logger.warning( + "Aliased output %s maps to input index %d, out of range for %d " + "delegate args; leaving the mutation undeclared.", + out_name, + 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: + if buf_fqn is None: + # A caller-supplied cache that is not a registered buffer: there is no + # buffer to declare a mutation of, and the engine still has the output + # binding, so the delegate ends up one arg short at execute. + logger.warning( + "Aliased output %s updates %s in place, but that input is not a " + "registered buffer, so no buffer mutation is declared for it.", + out_name, + in_name, + ) + continue + if buf_fqn in already_exposed: + # Already declared, by this pass for another engine sharing the buffer + # or by the exporter at transform time. Expected, not a fault. + logger.debug( + "Buffer %s already carries a mutation spec; skipping aliased " + "output %s.", + buf_fqn, + out_name, + ) continue oi = out_names.index(out_name) while len(val_list) <= oi: diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index f533eb2adc..3978a8160f 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -221,14 +221,26 @@ def _reorder_input_names_for_executorch( first arg lists its input nodes in binding order, so sort the names by each node's slot among the graph placeholders (its runtime delegate-arg position). - Only inputs need this. Outputs are also bound positionally by the runtime, - but they are ``getitem(engine_node, idx)`` nodes whose index order equals the - engine output-binding order. ExecuTorch lowering can reorder delegate outputs - (``arrange_graph_outputs`` moves buffer-mutation outputs ahead of user - outputs), but a TensorRT delegate partition is a functional inference engine - with no mutation outputs, so that pass is a no-op here and the output order is - preserved. If a TRT partition ever produced mutation outputs, outputs would - need the same node-identity reordering as inputs. + Only inputs need this. Outputs are also bound positionally, but they are + ``getitem(engine_node, idx)`` nodes whose index order equals the engine + output-binding order, and that order survives lowering -- though not because + the partition is mutation-free. With aliased-I/O (KV-cache) support a TensorRT + partition *does* produce mutation outputs, and ``arrange_graph_outputs`` does + move buffer-mutation outputs ahead of user outputs. It stays a no-op here + because ``_keep_mutated_buffers_above_delegate`` (``partitioner.py``) strips + the ``delegation_tag`` from mutated buffer placeholders, so they stay out of + the delegate's state dict and constants; ExecuTorch's ``_get_new_signature`` + then records the mutation as a plain ``USER_OUTPUT`` rather than a + ``BUFFER_MUTATION`` (it uses the latter only when the delegate itself consumes + the buffer). The lowered submodule therefore has no mutation specs, so + ``arrange_graph_outputs`` computes the identity permutation and the getitem + indices still line up with the engine's output bindings. + + That guarantee is conditional, not structural: if a mutated buffer is ever + tagged into a delegate its spec becomes ``BUFFER_MUTATION``, the delegate's + outputs are permuted, and they would need the same node-identity reordering as + the inputs below. ``_validate_output_binding_order`` checks that correspondence + on every preprocess, so it would fail loudly rather than mis-bind. """ input_nodes = list(engine_node.args[0]) if len(input_nodes) != len(input_names):