From aaa30847f717de2c864ff29917db423356967f69 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Mon, 17 Aug 2026 18:57:57 -0700 Subject: [PATCH 01/22] feat(executorch): let a TensorRT engine update its aliased KV buffer in place An engine with aliased I/O writes its aliased output *through* the aliased input's pointer, so by the time the engine returns the caller's buffer is already updated. Nothing in the ExecuTorch pipeline knows that, and it pays for the update twice on every execution: PropagateDevicePass hands the delegate an `_h2d_copy` staging copy instead of the buffer, and the aliased output is threaded back out so ExecuTorch can copy it into the buffer afterwards. For a KV cache that is two cache-sized copies per token. Add the two passes that remove them. They are separate functions because they run at different points and there is no single point that works: `rewire_aliased_mutations_to_buffers` runs on the exported program, before partitioning, because the partition boundary is what fixes the delegate's outputs. Pointing each aliased BUFFER_MUTATION at its own buffer placeholder says "the buffer is the result", so there is nothing to copy back and the aliased output -- now userless -- leaves the graph and the delegate with it. `unstage_aliased_buffers_pass` runs as a `to_out_var_pass`, the last hook in the window after PropagateDevicePass and before memory planning -- the window where the staging copies exist and the buffers' placement is still open. (`sym_shape_eval_pass` is a caller-supplied hook in that window too, but it runs first.) Which mutations are aliased is read from the engine's own `aliased_io`, not inferred from the graph. A copy-back mutation and a mutation of a buffer that is also an engine input are both a getitem off the engine node, indistinguishable by shape, and rewiring either would delete a real update with no error. Both silent-corruption shapes raise instead. Eliding every output of an engine would leave a delegate with no outputs, which nothing downstream reports: the runtime infers elision from a single argument count, and a delegate nothing reads is a pure node a later dead-code elimination can erase. And a marked buffer whose staging cannot be followed to CUDA would be written by the engine and then discarded, since its copy-back is already gone. Neither function is public: they are one feature and applying half of it is worse than applying none. The next commits reach them through an opt-in on export(). --- py/torch_tensorrt/executorch/_zero_copy.py | 364 ++++++++++++++++++ .../py/dynamo/executorch/test_zero_copy_kv.py | 323 ++++++++++++++++ 2 files changed, 687 insertions(+) create mode 100644 py/torch_tensorrt/executorch/_zero_copy.py create mode 100644 tests/py/dynamo/executorch/test_zero_copy_kv.py diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py new file mode 100644 index 00000000000..bac0fa59462 --- /dev/null +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -0,0 +1,364 @@ +"""Let a TensorRT engine update an aliased mutable buffer in place. + +An engine with aliased I/O (a KV cache) writes its aliased output *through* the +aliased input's pointer, so running the engine over the buffer already is the +update. Nothing in the ExecuTorch pipeline knows that, so by default the buffer +makes a full round trip on every execution: + +* ``PropagateDevicePass`` wraps every delegate input in ``et_copy._h2d_copy``, + so the delegate is handed a per-call staging copy rather than the caller's + buffer. The engine's in-place write lands in that copy. +* the aliased output is threaded back out as a delegate output, and ExecuTorch + copies it into the buffer afterwards to make the update stick. + +For a cache-sized buffer that is two copies per execution of something the +engine could have written directly. This module removes both, in the same +spirit as ``partitioner._keep_mutated_buffers_above_delegate``: let the upstream +pass run, then correct its output for the case Torch-TensorRT owns. + +The two halves are inseparable and run at different times: + +* :func:`rewire_aliased_mutations_to_buffers`, on the exported program before + partitioning, drops the copy-back by declaring that the buffer *is* the + mutation's result. The aliased output then has no user and disappears from the + partition. +* :func:`unstage_aliased_buffers_pass`, as a ``to_out_var_pass``, drops the + staging so the engine writes the caller's buffer rather than a copy. + +Applying only the first would leave the engine writing a discarded staging copy +with nothing to copy back -- the buffer would simply never update. So neither +half may be applied alone, and neither is part of the public API on its own: a +caller opts into both together or into neither. +""" + +import logging +import operator +from typing import Any, Dict, List, NamedTuple, Optional + +import torch +from torch.fx import Node + +_LOGGER = logging.getLogger(__name__) + + +def _engine_info_str(engine_info: List[Any], index: int) -> str: + if index < 0 or index >= len(engine_info) or engine_info[index] is None: + return "" + value = engine_info[index] + return value.decode("utf-8", "replace") if isinstance(value, bytes) else str(value) + + +def _aliased_inputs_by_output_index( + exported_program: Any, engine_node: Node +) -> Dict[int, Node]: + """Map each aliased output index of one engine to the input it writes in place. + + Reads the engine's own ``aliased_io`` rather than inferring aliasing from the + graph. The graph cannot tell the difference: an aliased KV mutation and a + copy-back mutation are both a ``getitem`` off the engine node whose buffer is + also an engine input, and rewiring a copy-back would silently drop a real + update. An output binding absent from the delegate's inputs is skipped + rather than reported -- ``_declare_aliased_kv_mutations_on_ep`` has already + warned about that same engine, and there is no mutation to rewire either way. + """ + 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 + + engine_info = _get_engine_info_for_node(exported_program, engine_node) + aliased_io = deserialize_aliased_io(_engine_info_str(engine_info, ALIASED_IO_IDX)) + if not aliased_io: + return {} + input_names = deserialize_binding_names( + _engine_info_str(engine_info, INPUT_BINDING_NAMES_IDX) + ) + output_names = deserialize_binding_names( + _engine_info_str(engine_info, OUTPUT_BINDING_NAMES_IDX) + ) + input_nodes = list(engine_node.args[0]) + + aliased: Dict[int, Node] = {} + for output_index, output_name in enumerate(output_names): + entry = aliased_io.get(output_name) + if entry is None: + continue + input_name = entry[0] + if input_name not in input_names: + continue + input_index = input_names.index(input_name) + if input_index >= len(input_nodes): + continue + aliased[output_index] = input_nodes[input_index] + return aliased + + +class _AliasedMutation(NamedTuple): + """One BUFFER_MUTATION an engine satisfies by writing the buffer in place.""" + + placeholder: Node # the buffer, as a graph input + aliased_output: Node # getitem(engine, i) currently standing in for it + engine: Node # the execute_engine call that performs the write + + +def _aliased_buffer_mutations( + exported_program: Any, +) -> Dict[int, _AliasedMutation]: + """Find the BUFFER_MUTATIONs an engine performs in place. + + Returns ``{index into graph_signature.output_specs: _AliasedMutation}``. + A mutation qualifies only when its value is ``getitem(engine_node, i)`` and + the engine declares output ``i`` as aliased onto that very buffer, so a + buffer mutated by an op outside the engine, or copied back out of one, is + left alone. + """ + from torch.export.graph_signature import OutputKind + + graph_module = exported_program.graph_module + signature = exported_program.graph_signature + execute_engine = torch.ops.tensorrt.execute_engine.default + + buffer_placeholders = { + fqn: node + for node in graph_module.graph.nodes + if node.op == "placeholder" + and (fqn := signature.inputs_to_buffers.get(node.name)) is not None + } + output_args = list(graph_module.graph.output_node().args[0]) + aliased_by_engine: Dict[Node, Dict[int, Node]] = {} + + mutations: Dict[int, _AliasedMutation] = {} + for spec_index, spec in enumerate(signature.output_specs): + if spec.kind != OutputKind.BUFFER_MUTATION or spec_index >= len(output_args): + continue + placeholder = buffer_placeholders.get(spec.target) + if placeholder is None: + continue + value = output_args[spec_index] + if ( + not isinstance(value, Node) + or value.op != "call_function" + or value.target is not operator.getitem + ): + continue + engine_node = value.args[0] + if ( + not isinstance(engine_node, Node) + or engine_node.op != "call_function" + or engine_node.target is not execute_engine + ): + continue + if engine_node not in aliased_by_engine: + aliased_by_engine[engine_node] = _aliased_inputs_by_output_index( + exported_program, engine_node + ) + if aliased_by_engine[engine_node].get(value.args[1]) is placeholder: + mutations[spec_index] = _AliasedMutation( + placeholder=placeholder, aliased_output=value, engine=engine_node + ) + return mutations + + +def rewire_aliased_mutations_to_buffers(exported_program: Any) -> int: + """Declare that an aliased buffer *is* its own mutation result. + + Export declares an aliased KV mutation as a ``getitem`` off the engine node: + the engine's aliased output, surfaced as a value. ExecuTorch implements that + mutation by copying the value back into the buffer, which is the copy this + removes. Repointing the mutation at the buffer placeholder leaves nothing to + copy, and with no other user the ``getitem`` dies -- so the aliased output + also leaves the partition and the delegate never receives an argument for it. + + This must run before partitioning, because it is the partition boundary that + freezes which outputs the delegate has. It must also run after export has + declared the aliased mutations, since it works from those declarations; each + placeholder it rewires is marked for + :func:`unstage_aliased_buffers_pass`, which cannot re-derive the aliasing + once lowering has turned the engine into an opaque blob. + + On its own this is not correct: ExecuTorch still stages the buffer, so the + engine's in-place write would land in per-call scratch and, with the + copy-back gone, be lost. It is only correct paired with the un-staging pass. + + Returns the number of mutations rewired. + """ + from torch.export.graph_signature import ( + ExportGraphSignature, + OutputKind, + OutputSpec, + TensorArgument, + ) + + graph_module = exported_program.graph_module + signature = exported_program.graph_signature + mutations = _aliased_buffer_mutations(exported_program) + if not mutations: + _LOGGER.debug("no aliased buffer mutations to rewire") + return 0 + + elided_by_engine: Dict[Node, List[Node]] = {} + output_node = graph_module.graph.output_node() + output_args = list(output_node.args[0]) + output_specs = list(signature.output_specs) + for spec_index, mutation in mutations.items(): + # Marked on the node rather than read back off the engine because the + # un-staging pass runs after lowering, where the engine's aliased_io is no + # longer reachable from the graph: it has become an opaque delegate blob. + mutation.placeholder.meta["_torch_tensorrt_aliased_buffer"] = True + output_args[spec_index] = mutation.placeholder + output_specs[spec_index] = OutputSpec( + OutputKind.BUFFER_MUTATION, + TensorArgument(name=mutation.placeholder.name), + output_specs[spec_index].target, + ) + elided_by_engine.setdefault(mutation.engine, []).append(mutation.aliased_output) + + output_node.args = (tuple(output_args),) + # Dropping every output of an engine would leave a delegate with no outputs. + # Nothing downstream reports that shape: the runtime infers elision from a + # single argument count, which a zero-output delegate satisfies, and a + # delegate nothing reads is a pure node that a later graph-wide dead-code + # elimination can erase, taking the computation with it. Stop here instead. + # The eliminate_dead_code() below does not erase this engine node: unlike + # the delegate, an execute_engine node is impure to FX and survives with no + # users. + for engine, elided in elided_by_engine.items(): + if all(user in elided and not user.users for user in engine.users): + raise RuntimeError( + "TensorRT zero-copy KV: every output of engine node " + f"'{engine.name}' is an aliased buffer written in place, so " + "eliding them would leave the delegate with no outputs at all. " + "This shape is not supported; export this method without " + "zero_copy_kv." + ) + graph_module.graph.eliminate_dead_code() + graph_module.graph.lint() + graph_module.recompile() + exported_program._graph_signature = ExportGraphSignature( + input_specs=list(signature.input_specs), output_specs=output_specs + ) + _LOGGER.debug("rewired %d aliased mutation(s) to their buffers", len(mutations)) + return len(mutations) + + +def _is_tensorrt_delegate(graph_module: torch.fx.GraphModule, node: Node) -> bool: + """True when ``node`` is a call_delegate dispatching to the TensorRT backend. + + Only a TensorRT engine promises the aliased-binding write; another backend's + delegate may legitimately need the staging copy. + """ + from executorch.exir.delegate import executorch_call_delegate + from torch_tensorrt.executorch.backend import TensorRTBackend + + if node.op != "call_function" or node.target is not executorch_call_delegate: + return False + lowered = node.args[0] if node.args else None + if not isinstance(lowered, Node) or lowered.op != "get_attr": + return False + module = getattr(graph_module, lowered.target, None) + return bool(getattr(module, "backend_id", None) == TensorRTBackend.__name__) + + +def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: + """Route TensorRT delegate inputs from their staging copy back to the buffer. + + A delegate input qualifies only when it is an ``_h2d_copy`` of a placeholder + carrying the mark left by :func:`rewire_aliased_mutations_to_buffers`. Every + other input keeps its staging, including a mutable buffer the engine does + not write in place. + + The placeholder's spec takes over the staging copy's device, so memory + planning puts the buffer in the delegate's device arena instead of a host + one -- which is what makes the engine's write land somewhere the caller can + still see afterwards. + + Raises when a marked buffer cannot be un-staged, because the alternative is + silence: its copy-back has already been removed, so leaving the staging in + place means the engine writes a scratch tensor and the buffer never updates. + + Returns the number of delegate inputs un-staged. + """ + from executorch.exir.schema import DeviceType + + h2d_copy = torch.ops.et_copy._h2d_copy.default + unstaged = 0 + for node in list(graph_module.graph.nodes): + if not _is_tensorrt_delegate(graph_module, node): + continue + new_args = list(node.args) + for i, arg in enumerate(node.args[1:], start=1): + if not isinstance(arg, Node) or arg.target is not h2d_copy: + continue + source = arg.args[0] + if not isinstance(source, Node) or source.op != "placeholder": + continue + if not source.meta.get("_torch_tensorrt_aliased_buffer"): + continue # not written in place; it needs its staging copy + staged_spec = arg.meta.get("spec") + source_spec = source.meta.get("spec") + if staged_spec is None or source_spec is None: + raise RuntimeError( + "TensorRT zero-copy KV: no TensorSpec on the staging copy of " + f"buffer '{source.name}', so it cannot be moved to the " + "delegate's device. The TensorRT engine writes this buffer in " + "place and its copy-back has already been removed, so the " + "update would be lost. This pass has to run as the " + "ExecutorchBackendConfig to_out_var_pass, which is where the " + "specs exist; torch_tensorrt.executorch.zero_copy_backend_config " + "installs it there." + ) + # spec.device is an exir schema DeviceType, not a torch.device. + if staged_spec.device != DeviceType.CUDA: + raise RuntimeError( + "TensorRT zero-copy KV: the staging copy of buffer " + f"'{source.name}' targets device {staged_spec.device!r}, not " + "CUDA, so moving the buffer there would put it where the " + "TensorRT engine cannot write it. The engine writes this " + "buffer in place and its copy-back has already been removed, " + "so the update would be lost." + ) + source_spec.device = staged_spec.device + source_spec.device_index = staged_spec.device_index + new_args[i] = source + unstaged += 1 + node.args = tuple(new_args) + if unstaged: + graph_module.graph.eliminate_dead_code() + graph_module.graph.lint() + graph_module.recompile() + return unstaged + + +def unstage_aliased_buffers_pass(inner_pass: Optional[Any] = None) -> Any: + """Build a ``to_out_var_pass`` that un-stages aliased buffers, then delegates. + + ``to_out_var_pass`` is the last hook that runs after ``PropagateDevicePass`` + and before memory planning -- the window in which the staging copies exist + and the buffers' placement is not yet fixed. (``sym_shape_eval_pass`` is a + caller-supplied hook in that window too, but it runs first.) + + ``inner_pass`` is the ``to_out_var_pass`` that would otherwise have run; it + runs after the un-staging. Omit it for ExecuTorch's default. + """ + from executorch.exir import ExecutorchBackendConfig + from executorch.exir.pass_base import PassBase + + inner = ( + inner_pass + if inner_pass is not None + else ExecutorchBackendConfig().to_out_var_pass + ) + + class _UnstageThenToOutVar(PassBase): # type: ignore[misc] + def call(self, graph_module: torch.fx.GraphModule) -> Any: + unstaged = _unstage_aliased_buffers(graph_module) + _LOGGER.debug("un-staged %d aliased delegate buffer(s)", unstaged) + return inner(graph_module) + + return _UnstageThenToOutVar() diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py new file mode 100644 index 00000000000..225c696fc08 --- /dev/null +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -0,0 +1,323 @@ +"""Export-side coverage for zero-copy aliased KV buffers. + +Two halves have to agree for a zero-copy ``.pte`` to be correct: + + * ``rewire_aliased_mutations_to_buffers`` declares the buffer to be its own + mutation result, which removes ExecuTorch's copy-back and takes the aliased + output out of the delegate. + * ``unstage_aliased_buffers_pass`` removes the host staging copy so the + engine's in-place write lands in the caller's buffer. + +The interesting failures are all silent -- a rewired mutation whose buffer is +still staged simply never updates -- so most of what is asserted here is which +mutations are left alone and which mis-shapes raise. +""" + +import operator +from types import SimpleNamespace + +import pytest + +pytest.importorskip("executorch.exir") + +import torch # noqa: E402 +import torch_tensorrt # noqa: E402 +from executorch.exir.delegate import executorch_call_delegate # noqa: E402 +from executorch.exir.schema import DeviceType # noqa: E402 +from torch.export.exported_program import ( # noqa: E402 + OutputKind, + OutputSpec, + TensorArgument, +) +from torch_tensorrt.executorch import _zero_copy as Z # noqa: E402 + +# The graphs below are built around torch.ops.tensorrt.execute_engine, which only +# exists once the Torch-TensorRT runtime operator library has loaded. +pytestmark = pytest.mark.skipif( + not torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime, + reason="Torch-TensorRT runtime operators are not available", +) + + +def _patch_engine_metadata(monkeypatch, *, aliased_io, input_names, output_names): + """Make every engine node report one fixed set of bindings and aliases.""" + import torch_tensorrt.dynamo.runtime._serialized_engine_layout as layout + import torch_tensorrt.dynamo.runtime._TorchTensorRTModule as trt_module + import torch_tensorrt.executorch.backend as backend + + info = ["x"] * (layout.ALIASED_IO_IDX + 1) + info[layout.INPUT_BINDING_NAMES_IDX] = "IN" + info[layout.OUTPUT_BINDING_NAMES_IDX] = "OUT" + monkeypatch.setattr(backend, "_get_engine_info_for_node", lambda ep, node: info) + monkeypatch.setattr(trt_module, "deserialize_aliased_io", lambda s: aliased_io) + monkeypatch.setattr( + layout, + "deserialize_binding_names", + lambda s: list(input_names) if s == "IN" else list(output_names), + ) + + +def _kv_program(*, mutation_value="aliased_getitem"): + """A one-engine program: engine(k_buffer, tokens) -> (logits, k_out). + + ``mutation_value`` picks what the KV buffer's BUFFER_MUTATION is bound to: + + * ``"aliased_getitem"``: the engine's aliased output (what export + declares for a caller-owned KV cache). + * ``"user_getitem"``: a non-aliased engine output, the shape a copy-back + mutation has. + * ``"external_op"``: a value produced outside the engine. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + tokens = graph.placeholder("tokens") + engine = graph.placeholder("engine") + engine_call = graph.call_function( + torch.ops.tensorrt.execute_engine.default, ([k_buffer, tokens], engine) + ) + logits = graph.call_function(operator.getitem, (engine_call, 0)) + k_out = graph.call_function(operator.getitem, (engine_call, 1)) + mutation = { + "aliased_getitem": k_out, + "user_getitem": logits, + "external_op": None, + }[mutation_value] + if mutation is None: + mutation = graph.call_function(torch.add, (k_buffer, k_buffer)) + graph.output((mutation, logits)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + signature = SimpleNamespace( + inputs_to_buffers={"b_k_0": "k_0"}, + input_specs=[], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(name=mutation.name), "k_0" + ), + OutputSpec(OutputKind.USER_OUTPUT, TensorArgument(name=logits.name), None), + ], + ) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=signature, + _graph_signature=signature, + ) + return program, k_buffer, k_out + + +@pytest.mark.unit +def test_rewire_points_the_mutation_at_its_buffer_and_marks_it(monkeypatch): + """The aliased output is replaced by the buffer itself and then dies. + + With the mutation bound to the placeholder there is nothing for ExecuTorch + to copy back, and with no other user the getitem leaves the graph -- which + is what takes the aliased output out of the delegate. + """ + program, k_buffer, k_out = _kv_program() + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in", "tokens"], + output_names=["logits", "out_k"], + ) + + assert Z.rewire_aliased_mutations_to_buffers(program) == 1 + + specs = program._graph_signature.output_specs + assert specs[0].kind == OutputKind.BUFFER_MUTATION + assert specs[0].target == "k_0" + assert specs[0].arg.name == k_buffer.name + output_node = program.graph_module.graph.output_node() + assert output_node.args[0][0] is k_buffer + assert k_out not in program.graph_module.graph.nodes + assert k_buffer.meta["_torch_tensorrt_aliased_buffer"] is True + + +@pytest.mark.unit +@pytest.mark.parametrize("mutation_value", ["user_getitem", "external_op"]) +def test_rewire_leaves_mutations_the_engine_does_not_alias(monkeypatch, mutation_value): + """Only a mutation the engine satisfies in place may be rewired. + + A copy-back mutation ("user_getitem") and a mutation computed outside the + engine ("external_op") both need their value copied into the buffer. Both + look exactly like an aliased mutation in the graph, so the discriminator has + to be the engine's own aliased_io -- rewiring either would delete a real + update with no error. + """ + program, k_buffer, _ = _kv_program(mutation_value=mutation_value) + original_spec = program._graph_signature.output_specs[0] + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in", "tokens"], + output_names=["logits", "out_k"], + ) + + assert Z.rewire_aliased_mutations_to_buffers(program) == 0 + assert program._graph_signature.output_specs[0] is original_spec + assert "_torch_tensorrt_aliased_buffer" not in k_buffer.meta + + +@pytest.mark.unit +def test_rewire_is_a_noop_without_aliased_io(monkeypatch): + program, k_buffer, _ = _kv_program() + _patch_engine_metadata( + monkeypatch, + aliased_io={}, + input_names=["k_in", "tokens"], + output_names=["logits", "out_k"], + ) + + assert Z.rewire_aliased_mutations_to_buffers(program) == 0 + assert "_torch_tensorrt_aliased_buffer" not in k_buffer.meta + + +@pytest.mark.unit +def test_rewire_rejects_an_engine_whose_every_output_is_aliased(monkeypatch): + """A delegate with no outputs at all is not a shape anything supports. + + Nothing downstream reports it: the runtime reads elision off a single + argument count, which a zero-output delegate satisfies, and the delegate + itself is a pure node a later graph-wide dead-code elimination can erase. + So the failure has to be raised here. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + engine = graph.placeholder("engine") + engine_call = graph.call_function( + torch.ops.tensorrt.execute_engine.default, ([k_buffer], engine) + ) + k_out = graph.call_function(operator.getitem, (engine_call, 0)) + graph.output((k_out,)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + signature = SimpleNamespace( + inputs_to_buffers={"b_k_0": "k_0"}, + input_specs=[], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(name=k_out.name), "k_0" + ) + ], + ) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=signature, + _graph_signature=signature, + ) + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in"], + output_names=["out_k"], + ) + + with pytest.raises(RuntimeError, match="no outputs at all"): + Z.rewire_aliased_mutations_to_buffers(program) + + +def _staged_delegate_graph(*, backend_id="TensorRTBackend", device=DeviceType.CUDA): + """A lowered graph: delegate(lowered, _h2d_copy(k_buffer), _h2d_copy(tokens)).""" + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + tokens = graph.placeholder("tokens") + lowered = graph.get_attr("lowered_module_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + staged_tokens = graph.call_function(h2d, (tokens,)) + delegate = graph.call_function( + executorch_call_delegate, (lowered, staged_k, staged_tokens) + ) + graph.output((delegate,)) + + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace(backend_id=backend_id) + graph_module = torch.fx.GraphModule(root, graph) + + for node, spec_device in ( + (k_buffer, DeviceType.CPU), + (tokens, DeviceType.CPU), + (staged_k, device), + (staged_tokens, device), + ): + node.meta["spec"] = SimpleNamespace(device=spec_device, device_index=3) + return graph_module, k_buffer, staged_k, delegate + + +@pytest.mark.unit +def test_unstage_feeds_the_buffer_straight_to_the_delegate(): + """The marked buffer replaces its staging copy and moves to the device. + + Moving the spec is not cosmetic: memory planning reads it, and a buffer left + in a host arena is somewhere the engine cannot write. + """ + graph_module, k_buffer, staged_k, delegate = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + assert Z._unstage_aliased_buffers(graph_module) == 1 + + assert delegate.args[1] is k_buffer + assert k_buffer.meta["spec"].device == DeviceType.CUDA + assert k_buffer.meta["spec"].device_index == 3 + # The other input is an ordinary one and keeps its staging copy. + assert delegate.args[2] is not None + assert delegate.args[2].target is torch.ops.et_copy._h2d_copy.default + + +@pytest.mark.unit +def test_unstage_keeps_staging_for_an_unmarked_buffer(): + graph_module, k_buffer, staged_k, delegate = _staged_delegate_graph() + + assert Z._unstage_aliased_buffers(graph_module) == 0 + assert delegate.args[1] is staged_k + assert k_buffer.meta["spec"].device == DeviceType.CPU + + +@pytest.mark.unit +def test_unstage_ignores_another_backends_delegate(): + """Only a TensorRT engine promises the in-place write.""" + graph_module, k_buffer, staged_k, delegate = _staged_delegate_graph( + backend_id="CudaBackend" + ) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + assert Z._unstage_aliased_buffers(graph_module) == 0 + assert delegate.args[1] is staged_k + + +@pytest.mark.unit +def test_unstage_raises_when_the_staging_copy_is_not_on_cuda(): + """Following the staging to the CPU would put the buffer out of the engine's + reach, and the copy-back that would have saved it is already gone.""" + graph_module, k_buffer, _, _ = _staged_delegate_graph(device=DeviceType.CPU) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="not.*CUDA"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_raises_when_the_staging_copy_has_no_spec(): + graph_module, k_buffer, staged_k, _ = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + del staged_k.meta["spec"] + + with pytest.raises(RuntimeError, match="no TensorSpec"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_pass_runs_the_inner_pass_after_unstaging(): + """A caller's own to_out_var_pass has to survive being composed with.""" + graph_module, k_buffer, _, delegate = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + seen = [] + + def inner(gm): + # The un-staging is already done by the time the inner pass sees the graph. + seen.append(delegate.args[1] is k_buffer) + return "inner-result" + + result = Z.unstage_aliased_buffers_pass(inner).call(graph_module) + + assert seen == [True] + assert result == "inner-result" From 53fd4e12255a656a37e3386f8def19960a42f277 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Mon, 17 Aug 2026 19:00:27 -0700 Subject: [PATCH 02/22] feat(executorch): accept a delegate whose in-place outputs are elided The output-binding validator assumed every engine output binding is a delegate output. With zero-copy KV that stops being true: the aliased outputs are gone from the delegate, because the engine's in-place write through the aliased input already is the buffer update. Accept that shape, but only when the caller asked for it. A delegate missing its aliased outputs because nothing declared them as buffer mutations looks exactly like a zero-copy one -- here and in the runtime, which reads elision off a single argument count -- and today that case is a loud export-time error. Relaxing the check unconditionally would turn it into a .pte that runs and quietly never updates its cache. So the permission travels down explicitly, on a CompileSpec that only export() sets, over the partitioner's DelegationSpec: the one channel from the export call to `preprocess`. Even with permission, a partial drop stays an error. It would be a buffer update lost without a word, and the argument count cannot express it in any case. --- py/torch_tensorrt/executorch/backend.py | 111 +++++++++++++++++++-- tests/py/dynamo/executorch/test_backend.py | 69 +++++++++++++ 2 files changed, 170 insertions(+), 10 deletions(-) diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index 655905bfe74..799732c06a5 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -1,7 +1,8 @@ # ExecuTorch TensorRT backend: serialize engines to a libtorch-free runtime blob. +import json import operator -from typing import Any, List, final +from typing import Any, Container, Iterable, List, Optional, Set, final import torch import torch.fx @@ -32,6 +33,16 @@ _BINDING_DELIM = "%" +# CompileSpec key by which export() tells this backend which aliased outputs it +# deliberately took out of the delegate. Its value is the JSON list of those +# engine output binding names (see _serialize_elided_output_names), NOT a bare +# flag: export only elides the aliased outputs backed by a registered buffer, so +# the backend must exempt exactly those and still reject a delegate that dropped +# any other binding. It travels on the partitioner's DelegationSpec, the only +# channel from the export call down to preprocess. Without it a delegate short of +# its aliased outputs is a bug, not a zero-copy program, and stays an error. +ZERO_COPY_KV_COMPILE_SPEC_KEY = "zero_copy_kv" + def _schema_name(target: Any) -> str: """Return the qualified op schema name for an OpOverload or EdgeOpOverload.""" @@ -274,7 +285,10 @@ def _reorder_input_names_for_executorch( def _validate_output_binding_order( - edge_program: ExportedProgram, engine_node: Any, output_names: List[str] + edge_program: ExportedProgram, + engine_node: Any, + output_names: List[str], + elidable_output_names: Optional[Container[str]] = None, ) -> None: """Check the delegate's outputs are the engine's output bindings, in order. @@ -286,19 +300,45 @@ def _validate_output_binding_order( delegate -- would swap the names silently. Inputs cannot rely on position at all and recover their order by node identity in ``_reorder_input_names_for_executorch``. + + ``elidable_output_names`` names the bindings the delegate is *allowed* to + have dropped, which zero-copy KV sets to exactly the aliased outputs export + rewired to write in place (never the whole aliased_io): the engine's in-place + write through the aliased input already is the buffer update, so no argument + is passed for them. Pass ``None`` (the default) when elision was not asked + for, and the delegate must carry every binding. + + A delegate that dropped its aliased outputs because nothing declared them as + mutations looks exactly like a zero-copy one, and the runtime reads elision off + a single argument count, so it cannot tell them apart either. That is also why + a partial drop stays an error: the count cannot express which bindings went. """ + elidable_names = elidable_output_names if elidable_output_names is not None else () + all_indices = list(range(len(output_names))) + unaliased_indices = [ + i for i in all_indices if output_names[i] not in elidable_names + ] + output_node = next( node for node in edge_program.graph_module.graph.nodes if node.op == "output" ) out_args = list(output_node.args[0]) # A single-output engine is returned directly rather than through a getitem, - # and one binding has no order to get wrong. + # and one binding has no order to get wrong. The same holds under elision + # when exactly one binding is left unaliased. if len(out_args) == 1 and out_args[0] is engine_node: - if len(output_names) != 1: + if len(all_indices) != 1 and len(unaliased_indices) != 1: + remaining = ( + f", {len(unaliased_indices)} of them after eliding the in-place " + "outputs" + if unaliased_indices != all_indices + else "" + ) raise ValueError( "TensorRT ExecuTorch backend: the delegate returns the engine node " f"directly but the engine declares {len(output_names)} output " - "bindings; only a single-output engine can be returned unwrapped." + f"bindings{remaining}; only a single-output engine can be returned " + "unwrapped." ) return indices: List[Any] = [] @@ -315,15 +355,61 @@ def _validate_output_binding_order( "node; cannot establish a reliable output binding order." ) indices.append(node.args[1]) - if indices != list(range(len(output_names))): + if indices not in (all_indices, unaliased_indices): + expected = ( + f"{all_indices}, or {unaliased_indices} with the in-place outputs elided" + if unaliased_indices != all_indices + else f"{all_indices}" + ) + # Outputs are missing and nothing exempted them. That is what an export + # asking for zero_copy_kv looks like when the aliased-buffer mark did not + # reach the partitioner, so no delegate was stamped and none is exempt -- + # a failure mode with no other symptom, hence naming it here. + unexempted_drop = elidable_output_names is None and len(indices) < len( + all_indices + ) raise ValueError( "TensorRT ExecuTorch backend: delegate outputs map to engine output " - f"indices {indices}, expected {list(range(len(output_names)))} -- the " - "runtime binds output i to output_binding_names[i], so a permuted or " - "incomplete output list would bind the wrong tensors." + f"indices {indices}, expected {expected} -- the runtime binds each " + "output it is given in binding order, so a permuted, incomplete, or " + "partially elided output list would bind the wrong tensors." + + ( + " No output was declared elidable for this delegate, so if the " + "export asked for zero_copy_kv the aliased-buffer mark did not " + "survive lowering." + if unexempted_drop + else "" + ) ) +def _serialize_elided_output_names(names: Iterable[str]) -> bytes: + """Encode the elided aliased-output binding names for the compile spec. + + JSON, not the ``%`` / ``@`` delimiters that separate ``engine_info``'s + binding-name and aliased_io fields, so a binding name containing one of those + cannot corrupt the record. + """ + return json.dumps(sorted(set(names))).encode("utf-8") + + +def _elided_output_names(compile_specs: List[CompileSpec]) -> Optional[Set[str]]: + """The aliased-output binding names export declared elidable, or ``None``. + + ``None`` when no zero-copy spec is present, which keeps a missing output an + error: only a caller who asked for zero-copy may drop the aliased outputs, + and then only exactly the ones export rewired to write in place. + """ + for spec in compile_specs: + if getattr(spec, "key", None) != ZERO_COPY_KV_COMPILE_SPEC_KEY: + continue + value = spec.value + if isinstance(value, (bytes, bytearray)): + value = bytes(value).decode("utf-8") + return set(json.loads(value)) + return None + + def _get_str(engine_info: List[Any], index: int, default: str = "") -> str: if index < 0 or index >= len(engine_info): return default @@ -375,7 +461,12 @@ 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) + _validate_output_binding_order( + edge_program, + engine_node, + output_names, + _elided_output_names(compile_specs), + ) 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 9c687bba6e1..604c58f7de6 100644 --- a/tests/py/dynamo/executorch/test_backend.py +++ b/tests/py/dynamo/executorch/test_backend.py @@ -437,3 +437,72 @@ def test_validate_output_binding_order_accepts_unwrapped_single_output(): g.output((engine,)) ep = SimpleNamespace(graph_module=torch.fx.GraphModule(torch.nn.Module(), g)) _validate_output_binding_order(ep, engine, ["out"]) + + +@pytest.mark.unit +def test_validate_output_binding_order_accepts_elided_aliased_outputs(): + """Zero-copy KV drops the in-place outputs from the delegate entirely. + + The engine still declares them as bindings -- the runtime binds them to + their aliased input's tensor -- but no delegate argument is passed for them. + """ + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([0]) + _validate_output_binding_order(ep, engine, ["out0", "kv0", "kv1"], {"kv0", "kv1"}) + + +@pytest.mark.unit +def test_validate_output_binding_order_rejects_elision_that_was_not_requested(): + """The same graph is a bug unless the caller asked for zero-copy. + + Aliased outputs missing because nothing declared them as buffer mutations + look identical to aliased outputs deliberately elided, here and at runtime. + So the default stays strict and only an explicit opt-in relaxes it. + + It is also the shape a zero-copy export produces when the aliased-buffer mark + never reaches the partitioner: nothing is stamped, so nothing is exempt. The + message names the feature, because that failure has no other symptom. + """ + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([0]) + with pytest.raises(ValueError, match="engine output indices") as excinfo: + _validate_output_binding_order(ep, engine, ["out0", "kv0", "kv1"]) + assert "zero_copy_kv" in str(excinfo.value) + + +@pytest.mark.unit +def test_validate_output_binding_order_still_accepts_aliased_outputs_threaded(): + """Not eliding is the pre-existing shape and stays legal: without zero-copy + each aliased output is a delegate output like any other.""" + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([0, 1, 2]) + _validate_output_binding_order(ep, engine, ["out0", "kv0", "kv1"], {"kv0", "kv1"}) + + +@pytest.mark.unit +def test_validate_output_binding_order_rejects_partially_elided_outputs(): + """Half-elision is a lost buffer update, and the runtime cannot express it: + it infers elision from one argument count, so it is all-or-nothing.""" + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([0, 1]) + with pytest.raises(ValueError, match="engine output indices") as excinfo: + _validate_output_binding_order( + ep, engine, ["out0", "kv0", "kv1"], {"kv0", "kv1"} + ) + # The mark plainly did survive -- these bindings were declared elidable -- so + # the lost-mark hint would be a wrong lead here. + assert "zero_copy_kv" not in str(excinfo.value) + + +@pytest.mark.unit +def test_validate_output_binding_order_rejects_permutation_of_elided_outputs(): + """Elision removes outputs; it does not license reordering the survivors.""" + from torch_tensorrt.executorch.backend import _validate_output_binding_order + + ep, engine = _engine_partition([2, 0]) + with pytest.raises(ValueError, match="engine output indices"): + _validate_output_binding_order(ep, engine, ["out0", "kv0", "out2"], {"kv0"}) From c67ab1ff0f842630caa4a57b88041c23b95bf616 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Mon, 17 Aug 2026 19:01:32 -0700 Subject: [PATCH 03/22] feat(executorch): run a delegate whose aliased outputs are elided With zero-copy KV the aliased outputs are not delegate arguments at all: the engine's write through the aliased input's pointer is the buffer update, so there is no mutation slot to fill and no copy to reflect. Detect that from the argument count rather than a serialized flag. Export elides either all of an engine's aliased outputs or none, so the two shapes differ in arity by exactly the aliased-output count the handle already knows -- and a .pte written before zero-copy existed keeps taking the threaded branch with no version check and no new blob field. (Elision is all-or-nothing only because every aliased output today is a `kv_cache_update` on a buffer; a future `kind="user"` alias beside an elided one would break the identity, and the arity check would reject the .pte at execute.) `setTensorAddress` has already pointed the binding at the caller's buffer by the point the branch is taken, so eliding only skips consuming an argument and recording a reflect. The end-to-end check for this path comes with the next commit, which adds the export-side option that can produce such a .pte. --- .../executorch/TensorRTBackend.cpp | 48 ++++++++++++++----- 1 file changed, 35 insertions(+), 13 deletions(-) diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 8408c13e881..785843f8db0 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -551,8 +551,10 @@ Result TensorRTBackend::init( // their addresses; no separate output allocation is required. // // Args layout (mirroring the Python exporter): -// args[0 .. num_inputs-1] – input EValues -// args[num_inputs .. num_inputs+num_outputs-1] – output EValues +// args[0 .. num_inputs-1] -- input EValues +// args[num_inputs .. num_inputs+num_delegate_outputs-1] -- output EValues +// num_delegate_outputs is num_outputs, less the aliased outputs when zero-copy KV +// elided them from the delegate; see the arity branch at the top of execute(). // --------------------------------------------------------------------------- Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* handle, Span args) const { (void)context; @@ -561,11 +563,19 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* const size_t num_inputs = engine->num_inputs; const size_t num_outputs = engine->num_outputs; - // Caller-owned KV: every input is a delegate arg, and each aliased output is - // threaded as a delegate output arg (the caller-owned mutable buffer's mutation - // slot), so all engine bindings map 1:1 to delegate args. - const size_t num_delegate_outputs = num_outputs; + // Caller-owned KV comes in two shapes. Either each aliased output is threaded as + // a delegate output arg (the caller-owned mutable buffer's mutation slot), so all + // engine bindings map 1:1 to delegate args; or -- zero-copy KV -- the aliased + // outputs are elided from the delegate entirely, because the engine's in-place + // write through the aliased input already IS the buffer update. A .pte written + // before zero-copy existed still takes the first branch, and export will not emit + // the shorter arity unless zero-copy was asked for, so a short argument list + // cannot instead mean "the aliased outputs were never declared". const size_t num_delegate_inputs = num_inputs; + const size_t num_aliased_outputs = engine->num_aliased_outputs; + const bool aliased_outputs_elided = + num_aliased_outputs > 0 && args.size() == num_delegate_inputs + num_outputs - num_aliased_outputs; + const size_t num_delegate_outputs = num_outputs - (aliased_outputs_elided ? num_aliased_outputs : 0); if (args.size() < num_delegate_inputs + num_delegate_outputs) { ET_LOG( Error, @@ -795,9 +805,15 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* ET_LOG(Error, "TensorRTBackend::execute: setTensorAddress failed for aliased output '%s'", name.c_str()); return Error::InvalidState; } - // The aliased output IS a delegate output arg (the caller-owned mutable - // buffer's mutation slot). Consume it and record a reflect so ExecuTorch's - // write-back copy_ sees the engine's in-place update. + // Elided: setTensorAddress above pointed this output binding at the caller's + // buffer, so the engine writes the buffer itself; nothing to reflect into. + if (aliased_outputs_elided) { + continue; + } + + // Otherwise the aliased output IS a delegate output arg (the caller-owned + // mutable buffer's mutation slot). Consume it and record a reflect so + // ExecuTorch's write-back copy_ sees the engine's in-place update. const size_t arg_i = arg_idx++; EValue* out_arg = args[arg_i]; TORCHTRT_ET_CHECK_NOT_NULL( @@ -945,10 +961,16 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // next execute() and the destructor wait before reusing/freeing exec_ctx. The D2H // copies live in the must_sync branch: an output staged to host always sets // output_staged_to_host, so outputs_needing_copy is empty on the skip path. - // An aliased reflect enqueues the engine's in-place update into the delegate - // output EValue on `stream`; ExecuTorch's buffer-mutation copy_ reads that EValue - // after execute() returns, so the reflect must complete first. A model with - // aliased outputs therefore always syncs here. + // A non-elided aliased reflect enqueues the engine's in-place update into the + // delegate output EValue on `stream`; ExecuTorch's buffer-mutation copy_ reads + // that EValue after execute() returns, so the reflect must complete first, and a + // model that threads its aliased outputs as delegate output args syncs here. + // Under zero-copy KV skipping the sync is correct, because the aliased buffer + // stays device-resident and its next reader is the following engine execute() -- + // on the same `stream`, provided the runner honours the single-shared-stream + // contract every coalesced .pte already depends on. A host reader that + // inspected it immediately after execute() returns would see stale data unless + // it synchronized `stream` itself; ExecuTorch's KV path never does such a read. const bool aliased_reflect_pending = !aliased_reflects.empty(); const bool must_sync = output_staged_to_host || input_staged_from_host || aliased_reflect_pending || !caller_stream_set; From 02498ef41a4442cf8000d87163812c79a69bed2a Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Mon, 17 Aug 2026 19:09:49 -0700 Subject: [PATCH 04/22] feat(executorch): add the zero_copy_kv opt-in and its finalization config Until now the only way to get zero-copy KV was to monkey-patch a private function, because the rewiring has to land between two things export() does internally -- after the aliased mutations are declared, before the program is staged -- and there was no hook there. Give it a real one. export(..., zero_copy_kv=True) edge.to_executorch(zero_copy_backend_config(config)) Opt-in rather than automatic: a .pte whose aliased outputs are elided needs a runtime that understands that shape, so producing one unasked would break a runner built before this feature, and the option changes nothing for existing caller-owned KV users. Two calls rather than one because to_executorch() is ExecuTorch's, not ours -- export() returns at the Edge boundary and never sees the config the program is finalized with. zero_copy_backend_config composes onto the caller's config rather than replacing it, so their memory planning and their own to_out_var_pass survive. It is the module's only public name; the two passes stay private, since applying one without the other is worse than applying neither. The backend's permission to accept an elided delegate is granted per method and only where a mutation was actually rewired, so a method that lost an output for some other reason is still rejected. That leaves one thing the caller must not forget, and it is documented as such along with the other two quiet contracts of this feature: finalizing without the config, running the delegates on separate CUDA streams, and expecting one cache to be shared across methods without a memory-planning pass that says so. All three produce wrong values rather than an error. The persistence check gets a zero-copy variant in CI: the existing runner asserts a decode step sees the KV the previous step wrote, which is exactly the property zero-copy has to preserve after removing the copy that used to provide it. The two-call contract is unavoidable on the direct export()+to_executorch() path, but not through torch_tensorrt.save(output_format="executorch"): that path owns both steps -- it calls export() and then to_executorch() itself -- so it can make zero-copy foolproof. Give save() a single `zero_copy_kv=` flag that threads the opt-in into export() and installs zero_copy_backend_config before to_executorch(), so forgetting the second call is not possible there. It is single-method only, like the rest of save(); multi-method stays on the direct export path. Wrapping preserves a caller-supplied backend_config, so passing both is fine. The one exclusion is handing save() a backend_config that already carries the pass: save() wraps it again and finalization raises. The two entry points are mutually exclusive. --- .../verify-executorch-reference-runner.sh | 64 +- .github/workflows/executorch-test-linux.yml | 6 +- .../runtime_performance/saving_models.rst | 104 +- .../export_kv_cache_decode.py | 107 ++- py/torch_tensorrt/_compile.py | 32 +- py/torch_tensorrt/executorch/__init__.py | 3 + py/torch_tensorrt/executorch/_export.py | 120 ++- py/torch_tensorrt/executorch/_zero_copy.py | 283 +++++- py/torch_tensorrt/executorch/backend.py | 19 +- py/torch_tensorrt/executorch/partitioner.py | 122 ++- tests/py/dynamo/executorch/test_backend.py | 105 +++ tests/py/dynamo/executorch/test_edge_cases.py | 5 +- tests/py/dynamo/executorch/test_export.py | 205 ++++ .../py/dynamo/executorch/test_zero_copy_kv.py | 891 +++++++++++++++++- 14 files changed, 1971 insertions(+), 95 deletions(-) diff --git a/.github/scripts/verify-executorch-reference-runner.sh b/.github/scripts/verify-executorch-reference-runner.sh index fa5881f5fc7..5b52744241b 100755 --- a/.github/scripts/verify-executorch-reference-runner.sh +++ b/.github/scripts/verify-executorch-reference-runner.sh @@ -14,16 +14,19 @@ set +x # First argument: path to an existing .pte model. # EXECUTORCH_SOURCE_DIR=/path/to/executorch # -# Optional second argument: path to a caller-owned KV-cache decode .pte (see -# examples/torchtrt_executorch_example/export_kv_cache_decode.py). When given, -# kv_cache_decode_check is built and run against it as well. +# Optional trailing arguments: one or more caller-owned KV-cache decode .pte +# files (see examples/torchtrt_executorch_example/export_kv_cache_decode.py). +# When given, kv_cache_decode_check is built and run against each of them, +# staged or zero-copy. # -# Optional third argument: path to a coalesced TensorRT + CUDA .pte (see +# Optional --coalesced=PATH: path to a coalesced TensorRT + CUDA .pte (see # examples/torchtrt_executorch_example/export_coalesced.py). When given, the # runner built from source here is run against it and its output is compared to # the eager reference that export script wrote next to the model. Only that # runner: the packaged binary links the TensorRT delegate alone, so it has no -# CUDA backend for the partition a coalesced program hands to one. +# CUDA backend for the partition a coalesced program hands to one. Named rather +# than positional because the KV-cache decode models are variadic, so a bare +# path after the first argument cannot be told apart from one of those. # # Optional: # TensorRT_ROOT=/path/to/extracted/TensorRT @@ -40,21 +43,35 @@ set +x repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" cd "${repo_root}" -if [[ $# -lt 1 || $# -gt 3 ]]; then - echo "Usage: $0 PATH_TO_MODEL.pte [PATH_TO_KV_CACHE_DECODE.pte [PATH_TO_COALESCED.pte]]" >&2 +if [[ $# -lt 1 ]]; then + echo "Usage: $0 PATH_TO_MODEL.pte [--coalesced=PATH_TO_COALESCED.pte]" \ + "[PATH_TO_KV_CACHE_DECODE.pte ...]" >&2 exit 1 fi model_path="$1" +shift if [[ ! -f "${model_path}" ]]; then echo "ExecuTorch model not found: ${model_path}" >&2 exit 1 fi -kv_model_path="${2:-}" -if [[ -n "${kv_model_path}" && ! -f "${kv_model_path}" ]]; then - echo "KV-cache decode model not found: ${kv_model_path}" >&2 - exit 1 -fi -coalesced_model_path="${3:-}" +coalesced_model_path="" +kv_model_paths=() +for arg in "$@"; do + case "${arg}" in + --coalesced=*) + coalesced_model_path="${arg#--coalesced=}" + ;; + *) + kv_model_paths+=("${arg}") + ;; + esac +done +for kv_model_path in "${kv_model_paths[@]:-}"; do + if [[ -n "${kv_model_path}" && ! -f "${kv_model_path}" ]]; then + echo "KV-cache decode model not found: ${kv_model_path}" >&2 + exit 1 + fi +done if [[ -n "${coalesced_model_path}" && ! -f "${coalesced_model_path}" ]]; then echo "Coalesced model not found: ${coalesced_model_path}" >&2 exit 1 @@ -318,7 +335,7 @@ fi cmake "${cmake_args[@]}" build_targets=(example_executorch_runner) -if [[ -n "${kv_model_path}" ]]; then +if [[ ${#kv_model_paths[@]} -gt 0 ]]; then build_targets+=(kv_cache_decode_check) fi @@ -551,11 +568,7 @@ for _log in "${runner_log}" "${packaged_runner_log}"; do assert_runner_output "${_log}" "[2,3,4,4]" "2.0000" 0 done -if [[ -n "${kv_model_path}" ]]; then - # kv_cache_decode_check exits non-zero when a decode step does not observe the KV - # the previous step wrote; the grep additionally pins the assertion itself, so - # weakening the check inside the binary cannot quietly turn this into a no-op. - kv_check_log="${verify_root}/kv_cache_decode_check.log" +if [[ ${#kv_model_paths[@]} -gt 0 ]]; then kv_check_path="${verify_root}/build-executorch-reference-runner/kv_cache_decode_check" if command -v ldd >/dev/null 2>&1 && ldd "${kv_check_path}" | @@ -564,8 +577,17 @@ if [[ -n "${kv_model_path}" ]]; then exit 1 fi - "${kv_check_path}" --model_path="${kv_model_path}" 2>&1 | tee "${kv_check_log}" - grep -q "PASS: decode at pos=1 observed the KV written at pos=0" "${kv_check_log}" + kv_index=0 + for kv_model_path in "${kv_model_paths[@]}"; do + # kv_cache_decode_check exits non-zero when a decode step does not observe the + # KV the previous step wrote; the 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_${kv_index}.log" + "${kv_check_path}" --model_path="${kv_model_path}" 2>&1 | tee "${kv_check_log}" + grep -q "PASS: decode at pos=1 observed the KV written at pos=0" "${kv_check_log}" + kv_index=$((kv_index + 1)) + done fi if [[ -n "${coalesced_model_path}" ]]; then diff --git a/.github/workflows/executorch-test-linux.yml b/.github/workflows/executorch-test-linux.yml index 9751f2bb54a..92be5830970 100644 --- a/.github/workflows/executorch-test-linux.yml +++ b/.github/workflows/executorch-test-linux.yml @@ -95,11 +95,15 @@ jobs: --model_path="${RUNNER_TEMP}/torchtrt-python.pte" python examples/torchtrt_executorch_example/export_kv_cache_decode.py \ --model_path="${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte" + python examples/torchtrt_executorch_example/export_kv_cache_decode.py \ + --model_path="${RUNNER_TEMP}/torchtrt-kv-cache-decode-zero-copy.pte" \ + --zero_copy python examples/torchtrt_executorch_example/export_coalesced.py \ --model_path="${RUNNER_TEMP}/torchtrt-coalesced.pte" .github/scripts/verify-executorch-reference-runner.sh \ "${RUNNER_TEMP}/torchtrt-python.pte" \ + --coalesced="${RUNNER_TEMP}/torchtrt-coalesced.pte" \ "${RUNNER_TEMP}/torchtrt-kv-cache-decode.pte" \ - "${RUNNER_TEMP}/torchtrt-coalesced.pte" + "${RUNNER_TEMP}/torchtrt-kv-cache-decode-zero-copy.pte" python examples/executorch_reference_runner/load_model.py \ --model_path="${RUNNER_TEMP}/torchtrt-python.pte" --num_runs=1 diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 6470afd72e4..0fcbff1112b 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -356,6 +356,77 @@ points but does not by itself give them shared mutable state. Neither case raises an error or a warning, so treat every shared payload as read-only. +.. _executorch_zero_copy_kv: + +**Zero-copy KV cache** + +When a TensorRT engine has aliased I/O -- a KV cache it updates through an +aliased binding -- running the engine over the cache already is the update. +ExecuTorch does not know that, so by default it pays for the update twice per +execution: it hands the delegate an ``_h2d_copy`` staging copy of the buffer +instead of the buffer itself, then copies the engine's aliased output back into +the buffer afterwards. For a KV cache both copies are cache-sized, per token. + +``zero_copy_kv=True`` removes them, so the engine writes the caller's buffer +directly. Through the two-step ``export`` + ``to_executorch`` path it takes two +calls, one at each end of the Edge boundary: + +.. code-block:: python + + from torch_tensorrt.executorch import export, zero_copy_backend_config + + edge = export( + {"prefill": prefill_program, "decode": decode_program}, + partitioners={"prefill": [CudaPartitioner([])], "decode": [CudaPartitioner([])]}, + zero_copy_kv=True, + ) + + # zero_copy_backend_config composes onto your own config; every other field + # (memory planning, passes) is preserved. + program = edge.to_executorch(zero_copy_backend_config(backend_config)) + +It is opt-in rather than automatic because the resulting ``.pte`` needs a +runtime that understands a delegate whose aliased outputs are elided. Producing +one silently would break a runner built before this feature. + +.. warning:: + + **Both calls are required.** Exporting with ``zero_copy_kv=True`` and then + finalizing without ``zero_copy_backend_config`` does not raise: the engine + writes a per-call staging copy that is discarded, and the cache never + updates. For a KV cache that is wrong output, not a crash. Nothing + downstream can detect the omission, so pairing the two is on the caller -- + unless ``torch_tensorrt.save`` is the one writing the ``.pte``, which owns + both ends and leaves nothing to pair. + +``torch_tensorrt.save`` finalizes the program itself, so a single +``zero_copy_kv=True`` covers both steps: + +.. code-block:: python + + torch_tensorrt.save( + trt_gm, "decode.pte", output_format="executorch", + arg_inputs=inputs, retrace=False, + zero_copy_kv=True, + ) + +It installs ``zero_copy_backend_config`` for you, so do not hand it one as +``backend_config`` as well: the pass would be installed twice and finalization +raises. The two entry points are alternatives, not a pair. + +Two further responsibilities are the caller's, and both fail quietly: + +* **One CUDA stream for every delegate**, if the ``.pte`` is coalesced -- and the + synchronization it calls for, which zero-copy makes load-bearing. See + :ref:`Running a coalesced .pte `. + +* **Sharing one cache between methods.** Zero-copy is per method: it makes each + method's engine write that method's buffer. Giving a prefill and a decode + method *the same* cache is a memory-planning question -- their mutable buffers + have to land at the same arena offsets -- which ExecuTorch's memory planner + owns and which is deployment-specific. Supply your own + ``memory_planning_pass`` for it; ``zero_copy_backend_config`` preserves it. + **Coalesced TensorRT + CUDA .pte** To run the ops TensorRT does not take on ExecuTorch's CUDA (AOTInductor) backend @@ -390,6 +461,8 @@ must be pointed at those data files to load them. ``.pte`` into the same directory overwrites the blob and the first ``.pte`` will fail to load. Save each coalesced model into its own directory. +.. _executorch_single_stream: + **Running a coalesced .pte: use a single CUDA stream** A coalesced ``.pte`` runs on more than one backend delegate (the TensorRT delegate @@ -404,20 +477,31 @@ illegal memory access. The runtime does not impose a shared stream across delegates, so it is the **runner's responsibility** to run all delegates on one CUDA stream. Create a single stream and, for the duration of execution, direct every backend to use it -(each backend exposes a caller-stream hook). All GPU work is then enqueued in order -and every cross-boundary dependency is satisfied, while execution stays -asynchronous. +(each backend exposes a caller-stream hook: scope both +``torch_tensorrt::executorch_backend::CudaStreamGuard`` and +``executorch::extension::cuda::CallerStreamGuard`` over that stream, since +installing one of them leaves the other backend on its own). All GPU work is then +enqueued in order and every cross-boundary dependency is satisfied, while +execution stays asynchronous. If the runner reads a delegate's outputs between calls (for example, an autoregressive decode loop), synchronize the shared stream before reading: the work may still be in flight when ``execute()`` returns, and a host-side copy on -the default stream will not wait for a non-blocking stream. +the default stream will not wait for a non-blocking stream. A model that threads +its aliased outputs through the delegate is insulated from this in practice: +reflecting each aliased output into its delegate output makes the delegate wait +for the engine before it returns. Under +:ref:`zero-copy KV ` those outputs are elided, so there +is nothing to reflect and the delegate returns with the engine still running -- +the synchronization is then the only thing making a host read see the new values. **ExecuTorch lowering options** -When ``output_format="executorch"``, ``torch_tensorrt.save`` forwards the following -keyword arguments to ExecuTorch's ``to_edge_transform_and_lower(...)``. They are -only consulted for the ``executorch`` format; passing them with any other -``output_format`` logs a warning and is otherwise ignored. +``torch_tensorrt.save`` takes these extra keyword arguments. They are only +consulted for the ``executorch`` format; passing them with any other +``output_format`` logs a warning and is otherwise ignored. ``constant_methods``, +``transform_passes`` and ``compile_config`` are forwarded to ExecuTorch's +``to_edge_transform_and_lower(...)``; the rest are consumed at other points in +``save``. * ``constant_methods`` — a ``dict`` of extra constant methods to embed in the ``.pte`` (e.g. ``{"get_max_seq_len": 2048}`` for an LLM runner). @@ -430,6 +514,10 @@ only consulted for the ``executorch`` format; passing them with any other your graph carries TensorRT engines, set ``_check_ir_validity=False`` explicitly. * ``backend_config`` — an ``ExecutorchBackendConfig`` forwarded to ``to_executorch(...)``. +* ``zero_copy_kv`` — a ``bool`` (default ``False``, single-method only). Lets the + TensorRT engine update an aliased KV cache in place. ``save`` owns both ends of + the Edge boundary, so this one argument covers what the two-call path spells + out; see :ref:`Zero-copy KV cache `. * ``generate_etrecord`` — a ``bool`` (default ``False``). When ``True``, an `ETRecord `_ is written next to the ``.pte`` as ``_etrecord.bin`` (e.g. ``trt.pte`` → diff --git a/examples/torchtrt_executorch_example/export_kv_cache_decode.py b/examples/torchtrt_executorch_example/export_kv_cache_decode.py index c8590d98b03..eedf5611592 100644 --- a/examples/torchtrt_executorch_example/export_kv_cache_decode.py +++ b/examples/torchtrt_executorch_example/export_kv_cache_decode.py @@ -14,6 +14,15 @@ ``.pte`` and asserts that a decode step observes the KV a previous step wrote (i.e. the cache is shared across ``execute()`` calls). +``--zero_copy`` exports the same model with the engine writing the cache buffer +directly, instead of ExecuTorch staging a copy for the delegate and copying the +result back. The persistence check is the same, and it is the check that matters +here: zero-copy removes the copy that was making the update stick, so if the +engine's in-place write is not reaching the caller's buffer the run fails. What +it cannot see is a ``--zero_copy`` export that degenerated into an ordinary +staged ``.pte`` -- the two are indistinguishable to it -- so ``_check_zero_copy`` +refuses to write one. + Prerequisites ------------- Install Torch-TensorRT with the ExecuTorch extra before running this example:: @@ -22,6 +31,8 @@ """ import argparse +import os +from typing import Any import torch import torch_tensorrt @@ -81,11 +92,90 @@ def split_heads(proj: torch.Tensor) -> torch.Tensor: return self.lm(self.o(out)) +def _check_zero_copy(program: Any) -> None: + """Refuse to write a ``--zero_copy`` .pte that is really an ordinary staged one. + + Neither half of zero-copy fails when it finds nothing to do: + ``zero_copy_kv=True`` logs a warning and carries on when no aliased buffer + mutation turns up, and the finalization pass, handed a program with nothing + marked, un-stages nothing and returns without a word. The ``.pte`` that comes + out still runs, and ``kv_cache_decode_check`` cannot tell it from the staged + model exported beside it -- so without this the lane would stay green while + covering none of the feature. + """ + from executorch.exir.delegate import executorch_call_delegate + + graph_module = program.exported_program().graph_module + marked = [ + node + for node in graph_module.graph.nodes + if node.op == "placeholder" and node.meta.get("_torch_tensorrt_aliased_buffer") + ] + if not marked: + raise RuntimeError( + "zero_copy_kv=True rewired no aliased buffer, so this .pte stages its " + "KV cache like any other." + ) + delegate_args = { + arg + for node in graph_module.graph.nodes + if node.op == "call_function" and node.target is executorch_call_delegate + for arg in node.args[1:] + } + staged = [node.name for node in marked if node not in delegate_args] + if staged: + raise RuntimeError( + f"buffer(s) {staged} still reach the delegate through a staging copy, " + "so the engine writes scratch that is thrown away and the cache never " + "updates." + ) + + +def _save_zero_copy(trt_gm: torch.fx.GraphModule, inputs: tuple, path: str) -> None: + """Save a .pte whose engine updates the KV cache in place. + + Zero-copy needs both ends of the Edge boundary: ``zero_copy_kv`` before + lowering, and ``zero_copy_backend_config`` on the config the program is + finalized with. Without the second the cache is still staged and its updates + are dropped, with no error. ``torch_tensorrt.save(output_format="executorch", + zero_copy_kv=True)`` owns both steps and is the shorter way to the same .pte; + this spells them out because both halves are shown, and because a program + that reaches ``to_executorch()`` by any other route has to install the config + itself. + """ + from torch_tensorrt.executorch import export, zero_copy_backend_config + + # retrace=True here, retrace=False for the plain save() below, so the two + # exporters are both covered. Which way round matters: the legacy exporter + # declares the aliased KV outputs while building the program, so on that lane + # export()'s declaration pass reads each engine's aliased_io only to find the + # mutations already declared. A retraced program arrives undeclared, so this + # is the lane where that read decides anything -- where an engine-aliased + # cache is told from an ordinary copy-back buffer. + edge = export(trt_gm, arg_inputs=inputs, retrace=True, zero_copy_kv=True) + program = edge.to_executorch(zero_copy_backend_config()) + _check_zero_copy(program) + with open(path, "wb") as output: + program.write_to_file(output) + if program._tensor_data: + # A delegate carrying external weights (the CUDA backend does) keeps them + # outside the .pte, and write_to_file does not persist them; without the + # .ptd next to it the .pte cannot load. This model has none, but a model + # built from this one may. + program.write_tensor_data_to_file(os.path.dirname(os.path.abspath(path))) + + def main() -> None: parser = argparse.ArgumentParser() parser.add_argument( "--model_path", default="kv_cache_decode.pte", help="Path to save the .pte" ) + parser.add_argument( + "--zero_copy", + action="store_true", + help="Let the engine write the KV cache in place (elides the aliased " + "delegate outputs; needs a runtime that supports them).", + ) args = parser.parse_args() with torch.no_grad(): @@ -101,13 +191,16 @@ def main() -> None: min_block_size=1, truncate_double=True, ) - torch_tensorrt.save( - trt_gm, - args.model_path, - output_format="executorch", - arg_inputs=(tokens, input_pos), - retrace=False, - ) + if args.zero_copy: + _save_zero_copy(trt_gm, (tokens, input_pos), args.model_path) + else: + torch_tensorrt.save( + trt_gm, + args.model_path, + output_format="executorch", + arg_inputs=(tokens, input_pos), + retrace=False, + ) print(f"Saved {args.model_path} successfully.") diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index fa70e13c461..d10692a785c 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -815,6 +815,14 @@ def save( resident. Requires the engine to have been compiled with ``enable_weight_streaming=True``. See :func:`torch_tensorrt.executorch.export` for the full description. + ``zero_copy_kv=`` (default ``False``, single-method only) opts a + decode method's KV cache into in-place updates: the TensorRT + engine writes the aliased KV buffer directly instead of receiving + a staging copy that ExecuTorch copies back afterward. Unlike the + direct ``executorch.export()`` + ``to_executorch()`` path -- where + producing zero-copy KV takes two paired calls the caller must not + forget -- ``save()`` owns both steps and installs the finalization + config itself, so a single ``zero_copy_kv=True`` is enough. """ if isinstance(module, CudaGraphsTorchTensorRTModule): module = module.compiled_module @@ -849,6 +857,7 @@ def save( executorch_weight_streaming_budget_per_engine = kwargs.pop( "weight_streaming_budget_per_engine", None ) + executorch_zero_copy_kv = kwargs.pop("zero_copy_kv", False) if output_format not in accepted_formats: raise ValueError( @@ -1028,6 +1037,11 @@ def _extract_tensor(obj: Any) -> Any: "output_format='executorch' and will be ignored for " f"output_format='{output_format}'." ) + if executorch_zero_copy_kv and output_format != "executorch": + logger.warning( + "zero_copy_kv= is only used with output_format='executorch' and will " + f"be ignored for output_format='{output_format}'." + ) if output_format == "aot_inductor" and platform.system() != "Linux": raise ValueError( f"The AOT Inductor format is only supported on Linux, {platform.system()} is not a supported platform for this format" @@ -1121,6 +1135,7 @@ def _extract_tensor(obj: Any) -> Any: compile_config=executorch_compile_config, generate_etrecord=executorch_generate_etrecord, weight_streaming_budget_per_engine=executorch_weight_streaming_budget_per_engine, + zero_copy_kv=executorch_zero_copy_kv, ) else: raise RuntimeError( @@ -1238,6 +1253,7 @@ def _extract_tensor(obj: Any) -> Any: compile_config=executorch_compile_config, generate_etrecord=executorch_generate_etrecord, weight_streaming_budget_per_engine=executorch_weight_streaming_budget_per_engine, + zero_copy_kv=executorch_zero_copy_kv, ) else: raise RuntimeError( @@ -1364,6 +1380,7 @@ def _extract_tensor(obj: Any) -> Any: compile_config=executorch_compile_config, generate_etrecord=executorch_generate_etrecord, weight_streaming_budget_per_engine=executorch_weight_streaming_budget_per_engine, + zero_copy_kv=executorch_zero_copy_kv, ) else: raise RuntimeError( @@ -1402,7 +1419,7 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None "(torch_tensorrt_runtime). Reinstall torch_tensorrt with the runtime extension." ) try: - from torch_tensorrt.executorch import export + from torch_tensorrt.executorch import export, zero_copy_backend_config except ImportError: raise ImportError( "ExecuTorch is not installed. Install with: pip install " @@ -1421,6 +1438,7 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None ) generate_etrecord = kwargs.get("generate_etrecord", False) + zero_copy_kv = kwargs.get("zero_copy_kv", False) # export() runs the TRT partitioner and to_edge_transform_and_lower itself; it # defaults compile_config to get_edge_compile_config() (_check_ir_validity=False, # since the TRT execute_engine placeholder graph fails edge IR validation) when a @@ -1436,8 +1454,18 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None weight_streaming_budget_per_engine=kwargs.get( "weight_streaming_budget_per_engine" ), + zero_copy_kv=zero_copy_kv, ) - executorch_program = edge_program.to_executorch(config=kwargs.get("backend_config")) + # Unlike the direct export()+to_executorch() path -- where the two steps + # belong to different owners and pairing them is the caller's job -- save() + # owns both, so it installs the finalization pass itself. Wrapping preserves + # every field of the caller's config. save() installs the pass once; a + # backend_config that already carries it is wrapped again here, and + # finalization then raises. + backend_config = kwargs.get("backend_config") + if zero_copy_kv: + backend_config = zero_copy_backend_config(backend_config) + executorch_program = edge_program.to_executorch(config=backend_config) with open(file_path, "wb") as f: executorch_program.write_to_file(f) _write_external_tensor_data(executorch_program, file_path) diff --git a/py/torch_tensorrt/executorch/__init__.py b/py/torch_tensorrt/executorch/__init__.py index fef0943ce7b..29c9c1e4efe 100644 --- a/py/torch_tensorrt/executorch/__init__.py +++ b/py/torch_tensorrt/executorch/__init__.py @@ -33,9 +33,11 @@ def __getattr__(name: str) -> NoReturn: "TensorRTPartitioner", "TensorRTBackend", "export", + "zero_copy_backend_config", ] else: from torch_tensorrt.executorch._export import export + from torch_tensorrt.executorch._zero_copy import zero_copy_backend_config from torch_tensorrt.executorch.backend import TensorRTBackend from torch_tensorrt.executorch.partitioner import TensorRTPartitioner @@ -50,4 +52,5 @@ def get_edge_compile_config() -> "EdgeCompileConfig": "TensorRTPartitioner", "TensorRTBackend", "export", + "zero_copy_backend_config", ] diff --git a/py/torch_tensorrt/executorch/_export.py b/py/torch_tensorrt/executorch/_export.py index e793e487519..7710110be42 100644 --- a/py/torch_tensorrt/executorch/_export.py +++ b/py/torch_tensorrt/executorch/_export.py @@ -296,6 +296,43 @@ def _apply_weight_streaming_budget( specs.append(CompileSpec(WEIGHT_STREAMING_BUDGET_COMPILE_SPEC_KEY, spec_value)) +def _apply_zero_copy_kv( + program_map: dict[str, ExportedProgram], +) -> dict[str, list[str]]: + """Hand each method's aliased buffers to the engine to update in place. + + Runs immediately after the mutations are declared and before anything + partitions the program: the rewiring works from those declarations and has to + land before the partition boundary fixes the delegate's outputs. It operates + on the staged programs, never the caller's, so a reused ExportedProgram is + left intact. + + Returns, per method that actually lost an output, the engine output binding + names it elided -- narrower than the methods the caller asked about, and + narrower than the engine's full aliased_io. Only these names may be exempted + from the backend's output-binding check, so a method that dropped an output + for some other reason, or an aliased output export never rewired, is still + caught. + """ + from torch_tensorrt.executorch._zero_copy import ( + rewire_aliased_mutations_to_buffers, + ) + + elided = { + name: rewire_aliased_mutations_to_buffers(program) + for name, program in program_map.items() + } + if not any(elided.values()): + logger.warning( + "zero_copy_kv=True, but no aliased buffer mutation was found in %s, " + "so zero-copy KV was not applied.", + ", ".join(sorted(program_map)), + ) + return {} + logger.debug("zero-copy KV: elided outputs per method: %s", elided) + return {name: names for name, names in elided.items() if names} + + def _per_method_values( value: Sequence[Any] | Mapping[str, Sequence[Any] | None] | None, method_names: tuple[str, ...], @@ -328,6 +365,42 @@ def _per_method_values( return {name: list(shared) for name in method_names} +def _reject_caller_set_zero_copy_spec( + method_compile_specs: dict[str, list[Any]], +) -> None: + """Reject a caller who sets the reserved zero-copy key by hand. + + export() sets this key itself, and only for the aliased outputs it actually + elided. A hand-set value would tell the backend that outputs were taken out + of the delegate that were not, so the backend would stop rejecting a delegate + that is genuinely short an output -- dropping a real KV update silently. + + Written for this one key rather than parametrized over a reserved key: the + wording is the safety property this key carries, which no other key shares. + ``_apply_weight_streaming_budget`` rejects its own reserved key inline rather + than through this function: it also *writes* a spec, so its rejection is one + branch of a larger operation, and the two keys fail for different reasons -- + that one has a supported argument to point the caller at, this one has a + safety property that no argument can restore. + + The message names the method whose specs carry the key, as + ``_apply_weight_streaming_budget`` does. A caller who gave one shared list + gets it fanned out to every method, so the name is then whichever method came + first; that imprecision is the sibling's too. + """ + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + for name, specs in method_compile_specs.items(): + for spec in specs: + if getattr(spec, "key", None) == ZERO_COPY_KV_COMPILE_SPEC_KEY: + raise ValueError( + f"compile_specs for {name!r} may not set the reserved key " + f"'{ZERO_COPY_KV_COMPILE_SPEC_KEY}'; export() sets it " + "itself when zero_copy_kv=True elides a method's aliased " + "outputs. Setting it by hand could drop a real KV update silently." + ) + + def export( source: ExportedProgram | torch.fx.GraphModule | Mapping[str, ExportedProgram], *, @@ -345,6 +418,7 @@ def export( ) = None, compile_config: "EdgeCompileConfig | None" = None, constant_methods: Mapping[str, Any] | None = None, + zero_copy_kv: bool = False, generate_etrecord: bool = False, weight_streaming_budget_per_engine: int | None = None, ) -> "EdgeProgramManager": @@ -372,6 +446,17 @@ def export( specs name no method is not rejected here, but a backend that reads its own method name from its specs, such as the CUDA backend, then raises during lowering. + ``zero_copy_kv=True`` lets a TensorRT engine update an aliased mutable buffer + -- a KV cache -- in place, instead of ExecuTorch handing the delegate a + staging copy and copying the engine's result back afterwards. It is opt-in + rather than automatic for two reasons: the resulting ``.pte`` needs a + runtime that understands a delegate with its aliased outputs elided, so + producing one silently would break an older runner; and it is only half the + change. Finalize such a program with + ``to_executorch(torch_tensorrt.executorch.zero_copy_backend_config(config))`` + -- without it the buffer is still staged and its updates are discarded, with + no error. + ``generate_etrecord=True`` is outside the payload sharing described above. It makes ExecuTorch deep copy the whole program, so peak memory grows by roughly the size of the program including engines. @@ -419,6 +504,13 @@ def export( constant_methods (Dict[str, Any]): Methods returning a constant, such as a vocab size. Keys must be valid Python identifiers and must not name a method of ``source``. + zero_copy_kv (bool): Let a TensorRT engine write its aliased mutable buffer -- + a KV cache -- in place, instead of receiving a staging copy that ExecuTorch + copies back. Requires finalizing the returned program with + :func:`torch_tensorrt.executorch.zero_copy_backend_config`; without that + second call the buffer is still staged and every update is discarded, with + no error. ``torch_tensorrt.save(output_format="executorch", + zero_copy_kv=True)`` owns both steps and needs only the one flag. generate_etrecord (bool): Ask ExecuTorch for an ETRecord for later debugging. This copies the whole program, engines included. weight_streaming_budget_per_engine (Optional[int]): Bytes of engine weights that @@ -456,6 +548,7 @@ def export( import torch_tensorrt.dynamo.runtime.meta_ops.register_meta_ops # noqa: F401 from executorch.exir import to_edge_transform_and_lower + from executorch.exir.backend.compile_spec_schema import CompileSpec as _CompileSpec from torch_tensorrt.dynamo._exporter import _declare_aliased_kv_mutations_on_ep from torch_tensorrt.executorch import TensorRTPartitioner, get_edge_compile_config from torch_tensorrt.executorch._export_utils import ( @@ -463,6 +556,10 @@ def export( stage_exported_program, validate_engine_program, ) + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names, + ) programs, method_names = _prepare_programs( source, @@ -501,6 +598,7 @@ def export( method_compile_specs = _per_method_values( compile_specs, method_names, "compile_specs" ) + _reject_caller_set_zero_copy_spec(method_compile_specs) _apply_weight_streaming_budget( method_compile_specs, weight_streaming_budget_per_engine ) @@ -568,6 +666,8 @@ def export( ) for name, program in program_map.items() } + zero_copy_methods = _apply_zero_copy_kv(staged_programs) if zero_copy_kv else {} + rewritten: dict[str, ExportedProgram] = {} method_partitioners: dict[str, list[Partitioner]] = {} for name, program in staged_programs.items(): @@ -589,8 +689,26 @@ def export( ) # Drop this method's engine payloads as soon as they are in the graph. rewritten[name] = replace_execute_engine(program, resolved_engines.pop(name)) + trt_compile_specs = list(method_compile_specs[name]) + if name in zero_copy_methods: + # Signal to TensorRTPartitioner that this method elided aliased + # outputs, and carry the method-wide binding names it elided. The + # partitioner does not apply this spec to every partition: it + # re-derives, per engine, exactly which of a delegate's aliased + # outputs were elided and stamps only those onto only that delegate + # (see TensorRTPartitioner._partition_elided_output_names), so a + # method that lowers to several TensorRT delegates marks only the KV + # one and a plain-compute delegate beside it carries no zero-copy + # spec. Without any spec the backend rejects a delegate short of its + # bindings, which keeps an accidentally dropped output an error. + trt_compile_specs.append( + _CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(zero_copy_methods[name]), + ) + ) method_partitioners[name] = [ - TensorRTPartitioner(compile_specs=method_compile_specs[name]), + TensorRTPartitioner(compile_specs=trt_compile_specs), *extra_partitioners[name], ] diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index bac0fa59462..c4b401dd818 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -27,17 +27,21 @@ Applying only the first would leave the engine writing a discarded staging copy with nothing to copy back -- the buffer would simply never update. So neither -half may be applied alone, and neither is part of the public API on its own: a -caller opts into both together or into neither. +pass is public on its own: the rewiring is reached only through +``export(..., zero_copy_kv=True)``, and the un-staging only through +:func:`zero_copy_backend_config`, which is this module's one exported name. """ import logging import operator -from typing import Any, Dict, List, NamedTuple, Optional +from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Set import torch from torch.fx import Node +if TYPE_CHECKING: + from executorch.exir import ExecutorchBackendConfig + _LOGGER = logging.getLogger(__name__) @@ -70,9 +74,12 @@ def _aliased_inputs_by_output_index( from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( deserialize_aliased_io, ) - from torch_tensorrt.executorch.backend import _get_engine_info_for_node + from torch_tensorrt.executorch._export_utils import _resolve_engine_info - engine_info = _get_engine_info_for_node(exported_program, engine_node) + # Only aliased_io and the binding names are read, never the engine itself. + engine_info = _resolve_engine_info( + exported_program, engine_node, metadata_only=True + ) aliased_io = deserialize_aliased_io(_engine_info_str(engine_info, ALIASED_IO_IDX)) if not aliased_io: return {} @@ -99,6 +106,29 @@ def _aliased_inputs_by_output_index( return aliased +def _engine_output_binding_names(exported_program: Any, engine_node: Node) -> List[str]: + """Return one engine's output binding names, in binding (index) order. + + Resolved metadata-only: reading the record without that costs a full + re-serialization of the engine through ``TRTEngine.__getstate__``, and only + the binding names are wanted here. Callers that read this repeatedly for the + same engine memoize it themselves -- ``_resolve_engine_info`` holds no cache. + """ + from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ( + OUTPUT_BINDING_NAMES_IDX, + deserialize_binding_names, + ) + from torch_tensorrt.executorch._export_utils import _resolve_engine_info + + engine_info = _resolve_engine_info( + exported_program, engine_node, metadata_only=True + ) + names: List[str] = deserialize_binding_names( + _engine_info_str(engine_info, OUTPUT_BINDING_NAMES_IDX) + ) + return names + + class _AliasedMutation(NamedTuple): """One BUFFER_MUTATION an engine satisfies by writing the buffer in place.""" @@ -165,7 +195,7 @@ def _aliased_buffer_mutations( return mutations -def rewire_aliased_mutations_to_buffers(exported_program: Any) -> int: +def rewire_aliased_mutations_to_buffers(exported_program: Any) -> List[str]: """Declare that an aliased buffer *is* its own mutation result. Export declares an aliased KV mutation as a ``getitem`` off the engine node: @@ -186,7 +216,12 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> int: engine's in-place write would land in per-call scratch and, with the copy-back gone, be lost. It is only correct paired with the un-staging pass. - Returns the number of mutations rewired. + Returns the engine output binding names of the aliased outputs it elided, + one per rewired mutation. Only these names may later be exempted from the + backend's output-binding check -- every *other* aliased output (a user alias + on a plain, non-buffer input, which export never rewired) must still be a + delegate output, so an engine mixing the two is caught rather than silently + dropping the un-rewired one's update into scratch. """ from torch.export.graph_signature import ( ExportGraphSignature, @@ -200,9 +235,11 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> int: mutations = _aliased_buffer_mutations(exported_program) if not mutations: _LOGGER.debug("no aliased buffer mutations to rewire") - return 0 + return [] elided_by_engine: Dict[Node, List[Node]] = {} + output_names_by_engine: Dict[Node, List[str]] = {} + elided_output_names: List[str] = [] output_node = graph_module.graph.output_node() output_args = list(output_node.args[0]) output_specs = list(signature.output_specs) @@ -218,6 +255,13 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> int: output_specs[spec_index].target, ) elided_by_engine.setdefault(mutation.engine, []).append(mutation.aliased_output) + names = output_names_by_engine.get(mutation.engine) + if names is None: + names = _engine_output_binding_names(exported_program, mutation.engine) + output_names_by_engine[mutation.engine] = names + output_index = mutation.aliased_output.args[1] + if 0 <= output_index < len(names): + elided_output_names.append(names[output_index]) output_node.args = (tuple(output_args),) # Dropping every output of an engine would leave a delegate with no outputs. @@ -240,11 +284,18 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> int: graph_module.graph.eliminate_dead_code() graph_module.graph.lint() graph_module.recompile() + # The signature is replaced in place rather than by rebuilding the program: + # the graph has already been edited in place, and every other field would be + # copied across unchanged. exported_program._graph_signature = ExportGraphSignature( input_specs=list(signature.input_specs), output_specs=output_specs ) - _LOGGER.debug("rewired %d aliased mutation(s) to their buffers", len(mutations)) - return len(mutations) + _LOGGER.debug( + "rewired %d aliased mutation(s) to their buffers, eliding outputs %s", + len(mutations), + elided_output_names, + ) + return elided_output_names def _is_tensorrt_delegate(graph_module: torch.fx.GraphModule, node: Node) -> bool: @@ -265,6 +316,75 @@ def _is_tensorrt_delegate(graph_module: torch.fx.GraphModule, node: Node) -> boo return bool(getattr(module, "backend_id", None) == TensorRTBackend.__name__) +def _delegate_declares_zero_copy( + graph_module: torch.fx.GraphModule, node: Node +) -> bool: + """True when a TensorRT delegate carries the zero-copy KV compile spec. + + ``TensorRTPartitioner`` stamps this spec per partition, onto only the delegate + whose own engine had an aliased output elided (derived per engine in + ``TensorRTPartitioner._partition_elided_output_names``), so a delegate that + declares it must have had a buffer un-staged here. A method that lowers to + several TensorRT delegates therefore marks only the KV one, never the plain + compute engines beside it -- which is what keeps this cross-check from + demanding an aliased buffer from a delegate that never had one. A delegate + that declares it but un-staged nothing is a lost KV update -- the mark that + would have driven the un-staging did not survive to this pass -- and is caught + in :func:`_unstage_aliased_buffers`. + """ + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + lowered = node.args[0] if node.args else None + if not isinstance(lowered, Node) or lowered.op != "get_attr": + return False + module = getattr(graph_module, lowered.target, None) + return any( + getattr(spec, "key", None) == ZERO_COPY_KV_COMPILE_SPEC_KEY + for spec in (getattr(module, "compile_specs", None) or []) + ) + + +def _device_move_is_safe( + source: Node, h2d_copy: Any, target_device: Any, target_device_index: Any +) -> bool: + """True when moving ``source``'s spec device disturbs no other consumer. + + A placeholder's device is shared by every user, so it can only be retargeted + to the delegate's device when nothing *reads* it on another one. ExecuTorch + guards the same hazard, more strictly and only under its opt-in + ``skip_h2d_for_method_inputs``: it demands the placeholder have exactly one + user. The rule here is looser because two kinds of user impose no such + constraint and are allowed: the graph ``output`` node -- the buffer is its + own BUFFER_MUTATION result, which is exactly what zero-copy sets up and + which carries no device of its own -- and another ``_h2d_copy`` staging to + the same GPU, superseded by the un-staging when it feeds a TensorRT + delegate, and otherwise left in place, still reading a buffer that now lives + on the GPU it was copying to. Any other reader (a compute op, or a staging + to a different device) makes the move unsafe. + + The index is compared as well as the type, because ``spec.device`` is only + ``CUDA``/``CPU``: two engines resolved to ``cuda:0`` and ``cuda:1`` stage the + same buffer to different GPUs, and un-staging both would leave whichever ran + last owning the buffer while the other engine writes an address on the wrong + device. + """ + for user in source.users: + if user.op == "output": + continue + if not ( + isinstance(user, Node) + and user.op == "call_function" + and user.target is h2d_copy + ): + return False + spec = user.meta.get("spec") + if spec is None or spec.device != target_device: + return False + if spec.device_index != target_device_index: + return False + return True + + def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: """Route TensorRT delegate inputs from their staging copy back to the buffer. @@ -274,23 +394,43 @@ def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: not write in place. The placeholder's spec takes over the staging copy's device, so memory - planning puts the buffer in the delegate's device arena instead of a host - one -- which is what makes the engine's write land somewhere the caller can - still see afterwards. - - Raises when a marked buffer cannot be un-staged, because the alternative is - silence: its copy-back has already been removed, so leaving the staging in - place means the engine writes a scratch tensor and the buffer never updates. + planning puts the buffer in the delegate's device arena rather than a host + one. That is what makes handing the buffer straight to the engine valid at + all: a host-arena pointer is not something the engine can write. That move is + refused when the buffer has another consumer (see + :func:`_device_move_is_safe`), which would otherwise have its device silently + changed too. + + A failure here is a lost KV update -- unless the program has already been + through this pass, the one case where nothing is lost -- so it is raised + rather than logged: export has already removed the copy-back, so a marked + buffer left staged has the engine write per-call scratch that is then + discarded and the buffer never updates. It raises when the staging copy has + no spec or is not on CUDA, when the device move is unsafe, and -- so a + discovery miss cannot pass silently -- after the loop when any marked buffer + was never un-staged, cross-checked against each delegate's own + ``zero_copy_kv`` spec: a TensorRT delegate that declares zero-copy but + un-staged nothing is broken and names the buffer. Returns the number of delegate inputs un-staged. """ from executorch.exir.schema import DeviceType + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY h2d_copy = torch.ops.et_copy._h2d_copy.default unstaged = 0 + unstaged_placeholders: Set[Node] = set() + orphaned_stagings: List[Node] = [] + zero_copy_delegates: List[Node] = [] + unstaged_per_delegate: Dict[Node, int] = {} + for node in list(graph_module.graph.nodes): if not _is_tensorrt_delegate(graph_module, node): continue + declares_zero_copy = _delegate_declares_zero_copy(graph_module, node) + if declares_zero_copy: + zero_copy_delegates.append(node) + unstaged_per_delegate[node] = 0 new_args = list(node.args) for i, arg in enumerate(node.args[1:], start=1): if not isinstance(arg, Node) or arg.target is not h2d_copy: @@ -323,13 +463,79 @@ def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: "buffer in place and its copy-back has already been removed, " "so the update would be lost." ) + already_placed = (source_spec.device, source_spec.device_index) == ( + staged_spec.device, + staged_spec.device_index, + ) + if not already_placed and not _device_move_is_safe( + source, h2d_copy, staged_spec.device, staged_spec.device_index + ): + raise RuntimeError( + "TensorRT zero-copy KV: buffer " + f"'{source.name}' is read by a consumer the move would " + "disturb -- an op outside the delegate, or a staging copy " + "bound for a different GPU -- so placing it on this engine's " + "device would silently change that consumer's device too. " + "Export this method without zero_copy_kv, or stop sharing the " + "aliased buffer." + ) source_spec.device = staged_spec.device source_spec.device_index = staged_spec.device_index new_args[i] = source unstaged += 1 + unstaged_placeholders.add(source) + orphaned_stagings.append(arg) + if declares_zero_copy: + unstaged_per_delegate[node] += 1 node.args = tuple(new_args) + + marked_but_unstaged = [ + node + for node in graph_module.graph.nodes + if node.op == "placeholder" + and node.meta.get("_torch_tensorrt_aliased_buffer") + and node not in unstaged_placeholders + ] + if marked_but_unstaged: + names = ", ".join(repr(node.name) for node in marked_but_unstaged) + raise RuntimeError( + "TensorRT zero-copy KV: buffer(s) " + f"{names} were marked for in-place update but no TensorRT delegate " + "staging was found to un-stage. Either they never reached a " + "TensorRT delegate, which is a broken zero-copy program -- export " + "removed their copy-back, so leaving them staged has the engine " + "write per-call scratch that is discarded and the buffer never " + "updates -- or this pass has already run over the program and they " + "are wired straight to the engine already, which happens whenever " + "it is installed twice: nesting zero_copy_backend_config, " + "finalizing the same program twice, or passing " + "save(zero_copy_kv=True) a config that already carries the pass. " + "Install it once." + ) + for delegate in zero_copy_delegates: + if unstaged_per_delegate[delegate] == 0: + staged_inputs = [ + arg.args[0].name + for arg in delegate.args[1:] + if isinstance(arg, Node) + and arg.target is h2d_copy + and isinstance(arg.args[0], Node) + ] + raise RuntimeError( + "TensorRT zero-copy KV: delegate " + f"'{delegate.name}' declares zero-copy KV " + f"(compile spec '{ZERO_COPY_KV_COMPILE_SPEC_KEY}') but no aliased " + f"buffer was un-staged for it (staged inputs: {staged_inputs}). Export " + "elided its aliased outputs, so the engine now writes per-call " + "scratch that is discarded and the cache never updates." + ) + if unstaged: - graph_module.graph.eliminate_dead_code() + # Erase only the stagings we orphaned. A graph-wide eliminate_dead_code() + # in a to_out_var_pass could delete another backend's unused delegate. + for staging in dict.fromkeys(orphaned_stagings): + if not staging.users: + graph_module.graph.erase_node(staging) graph_module.graph.lint() graph_module.recompile() return unstaged @@ -362,3 +568,44 @@ def call(self, graph_module: torch.fx.GraphModule) -> Any: return inner(graph_module) return _UnstageThenToOutVar() + + +def zero_copy_backend_config( + config: Optional["ExecutorchBackendConfig"] = None, +) -> "ExecutorchBackendConfig": + """Build the ``ExecutorchBackendConfig`` a zero-copy KV program needs. + + This is the second half of ``export(..., zero_copy_kv=True)``. Export has + already removed ExecuTorch's copy-back of the aliased buffers; this installs + the pass that removes their staging, so the engine writes the caller's + buffer instead of a scratch copy that is thrown away. + + The feature is split across two calls because ``to_executorch()`` belongs to + ExecuTorch, not to Torch-TensorRT: ``export()`` hands back an + ``EdgeProgramManager`` at the Edge boundary and never sees the config the + program is finalized with. + + ``config`` is your own configuration -- every field is preserved, and a + ``to_out_var_pass`` you already set runs after the un-staging. Omit it to + start from ExecuTorch's defaults. + + .. warning:: + Finalizing a ``zero_copy_kv=True`` program *without* this config does + not raise. The engine writes a per-call staging copy that is then + discarded and the buffer never updates, which for a KV cache is wrong + output rather than a crash. Nothing downstream can detect the omission, + so pairing the two is the caller's responsibility. + + The opposite mistake does raise. ``save(..., zero_copy_kv=True)`` + installs this pass itself, so handing it the result of this function as + ``backend_config`` applies the pass twice and finalization fails. The + two entry points are mutually exclusive: use one or the other. + """ + from dataclasses import replace + + from executorch.exir import ExecutorchBackendConfig + + base = config if config is not None else ExecutorchBackendConfig() + return replace( + base, to_out_var_pass=unstage_aliased_buffers_pass(base.to_out_var_pass) + ) diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index 799732c06a5..c38091934d0 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -33,14 +33,17 @@ _BINDING_DELIM = "%" -# CompileSpec key by which export() tells this backend which aliased outputs it -# deliberately took out of the delegate. Its value is the JSON list of those -# engine output binding names (see _serialize_elided_output_names), NOT a bare -# flag: export only elides the aliased outputs backed by a registered buffer, so -# the backend must exempt exactly those and still reject a delegate that dropped -# any other binding. It travels on the partitioner's DelegationSpec, the only -# channel from the export call down to preprocess. Without it a delegate short of -# its aliased outputs is a bug, not a zero-copy program, and stays an error. +# CompileSpec key naming the aliased outputs a delegate deliberately does not +# carry. Its value is the JSON list of those engine output binding names (see +# _serialize_elided_output_names), NOT a bare flag: only the aliased outputs +# backed by a registered buffer are elided, so the backend must exempt exactly +# those and still reject a delegate that dropped any other binding. export() +# appends a method-wide instance to signal the opt-in; TensorRTPartitioner strips +# that one and re-derives the per-engine value it puts on each delegate's +# DelegationSpec, the only channel that reaches preprocess (see +# TensorRTPartitioner._partition_elided_output_names). Without it a delegate +# short of its aliased outputs is a bug, not a zero-copy program, and stays an +# error. ZERO_COPY_KV_COMPILE_SPEC_KEY = "zero_copy_kv" diff --git a/py/torch_tensorrt/executorch/partitioner.py b/py/torch_tensorrt/executorch/partitioner.py index 559eb271de5..cc9dc0c16f8 100644 --- a/py/torch_tensorrt/executorch/partitioner.py +++ b/py/torch_tensorrt/executorch/partitioner.py @@ -1,7 +1,7 @@ # ExecuTorch partitioner: partition by execute_engine nodes. import logging -from typing import Callable, Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Set, Tuple import torch from executorch.exir.backend.compile_spec_schema import CompileSpec @@ -13,12 +13,18 @@ from executorch.exir.backend.utils import tag_constant_data from torch.export import ExportedProgram from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition +from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ( + OUTPUT_BINDING_NAMES_IDX, + deserialize_binding_names, +) from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import DEVICE_IDX from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, TensorRTBackend, _get_engine_info_for_node, _get_engine_nodes_in, _parse_device_id, + _serialize_elided_output_names, ) from torch_tensorrt.executorch.operator_support import TensorRTOperatorSupport @@ -127,6 +133,20 @@ def __init__( ) -> None: super().__init__() self.compile_specs = list(compile_specs) if compile_specs else [] + # The zero-copy KV spec is stamped per-partition in partition(), never + # applied to every partition like the rest of compile_specs. Its presence + # here only records that this method asked for zero-copy; the actual + # elided binding names are derived per engine at partition time, so a + # method that lowers to several TensorRT delegates marks only the one + # whose aliased outputs were elided. It has to stay out of the shared + # list: a plain-compute delegate carrying the spec would make the + # un-staging cross-check demand an aliased buffer it never had. + self._zero_copy_requested = any( + s.key == ZERO_COPY_KV_COMPILE_SPEC_KEY for s in self.compile_specs + ) + self._base_compile_specs = [ + s for s in self.compile_specs if s.key != ZERO_COPY_KV_COMPILE_SPEC_KEY + ] # Mirror CudaPartitioner: a target_device CompileSpec drives ExecuTorch's # PropagateDevicePass, which tags delegate I/O TensorSpecs with the device # and serializes it into the .pte's extra_tensor_info. When the caller pins @@ -134,8 +154,12 @@ def __init__( # its own engine node in partition() (engine nodes are not available here) # so a cuda:N engine is not mislabeled cuda:0. self._has_explicit_target_device = any( - s.key == _TARGET_DEVICE_COMPILE_SPEC_KEY for s in self.compile_specs + s.key == _TARGET_DEVICE_COMPILE_SPEC_KEY for s in self._base_compile_specs ) + # ExecuTorch partitioners conventionally hold a delegation_spec. partition() + # builds a fresh DelegationSpec per partition and never reads this one; the + # only reader is _export._declared_method_name, on a partitioner the caller + # passes to export(). self.delegation_spec = DelegationSpec( backend_id=TensorRTBackend.__name__, compile_specs=self.compile_specs, @@ -176,6 +200,64 @@ def _resolve_target_device_for_partition( ) return b"cuda:0" + def _partition_elided_output_names( + self, exported_program: ExportedProgram, partition: Partition + ) -> Set[str]: + """Engine output binding names this partition's delegate legitimately drops. + + Zero-copy KV elides an engine's aliased output when its aliased input is a + buffer export rewired to be written in place -- marked on the placeholder + with ``_torch_tensorrt_aliased_buffer``. This is derived from THIS + partition's own engine (its ``aliased_io`` paired with the marks on its own + input placeholders), never from a method-wide name list, so a second engine + that happens to share an output binding name is never told it may drop that + binding. That is what lets a real lost output on the plain delegate still + raise while the KV delegate's genuine elision is exempted. + + Any extraction failure returns an empty set: the delegate then carries every + binding and a genuinely missing aliased output stays an error in the + backend's ``_validate_output_binding_order``. + """ + from torch_tensorrt.executorch._zero_copy import _aliased_inputs_by_output_index + + try: + engine_nodes = _get_engine_nodes_in(partition.nodes) + if len(engine_nodes) != 1: + return set() + engine = engine_nodes[0] + aliased = _aliased_inputs_by_output_index(exported_program, engine) + if not aliased: + return set() + # Only OUTPUT_BINDING_NAMES_IDX is read, never the engine itself. + engine_info = _get_engine_info_for_node( + exported_program, engine, metadata_only=True + ) + raw = engine_info[OUTPUT_BINDING_NAMES_IDX] + if isinstance(raw, bytes): + raw = raw.decode("utf-8", "replace") + output_names = deserialize_binding_names(str(raw or "")) + elided: Set[str] = set() + for output_index, input_node in aliased.items(): + if ( + isinstance(input_node, torch.fx.Node) + and input_node.meta.get("_torch_tensorrt_aliased_buffer") + and 0 <= output_index < len(output_names) + ): + elided.add(output_names[output_index]) + return elided + except Exception as e: + # Broad by design, mirroring _resolve_target_device_for_partition: any + # extraction failure must not abort the export. It degrades safely -- + # the delegate keeps every binding, so a truly-elided output is caught + # downstream rather than dropped. + logger.warning( + "zero-copy KV: could not resolve elided outputs for partition %s " + "(%s); the delegate will carry every binding.", + getattr(partition, "id", "?"), + e, + ) + return set() + def partition(self, exported_program: ExportedProgram) -> PartitionResult: capability_partitioner = CapabilityBasedPartitioner( exported_program.graph_module, @@ -189,21 +271,31 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: tag = f"tensorrt_{partition.id}" for node in partition.nodes: node.meta["delegation_tag"] = tag - if self._has_explicit_target_device: - partition_tags[tag] = self.delegation_spec - else: - partition_tags[tag] = DelegationSpec( - backend_id=TensorRTBackend.__name__, - compile_specs=self.compile_specs - + [ + specs = list(self._base_compile_specs) + if not self._has_explicit_target_device: + specs.append( + CompileSpec( + _TARGET_DEVICE_COMPILE_SPEC_KEY, + self._resolve_target_device_for_partition( + exported_program, partition + ), + ) + ) + if self._zero_copy_requested: + elided = self._partition_elided_output_names( + exported_program, partition + ) + if elided: + specs.append( CompileSpec( - _TARGET_DEVICE_COMPILE_SPEC_KEY, - self._resolve_target_device_for_partition( - exported_program, partition - ), + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(elided), ) - ], - ) + ) + partition_tags[tag] = DelegationSpec( + backend_id=TensorRTBackend.__name__, + compile_specs=specs, + ) tag_constant_data(exported_program) _keep_mutated_buffers_above_delegate(exported_program) diff --git a/tests/py/dynamo/executorch/test_backend.py b/tests/py/dynamo/executorch/test_backend.py index 604c58f7de6..c4eed938cc9 100644 --- a/tests/py/dynamo/executorch/test_backend.py +++ b/tests/py/dynamo/executorch/test_backend.py @@ -9,6 +9,7 @@ executorch = pytest.importorskip("executorch.exir") import torch # noqa: E402 +from executorch.exir.backend.compile_spec_schema import CompileSpec # noqa: E402 from torch.export.graph_signature import InputKind # noqa: E402 from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( # noqa: E402 ALIASED_IO_IDX, @@ -506,3 +507,107 @@ def test_validate_output_binding_order_rejects_permutation_of_elided_outputs(): ep, engine = _engine_partition([2, 0]) with pytest.raises(ValueError, match="engine output indices"): _validate_output_binding_order(ep, engine, ["out0", "kv0", "out2"], {"kv0"}) + + +def _aliased_edge_program(present_indices, out_names, aliased_io_map): + """A one-engine partition declaring aliased_io whose delegate keeps only + ``present_indices`` of its output bindings (the rest elided, as zero-copy + export produces). One input, ``tokens``, which the aliases point at.""" + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import serialize_aliased_io + + engine_info = [""] * SERIALIZATION_LEN + engine_info[ENGINE_IDX] = _engine_tensor(b"engine-bytes") + engine_info[INPUT_BINDING_NAMES_IDX] = "tokens" + engine_info[OUTPUT_BINDING_NAMES_IDX] = "%".join(out_names) + engine_info[ALIASED_IO_IDX] = serialize_aliased_io(aliased_io_map) + + graph = torch.fx.Graph() + tokens = graph.placeholder("tokens") + engine_node = graph.call_function(_ENGINE_OP, ([tokens], *engine_info)) + graph.output( + tuple( + graph.call_function(operator.getitem, (engine_node, i)) + for i in present_indices + ) + ) + return SimpleNamespace( + graph_module=SimpleNamespace(graph=graph), + graph_signature=SimpleNamespace( + input_specs=[ + SimpleNamespace( + kind=InputKind.USER_INPUT, arg=SimpleNamespace(name="tokens") + ) + ] + ), + constants={}, + ) + + +@pytest.mark.unit +def test_preprocess_rejects_elided_aliased_output_without_zero_copy_spec(): + """The gate at preprocess: a delegate that dropped its aliased output is a + bug unless a zero-copy compile spec says the drop was deliberate. Deleting + the gate (treating the drop as always allowed) would let this pass silently. + """ + from torch_tensorrt.executorch.backend import TensorRTBackend + + edge_program = _aliased_edge_program( + present_indices=[0], + out_names=["logits", "out_k"], + aliased_io_map={"out_k": ("tokens", "kv_cache_update")}, + ) + with pytest.raises(ValueError, match="engine output indices"): + TensorRTBackend.preprocess(edge_program, []) + + +@pytest.mark.unit +def test_preprocess_accepts_elided_aliased_output_with_zero_copy_spec(): + """The same delegate is accepted once the compile spec names the elided + binding, and only then.""" + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + TensorRTBackend, + _serialize_elided_output_names, + ) + + edge_program = _aliased_edge_program( + present_indices=[0], + out_names=["logits", "out_k"], + aliased_io_map={"out_k": ("tokens", "kv_cache_update")}, + ) + spec = CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, _serialize_elided_output_names(["out_k"]) + ) + # Does not raise, and produces a valid engine blob (aliased_io present bumps + # the blob format magic to TR02, so match the family, not the exact base). + result = TensorRTBackend.preprocess(edge_program, [spec]) + assert isinstance(result.processed_bytes, bytes) + assert result.processed_bytes[:2] == TENSORRT_MAGIC[:2] + + +@pytest.mark.unit +def test_preprocess_rejects_a_non_buffer_alias_elided_alongside_a_buffer_alias(): + """An engine mixing a buffer-backed aliased output (export rewired it and + named it in the spec) with a non-buffer-backed one (never rewired) must not + have the non-buffer one silently elided. The spec names only ``out_k``, so a + delegate that also dropped ``out_v`` is rejected -- exactly as it would be + without zero-copy.""" + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + TensorRTBackend, + _serialize_elided_output_names, + ) + + edge_program = _aliased_edge_program( + present_indices=[0], # both out_k (buffer) and out_v (non-buffer) dropped + out_names=["logits", "out_k", "out_v"], + aliased_io_map={ + "out_k": ("tokens", "kv_cache_update"), + "out_v": ("tokens", "kv_cache_update"), + }, + ) + spec = CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, _serialize_elided_output_names(["out_k"]) + ) + with pytest.raises(ValueError, match="engine output indices"): + TensorRTBackend.preprocess(edge_program, [spec]) diff --git a/tests/py/dynamo/executorch/test_edge_cases.py b/tests/py/dynamo/executorch/test_edge_cases.py index 4e5d7323d0d..f5cf9b417c9 100644 --- a/tests/py/dynamo/executorch/test_edge_cases.py +++ b/tests/py/dynamo/executorch/test_edge_cases.py @@ -56,7 +56,7 @@ def test_save_as_executorch_uses_public_lowering_and_persists_data( backend_config=backend_config, ) - # The complete set of lowering options _save_as_executorch forwards. The five this + # The complete set of lowering options _save_as_executorch forwards. The six this # test does not pass are still forwarded explicitly, as None or False rather than # left out. backend_config is absent by design -- it is not a lowering option and is # routed to to_executorch() below. @@ -69,7 +69,10 @@ def test_save_as_executorch_uses_public_lowering_and_persists_data( constant_methods=None, generate_etrecord=False, weight_streaming_budget_per_engine=None, + zero_copy_kv=False, ) + # zero_copy_kv defaults off, so the backend_config reaches to_executorch() + # unwrapped and the KV buffer keeps its staging and its copy-back. edge.to_executorch.assert_called_once_with(config=backend_config) program.write_to_file.assert_called_once() program.write_tensor_data_to_file.assert_called_once_with(str(tmp_path)) diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py index 3801d9a2cab..f10c12ea243 100644 --- a/tests/py/dynamo/executorch/test_export.py +++ b/tests/py/dynamo/executorch/test_export.py @@ -500,6 +500,167 @@ def test_export_returns_edge_and_forwards_all_options(monkeypatch): assert compile_specs == [compile_spec] +def _patch_declare(monkeypatch): + """Record the programs the declaration pass sees, and hand each one back. + + export() imports the symbol inside its own body, so the patch has to land on + the module that owns it. The stub takes ``**kw`` because export() passes + ``copyback_buffers=``. + """ + import torch_tensorrt.dynamo._exporter as dynamo_exporter + + seen = [] + + def _declare(program, **kw): + seen.append(program) + return program + + monkeypatch.setattr( + dynamo_exporter, "_declare_aliased_kv_mutations_on_ep", _declare + ) + return seen + + +def _patch_rewire(monkeypatch, elided_names=("kv",)): + import torch_tensorrt.executorch._zero_copy as zero_copy + + seen = [] + + def _rewire(program): + seen.append(program) + return list(elided_names) + + monkeypatch.setattr(zero_copy, "rewire_aliased_mutations_to_buffers", _rewire) + return seen + + +@pytest.mark.unit +def test_export_zero_copy_kv_rewires_every_method(monkeypatch): + """The opt-in is what makes the aliased buffers zero-copy; nothing else does. + + It has to run per method and after the declaration, since it works from the + mutations that declaration produced. + """ + export_module, lower = _patch_lowering(monkeypatch) + declared = _patch_declare(monkeypatch) + rewired = _patch_rewire(monkeypatch) + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + + export_module.export( + {"prefill": prefill, "decode": decode}, + partitioners={"prefill": [object()], "decode": [object()]}, + zero_copy_kv=True, + ) + + assert declared == [prefill, decode] + assert rewired == [prefill, decode] + # The backend rejects a delegate missing its aliased outputs unless it is + # told the omission was deliberate, and this spec is the only channel. + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _elided_output_names, + ) + + for pipeline in lower.call_args.kwargs["partitioner"].values(): + assert any( + spec.key == ZERO_COPY_KV_COMPILE_SPEC_KEY + for spec in pipeline[0].compile_specs + ) + # The spec carries the elided binding names, not a bare flag. + assert _elided_output_names(pipeline[0].compile_specs) == {"kv"} + + +@pytest.mark.unit +def test_export_does_not_exempt_a_method_that_kept_all_its_outputs(monkeypatch): + """The exemption is per method and only where an output was actually elided. + + A method that lost an output for some other reason must still be caught by + the backend's output-binding check. + """ + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + export_module, lower = _patch_lowering(monkeypatch) + import torch_tensorrt.executorch._zero_copy as zero_copy + + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + elided = {prefill: [], decode: ["k0", "k1"]} + monkeypatch.setattr( + zero_copy, "rewire_aliased_mutations_to_buffers", lambda p: elided[p] + ) + + export_module.export( + {"prefill": prefill, "decode": decode}, + partitioners={"prefill": [object()], "decode": [object()]}, + zero_copy_kv=True, + ) + + exempt = { + name: any( + spec.key == ZERO_COPY_KV_COMPILE_SPEC_KEY + for spec in pipeline[0].compile_specs + ) + for name, pipeline in lower.call_args.kwargs["partitioner"].items() + } + assert exempt == {"prefill": False, "decode": True} + + +@pytest.mark.unit +def test_export_zero_copy_kv_keeps_the_weight_streaming_spec(monkeypatch): + """Both options stamp the same partitioner, and one must not displace the other. + + The zero-copy spec is appended to a copy of the method's compile specs, so a + budget baked in earlier has to still reach the delegate. + """ + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + from torch_tensorrt.executorch.partitioner import ( + WEIGHT_STREAMING_BUDGET_COMPILE_SPEC_KEY, + ) + + export_module, lower = _patch_lowering(monkeypatch) + _patch_rewire(monkeypatch) + + export_module.export( + FakeExportedProgram(), + zero_copy_kv=True, + weight_streaming_budget_per_engine=1 << 20, + ) + + keys = {spec.key for spec in lower.call_args.kwargs["partitioner"][0].compile_specs} + assert keys == { + WEIGHT_STREAMING_BUDGET_COMPILE_SPEC_KEY, + ZERO_COPY_KV_COMPILE_SPEC_KEY, + } + + +@pytest.mark.unit +def test_export_leaves_kv_buffers_staged_by_default(monkeypatch): + """Zero-copy is opt-in: a .pte that elides its aliased outputs cannot be run + by a runtime that predates the feature, so export must not produce one + unasked.""" + export_module, lower = _patch_lowering(monkeypatch) + rewired = _patch_rewire(monkeypatch) + + export_module.export(FakeExportedProgram()) + + assert rewired == [] + + +@pytest.mark.unit +def test_export_warns_when_zero_copy_kv_has_nothing_to_do(monkeypatch, caplog): + """Asking for zero-copy on a model with no engine-aliased buffer is not an + error, but silently doing nothing would leave the caller expecting a speedup + that is not coming.""" + export_module, lower = _patch_lowering(monkeypatch) + _patch_rewire(monkeypatch, elided_names=()) + + with caplog.at_level("WARNING", logger=export_module.logger.name): + export_module.export(FakeExportedProgram(), zero_copy_kv=True) + + assert "no aliased buffer mutation was found" in caplog.text + + @pytest.mark.unit def test_export_preserves_independent_method_mapping(monkeypatch): prefill = FakeExportedProgram() @@ -1677,3 +1838,47 @@ def counting_get_engine_info_from_state(engine_obj, *, metadata_only=False): # is the tensor path fetching the bytes, which it re-resolves rather than trusting # the metadata-only record. Without the handoff the first entry would repeat. assert calls == [True, False] + + +@pytest.mark.unit +def test_export_rejects_caller_set_zero_copy_compile_spec(monkeypatch): + """The zero-copy key is reserved. export() sets it itself and only for the + outputs it elided; a hand-set value would exempt the delegate from the + output-binding check and could drop a real KV update silently. + + A bare list of compile specs is one unnamed list, but export() still fans it + out per method, so the message names the method it landed on. + """ + from executorch.exir.backend.compile_spec_schema import CompileSpec + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + export_module, lower = _patch_lowering(monkeypatch) + with pytest.raises(ValueError, match="reserved key") as excinfo: + export_module.export( + FakeExportedProgram(), + compile_specs=[CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")], + ) + assert "compile_specs for 'forward'" in str(excinfo.value) + + +@pytest.mark.unit +def test_export_rejects_caller_set_zero_copy_compile_spec_per_method(monkeypatch): + """The reserved-key rejection also covers the per-method mapping form, and + names the one method whose list carries the key rather than the method the + mapping happens to start with.""" + from executorch.exir.backend.compile_spec_schema import CompileSpec + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + export_module, lower = _patch_lowering(monkeypatch) + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + with pytest.raises(ValueError, match="reserved key") as excinfo: + export_module.export( + {"prefill": prefill, "decode": decode}, + partitioners={"prefill": [object()], "decode": [object()]}, + compile_specs={ + "prefill": [], + "decode": [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")], + }, + ) + assert "compile_specs for 'decode'" in str(excinfo.value) diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index 225c696fc08..10d017eafd1 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -10,7 +10,8 @@ The interesting failures are all silent -- a rewired mutation whose buffer is still staged simply never updates -- so most of what is asserted here is which -mutations are left alone and which mis-shapes raise. +mutations are left alone, which mis-shapes raise, and that a marked buffer that +is never un-staged is caught rather than dropped. """ import operator @@ -22,6 +23,7 @@ import torch # noqa: E402 import torch_tensorrt # noqa: E402 +from executorch.exir.backend.compile_spec_schema import CompileSpec # noqa: E402 from executorch.exir.delegate import executorch_call_delegate # noqa: E402 from executorch.exir.schema import DeviceType # noqa: E402 from torch.export.exported_program import ( # noqa: E402 @@ -30,6 +32,9 @@ TensorArgument, ) from torch_tensorrt.executorch import _zero_copy as Z # noqa: E402 +from torch_tensorrt.executorch.backend import ( # noqa: E402 + ZERO_COPY_KV_COMPILE_SPEC_KEY, +) # The graphs below are built around torch.ops.tensorrt.execute_engine, which only # exists once the Torch-TensorRT runtime operator library has loaded. @@ -43,12 +48,22 @@ def _patch_engine_metadata(monkeypatch, *, aliased_io, input_names, output_names """Make every engine node report one fixed set of bindings and aliases.""" import torch_tensorrt.dynamo.runtime._serialized_engine_layout as layout import torch_tensorrt.dynamo.runtime._TorchTensorRTModule as trt_module - import torch_tensorrt.executorch.backend as backend + import torch_tensorrt.executorch._export_utils as export_utils info = ["x"] * (layout.ALIASED_IO_IDX + 1) info[layout.INPUT_BINDING_NAMES_IDX] = "IN" info[layout.OUTPUT_BINDING_NAMES_IDX] = "OUT" - monkeypatch.setattr(backend, "_get_engine_info_for_node", lambda ep, node: info) + + # The rewiring resolves engine info through _resolve_engine_info (the node is + # still an execute_engine at this stage), so that is what to fake. The stub + # requires metadata_only: without it the read goes through + # TRTEngine.__getstate__ and re-serializes the whole engine to recover the + # binding names and aliased_io, which are the only fields wanted here. + def _fake_resolve(ep, node, *, metadata_only=False): + assert metadata_only, "zero-copy reads binding metadata, not the engine" + return info + + monkeypatch.setattr(export_utils, "_resolve_engine_info", _fake_resolve) monkeypatch.setattr(trt_module, "deserialize_aliased_io", lambda s: aliased_io) monkeypatch.setattr( layout, @@ -111,7 +126,8 @@ def test_rewire_points_the_mutation_at_its_buffer_and_marks_it(monkeypatch): With the mutation bound to the placeholder there is nothing for ExecuTorch to copy back, and with no other user the getitem leaves the graph -- which - is what takes the aliased output out of the delegate. + is what takes the aliased output out of the delegate. The elided output's + binding name is returned so the backend can exempt exactly that one. """ program, k_buffer, k_out = _kv_program() _patch_engine_metadata( @@ -121,7 +137,7 @@ def test_rewire_points_the_mutation_at_its_buffer_and_marks_it(monkeypatch): output_names=["logits", "out_k"], ) - assert Z.rewire_aliased_mutations_to_buffers(program) == 1 + assert Z.rewire_aliased_mutations_to_buffers(program) == ["out_k"] specs = program._graph_signature.output_specs assert specs[0].kind == OutputKind.BUFFER_MUTATION @@ -153,7 +169,7 @@ def test_rewire_leaves_mutations_the_engine_does_not_alias(monkeypatch, mutation output_names=["logits", "out_k"], ) - assert Z.rewire_aliased_mutations_to_buffers(program) == 0 + assert Z.rewire_aliased_mutations_to_buffers(program) == [] assert program._graph_signature.output_specs[0] is original_spec assert "_torch_tensorrt_aliased_buffer" not in k_buffer.meta @@ -168,7 +184,7 @@ def test_rewire_is_a_noop_without_aliased_io(monkeypatch): output_names=["logits", "out_k"], ) - assert Z.rewire_aliased_mutations_to_buffers(program) == 0 + assert Z.rewire_aliased_mutations_to_buffers(program) == [] assert "_torch_tensorrt_aliased_buffer" not in k_buffer.meta @@ -215,7 +231,9 @@ def test_rewire_rejects_an_engine_whose_every_output_is_aliased(monkeypatch): Z.rewire_aliased_mutations_to_buffers(program) -def _staged_delegate_graph(*, backend_id="TensorRTBackend", device=DeviceType.CUDA): +def _staged_delegate_graph( + *, backend_id="TensorRTBackend", device=DeviceType.CUDA, compile_specs=None +): """A lowered graph: delegate(lowered, _h2d_copy(k_buffer), _h2d_copy(tokens)).""" graph = torch.fx.Graph() k_buffer = graph.placeholder("b_k_0") @@ -230,7 +248,9 @@ def _staged_delegate_graph(*, backend_id="TensorRTBackend", device=DeviceType.CU graph.output((delegate,)) root = torch.nn.Module() - root.lowered_module_0 = SimpleNamespace(backend_id=backend_id) + root.lowered_module_0 = SimpleNamespace( + backend_id=backend_id, compile_specs=compile_specs + ) graph_module = torch.fx.GraphModule(root, graph) for node, spec_device in ( @@ -261,6 +281,8 @@ def test_unstage_feeds_the_buffer_straight_to_the_delegate(): # The other input is an ordinary one and keeps its staging copy. assert delegate.args[2] is not None assert delegate.args[2].target is torch.ops.et_copy._h2d_copy.default + # The orphaned staging is erased, but only it -- the other staging survives. + assert staged_k not in graph_module.graph.nodes @pytest.mark.unit @@ -273,15 +295,97 @@ def test_unstage_keeps_staging_for_an_unmarked_buffer(): @pytest.mark.unit -def test_unstage_ignores_another_backends_delegate(): - """Only a TensorRT engine promises the in-place write.""" +def test_unstage_raises_for_a_marked_buffer_on_another_backends_delegate(): + """Only a TensorRT engine promises the in-place write, so a marked buffer + routed to another backend's delegate is never un-staged -- and because + export has already dropped its copy-back, that is a broken program, not a + silent no-op.""" graph_module, k_buffer, staged_k, delegate = _staged_delegate_graph( backend_id="CudaBackend" ) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - assert Z._unstage_aliased_buffers(graph_module) == 0 - assert delegate.args[1] is staged_k + with pytest.raises(RuntimeError, match="no TensorRT delegate staging"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_raises_when_a_marked_buffer_is_never_unstaged(): + """A marked buffer that reaches no TensorRT delegate at all must raise, not + return 0. Its copy-back is already gone, so leaving it staged would silently + discard every update.""" + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + graph.output((k_buffer,)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="no TensorRT delegate staging"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_raises_when_a_zero_copy_delegate_unstaged_nothing(): + """A TensorRT delegate that declares zero-copy but had no buffer un-staged + (its mark did not survive to this pass) is unambiguously broken and must + raise, naming the delegate.""" + graph_module, k_buffer, staged_k, delegate = _staged_delegate_graph( + compile_specs=[CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")] + ) + # k_buffer deliberately left unmarked: nothing gets un-staged for the delegate. + + with pytest.raises(RuntimeError, match="declares zero-copy KV"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_leaves_another_backends_same_gpu_staging_in_place(): + """A marked buffer read by a second backend on the *same* GPU is still moved. + + This is the accept side of the ``_h2d_copy`` allowance in + ``_device_move_is_safe``: the other backend's delegate keeps its staging copy, + which after the move reads a buffer already resident on the GPU it was copying + to, so nothing it sees changes. The function's other accept branch, the graph + ``output`` node, is pinned by + ``test_unstage_allows_a_buffer_that_is_also_its_mutation_output``; every other + unit test that gives the buffer a second ``_h2d_copy`` pins a refusal, so a + rule that allowed no second staging at all would still pass all of those. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_trt = graph.call_function(h2d, (k_buffer,)) + staged_other = graph.call_function(h2d, (k_buffer,)) + lowered_trt = graph.get_attr("lowered_module_0") + lowered_other = graph.get_attr("lowered_module_1") + delegate_trt = graph.call_function( + executorch_call_delegate, (lowered_trt, staged_trt) + ) + delegate_other = graph.call_function( + executorch_call_delegate, (lowered_other, staged_other) + ) + graph.output((k_buffer, delegate_trt, delegate_other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="CudaBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + staged_trt.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_other.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + assert Z._unstage_aliased_buffers(graph_module) == 1 + + assert delegate_trt.args[1] is k_buffer + assert delegate_other.args[1] is staged_other + assert staged_other in graph_module.graph.nodes + assert k_buffer.meta["spec"].device == DeviceType.CUDA + assert k_buffer.meta["spec"].device_index == 0 @pytest.mark.unit @@ -305,6 +409,227 @@ def test_unstage_raises_when_the_staging_copy_has_no_spec(): Z._unstage_aliased_buffers(graph_module) +@pytest.mark.unit +def test_unstage_allows_a_buffer_that_is_also_its_mutation_output(): + """The zero-copy shape itself: the marked buffer is both the delegate's + staged input and its own BUFFER_MUTATION graph output. The output-node + reference carries no device of its own, so the device move must be allowed -- + the real lowered KV graph looks exactly like this.""" + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + lowered = graph.get_attr("lowered_module_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + delegate = graph.call_function(executorch_call_delegate, (lowered, staged_k)) + graph.output((k_buffer, delegate)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=3) + staged_k.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=3) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + assert Z._unstage_aliased_buffers(graph_module) == 1 + assert delegate.args[1] is k_buffer + assert k_buffer.meta["spec"].device == DeviceType.CUDA + + +@pytest.mark.unit +def test_unstage_refuses_to_move_a_shared_buffer(): + """A buffer read by a consumer other than its TensorRT delegate staging + cannot have its device moved -- that would silently retarget the other + consumer too, exactly what ExecuTorch's PropagateDevicePass rejects. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + lowered = graph.get_attr("lowered_module_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + other = graph.call_function(torch.add, (k_buffer, k_buffer)) + delegate = graph.call_function(executorch_call_delegate, (lowered, staged_k)) + graph.output((delegate, other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=3) + staged_k.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=3) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer the move would disturb"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus(): + """One buffer staged to two TensorRT delegates on *different* GPUs cannot be + un-staged for either: a spec carries one device index, so whichever engine + lost the race would be handed an address on the other's GPU. ``spec.device`` + is only CUDA/CPU, so it is the device-index comparison in + ``_device_move_is_safe`` that refuses the first delegate here. That is what + separates this from the supported two-delegates-one-GPU shape, but it is not + what this test pins: two branches of that function refuse the shape in + sequence, and were the index comparison gone, delegate 0 would be un-staged + and the direct-consumer branch would refuse delegate 1 instead, leaving this + test green. The index comparison itself is pinned by + ``test_unstage_refuses_a_buffer_another_backend_stages_to_a_different_gpu``. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_0 = graph.call_function(h2d, (k_buffer,)) + staged_1 = graph.call_function(h2d, (k_buffer,)) + lowered_0 = graph.get_attr("lowered_module_0") + lowered_1 = graph.get_attr("lowered_module_1") + delegate_0 = graph.call_function(executorch_call_delegate, (lowered_0, staged_0)) + delegate_1 = graph.call_function(executorch_call_delegate, (lowered_1, staged_1)) + graph.output((k_buffer, delegate_0, delegate_1)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + staged_0.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_1.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=1) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer the move would disturb"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_refuses_a_buffer_another_backend_stages_to_a_different_gpu(): + """A marked buffer staged to a TensorRT delegate on cuda:0 and to a + *non*-TensorRT delegate on cuda:1 cannot be moved either. + + Un-staging skips the other backend's delegate, so its staging copy keeps + reading the buffer while staging it to cuda:1, and re-homing the buffer onto + the TensorRT engine's cuda:0 would move the source of that read to the wrong + GPU. The device-index comparison is the only thing that refuses this shape: + the other staging is an ``_h2d_copy`` like every supported one, and because + it is never un-staged its delegate never becomes the kind of direct consumer + the shared-buffer rule catches. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_trt = graph.call_function(h2d, (k_buffer,)) + staged_other = graph.call_function(h2d, (k_buffer,)) + lowered_trt = graph.get_attr("lowered_module_0") + lowered_other = graph.get_attr("lowered_module_1") + delegate_trt = graph.call_function( + executorch_call_delegate, (lowered_trt, staged_trt) + ) + delegate_other = graph.call_function( + executorch_call_delegate, (lowered_other, staged_other) + ) + graph.output((k_buffer, delegate_trt, delegate_other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="CudaBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + staged_trt.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_other.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=1) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer the move would disturb"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_refuses_to_rehome_a_buffer_already_on_another_gpu(): + """A buffer already resident on cuda:0 is not re-homed to a second TensorRT + delegate's cuda:1. + + Whether the move needs checking at all is decided by comparing the buffer's + device *and index* against the staging copy's. Comparing the device alone + would call this buffer already placed -- both ends are CUDA -- skip the + check, and overwrite the index with the second engine's, leaving the first + engine holding an address on the other GPU. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_0 = graph.call_function(h2d, (k_buffer,)) + staged_1 = graph.call_function(h2d, (k_buffer,)) + lowered_0 = graph.get_attr("lowered_module_0") + lowered_1 = graph.get_attr("lowered_module_1") + delegate_0 = graph.call_function(executorch_call_delegate, (lowered_0, staged_0)) + delegate_1 = graph.call_function(executorch_call_delegate, (lowered_1, staged_1)) + graph.output((k_buffer, delegate_0, delegate_1)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_0.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_1.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=1) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer the move would disturb"): + Z._unstage_aliased_buffers(graph_module) + assert k_buffer.meta["spec"].device_index == 0 + + +@pytest.mark.unit +def test_zero_copy_backend_config_keeps_the_callers_config(): + """It composes onto a config rather than replacing one: a caller finalizing + a zero-copy program still needs their own memory planning and passes.""" + from executorch.exir import ExecutorchBackendConfig + + inner = object() + base = ExecutorchBackendConfig(to_out_var_pass=inner, emit_stacktrace=True) + + config = torch_tensorrt.executorch.zero_copy_backend_config(base) + + assert config.emit_stacktrace is True + assert config.memory_planning_pass is base.memory_planning_pass + assert config.to_out_var_pass is not inner + # The caller's to_out_var_pass is not dropped, it is run after the un-staging. + graph_module, k_buffer, _, delegate = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + seen = [] + base = ExecutorchBackendConfig(to_out_var_pass=lambda gm: seen.append(gm)) + torch_tensorrt.executorch.zero_copy_backend_config(base).to_out_var_pass.call( + graph_module + ) + assert seen == [graph_module] + assert delegate.args[1] is k_buffer + + +@pytest.mark.unit +def test_zero_copy_backend_config_defaults_to_executorch_defaults(): + """Called with no config it starts from ExecuTorch's defaults, and the one + field it replaces is to_out_var_pass, wrapped in the un-staging pass.""" + from executorch.exir import ExecutorchBackendConfig + + config = torch_tensorrt.executorch.zero_copy_backend_config() + + defaults = ExecutorchBackendConfig() + # The un-staging pass specifically, not merely "some object that is not the + # default" -- which is all any wrapper would have to be. + assert type(config.to_out_var_pass).__name__ == "_UnstageThenToOutVar" + assert type(config.memory_planning_pass) is type(defaults.memory_planning_pass) + assert type(config.sym_shape_eval_pass) is type(defaults.sym_shape_eval_pass) + assert config.emit_stacktrace == defaults.emit_stacktrace + + @pytest.mark.unit def test_unstage_pass_runs_the_inner_pass_after_unstaging(): """A caller's own to_out_var_pass has to survive being composed with.""" @@ -321,3 +646,543 @@ def inner(gm): assert seen == [True] assert result == "inner-result" + + +# -------------------------------------------------------------------------- +# Multi-delegate: a method that lowers to two TensorRT engines -- one with an +# aliased+elided KV buffer, one plain-compute engine with none. The zero-copy +# CompileSpec is appended once, to a single TensorRTPartitioner, and the +# partitioner must stamp it onto ONLY the delegate whose own engine had an +# aliased output elided. Stamped partition-wide instead, the plain delegate +# declares zero-copy while un-staging nothing, and the un-staging cross-check +# then rejects an otherwise-correct program. +# -------------------------------------------------------------------------- + + +def _no_op_engine_node(graph, input_nodes, *, aliased_io, input_names, output_names): + """A no_op_placeholder_for_execute_engine node with inlined engine info. + + Mirrors what replace_execute_engine() produces before partitioning: args are + ``(input_list, *engine_info)`` with the binding names and aliased_io in their + serialized wire form, so the partitioner's real per-engine resolution + (_resolve_engine_info / _aliased_inputs_by_output_index) runs unmocked. + """ + from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ( + ALIASED_IO_IDX, + DEVICE_IDX, + ENGINE_IDX, + INPUT_BINDING_NAMES_IDX, + OUTPUT_BINDING_NAMES_IDX, + SERIALIZATION_LEN, + SERIALIZED_ENGINE_BINDING_DELIM, + ) + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import serialize_aliased_io + + info = [""] * SERIALIZATION_LEN + info[ENGINE_IDX] = "" # not read by the partitioner's elision resolution + info[DEVICE_IDX] = "0" + info[INPUT_BINDING_NAMES_IDX] = SERIALIZED_ENGINE_BINDING_DELIM.join(input_names) + info[OUTPUT_BINDING_NAMES_IDX] = SERIALIZED_ENGINE_BINDING_DELIM.join(output_names) + info[ALIASED_IO_IDX] = serialize_aliased_io(aliased_io) + return graph.call_function( + torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default, + (list(input_nodes), *info), + ) + + +def _two_engine_program(): + """engine_a(k_buffer, tokens) -> (logits, out_k[aliased]); engine_b(x) -> (y). + + ``k_buffer`` carries ``_torch_tensorrt_aliased_buffer`` (rewiring already + ran); engine_a aliases its ``out_k`` output onto it, engine_b aliases nothing. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + tokens = graph.placeholder("tokens") + x = graph.placeholder("x") + engine_a = _no_op_engine_node( + graph, + [k_buffer, tokens], + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in", "tokens"], + output_names=["logits", "out_k"], + ) + engine_b = _no_op_engine_node( + graph, + [x], + aliased_io={}, + input_names=["x_in"], + output_names=["y"], + ) + graph.output((engine_a, engine_b)) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=SimpleNamespace(buffers_to_mutate={}, inputs_to_buffers={}), + constants={}, + ) + return program, engine_a, engine_b + + +def _partition_two_engines(program, engine_a, engine_b, monkeypatch): + """Run the real TensorRTPartitioner, one partition per engine node.""" + from torch_tensorrt.executorch.backend import _serialize_elided_output_names + from torch_tensorrt.executorch.partitioner import TensorRTPartitioner + + class _FakeCap: + def __init__(self, graph_module, *args, **kwargs): + self._engines = [engine_a, engine_b] + + def propose_partitions(self): + return [ + SimpleNamespace(id=i, nodes=[node]) + for i, node in enumerate(self._engines) + ] + + monkeypatch.setattr( + "torch_tensorrt.executorch.partitioner.CapabilityBasedPartitioner", _FakeCap + ) + monkeypatch.setattr( + "torch_tensorrt.executorch.partitioner.tag_constant_data", + lambda exported_program: None, + ) + # Appended once, method-wide -- exactly how export() builds the partitioner. + partitioner = TensorRTPartitioner( + compile_specs=[ + CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(["out_k"]), + ) + ] + ) + return partitioner.partition(program) + + +def _zero_copy_names(compile_specs): + from torch_tensorrt.executorch.backend import _elided_output_names + + return _elided_output_names(compile_specs) + + +@pytest.mark.unit +def test_partition_stamps_zero_copy_only_on_the_kv_delegate(monkeypatch): + """The KV delegate carries the zero-copy spec naming its own elided binding, + and the plain-compute delegate carries no zero-copy spec at all. + + The method-wide spec the partitioner is constructed with must not reach every + partition: the names it holds are the method's, and only this engine's own + aliased_io says which of them are its. + """ + program, engine_a, engine_b = _two_engine_program() + result = _partition_two_engines(program, engine_a, engine_b, monkeypatch) + + kv_specs = result.partition_tags["tensorrt_0"].compile_specs + plain_specs = result.partition_tags["tensorrt_1"].compile_specs + assert _zero_copy_names(kv_specs) == {"out_k"} + assert _zero_copy_names(plain_specs) is None + + +@pytest.mark.unit +def test_multi_delegate_zero_copy_lowers_without_false_raise(monkeypatch): + """A correct two-delegate zero-copy program survives the whole pipeline: run + the real partitioner, build the lowered two-delegate graph from the specs it + produced, and un-stage. + + The KV buffer is un-staged and the plain delegate is left alone. A plain + delegate stamped zero-copy would instead make _unstage_aliased_buffers raise + "declares zero-copy KV ... but no aliased buffer was un-staged for it" over a + program that is correct. + """ + program, engine_a, engine_b = _two_engine_program() + result = _partition_two_engines(program, engine_a, engine_b, monkeypatch) + kv_specs = result.partition_tags["tensorrt_0"].compile_specs + plain_specs = result.partition_tags["tensorrt_1"].compile_specs + + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + x = graph.placeholder("x") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + staged_x = graph.call_function(h2d, (x,)) + kv_lowered = graph.get_attr("lowered_module_0") + plain_lowered = graph.get_attr("lowered_module_1") + kv_delegate = graph.call_function(executorch_call_delegate, (kv_lowered, staged_k)) + plain_delegate = graph.call_function( + executorch_call_delegate, (plain_lowered, staged_x) + ) + graph.output((k_buffer, kv_delegate, plain_delegate)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=kv_specs + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=plain_specs + ) + graph_module = torch.fx.GraphModule(root, graph) + for node, dev in ( + (k_buffer, DeviceType.CPU), + (x, DeviceType.CPU), + (staged_k, DeviceType.CUDA), + (staged_x, DeviceType.CUDA), + ): + node.meta["spec"] = SimpleNamespace(device=dev, device_index=3) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + assert Z._unstage_aliased_buffers(graph_module) == 1 + assert kv_delegate.args[1] is k_buffer + # The plain delegate keeps its staging and is never demanded to un-stage. + assert plain_delegate.args[1] is staged_x + + +@pytest.mark.unit +def test_unstage_raises_when_the_plain_delegate_is_wrongly_stamped(): + """The other side of the per-partition stamping, in isolation: a delegate + that carries the zero-copy spec and un-stages nothing must raise, whatever + put the spec there. Narrowing which delegates get stamped must not weaken + this -- it is the lost-update guard for the KV delegate too. + """ + zero_copy_spec = [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b'["out_k"]')] + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + x = graph.placeholder("x") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + staged_x = graph.call_function(h2d, (x,)) + kv_lowered = graph.get_attr("lowered_module_0") + plain_lowered = graph.get_attr("lowered_module_1") + kv_delegate = graph.call_function(executorch_call_delegate, (kv_lowered, staged_k)) + plain_delegate = graph.call_function( + executorch_call_delegate, (plain_lowered, staged_x) + ) + graph.output((k_buffer, kv_delegate, plain_delegate)) + root = torch.nn.Module() + # Both delegates wrongly carry the spec -- the shape that per-engine stamping + # in TensorRTPartitioner exists to prevent. + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=list(zero_copy_spec) + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=list(zero_copy_spec) + ) + graph_module = torch.fx.GraphModule(root, graph) + for node, dev in ( + (k_buffer, DeviceType.CPU), + (x, DeviceType.CPU), + (staged_k, DeviceType.CUDA), + (staged_x, DeviceType.CUDA), + ): + node.meta["spec"] = SimpleNamespace(device=dev, device_index=3) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="declares zero-copy KV"): + Z._unstage_aliased_buffers(graph_module) + + +# -------------------------------------------------------------------------- +# Single-engine, on the same engine-node helper: the same per-engine derivation +# also has to narrow *within* one engine, from every aliased output down to the +# ones whose aliased input is a buffer export rewired. +# -------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_partition_elides_only_the_outputs_aliased_onto_a_marked_buffer(): + """One engine, two aliased outputs, one marked buffer: only the marked one is + named elidable. + + An aliased output whose input is not a buffer export rewired -- a user alias, + which nothing rewires and whose placeholder therefore carries no + ``_torch_tensorrt_aliased_buffer`` -- is still a delegate output. Deriving the + elidable set from the engine's aliased_io alone would exempt it too, and the + backend would then accept a delegate that dropped a mutation nothing writes + back. + """ + from torch_tensorrt.executorch._zero_copy import _aliased_inputs_by_output_index + from torch_tensorrt.executorch.partitioner import TensorRTPartitioner + + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + user_alias = graph.placeholder("u") + engine = _no_op_engine_node( + graph, + [k_buffer, user_alias], + aliased_io={ + "out_k": ("k_in", "kv_cache_update"), + "out_u": ("u_in", "user"), + }, + input_names=["k_in", "u_in"], + output_names=["out_k", "out_u"], + ) + graph.output((engine,)) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=SimpleNamespace(buffers_to_mutate={}, inputs_to_buffers={}), + constants={}, + ) + partition = SimpleNamespace(id=0, nodes=[engine]) + + # Both outputs are aliased, or the narrowing below would have nothing to do. + assert set(_aliased_inputs_by_output_index(program, engine)) == {0, 1} + + # The spec deliberately names the output the derivation must NOT pick, so a + # result of {"out_k"} can only have come from the engine's aliased_io and the + # marks on its inputs. + partitioner = TensorRTPartitioner( + compile_specs=[CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b'["out_u"]')] + ) + assert partitioner._partition_elided_output_names(program, partition) == {"out_k"} + + +# -------------------------------------------------------------------------- +# GPU integration: the mark set during rewiring must survive real lowering, or +# the un-staging pass has nothing to act on and every KV update is lost. Only a +# real export exercises that -- the stub graphs above set the mark by hand. +# -------------------------------------------------------------------------- +VOCAB = 64 +DIM = 32 +HEADS = 2 +HEAD_DIM = 16 +MAX_LEN = 16 + + +class _KVDecodeStep(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.embed = torch.nn.Embedding(VOCAB, DIM) + self.pos_embed = torch.nn.Embedding(MAX_LEN, DIM) + self.q = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.k = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.v = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.o = torch.nn.Linear(HEADS * HEAD_DIM, DIM, bias=False) + self.lm = torch.nn.Linear(DIM, VOCAB, bias=False) + self.register_buffer("k_cache", torch.zeros(1, HEADS, MAX_LEN, HEAD_DIM)) + self.register_buffer("v_cache", torch.zeros(1, HEADS, MAX_LEN, HEAD_DIM)) + + def forward(self, tokens: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor: + pos_idx = input_pos.reshape(-1) + pos = input_pos.reshape(()) + x = self.embed(tokens) + self.pos_embed(input_pos.reshape(1, 1)) + + def split_heads(proj: torch.Tensor) -> torch.Tensor: + return proj.view(1, 1, HEADS, HEAD_DIM).transpose(1, 2) + + q = split_heads(self.q(x)) + k = split_heads(self.k(x)) + v = split_heads(self.v(x)) + self.k_cache.index_copy_(2, pos_idx, k) + self.v_cache.index_copy_(2, pos_idx, v) + scores = (q @ self.k_cache.transpose(-1, -2)) / (HEAD_DIM**0.5) + allowed = torch.arange(MAX_LEN, device=x.device) <= pos + bias = torch.where( + allowed, + torch.zeros((), dtype=x.dtype, device=x.device), + torch.full((), torch.finfo(x.dtype).min, dtype=x.dtype, device=x.device), + ) + attn = torch.softmax(scores + bias.view(1, 1, 1, MAX_LEN), dim=-1) + out = (attn @ self.v_cache).transpose(1, 2).reshape(1, 1, HEADS * HEAD_DIM) + return self.lm(self.o(out)) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA + TensorRT for a real engine" +) +@pytest.mark.parametrize("generate_etrecord", [False, True], ids=["plain", "etrecord"]) +def test_aliased_buffer_mark_survives_real_lowering(generate_etrecord): + """After a real export(..., zero_copy_kv=True), the KV buffer placeholder in + the lowered edge program still carries ``_torch_tensorrt_aliased_buffer`` -- + the token the to_out_var_pass keys the un-staging on. + + ``generate_etrecord=True`` is covered because it makes ExecuTorch deep copy + the whole program, and the mark rides on node meta. Losing it there would not + raise here: it surfaces later as the un-staging pass finding a marked buffer + it never un-staged, by which point the connection to this option is gone. + """ + with torch.no_grad(): + torch.manual_seed(0) + model = _KVDecodeStep().eval().cuda() + tokens = torch.zeros(1, 1, dtype=torch.long).cuda() + input_pos = torch.tensor([0], dtype=torch.long).cuda() + + exported_program = torch.export.export(model, (tokens, input_pos)) + trt_gm = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=(tokens, input_pos), + min_block_size=1, + truncate_double=True, + ) + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=(tokens, input_pos), + retrace=False, + zero_copy_kv=True, + generate_etrecord=generate_etrecord, + ) + + ep = edge.exported_program() + marked = [ + node + for node in ep.graph_module.graph.nodes + if node.op == "placeholder" and node.meta.get("_torch_tensorrt_aliased_buffer") + ] + assert marked, "no placeholder kept _torch_tensorrt_aliased_buffer through lowering" + + +# -------------------------------------------------------------------------- +# save() path: unlike the direct export()+to_executorch() contract -- two paired +# calls the caller must not forget -- torch_tensorrt.save() owns both steps, so a +# single zero_copy_kv=True must both hand export() the opt-in and install the +# finalization config before to_executorch(). These are CPU-only: the TensorRT +# lowering and the ExecuTorch finalization are stubbed so the wiring is checked +# without a GPU. The passes they invoke have their own coverage above; a real +# end-to-end run is exercised by kv_cache_decode_check on GPU. +# -------------------------------------------------------------------------- +def _trivial_exported_program(): + """A tiny CPU ExportedProgram -- enough for save() to reach _save_as_executorch. + + It carries no execute_engine node, so the retrace=True KV-declaration pass is + a no-op on it and the stubs below stand in for the real lowering. + """ + + class _Add(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + 1 + + return torch.export.export(_Add(), (torch.randn(3),)) + + +def _install_save_stubs(monkeypatch, *, wrap_config=True): + """Stub the executorch lowering that save() drives and record how it is called. + + Returns a namespace capturing the kwargs export() received, the arguments + zero_copy_backend_config() was wrapped with, and the config finally handed to + to_executorch(). When ``wrap_config`` is False the real + zero_copy_backend_config runs, so the recorded config is the genuine one. + """ + import torch_tensorrt._compile as compile_module + import torch_tensorrt.executorch as executorch_api + + monkeypatch.setattr( + compile_module, + "ENABLED_FEATURES", + compile_module.ENABLED_FEATURES._replace(torch_tensorrt_runtime=True), + ) + + calls = SimpleNamespace( + export_kwargs=None, wrap_args=[], to_executorch_config="unset" + ) + + def _to_executorch(config=None): + calls.to_executorch_config = config + return SimpleNamespace( + _tensor_data=None, write_to_file=lambda f: f.write(b"stub-pte") + ) + + edge = SimpleNamespace(to_executorch=_to_executorch) + + def _export(exp_program, **kwargs): + calls.export_kwargs = kwargs + return edge + + monkeypatch.setattr(executorch_api, "export", _export) + + if wrap_config: + wrapped = object() + + def _wrap(config=None): + calls.wrap_args.append(config) + return wrapped + + monkeypatch.setattr(executorch_api, "zero_copy_backend_config", _wrap) + calls.wrapped_sentinel = wrapped + else: + real_wrap = executorch_api.zero_copy_backend_config + + def _wrap(config=None): + calls.wrap_args.append(config) + return real_wrap(config) + + monkeypatch.setattr(executorch_api, "zero_copy_backend_config", _wrap) + + return calls + + +@pytest.mark.unit +def test_save_zero_copy_kv_true_threads_flag_and_installs_config(monkeypatch, tmp_path): + """save(zero_copy_kv=True, backend_config=cfg) opts export() in and wraps the + caller's config exactly once, forwarding the wrapped one to to_executorch().""" + calls = _install_save_stubs(monkeypatch) + user_cfg = object() + + torch_tensorrt.save( + _trivial_exported_program(), + str(tmp_path / "model.pte"), + output_format="executorch", + zero_copy_kv=True, + backend_config=user_cfg, + ) + + assert calls.export_kwargs["zero_copy_kv"] is True + # The user's config is wrapped once (preserving their fields), not double-wrapped. + assert calls.wrap_args == [user_cfg] + assert calls.to_executorch_config is calls.wrapped_sentinel + + +@pytest.mark.unit +def test_save_zero_copy_kv_true_wraps_defaults_without_a_config(monkeypatch, tmp_path): + """With no backend_config, zero_copy_backend_config(None) starts from ET + defaults; the finalization config is still installed.""" + calls = _install_save_stubs(monkeypatch) + + torch_tensorrt.save( + _trivial_exported_program(), + str(tmp_path / "model.pte"), + output_format="executorch", + zero_copy_kv=True, + ) + + assert calls.export_kwargs["zero_copy_kv"] is True + assert calls.wrap_args == [None] + assert calls.to_executorch_config is calls.wrapped_sentinel + + +@pytest.mark.unit +def test_save_zero_copy_kv_true_installs_the_real_unstaging_pass(monkeypatch, tmp_path): + """End of the wiring with the real config builder: the config reaching + to_executorch() carries the un-staging to_out_var_pass, not ET's default.""" + calls = _install_save_stubs(monkeypatch, wrap_config=False) + + torch_tensorrt.save( + _trivial_exported_program(), + str(tmp_path / "model.pte"), + output_format="executorch", + zero_copy_kv=True, + ) + + assert calls.wrap_args == [None] + assert ( + type(calls.to_executorch_config.to_out_var_pass).__name__ + == "_UnstageThenToOutVar" + ) + + +@pytest.mark.unit +def test_save_defaults_leave_kv_staged(monkeypatch, tmp_path): + """Default save() (zero_copy_kv omitted) never wraps the config, so the KV + buffer keeps its staging and its copy-back: the caller's config reaches + to_executorch() untouched and export() is told zero_copy_kv=False.""" + calls = _install_save_stubs(monkeypatch) + user_cfg = object() + + torch_tensorrt.save( + _trivial_exported_program(), + str(tmp_path / "model.pte"), + output_format="executorch", + backend_config=user_cfg, + ) + + assert calls.export_kwargs["zero_copy_kv"] is False + assert calls.wrap_args == [] + assert calls.to_executorch_config is user_cfg From 021899e67a6e107ff8af974431925aa4647f0c86 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Mon, 24 Aug 2026 21:21:24 -0700 Subject: [PATCH 05/22] feat(executorch): support zero_copy_kv beside a copy-back buffer A method may hold both kinds of mutable buffer: a KV cache the engine writes through an aliased binding, and a non-KV buffer -- a convolution state -- whose new value comes back as a trailing delegate output for ExecuTorch to copy back. The copy-back path predates `zero_copy_kv`, but the combination of the two was never decided or covered, only mechanically tolerated. Treat it as supported. `_aliased_buffer_mutations` already discriminates the two by reading the engine's own `aliased_io` rather than the graph, in which they are identical -- both a `getitem` off the engine node whose buffer is also an engine input. The aliased caches go zero-copy; the copy-back buffer keeps its staging copy and the output that writes it. Refusing the combination instead would give up the feature for every model carrying one non-KV mutable buffer beside its cache, and rewiring the copy-back would delete a real update with no error. Neither is acceptable, and the discriminator that avoids both is already load-bearing, so pin it: a stub-level test on one method holding both mutations, and a real-engine export of a decode step with `k_cache`/`v_cache` beside a ring-shifted `conv_state`, asserting the caches end up bound to their own placeholders while `conv_state` stays bound to a delegate output. That real-engine export runs on both exporters, because they reach the discriminator by different routes. The legacy exporter (`retrace=False`) declares all three mutations as it builds the program, leaving `_declare_aliased_kv_mutations_on_ep` nothing to do. Under `retrace=True`, which is `save()`'s default, the retraced program arrives with no mutation declared at all -- `torch.export` drops the aliased outputs at the fx boundary and leaves the copy-back value as a plain return -- so that post-export pass is what separates the two kinds, and it is exercised only on that parameter. --- py/torch_tensorrt/executorch/_export.py | 6 + .../py/dynamo/executorch/test_zero_copy_kv.py | 258 ++++++++++++++++++ 2 files changed, 264 insertions(+) diff --git a/py/torch_tensorrt/executorch/_export.py b/py/torch_tensorrt/executorch/_export.py index 7710110be42..04c4a64e703 100644 --- a/py/torch_tensorrt/executorch/_export.py +++ b/py/torch_tensorrt/executorch/_export.py @@ -457,6 +457,12 @@ def export( -- without it the buffer is still staged and its updates are discarded, with no error. + Only a buffer the engine declares aliased is affected. A method may hold both + kinds at once: a mutable buffer with no aliasing available -- a convolution + state, say -- keeps its staging copy and the copy-back that writes it, while + the aliased caches beside it go zero-copy. The two are told apart by the + engine's own ``aliased_io``, not by the graph, in which they look identical. + ``generate_etrecord=True`` is outside the payload sharing described above. It makes ExecuTorch deep copy the whole program, so peak memory grows by roughly the size of the program including engines. diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index 10d017eafd1..102cfc5b8d2 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -174,6 +174,91 @@ def test_rewire_leaves_mutations_the_engine_does_not_alias(monkeypatch, mutation assert "_torch_tensorrt_aliased_buffer" not in k_buffer.meta +def _mixed_program(): + """One engine, one method, both kinds of mutation at once. + + ``engine(b_k_0, b_state_0, tokens) -> (logits, out_k, out_state)`` where the + engine aliases only ``out_k`` onto ``b_k_0``. ``b_state_0`` is the #4459 + shape: a mutable buffer with no aliasing available, whose new value + ``lift_mutated_buffers`` appended as a trailing output for ExecuTorch to copy + back. In the graph the two mutations are indistinguishable -- each is a + ``getitem`` off the engine node whose buffer is also an engine input. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + state_buffer = graph.placeholder("b_state_0") + tokens = graph.placeholder("tokens") + engine = graph.placeholder("engine") + engine_call = graph.call_function( + torch.ops.tensorrt.execute_engine.default, + ([k_buffer, state_buffer, tokens], engine), + ) + logits = graph.call_function(operator.getitem, (engine_call, 0)) + k_out = graph.call_function(operator.getitem, (engine_call, 1)) + state_out = graph.call_function(operator.getitem, (engine_call, 2)) + graph.output((k_out, state_out, logits)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + signature = SimpleNamespace( + inputs_to_buffers={"b_k_0": "k_0", "b_state_0": "state_0"}, + input_specs=[], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(name=k_out.name), "k_0" + ), + OutputSpec( + OutputKind.BUFFER_MUTATION, + TensorArgument(name=state_out.name), + "state_0", + ), + OutputSpec(OutputKind.USER_OUTPUT, TensorArgument(name=logits.name), None), + ], + ) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=signature, + _graph_signature=signature, + ) + return program, k_buffer, state_buffer, k_out, state_out + + +@pytest.mark.unit +def test_rewire_keeps_the_copyback_in_a_method_that_also_has_an_aliased_kv(monkeypatch): + """Zero-copy and a copy-back buffer may share one method, and must not mix. + + Rewiring the copy-back would delete a real update with no error, and refusing + the aliased one would give up the whole feature for any model carrying a + non-KV mutable buffer beside its cache. The engine's own aliased_io is what + separates them: only ``out_k`` is listed, so only ``b_k_0`` is rewired and + only its binding name is offered to the backend as elided. ``b_state_0`` + keeps its delegate output, which is the value ExecuTorch copies back. + """ + program, k_buffer, state_buffer, k_out, state_out = _mixed_program() + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in", "state_in", "tokens"], + output_names=["logits", "out_k", "out_state"], + ) + + assert Z.rewire_aliased_mutations_to_buffers(program) == ["out_k"] + + kv_spec, state_spec, _ = program._graph_signature.output_specs + assert kv_spec.arg.name == k_buffer.name + assert k_buffer.meta["_torch_tensorrt_aliased_buffer"] is True + assert k_out not in program.graph_module.graph.nodes + + assert state_spec.kind == OutputKind.BUFFER_MUTATION + assert state_spec.target == "state_0" + assert state_spec.arg.name == state_out.name + assert state_out in program.graph_module.graph.nodes + assert program.graph_module.graph.output_node().args[0][1] is state_out + # Un-staging keys on this mark, so leaving it off b_state_0 is what keeps the + # copy-back buffer's staging copy -- the engine writes that copy and + # ExecuTorch copies it back, exactly as without zero-copy. + assert "_torch_tensorrt_aliased_buffer" not in state_buffer.meta + + @pytest.mark.unit def test_rewire_is_a_noop_without_aliased_io(monkeypatch): program, k_buffer, _ = _kv_program() @@ -1030,6 +1115,179 @@ def test_aliased_buffer_mark_survives_real_lowering(generate_etrecord): assert marked, "no placeholder kept _torch_tensorrt_aliased_buffer through lowering" +class _MixedDecodeStep(torch.nn.Module): + """A decode step with an engine-aliased KV cache and a copy-back buffer. + + ``k_cache``/``v_cache`` are written by ``index_copy_`` on the sequence axis, + which the converter turns into an aliased engine binding. ``conv_state`` is a + ring shift -- a whole-buffer rewrite with no position to alias on -- so + ``lift_mutated_buffers`` records it in ``_copyback_mutation_buffers`` and its + new value comes back as a trailing delegate output instead. + """ + + def __init__(self) -> None: + super().__init__() + self.embed = torch.nn.Embedding(VOCAB, DIM) + self.q = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.k = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.v = torch.nn.Linear(DIM, HEADS * HEAD_DIM, bias=False) + self.o = torch.nn.Linear(HEADS * HEAD_DIM, DIM, bias=False) + self.lm = torch.nn.Linear(DIM, VOCAB, bias=False) + self.register_buffer("k_cache", torch.zeros(1, HEADS, MAX_LEN, HEAD_DIM)) + self.register_buffer("v_cache", torch.zeros(1, HEADS, MAX_LEN, HEAD_DIM)) + self.register_buffer("conv_state", torch.zeros(1, DIM, 4)) + + def forward(self, tokens: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor: + pos_idx = input_pos.reshape(-1) + pos = input_pos.reshape(()) + x = self.embed(tokens) + + shifted = torch.cat([self.conv_state[:, :, 1:], x.reshape(1, DIM, 1)], dim=2) + self.conv_state.copy_(shifted) + x = x + self.conv_state.sum(dim=2).reshape(1, 1, DIM) + + def split_heads(proj: torch.Tensor) -> torch.Tensor: + return proj.view(1, 1, HEADS, HEAD_DIM).transpose(1, 2) + + q = split_heads(self.q(x)) + k = split_heads(self.k(x)) + v = split_heads(self.v(x)) + self.k_cache.index_copy_(2, pos_idx, k) + self.v_cache.index_copy_(2, pos_idx, v) + scores = (q @ self.k_cache.transpose(-1, -2)) / (HEAD_DIM**0.5) + allowed = torch.arange(MAX_LEN, device=x.device) <= pos + bias = torch.where( + allowed, + torch.zeros((), dtype=x.dtype, device=x.device), + torch.full((), torch.finfo(x.dtype).min, dtype=x.dtype, device=x.device), + ) + attn = torch.softmax(scores + bias.view(1, 1, 1, MAX_LEN), dim=-1) + out = (attn @ self.v_cache).transpose(1, 2).reshape(1, 1, HEADS * HEAD_DIM) + return self.lm(self.o(out)) + + +def _real_delegates(graph_module): + return [ + node + for node in graph_module.graph.nodes + if node.op == "call_function" and node.target is executorch_call_delegate + ] + + +def _lowered_module(graph_module, delegate): + return getattr(graph_module, delegate.args[0].target) + + +def _assert_marked_buffers_reach_the_engine_unstaged(graph_module): + """Every marked buffer is a direct argument of a TensorRT delegate. + + Finalizing a zero-copy program without raising is a weak signal, because both + of the completeness raises at the end of ``_unstage_aliased_buffers`` fire off + its own bookkeeping: a pass that records each un-staging and then leaves the + argument pointing at the staging copy still satisfies them. Only the graph says + whether the rewiring happened, and getting it wrong is silent -- the engine + writes per-call scratch that is discarded and the cache never updates. + """ + marked = [ + node + for node in graph_module.graph.nodes + if node.op == "placeholder" and node.meta.get("_torch_tensorrt_aliased_buffer") + ] + assert marked, "no buffer was marked for in-place update" + reached = { + arg + for node in _real_delegates(graph_module) + if _lowered_module(graph_module, node).backend_id == "TensorRTBackend" + for arg in node.args[1:] + if isinstance(arg, torch.fx.Node) + } + for node in marked: + assert node in reached, ( + f"buffer '{node.name}' is marked for in-place update but is not a " + "direct argument of any TensorRT delegate -- it either still reaches " + "one through a staging copy, or reaches none at all. Either way " + "nothing writes the caller's buffer and the cache never updates" + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA + TensorRT for a real engine" +) +@pytest.mark.parametrize("retrace", [False, True], ids=["legacy", "retrace"]) +def test_zero_copy_kv_keeps_a_copyback_buffer_in_the_same_method(retrace): + """A real method holding both kinds of mutable buffer exports and keeps both. + + The KV caches end up bound to their own placeholders -- no value for + ExecuTorch to copy, which is the zero copy -- while ``conv_state`` stays bound + to a delegate output, which is the value ExecuTorch copies back into it. + Losing that distinction in either direction is silent wrong output, so it is + pinned on a real engine rather than a stub: the aliased_io the discriminator + reads is produced by the converter, not by this test. + + Both exporters are covered because they reach that distinction by different + routes, and only one of them is the ``save()`` default. The legacy exporter + declares all three mutations while it inlines the engines, so + ``_declare_aliased_kv_mutations_on_ep`` finds nothing left to do and the + discriminator never runs. Under ``retrace=True`` the retraced program arrives + with no mutations declared at all -- torch.export drops the aliased outputs at + the fx boundary and leaves the copy-back value as a plain return -- so that + pass is what separates the two kinds, by reading each engine's ``aliased_io``. + """ + with torch.no_grad(): + torch.manual_seed(0) + model = _MixedDecodeStep().eval().cuda() + tokens = torch.zeros(1, 1, dtype=torch.long).cuda() + input_pos = torch.tensor([0], dtype=torch.long).cuda() + + exported_program = torch.export.export(model, (tokens, input_pos)) + trt_gm = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=(tokens, input_pos), + min_block_size=1, + truncate_double=True, + ) + assert trt_gm.meta.get("_copyback_mutation_buffers") == ["conv_state"], ( + "the model no longer produces a copy-back buffer, so this test would " + "pass without exercising the combination it exists for" + ) + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=(tokens, input_pos), + retrace=retrace, + zero_copy_kv=True, + ) + + ep = edge.exported_program() + output_args = list(ep.graph_module.graph.output_node().args[0]) + bound = { + spec.target: value + for spec, value in zip(ep.graph_signature.output_specs, output_args) + if spec.kind == OutputKind.BUFFER_MUTATION + } + assert set(bound) == {"k_cache", "v_cache", "conv_state"} + for name in ("k_cache", "v_cache"): + assert bound[name].op == "placeholder", ( + f"{name} is still satisfied by a delegate output, so ExecuTorch will " + "copy it back and zero-copy bought nothing" + ) + assert bound[name].meta.get("_torch_tensorrt_aliased_buffer") is True + assert bound["conv_state"].op == "call_function", ( + "conv_state was rewired to its own placeholder, which deletes the " + "copy-back of a buffer no engine writes in place -- a lost update" + ) + assert "_torch_tensorrt_aliased_buffer" not in bound["conv_state"].meta + + # Everything above is the export half. The staging the other half removes does + # not exist until PropagateDevicePass runs inside to_executorch, so this is the + # earliest point at which the caches can be seen reaching the engine directly. + program = edge.to_executorch( + config=torch_tensorrt.executorch.zero_copy_backend_config() + ) + _assert_marked_buffers_reach_the_engine_unstaged( + program.exported_program().graph_module + ) + + # -------------------------------------------------------------------------- # save() path: unlike the direct export()+to_executorch() contract -- two paired # calls the caller must not forget -- torch_tensorrt.save() owns both steps, so a From afadb797b93944bc8c08d10f5e0284c6bf80bec7 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Tue, 25 Aug 2026 12:13:26 -0700 Subject: [PATCH 06/22] test(executorch): cover zero_copy_kv on the two multi-delegate shapes `zero_copy_kv` was pinned on one real engine only: a single method, a single TensorRT delegate, the aliased caches and a `conv_state` copy-back side by side. Two shapes that shape does not reach are covered here, both on real engines. A method whose copy-back rides on a *different* TensorRT delegate than the aliased caches. `TensorRTPartitioner` derives the elided binding names per engine and stamps `zero_copy_kv` only on the delegate that lost an output, so the plain compute delegate beside it must carry no spec; if it did, `_unstage_aliased_buffers`'s cross-check would demand an aliased buffer it never had and the export would die. The test asserts the stamping directly and then finalizes, which is where that cross-check runs. A method that also holds an ExecuTorch CUDA (AOTI) delegate. `erfinv` has no TensorRT converter, so with a `CudaPartitioner` catch-all the method lowers to TensorRT, CudaBackend and TensorRT in sequence. Un-staging must reach the TensorRT delegate only -- `_is_tensorrt_delegate` gates it, because no other backend promises the in-place write through an aliased binding. Both models assert their own shape before asserting the behaviour: a partitioner change that collapsed either back to one delegate would otherwise leave the test passing while covering nothing. The trailing `Linear` in the CUDA model exists for that reason -- ending on `erfinv` leaves exactly one TensorRT delegate, and "only the KV delegate is stamped" is then true no matter what the partitioner does. Both of these run on both exporters. The split-delegate one is the only shape where `_declare_aliased_kv_mutations_on_ep` has to pick the aliased engine out of several -- it scans every `execute_engine` node and skips the ones whose `aliased_io` is empty, and the copy-back value it detaches comes off a different engine than the caches it declares. The legacy exporter declares all of that as it builds the program, so that scan runs only under `retrace=True`, which is `save()`'s default and had no real-engine `zero_copy_kv` exercise before. --- .../py/dynamo/executorch/test_zero_copy_kv.py | 256 ++++++++++++++++++ 1 file changed, 256 insertions(+) diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index 102cfc5b8d2..199897414e2 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -1288,6 +1288,262 @@ def test_zero_copy_kv_keeps_a_copyback_buffer_in_the_same_method(retrace): ) +class _SplitRolesDecodeStep(_MixedDecodeStep): + """The same two buffer kinds, but on two different TensorRT engines. + + ``torch.sinh`` is pinned out of TensorRT by the test, so the attention half + -- which holds the engine-aliased caches -- and the ``conv_state`` half end up + in separate partitions. Only the first engine has aliased outputs, and the + copy-back output rides on the second. + """ + + def forward(self, tokens: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor: + pos_idx = input_pos.reshape(-1) + pos = input_pos.reshape(()) + x = self.embed(tokens) + + def split_heads(proj: torch.Tensor) -> torch.Tensor: + return proj.view(1, 1, HEADS, HEAD_DIM).transpose(1, 2) + + q = split_heads(self.q(x)) + k = split_heads(self.k(x)) + v = split_heads(self.v(x)) + self.k_cache.index_copy_(2, pos_idx, k) + self.v_cache.index_copy_(2, pos_idx, v) + scores = (q @ self.k_cache.transpose(-1, -2)) / (HEAD_DIM**0.5) + allowed = torch.arange(MAX_LEN, device=x.device) <= pos + bias = torch.where( + allowed, + torch.zeros((), dtype=x.dtype, device=x.device), + torch.full((), torch.finfo(x.dtype).min, dtype=x.dtype, device=x.device), + ) + attn = torch.softmax(scores + bias.view(1, 1, 1, MAX_LEN), dim=-1) + out = (attn @ self.v_cache).transpose(1, 2).reshape(1, 1, HEADS * HEAD_DIM) + + h = torch.sinh(self.o(out) + x) + shifted = torch.cat([self.conv_state[:, :, 1:], h.reshape(1, DIM, 1)], dim=2) + self.conv_state.copy_(shifted) + return self.lm(h + self.conv_state.sum(dim=2).reshape(1, 1, DIM)) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA + TensorRT for a real engine" +) +@pytest.mark.parametrize("retrace", [False, True], ids=["legacy", "retrace"]) +def test_zero_copy_kv_with_the_copyback_on_a_second_delegate(retrace): + """Two TensorRT delegates, one with the aliased caches and one with the copy-back. + + This is the shape ``_delegate_declares_zero_copy`` reasons about: the + partitioner must stamp ``zero_copy_kv`` on the KV delegate only, or + ``_unstage_aliased_buffers``'s cross-check demands an aliased buffer from the + plain compute delegate and the export dies. Finalizing here is the assertion: + a wrongly stamped delegate raises inside ``to_executorch``. + + Under ``retrace=True`` this is also the only shape where + ``_declare_aliased_kv_mutations_on_ep`` has to pick the aliased engine out of + several: it scans every ``execute_engine`` node and skips the ones whose + ``aliased_io`` is empty, and the copy-back value it detaches comes off a + different engine than the caches it declares. The legacy exporter declares all + of that while inlining, so that scan runs only on this parameter. + """ + with torch.no_grad(): + torch.manual_seed(0) + model = _SplitRolesDecodeStep().eval().cuda() + tokens = torch.zeros(1, 1, dtype=torch.long).cuda() + input_pos = torch.tensor([0], dtype=torch.long).cuda() + + exported_program = torch.export.export(model, (tokens, input_pos)) + trt_gm = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=(tokens, input_pos), + min_block_size=1, + truncate_double=True, + torch_executed_ops={"torch.ops.aten.sinh.default"}, + ) + aliased_per_engine = [ + bool(getattr(sub, "aliased_io", None)) for _, sub in trt_gm.named_children() + ] + assert len(aliased_per_engine) > 1 and sum(aliased_per_engine) == 1, ( + "the model no longer lowers to several engines with the aliasing on " + f"exactly one of them ({aliased_per_engine}), so it does not exercise " + "the multi-delegate split this test exists for" + ) + assert trt_gm.meta.get("_copyback_mutation_buffers") == ["conv_state"] + + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=(tokens, input_pos), + retrace=retrace, + zero_copy_kv=True, + ) + + ep = edge.exported_program() + graph_module = ep.graph_module + output_args = list(graph_module.graph.output_node().args[0]) + bound = { + spec.target: value + for spec, value in zip(ep.graph_signature.output_specs, output_args) + if spec.kind == OutputKind.BUFFER_MUTATION + } + assert set(bound) == {"k_cache", "v_cache", "conv_state"} + for name in ("k_cache", "v_cache"): + assert bound[name].op == "placeholder" + assert bound[name].meta.get("_torch_tensorrt_aliased_buffer") is True + assert bound["conv_state"].target is operator.getitem + + delegates = _real_delegates(graph_module) + assert len(delegates) > 1 + kv_delegate = next( + node + for node in delegates + if any( + isinstance(arg, torch.fx.Node) + and arg.meta.get("_torch_tensorrt_aliased_buffer") + for arg in node.args[1:] + ) + ) + copyback_delegate = bound["conv_state"].args[0] + assert copyback_delegate in delegates + assert copyback_delegate is not kv_delegate, ( + "the copy-back landed on the same delegate as the aliased caches, so this " + "test is running the single-delegate shape again" + ) + + stamped = [ + node + for node in delegates + if any( + spec.key == ZERO_COPY_KV_COMPILE_SPEC_KEY + for spec in _lowered_module(graph_module, node).compile_specs + ) + ] + assert stamped == [kv_delegate], ( + "the zero-copy spec must sit on the delegate whose engine lost an output " + "and on no other; a plain compute delegate carrying it is asked for an " + "aliased buffer it never had" + ) + + # The un-staging cross-check runs here, not above -- and so does the + # un-staging itself, which only the finalized graph shows. + program = edge.to_executorch( + config=torch_tensorrt.executorch.zero_copy_backend_config() + ) + _assert_marked_buffers_reach_the_engine_unstaged( + program.exported_program().graph_module + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA + TensorRT for a real engine" +) +def test_zero_copy_kv_beside_an_executorch_cuda_delegate(): + """An aliased KV cache in a method that also holds an ExecuTorch CUDA delegate. + + ``erfinv`` has no TensorRT converter, so with a ``CudaPartitioner`` catch-all + the method lowers to TensorRT, CudaBackend and TensorRT delegates in sequence. + The un-staging must reach into the TensorRT delegate only. What is asserted + here is the reachable half -- that only the KV TensorRT delegate is stamped, + and that the caches reach it un-staged with a CUDA delegate in the middle. + That the gate itself refuses a marked buffer on another backend is pinned by + ``test_unstage_raises_for_a_marked_buffer_on_another_backends_delegate``. + """ + cuda_backend = pytest.importorskip("executorch.backends.cuda.cuda_backend") + cuda_partitioner = pytest.importorskip("executorch.backends.cuda.cuda_partitioner") + + class _CudaNeighbourDecodeStep(_KVDecodeStep): + def __init__(self): + super().__init__() + # A TensorRT-supported op AFTER erfinv, so the CUDA delegate is + # sandwiched between two TensorRT ones. Ending on erfinv would leave + # the method with a single TensorRT delegate, and the assertion below + # that only the KV delegate carries the zero-copy spec would then hold + # whatever the partitioner did. + self.tail = torch.nn.Linear(VOCAB, VOCAB, bias=False) + + def forward(self, tokens, input_pos): + h = super().forward(tokens, input_pos) + return self.tail(torch.erfinv(torch.tanh(h))) + + with torch.no_grad(): + torch.manual_seed(0) + model = _CudaNeighbourDecodeStep().eval().cuda() + tokens = torch.zeros(1, 1, dtype=torch.long).cuda() + input_pos = torch.tensor([0], dtype=torch.long).cuda() + + exported_program = torch.export.export(model, (tokens, input_pos)) + trt_gm = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=(tokens, input_pos), + min_block_size=1, + truncate_double=True, + ) + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=(tokens, input_pos), + retrace=False, + zero_copy_kv=True, + partitioners=[ + cuda_partitioner.CudaPartitioner( + [ + cuda_backend.CudaBackend.generate_method_name_compile_spec( + "forward" + ) + ] + ) + ], + ) + + graph_module = edge.exported_program().graph_module + delegates = _real_delegates(graph_module) + backends = { + node: _lowered_module(graph_module, node).backend_id for node in delegates + } + assert sorted(backends.values()) == [ + "CudaBackend", + "TensorRTBackend", + "TensorRTBackend", + ], ( + f"the method no longer lowers to TensorRT/CudaBackend/TensorRT " + f"({sorted(backends.values())}), so it does not cover the sandwiched " + "CUDA delegate this test exists for" + ) + + def _marked_args(node): + return [ + arg.name + for arg in node.args[1:] + if isinstance(arg, torch.fx.Node) + and arg.meta.get("_torch_tensorrt_aliased_buffer") + ] + + kv_delegates = [node for node in delegates if _marked_args(node)] + assert len(kv_delegates) == 1 and backends[kv_delegates[0]] == "TensorRTBackend", ( + "the aliased buffers must reach exactly one TensorRT delegate; " + f"got {[(backends[n], _marked_args(n)) for n in kv_delegates]}" + ) + + stamped = [ + node + for node in delegates + if any( + spec.key == ZERO_COPY_KV_COMPILE_SPEC_KEY + for spec in _lowered_module(graph_module, node).compile_specs + ) + ] + assert stamped == kv_delegates, ( + "only the delegate whose engine lost an aliased output may carry the " + "zero-copy spec; the CUDA delegate and the trailing TensorRT one had no " + f"aliased buffer, yet {[backends[n] for n in stamped]} are stamped" + ) + + program = edge.to_executorch( + config=torch_tensorrt.executorch.zero_copy_backend_config() + ) + _assert_marked_buffers_reach_the_engine_unstaged( + program.exported_program().graph_module + ) + + # -------------------------------------------------------------------------- # save() path: unlike the direct export()+to_executorch() contract -- two paired # calls the caller must not forget -- torch_tensorrt.save() owns both steps, so a From 6fdb0b500848c9a34fa3b2c4e3c5060936b2582b Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Wed, 2 Sep 2026 23:13:35 -0700 Subject: [PATCH 07/22] docs: name only the ExecuTorch caller-stream guard for a coalesced .pte The single-stream section told runners to scope `torch_tensorrt::executorch_backend::CudaStreamGuard` alongside `executorch::extension::cuda::CallerStreamGuard`. That class no longer exists -- `cpp/src/torch_tensorrt/executorch/README.md` records its removal, deliberately with no deprecated alias -- so the snippet named a type that does not compile. The reason given for scoping both was backwards as well. One caller-stream guard already reaches every CUDA-capable delegate, because they resolve a single shared `libextension_cuda` and so read the same caller-stream storage; that is why the old class could be dropped rather than aliased. Say that instead. --- .../user_guide/runtime_performance/saving_models.rst | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 0fcbff1112b..613a8f20e2b 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -476,12 +476,11 @@ illegal memory access. The runtime does not impose a shared stream across delegates, so it is the **runner's responsibility** to run all delegates on one CUDA stream. Create a -single stream and, for the duration of execution, direct every backend to use it -(each backend exposes a caller-stream hook: scope both -``torch_tensorrt::executorch_backend::CudaStreamGuard`` and -``executorch::extension::cuda::CallerStreamGuard`` over that stream, since -installing one of them leaves the other backend on its own). All GPU work is then -enqueued in order and every cross-boundary dependency is satisfied, while +single stream and scope ``executorch::extension::cuda::CallerStreamGuard`` over it +for the duration of execution. That one guard reaches every CUDA-capable +delegate: they resolve a single shared ``libextension_cuda``, so the TensorRT +backend and the CUDA backend read the same caller-stream storage. All GPU work is +then enqueued in order and every cross-boundary dependency is satisfied, while execution stays asynchronous. If the runner reads a delegate's outputs between calls (for example, an From 232176fbfa1bc2d9bcca28f94725f4787ffab97f Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Wed, 2 Sep 2026 23:20:53 -0700 Subject: [PATCH 08/22] feat(executorch): check a finalized zero-copy program in the library The warning on `zero_copy_backend_config` said nothing downstream could detect a program exported with `zero_copy_kv=True` and then finalized without it. The example in this same change detects exactly that: it walks the finalized program for marked buffers that no longer reach a delegate directly and refuses to write the `.pte`. So the claim was false, and the check was written twice -- in the example and in the tests -- while the one path that owned both ends, `save()`, held the finalized program and did not run it. Move it into the library as `torch_tensorrt.executorch.check_zero_copy_kv`, call it from `save(..., zero_copy_kv=True)` before the `.pte` is written, and have the example call the library version. It refuses two shapes: a marked buffer that is not a direct argument of a TensorRT delegate (finalized without the config, so the engine writes a staging copy that is discarded), and nothing marked for in-place update at all (`zero_copy_kv=True` only warns when it finds no aliased buffer mutation). Where the un-staging pass does run it has already raised on the first shape; this catches the case where it never ran. Only a TensorRT delegate counts as that direct argument, which is the same filter the un-staging pass applies. A buffer is marked because a TensorRT engine writes it in place, so another backend's delegate holding it says nothing about whether the engine got it un-staged: a program can hand the buffer straight to a `CudaBackend` delegate while the TensorRT engine beside it still reads a staging copy that is thrown away. Every method is read, not only `forward`. `ExecutorchProgramManager.exported_program()` defaults to `forward`, and `export()` rewires each method separately, so a check that took the default would raise a bare `KeyError: 'forward'` on the prefill/decode program the user guide's own zero-copy example builds, and on a program that does have a `forward` beside other methods it would pass one whose decode had degenerated to staged. The failure names the method the buffer is in. The "nothing marked" refusal stays about the program rather than about each method, matching the warning `export()` emits for the same condition: a method with no aliased buffer mutation of its own is not an error, so a model that rewires only its decode step is accepted, and only a program where no method rewired anything is refused. The docs and the docstring now point at it instead of asserting the mistake is undetectable. The three real-engine tests hand it their finalized program before running their own stricter graph assertion. Every other test of it builds the program itself -- a `SimpleNamespace` for the check's own cases, a monkeypatch for `save()`'s -- so nothing else puts `methods` or `exported_program(name)` in front of a real `ExecutorchProgramManager`, and an upstream rename would leave all of them green while `save(zero_copy_kv=True)` raised `AttributeError` for every caller. --- .../runtime_performance/saving_models.rst | 18 +- .../export_kv_cache_decode.py | 52 +---- py/torch_tensorrt/_compile.py | 13 +- py/torch_tensorrt/executorch/__init__.py | 7 +- py/torch_tensorrt/executorch/_zero_copy.py | 97 +++++++++- .../py/dynamo/executorch/test_zero_copy_kv.py | 178 ++++++++++++++++-- 6 files changed, 293 insertions(+), 72 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 613a8f20e2b..05d7fbf6005 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -392,12 +392,18 @@ one silently would break a runner built before this feature. .. warning:: **Both calls are required.** Exporting with ``zero_copy_kv=True`` and then - finalizing without ``zero_copy_backend_config`` does not raise: the engine - writes a per-call staging copy that is discarded, and the cache never - updates. For a KV cache that is wrong output, not a crash. Nothing - downstream can detect the omission, so pairing the two is on the caller -- - unless ``torch_tensorrt.save`` is the one writing the ``.pte``, which owns - both ends and leaves nothing to pair. + finalizing without ``zero_copy_backend_config`` does not raise on its own: + the engine writes a per-call staging copy that is discarded, and the cache + never updates. For a KV cache that is wrong output, not a crash. Pass the + finalized program to ``torch_tensorrt.executorch.check_zero_copy_kv``, which + reads the graph back and refuses one whose caches are still staged, before + writing the ``.pte``:: + + program = edge.to_executorch(zero_copy_backend_config(backend_config)) + torch_tensorrt.executorch.check_zero_copy_kv(program) + + ``torch_tensorrt.save`` owns both ends and runs that check itself, so there + is nothing to pair on that path. ``torch_tensorrt.save`` finalizes the program itself, so a single ``zero_copy_kv=True`` covers both steps: diff --git a/examples/torchtrt_executorch_example/export_kv_cache_decode.py b/examples/torchtrt_executorch_example/export_kv_cache_decode.py index eedf5611592..5c1c0d8b3ce 100644 --- a/examples/torchtrt_executorch_example/export_kv_cache_decode.py +++ b/examples/torchtrt_executorch_example/export_kv_cache_decode.py @@ -20,8 +20,8 @@ here: zero-copy removes the copy that was making the update stick, so if the engine's in-place write is not reaching the caller's buffer the run fails. What it cannot see is a ``--zero_copy`` export that degenerated into an ordinary -staged ``.pte`` -- the two are indistinguishable to it -- so ``_check_zero_copy`` -refuses to write one. +staged ``.pte`` -- the two are indistinguishable to it -- so +``check_zero_copy_kv`` refuses to write one. Prerequisites ------------- @@ -32,7 +32,6 @@ import argparse import os -from typing import Any import torch import torch_tensorrt @@ -92,45 +91,6 @@ def split_heads(proj: torch.Tensor) -> torch.Tensor: return self.lm(self.o(out)) -def _check_zero_copy(program: Any) -> None: - """Refuse to write a ``--zero_copy`` .pte that is really an ordinary staged one. - - Neither half of zero-copy fails when it finds nothing to do: - ``zero_copy_kv=True`` logs a warning and carries on when no aliased buffer - mutation turns up, and the finalization pass, handed a program with nothing - marked, un-stages nothing and returns without a word. The ``.pte`` that comes - out still runs, and ``kv_cache_decode_check`` cannot tell it from the staged - model exported beside it -- so without this the lane would stay green while - covering none of the feature. - """ - from executorch.exir.delegate import executorch_call_delegate - - graph_module = program.exported_program().graph_module - marked = [ - node - for node in graph_module.graph.nodes - if node.op == "placeholder" and node.meta.get("_torch_tensorrt_aliased_buffer") - ] - if not marked: - raise RuntimeError( - "zero_copy_kv=True rewired no aliased buffer, so this .pte stages its " - "KV cache like any other." - ) - delegate_args = { - arg - for node in graph_module.graph.nodes - if node.op == "call_function" and node.target is executorch_call_delegate - for arg in node.args[1:] - } - staged = [node.name for node in marked if node not in delegate_args] - if staged: - raise RuntimeError( - f"buffer(s) {staged} still reach the delegate through a staging copy, " - "so the engine writes scratch that is thrown away and the cache never " - "updates." - ) - - def _save_zero_copy(trt_gm: torch.fx.GraphModule, inputs: tuple, path: str) -> None: """Save a .pte whose engine updates the KV cache in place. @@ -143,7 +103,11 @@ def _save_zero_copy(trt_gm: torch.fx.GraphModule, inputs: tuple, path: str) -> N that reaches ``to_executorch()`` by any other route has to install the config itself. """ - from torch_tensorrt.executorch import export, zero_copy_backend_config + from torch_tensorrt.executorch import ( + check_zero_copy_kv, + export, + zero_copy_backend_config, + ) # retrace=True here, retrace=False for the plain save() below, so the two # exporters are both covered. Which way round matters: the legacy exporter @@ -154,7 +118,7 @@ def _save_zero_copy(trt_gm: torch.fx.GraphModule, inputs: tuple, path: str) -> N # cache is told from an ordinary copy-back buffer. edge = export(trt_gm, arg_inputs=inputs, retrace=True, zero_copy_kv=True) program = edge.to_executorch(zero_copy_backend_config()) - _check_zero_copy(program) + check_zero_copy_kv(program) with open(path, "wb") as output: program.write_to_file(output) if program._tensor_data: diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index d10692a785c..93461b0e64b 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -1419,7 +1419,11 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None "(torch_tensorrt_runtime). Reinstall torch_tensorrt with the runtime extension." ) try: - from torch_tensorrt.executorch import export, zero_copy_backend_config + from torch_tensorrt.executorch import ( + check_zero_copy_kv, + export, + zero_copy_backend_config, + ) except ImportError: raise ImportError( "ExecuTorch is not installed. Install with: pip install " @@ -1466,6 +1470,13 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None if zero_copy_kv: backend_config = zero_copy_backend_config(backend_config) executorch_program = edge_program.to_executorch(config=backend_config) + if zero_copy_kv: + # save() holds the finalized program here, which is the only place the + # graph shows whether the caches actually reach the engine un-staged. + # Both halves of zero-copy no-op quietly when they find nothing to do, so + # without this a save() that asked for zero-copy could still write an + # ordinary staged .pte. + check_zero_copy_kv(executorch_program) with open(file_path, "wb") as f: executorch_program.write_to_file(f) _write_external_tensor_data(executorch_program, file_path) diff --git a/py/torch_tensorrt/executorch/__init__.py b/py/torch_tensorrt/executorch/__init__.py index 29c9c1e4efe..7af29ebfb83 100644 --- a/py/torch_tensorrt/executorch/__init__.py +++ b/py/torch_tensorrt/executorch/__init__.py @@ -34,10 +34,14 @@ def __getattr__(name: str) -> NoReturn: "TensorRTBackend", "export", "zero_copy_backend_config", + "check_zero_copy_kv", ] else: from torch_tensorrt.executorch._export import export - from torch_tensorrt.executorch._zero_copy import zero_copy_backend_config + from torch_tensorrt.executorch._zero_copy import ( + check_zero_copy_kv, + zero_copy_backend_config, + ) from torch_tensorrt.executorch.backend import TensorRTBackend from torch_tensorrt.executorch.partitioner import TensorRTPartitioner @@ -53,4 +57,5 @@ def get_edge_compile_config() -> "EdgeCompileConfig": "TensorRTBackend", "export", "zero_copy_backend_config", + "check_zero_copy_kv", ] diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index c4b401dd818..ec8908cc183 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -29,7 +29,9 @@ with nothing to copy back -- the buffer would simply never update. So neither pass is public on its own: the rewiring is reached only through ``export(..., zero_copy_kv=True)``, and the un-staging only through -:func:`zero_copy_backend_config`, which is this module's one exported name. +:func:`zero_copy_backend_config`. That, plus :func:`check_zero_copy_kv` -- which +reads a finalized program back and refuses one where the pairing did not happen +-- is what this module exports. """ import logging @@ -570,6 +572,89 @@ def call(self, graph_module: torch.fx.GraphModule) -> Any: return _UnstageThenToOutVar() +def check_zero_copy_kv(program: Any) -> None: + """Raise unless a finalized program really updates its KV buffers in place. + + ``program`` is what ``to_executorch()`` returns. Both halves of zero-copy do + nothing quietly when they find nothing to do: ``zero_copy_kv=True`` warns and + carries on when the model holds no aliased buffer mutation, and + :func:`unstage_aliased_buffers_pass`, handed a program with nothing marked, + un-stages nothing and returns. Either way the ``.pte`` that comes out runs + and stages its cache like any other, which for a KV cache is wrong output + rather than a crash. + + Two shapes are refused: a marked buffer that is not a direct argument of any + *TensorRT* delegate -- it still reaches one through an ``_h2d_copy`` staging, + or reaches none -- and a program with no marked buffer in any method. The + first is what finalizing without :func:`zero_copy_backend_config` leaves + behind; when that config *is* installed, ``_unstage_aliased_buffers`` has + already raised on the same condition. + + Only a TensorRT delegate counts, matching what that pass un-stages. The mark + is put on a buffer because a TensorRT engine writes it in place, so another + backend's delegate taking the buffer directly says nothing about whether the + engine did: it can be handed the buffer while the engine beside it still + reads a staging copy whose contents are discarded. + + Every method is read, not only ``forward``. ``export()`` rewires each method + on its own, so a check that stopped at ``forward`` would pass a program whose + decode had degenerated to staged -- and on the prefill/decode pair the user + guide's zero-copy example exports it would not get that far, since a + multi-method program need not have a ``forward`` at all. The second refusal is + about the program rather than about one method, matching the warning + ``export()`` emits: a method with no aliased buffer mutation of its own is + not an error, so a model that rewires only its decode step is accepted. + + This reads the graph, so it says what the program does rather than what the + passes recorded. It says nothing about whether the engine's write is correct, + only that the buffer it writes is the caller's. + """ + method_names = sorted(program.methods) + staged_by_method: Dict[str, List[str]] = {} + marked_anywhere = False + for method_name in method_names: + graph_module = program.exported_program(method_name).graph_module + marked = [ + node + for node in graph_module.graph.nodes + if node.op == "placeholder" + and node.meta.get("_torch_tensorrt_aliased_buffer") + ] + if not marked: + continue + marked_anywhere = True + delegate_args = { + arg + for node in graph_module.graph.nodes + if _is_tensorrt_delegate(graph_module, node) + for arg in node.args[1:] + } + staged = [node.name for node in marked if node not in delegate_args] + if staged: + staged_by_method[method_name] = staged + if not marked_anywhere: + raise RuntimeError( + "TensorRT zero-copy KV: no buffer in this program is marked for " + f"in-place update, in any of its methods ({', '.join(method_names)}), " + "so it stages its caches like any other .pte. Either it was not " + "exported with zero_copy_kv=True, or it was and no aliased buffer " + "mutation was found -- export logs a warning for that case." + ) + if staged_by_method: + detail = ", ".join( + f"'{name}' in method '{method}'" + for method, names in staged_by_method.items() + for name in names + ) + raise RuntimeError( + f"TensorRT zero-copy KV: buffer(s) {detail} are marked for in-place " + "update but do not reach a TensorRT delegate directly, so the engine " + "writes a staging copy that is discarded and the cache never updates. Export " + "removed their copy-back, so nothing else would restore it. Finalize " + "with torch_tensorrt.executorch.zero_copy_backend_config()." + ) + + def zero_copy_backend_config( config: Optional["ExecutorchBackendConfig"] = None, ) -> "ExecutorchBackendConfig": @@ -591,10 +676,12 @@ def zero_copy_backend_config( .. warning:: Finalizing a ``zero_copy_kv=True`` program *without* this config does - not raise. The engine writes a per-call staging copy that is then - discarded and the buffer never updates, which for a KV cache is wrong - output rather than a crash. Nothing downstream can detect the omission, - so pairing the two is the caller's responsibility. + not raise on its own. The engine writes a per-call staging copy that is + then discarded and the buffer never updates, which for a KV cache is + wrong output rather than a crash. Hand the finalized program to + :func:`check_zero_copy_kv` before writing the ``.pte`` and that mistake + becomes an error; ``torch_tensorrt.save(..., zero_copy_kv=True)`` runs + the check for you. The opposite mistake does raise. ``save(..., zero_copy_kv=True)`` installs this pass itself, so handing it the result of this function as diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index 199897414e2..723df7cd6ac 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -698,6 +698,135 @@ def test_zero_copy_backend_config_keeps_the_callers_config(): assert delegate.args[1] is k_buffer +def _finalized_program(forward=None, **methods): + """The shape ``check_zero_copy_kv`` reads: to_executorch()'s return value. + + One positional graph module makes a single-method ``forward`` program; the + keywords name a method each. ``exported_program`` defaults to ``forward`` + and raises ``KeyError`` on a method the program does not have, like + ``ExecutorchProgramManager``'s -- which is what a program with no ``forward`` + does to a caller that never asked for one. + """ + if forward is not None: + methods = {"forward": forward, **methods} + return SimpleNamespace( + methods=set(methods), + exported_program=lambda method_name="forward": SimpleNamespace( + graph_module=methods[method_name] + ), + ) + + +def _unstaged_graph(): + """A graph whose marked buffer already reaches its delegate directly.""" + graph_module, k_buffer, _, _ = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + Z._unstage_aliased_buffers(graph_module) + return graph_module + + +@pytest.mark.unit +def test_check_zero_copy_kv_accepts_an_unstaged_buffer(): + graph_module, k_buffer, _, delegate = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + Z._unstage_aliased_buffers(graph_module) + + Z.check_zero_copy_kv(_finalized_program(graph_module)) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_still_staged_buffer(): + """The shape a program finalized without zero_copy_backend_config has: the + buffer is marked, so export dropped its copy-back, but it still reaches the + delegate through a staging copy the engine's write is thrown away with.""" + graph_module, k_buffer, _, _ = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="do not reach a TensorRT delegate"): + Z.check_zero_copy_kv(_finalized_program(graph_module)) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_buffer_only_another_backend_takes(): + """Another backend's delegate taking the buffer directly is not zero-copy. + + The mark is on this buffer because a TensorRT engine writes it in place, and + that engine here is still reading a staging copy whose contents are thrown + away. Counting any backend's delegate would pass this program. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + lowered_trt = graph.get_attr("lowered_module_0") + lowered_other = graph.get_attr("lowered_module_1") + delegate_trt = graph.call_function( + executorch_call_delegate, (lowered_trt, staged_k) + ) + delegate_other = graph.call_function( + executorch_call_delegate, (lowered_other, k_buffer) + ) + graph.output((k_buffer, delegate_trt, delegate_other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="CudaBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="do not reach a TensorRT delegate"): + Z.check_zero_copy_kv(_finalized_program(graph_module)) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_program_with_nothing_marked(): + """zero_copy_kv=True on a model with no engine-aliased buffer only warns, so + the .pte that comes out is an ordinary staged one. Refuse it rather than let + a caller who asked for zero-copy ship a program that never got it.""" + graph_module, _, _, _ = _staged_delegate_graph() + + with pytest.raises(RuntimeError, match="marked for in-place update"): + Z.check_zero_copy_kv(_finalized_program(graph_module)) + + +@pytest.mark.unit +def test_check_zero_copy_kv_accepts_a_program_with_no_forward_method(): + """The shape the user guide's zero-copy example exports: prefill and decode, + no ``forward``. Reading the default method would raise KeyError naming a + method the caller never asked for. A method that rewired nothing of its own + is not an error either, so only ``decode`` here carries a marked buffer.""" + unmarked, _, _, _ = _staged_delegate_graph() + + Z.check_zero_copy_kv(_finalized_program(prefill=unmarked, decode=_unstaged_graph())) + + +@pytest.mark.unit +def test_check_zero_copy_kv_catches_a_method_other_than_forward(): + """The silent case: ``forward`` got zero-copy and ``decode`` degenerated to + staged. Stopping at ``forward`` would write a .pte whose decode cache never + updates, so the failure has to name the method that lost it.""" + staged, k_buffer, _, _ = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="in method 'decode'"): + Z.check_zero_copy_kv(_finalized_program(_unstaged_graph(), decode=staged)) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_multi_method_program_with_nothing_marked(): + """Nothing marked anywhere is about the program, not about one method: a + method with no aliased buffer mutation is an error only when no other method + has one, and the failure lists every method it looked in.""" + first, _, _, _ = _staged_delegate_graph() + second, _, _, _ = _staged_delegate_graph() + + with pytest.raises(RuntimeError, match=r"\(decode, prefill\)"): + Z.check_zero_copy_kv(_finalized_program(prefill=first, decode=second)) + + @pytest.mark.unit def test_zero_copy_backend_config_defaults_to_executorch_defaults(): """Called with no config it starts from ExecuTorch's defaults, and the one @@ -1178,7 +1307,7 @@ def _lowered_module(graph_module, delegate): return getattr(graph_module, delegate.args[0].target) -def _assert_marked_buffers_reach_the_engine_unstaged(graph_module): +def _assert_marked_buffers_reach_the_engine_unstaged(program): """Every marked buffer is a direct argument of a TensorRT delegate. Finalizing a zero-copy program without raising is a weak signal, because both @@ -1187,7 +1316,16 @@ def _assert_marked_buffers_reach_the_engine_unstaged(graph_module): argument pointing at the staging copy still satisfies them. Only the graph says whether the rewiring happened, and getting it wrong is silent -- the engine writes per-call scratch that is discarded and the cache never updates. + + The library's own check runs first, on the whole program. It is weaker than + what follows -- it accepts a marked buffer reaching any backend's delegate -- + but it is the only place the container API it reads, ``methods`` and + ``exported_program(name)``, meets a real ``ExecutorchProgramManager``: every + other test of it builds the program itself, so an upstream rename would + leave those green and break ``save(zero_copy_kv=True)`` for every caller. """ + torch_tensorrt.executorch.check_zero_copy_kv(program) + graph_module = program.exported_program().graph_module marked = [ node for node in graph_module.graph.nodes @@ -1283,9 +1421,7 @@ def test_zero_copy_kv_keeps_a_copyback_buffer_in_the_same_method(retrace): program = edge.to_executorch( config=torch_tensorrt.executorch.zero_copy_backend_config() ) - _assert_marked_buffers_reach_the_engine_unstaged( - program.exported_program().graph_module - ) + _assert_marked_buffers_reach_the_engine_unstaged(program) class _SplitRolesDecodeStep(_MixedDecodeStep): @@ -1428,9 +1564,7 @@ def test_zero_copy_kv_with_the_copyback_on_a_second_delegate(retrace): program = edge.to_executorch( config=torch_tensorrt.executorch.zero_copy_backend_config() ) - _assert_marked_buffers_reach_the_engine_unstaged( - program.exported_program().graph_module - ) + _assert_marked_buffers_reach_the_engine_unstaged(program) @pytest.mark.skipif( @@ -1539,9 +1673,7 @@ def _marked_args(node): program = edge.to_executorch( config=torch_tensorrt.executorch.zero_copy_backend_config() ) - _assert_marked_buffers_reach_the_engine_unstaged( - program.exported_program().graph_module - ) + _assert_marked_buffers_reach_the_engine_unstaged(program) # -------------------------------------------------------------------------- @@ -1571,9 +1703,10 @@ def _install_save_stubs(monkeypatch, *, wrap_config=True): """Stub the executorch lowering that save() drives and record how it is called. Returns a namespace capturing the kwargs export() received, the arguments - zero_copy_backend_config() was wrapped with, and the config finally handed to - to_executorch(). When ``wrap_config`` is False the real - zero_copy_backend_config runs, so the recorded config is the genuine one. + zero_copy_backend_config() was wrapped with, the config finally handed to + to_executorch(), and the programs check_zero_copy_kv() was given. When + ``wrap_config`` is False the real zero_copy_backend_config runs, so the + recorded config is the genuine one. """ import torch_tensorrt._compile as compile_module import torch_tensorrt.executorch as executorch_api @@ -1585,17 +1718,30 @@ def _install_save_stubs(monkeypatch, *, wrap_config=True): ) calls = SimpleNamespace( - export_kwargs=None, wrap_args=[], to_executorch_config="unset" + export_kwargs=None, + wrap_args=[], + to_executorch_config="unset", + program=None, + checked=[], ) def _to_executorch(config=None): calls.to_executorch_config = config - return SimpleNamespace( + calls.program = SimpleNamespace( _tensor_data=None, write_to_file=lambda f: f.write(b"stub-pte") ) + return calls.program edge = SimpleNamespace(to_executorch=_to_executorch) + # The stub program has no graph, so the real check cannot read it; what these + # tests pin is that save() runs it, on the finalized program, before writing. + monkeypatch.setattr( + executorch_api, + "check_zero_copy_kv", + lambda program: calls.checked.append(program), + ) + def _export(exp_program, **kwargs): calls.export_kwargs = kwargs return edge @@ -1642,6 +1788,7 @@ def test_save_zero_copy_kv_true_threads_flag_and_installs_config(monkeypatch, tm # The user's config is wrapped once (preserving their fields), not double-wrapped. assert calls.wrap_args == [user_cfg] assert calls.to_executorch_config is calls.wrapped_sentinel + assert calls.checked == [calls.program] @pytest.mark.unit @@ -1700,3 +1847,4 @@ def test_save_defaults_leave_kv_staged(monkeypatch, tmp_path): assert calls.export_kwargs["zero_copy_kv"] is False assert calls.wrap_args == [] assert calls.to_executorch_config is user_cfg + assert calls.checked == [] From 84c55d60ba700d5af187cfe1a6cb4212c8a186c2 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Wed, 2 Sep 2026 23:24:20 -0700 Subject: [PATCH 09/22] fix(executorch): refuse a zero-output engine after the dead code is eliminated The guard against eliding every output of an engine fired only when *every* user of the engine was one of the elided getitems. A non-elided getitem that is itself dead defeated it: the check saw an output, the `eliminate_dead_code()` right below erased that getitem anyway, and the partition got the zero-output delegate the guard exists to refuse. Counting only the users that something reads does not settle it either, because it looks one step past the engine and a dead chain can be longer than that. The first link of a two-node dead chain does have a user, so the guard stays silent and the elimination then erases the whole chain. Measured on hand-built graphs with chains of 1, 2 and 3 nodes: only the one-node chain raised, and at 2 and 3 the rewiring returned normally leaving the engine node with no users at all. So run the elimination first and then ask whether the engine still has a user: what survives it is what the delegate will have, whatever the chain length. The test is parametrized over a one- and a two-node chain, the second being the length a rule reading only the engine's immediate users misses. The comment above it claimed an `execute_engine` node is impure to FX and survives DCE with no users. That is true, and it now says why, because the reason is not local to this repository: PyTorch defaults any operator taking a ScriptObject argument to an ORDERED effect (`torch._library.effects.EffectHolder._set_default_effect`), and `execute_engine` takes the engine as one. Confirmed on the op this code builds against -- `_get_effect(torch.ops.tensorrt.execute_engine.default)` is `EffectType.ORDERED`, `Node.is_impure()` is `True`, and a userless engine node survives `Graph.eliminate_dead_code()`; the delegate that replaces it after lowering has no effect registered and does not. The comment also no longer leaves the reader to infer that the guard, and not the DCE, is what stops the bad shape. --- py/torch_tensorrt/executorch/_zero_copy.py | 33 +++++++----- .../py/dynamo/executorch/test_zero_copy_kv.py | 51 +++++++++++++++++++ 2 files changed, 70 insertions(+), 14 deletions(-) diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index ec8908cc183..f96cb028e48 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -239,7 +239,7 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> List[str]: _LOGGER.debug("no aliased buffer mutations to rewire") return [] - elided_by_engine: Dict[Node, List[Node]] = {} + engines_with_elided_outputs: Set[Node] = set() output_names_by_engine: Dict[Node, List[str]] = {} elided_output_names: List[str] = [] output_node = graph_module.graph.output_node() @@ -256,7 +256,7 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> List[str]: TensorArgument(name=mutation.placeholder.name), output_specs[spec_index].target, ) - elided_by_engine.setdefault(mutation.engine, []).append(mutation.aliased_output) + engines_with_elided_outputs.add(mutation.engine) names = output_names_by_engine.get(mutation.engine) if names is None: names = _engine_output_binding_names(exported_program, mutation.engine) @@ -266,24 +266,29 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> List[str]: elided_output_names.append(names[output_index]) output_node.args = (tuple(output_args),) - # Dropping every output of an engine would leave a delegate with no outputs. + graph_module.graph.eliminate_dead_code() + # Leaving an engine with no output would leave its delegate with no outputs. # Nothing downstream reports that shape: the runtime infers elision from a # single argument count, which a zero-output delegate satisfies, and a # delegate nothing reads is a pure node that a later graph-wide dead-code - # elimination can erase, taking the computation with it. Stop here instead. - # The eliminate_dead_code() below does not erase this engine node: unlike - # the delegate, an execute_engine node is impure to FX and survives with no - # users. - for engine, elided in elided_by_engine.items(): - if all(user in elided and not user.users for user in engine.users): + # elimination can erase, taking the computation with it. This raise is what + # stops that. It reads the graph after the elimination above rather than + # before, so that an output kept alive only by a chain that is itself dead + # does not count: the elimination erases such a chain however long it is, + # and what survives it is what the delegate will really have. The engine + # node itself survives even with no users -- PyTorch defaults an operator + # taking a ScriptObject argument to an ORDERED effect + # (torch._library.effects), and execute_engine takes the engine as one, so + # FX reads it as impure where the delegate is not. + for engine in engines_with_elided_outputs: + if not engine.users: raise RuntimeError( - "TensorRT zero-copy KV: every output of engine node " - f"'{engine.name}' is an aliased buffer written in place, so " - "eliding them would leave the delegate with no outputs at all. " - "This shape is not supported; export this method without " + "TensorRT zero-copy KV: eliding the aliased buffers engine node " + f"'{engine.name}' writes in place leaves it with no output any " + "node reads, so the delegate would have no outputs at all. This " + "shape is not supported; export this method without " "zero_copy_kv." ) - graph_module.graph.eliminate_dead_code() graph_module.graph.lint() graph_module.recompile() # The signature is replaced in place rather than by rebuilding the program: diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index 723df7cd6ac..f97e9b7dbbc 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -316,6 +316,57 @@ def test_rewire_rejects_an_engine_whose_every_output_is_aliased(monkeypatch): Z.rewire_aliased_mutations_to_buffers(program) +@pytest.mark.unit +@pytest.mark.parametrize("dead_chain_length", [1, 2]) +def test_rewire_rejects_an_engine_whose_only_other_output_is_dead( + monkeypatch, dead_chain_length +): + """A surviving-but-unread output does not keep the engine's delegate alive. + + The engine has two outputs: an aliased buffer this elides, and a second one + whose consumers end in nothing. Counting the second as an output would let + the check pass, and the dead-code elimination would erase the whole chain, + leaving exactly the zero-output delegate the check exists to refuse. The + two-link chain is the case a check reading only the engine's immediate users + misses: that first link does have a user, so it reads as live. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + engine = graph.placeholder("engine") + engine_call = graph.call_function( + torch.ops.tensorrt.execute_engine.default, ([k_buffer], engine) + ) + k_out = graph.call_function(operator.getitem, (engine_call, 0)) + dead = graph.call_function(operator.getitem, (engine_call, 1)) + for _ in range(dead_chain_length - 1): + dead = graph.call_function(torch.add, (dead, dead)) + graph.output((k_out,)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + signature = SimpleNamespace( + inputs_to_buffers={"b_k_0": "k_0"}, + input_specs=[], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(name=k_out.name), "k_0" + ) + ], + ) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=signature, + _graph_signature=signature, + ) + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in"], + output_names=["out_k", "out_dead"], + ) + + with pytest.raises(RuntimeError, match="no outputs at all"): + Z.rewire_aliased_mutations_to_buffers(program) + + def _staged_delegate_graph( *, backend_id="TensorRTBackend", device=DeviceType.CUDA, compile_specs=None ): From bc0ecda01f062469968f381648d20058c6c03233 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Wed, 2 Sep 2026 23:25:46 -0700 Subject: [PATCH 10/22] fix(executorch): reject a repeated aliased_io entry when the blob is parsed `init` incremented the aliased-output count once per `aliased_io` entry rather than once per output binding it claimed. A blob listing the same entry twice passed every check -- both copies resolve to the same output and the same engine alias -- and left the count at two for one aliased output. `execute` reads that count to decide the delegate arity, so an inflated one is not just a log line. With the count above the number of flagged outputs, the elision test can be true while the output loop consumes more arguments than the span holds, and nothing bounds its index into `args`. With the count above the number of output bindings, `num_outputs - num_aliased_outputs` underflows, the sum in the length check wraps, and the check passes for any arity. A repeat is malformed for every reader of the blob, not only for `init`, and it is visible from the bytes alone -- so refuse it in `TensorRTBlobHeader::parse`, where the blob-header unit tests reach it without a GPU or a real engine. It is refused the way the parser refuses any other malformed metadata, which means the diagnostic is `init`'s generic parse failure rather than a message naming the output. `execute` keeps the `num_aliased_outputs <= num_outputs` bound. The parser cannot establish it: a blob may carry an empty `io_bindings` array, and then the output bindings are inferred from the deserialized engine, which the parser has not seen. Both subtractions there are unsigned and the length check is the only thing bounding how far the loops index into `args`, so the bound is still worth checking for a header that reached the backend some other way. Our own exporter cannot emit a duplicate; this needs a crafted or corrupt file. --- .../executorch/TensorRTBackend.cpp | 19 ++++++++++++ .../executorch/TensorRTBlobHeader.cpp | 11 +++++++ .../test_executorch_blob_header.cpp | 31 +++++++++++++++++++ 3 files changed, 61 insertions(+) diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 785843f8db0..5c4b0085a9a 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -440,6 +440,9 @@ Result TensorRTBackend::init( // Map each aliased output binding to the index of the input it aliases so // execute() can bind it to that input's device pointer (in-place). // Non-aliased models have an empty header.aliased_io -> all -1, unchanged path. + // The parser has already refused a blob claiming one output twice, so the count + // built below is one per distinct output binding, which is what execute() + // subtracts on. 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) { @@ -573,6 +576,22 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // cannot instead mean "the aliased outputs were never declared". const size_t num_delegate_inputs = num_inputs; const size_t num_aliased_outputs = engine->num_aliased_outputs; + // The blob parser refuses a second aliased_io entry for an output already + // claimed, and init refuses an entry whose output is not one of the recorded + // output bindings, so this holds for any header that reached here. It is + // checked anyway because both subtractions below are unsigned: on a header + // built some other way they wrap and the length check then accepts any + // argument count. That wrap is all this catches. A duplicate entry inflates + // the count while staying within num_outputs, passes both checks, and still + // indexes one past the end of args; the parser's refusal is what stops that. + if (num_aliased_outputs > num_outputs) { + ET_LOG( + Error, + "TensorRTBackend::execute: %zu aliased output(s) recorded for %zu output binding(s)", + num_aliased_outputs, + num_outputs); + return Error::InvalidProgram; + } const bool aliased_outputs_elided = num_aliased_outputs > 0 && args.size() == num_delegate_inputs + num_outputs - num_aliased_outputs; const size_t num_delegate_outputs = num_outputs - (aliased_outputs_elided ? num_aliased_outputs : 0); diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp index 9fc5600c35a..1eaeae996f0 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp @@ -4,6 +4,7 @@ #include #include #include +#include namespace torch_tensorrt { namespace executorch_backend { @@ -246,6 +247,7 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { return false; } ++apos; + std::unordered_set claimed_outputs; while (true) { apos = skip_ws(json, apos); if (apos >= json.size()) { @@ -302,6 +304,15 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { } } if (!ab.output.empty() && !ab.input.empty()) { + // An output binding may be claimed by at most one entry. A second entry + // for the same output names the same binding, so nothing that resolves + // the names can tell the two apart -- but a reader that counts aliased + // outputs per entry, as TensorRTBackend does to size the delegate + // argument list, counts one output twice. Refuse the blob here, where + // the repeat is visible from the bytes alone. + if (!claimed_outputs.insert(ab.output).second) { + return false; + } // The current Python serializer always writes "kind" (serialization.py), // and older blobs carry no aliased_io array at all, so this default is // defensive: it only fires for a blob that has an aliased_io entry but diff --git a/tests/cpp/executorch/test_executorch_blob_header.cpp b/tests/cpp/executorch/test_executorch_blob_header.cpp index 6a7e5cd91dc..90782ff4a3a 100644 --- a/tests/cpp/executorch/test_executorch_blob_header.cpp +++ b/tests/cpp/executorch/test_executorch_blob_header.cpp @@ -177,6 +177,37 @@ TEST(ExecuTorchTensorRTBlobHeader, InputNamedAliasedIoWithNoAliasesStillParses) EXPECT_TRUE(header.aliased_io.empty()); } +TEST(ExecuTorchTensorRTBlobHeader, RejectsRepeatedAliasedIoOutput) { + // Two entries claiming out_k. They resolve to the same binding and the same + // engine alias, so nothing that looks the names up can tell them from one + // entry -- but a reader counting aliased outputs per entry counts out_k twice, + // which is how the backend sizes the delegate argument list. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"out_v","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"},)" + R"({"output":"out_k","input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesAliasedIoEntriesForDistinctOutputs) { + // The minimal pair for the test above: the same blob with the second entry + // claiming its own output. A second entry is not itself the defect. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"out_v","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"},)" + R"({"output":"out_v","input":"in_v","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + ASSERT_EQ(header.aliased_io.size(), 2u); + EXPECT_EQ(header.aliased_io[0].output, "out_k"); + EXPECT_EQ(header.aliased_io[1].output, "out_v"); +} + TEST(ExecuTorchTensorRTBlobHeader, RejectsUnknownFutureMagic) { constexpr char kFutureMagic[4] = {'T', 'R', '0', '3'}; const std::string metadata = R"({"io_bindings":[{"name":"x","is_input":true}]})"; From 24ca7d4c276e716a448be7541720662ff462c528 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Wed, 2 Sep 2026 23:32:03 -0700 Subject: [PATCH 11/22] test(executorch): run the KV decode check with and without a caller stream The KV persistence check never installed a caller stream guard, so `getCallerStream()` was always empty and the backend always took the synchronizing branch at the end of `execute()`. The skip-the-sync path zero-copy KV depends on -- no host staging, no aliased reflect, a caller stream set -- had therefore never run under this check, and neither had the machinery that exists to make it safe: the `inflight_event` record, the wait on it at the top of the next `execute()`, and the drain in `~EngineHandle`. Run both scenarios twice, once with no caller stream and once with a `CallerStreamGuard` scoped over the decode loop on a stream the check owns, the way the main runner does. The guard is constructed only in the second mode: an explicitly null selection is still a selection (see `tests/cpp/executorch/test_caller_stream.cpp`), so an unconditional guard over a null stream would cover one branch twice. The guarded run synchronizes and destroys its stream before reading the logits. What that buys is branch coverage, and only on the zero-copy `.pte`. Instrumenting the branch decision shows the zero-copy model in the guarded mode taking the skip path on all three of its `execute()` calls, waiting on the previous enqueue at the top of the second `execute()` of the two-step scenario, and draining a pending enqueue at each of the two teardowns; the unguarded mode reproduces the old always-synchronize behaviour in the same binary. The staged `.pte` still synchronizes in both modes, because its aliased outputs are delegate output args and so an aliased reflect is always pending -- running it under the guard moves the engine onto a caller-supplied stream and nothing else. What it does not buy is a check that can fail when a synchronization is missing. The stream comes from `cudaStreamCreate`, which is a blocking stream, and the logits are read with a synchronous `cudaMemcpy` on the legacy default stream, which implicitly waits for every blocking stream in the context. Two controls say so on this model: deleting the check's own `cudaStreamSynchronize` still passes, and a `cudaStreamQuery` placed where that sync was reports the stream already complete every time. So repeated runs agreeing to the last bit are not evidence that the ordering is right. The check exercises the path and would fail on a hard error in it -- a rejected event record, a context reconfigured under a live enqueue, a mis-elided argument -- but it is not a race detector. Making it one would mean a `cudaStreamNonBlocking` stream and dropping the explicit sync that is redundant with the default-stream copy. `verify-executorch-reference-runner.sh` now greps for both modes by name, so dropping one cannot leave the lane green with the branch uncovered. Also: where the shared-stream contract is written down, the zero-copy section described getting it wrong as failing quietly. The coalesced-`.pte` section, a few paragraphs down, already calls it a race that can surface as wrong results or an illegal memory access. Say it at that strength in both places. --- .../verify-executorch-reference-runner.sh | 13 +- .../runtime_performance/saving_models.rst | 6 +- .../executorch_reference_runner/README.md | 7 + .../kv_cache_decode_check.cpp | 131 +++++++++++++----- 4 files changed, 119 insertions(+), 38 deletions(-) diff --git a/.github/scripts/verify-executorch-reference-runner.sh b/.github/scripts/verify-executorch-reference-runner.sh index 5b52744241b..a4b96146927 100755 --- a/.github/scripts/verify-executorch-reference-runner.sh +++ b/.github/scripts/verify-executorch-reference-runner.sh @@ -580,12 +580,19 @@ if [[ ${#kv_model_paths[@]} -gt 0 ]]; then kv_index=0 for kv_model_path in "${kv_model_paths[@]}"; do # kv_cache_decode_check exits non-zero when a decode step does not observe the - # KV the previous step wrote; the grep additionally pins the assertion itself, + # KV the previous step wrote; the greps additionally pin the assertion itself, # so weakening the check inside the binary cannot quietly turn this into a - # no-op. + # no-op. Both caller-stream modes are pinned by name: the backend skips its + # end-of-execute synchronization only when a caller stream is set, so dropping + # the "own" run would leave the branch zero-copy KV relies on uncovered while + # the lane stayed green. kv_check_log="${verify_root}/kv_cache_decode_check_${kv_index}.log" "${kv_check_path}" --model_path="${kv_model_path}" 2>&1 | tee "${kv_check_log}" - grep -q "PASS: decode at pos=1 observed the KV written at pos=0" "${kv_check_log}" + for kv_stream_mode in none own; do + grep -q \ + "PASS: decode at pos=1 observed the KV written at pos=0 across execute() calls (caller stream: ${kv_stream_mode})" \ + "${kv_check_log}" + done kv_index=$((kv_index + 1)) done fi diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 05d7fbf6005..34b56af8543 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -420,10 +420,12 @@ It installs ``zero_copy_backend_config`` for you, so do not hand it one as ``backend_config`` as well: the pass would be installed twice and finalization raises. The two entry points are alternatives, not a pair. -Two further responsibilities are the caller's, and both fail quietly: +Two further responsibilities are the caller's, and neither raises: * **One CUDA stream for every delegate**, if the ``.pte`` is coalesced -- and the - synchronization it calls for, which zero-copy makes load-bearing. See + synchronization it calls for, which zero-copy makes load-bearing. Getting this + wrong is a race, not a deterministic error: it is intermittent and can surface + as wrong results *or* as an illegal memory access. See :ref:`Running a coalesced .pte `. * **Sharing one cache between methods.** Zero-copy is per method: it makes each diff --git a/examples/executorch_reference_runner/README.md b/examples/executorch_reference_runner/README.md index 1b1ccba4544..3ec2f61fbb6 100644 --- a/examples/executorch_reference_runner/README.md +++ b/examples/executorch_reference_runner/README.md @@ -207,3 +207,10 @@ Because the causal attention at position 1 covers positions 0..1, the two logits differ only if the KV written at position 0 persisted across `execute()` calls. The runner prints `[kv-check] PASS` and returns 0 on success, or fails if the two are identical (the update did not persist). It requires a CUDA device. + +That pair runs twice over, printing `caller stream: none` and then `caller +stream: own`. The second scopes a `CallerStreamGuard` over the decode loop on a +stream the runner creates, which is what lets the backend return from +`execute()` with the enqueue still in flight; the first leaves the caller stream +unset, so every `execute()` synchronizes before it returns. A zero-copy `.pte` +reaches the skip-the-sync path only in the second, so both have to pass. diff --git a/examples/executorch_reference_runner/kv_cache_decode_check.cpp b/examples/executorch_reference_runner/kv_cache_decode_check.cpp index 2be3ae1a263..bda398126e5 100644 --- a/examples/executorch_reference_runner/kv_cache_decode_check.cpp +++ b/examples/executorch_reference_runner/kv_cache_decode_check.cpp @@ -22,6 +22,14 @@ * position-0 slot is still zero). Equal logits mean the update did not persist * (cache reset per call, or the aliased output bound to scratch), so we fail. * + * Both scenarios are run twice, once with no caller stream and once with a + * CallerStreamGuard scoped over the decode loop on a stream this runner owns. + * The backend takes a different path in each: with no caller stream it + * synchronizes at the end of every execute(), and with one -- and nothing to + * stage or reflect, which is what a zero-copy .pte leaves -- it may return with + * the enqueue still in flight. Only the second exercises that skip, and only + * there does the runner owe the synchronization before it reads the logits. + * * Usage: * kv_cache_decode_check --model_path=kv_cache_decode.pte [--tol=1e-3] */ @@ -33,10 +41,12 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -72,8 +82,14 @@ static const char* get_flag(int argc, char** argv, const char* flag, const char* // Load a FRESH method (zeroed caller-owned buffers), run one decode step per // entry in `positions` (token id fixed to 1, input_pos = the position), and -// return the final step's first output as host floats. -static std::vector run_decode(Program& program, const char* method_name, const std::vector& positions) { +// return the final step's first output as host floats. With `use_caller_stream` +// the decode loop runs under a CallerStreamGuard on a stream owned here, the way +// the main runner scopes one; without it the backend sees no caller stream. +static std::vector run_decode( + Program& program, + const char* method_name, + const std::vector& positions, + bool use_caller_stream) { Result method_meta = program.method_meta(method_name); ET_CHECK_MSG(method_meta.ok(), "method_meta failed: 0x%" PRIx32, static_cast(method_meta.error())); @@ -148,6 +164,17 @@ static std::vector run_decode(Program& program, const char* method_name, exec_aten::ScalarType::Long, nd, sizes[i].data(), data[i].data(), dim_order[i].data(), strides[i].data()); } + cudaStream_t caller_stream = nullptr; + // Optional, not a guard over a null stream: an explicitly null selection is + // still a selection, so constructing one unconditionally would leave both + // modes with a caller stream set and cover the same branch twice. The guard + // scopes the whole loop rather than each step, so consecutive decodes order on + // one stream instead of on an end-of-execute sync. + std::optional caller_stream_guard; + if (use_caller_stream) { + ET_CHECK_MSG(cudaStreamCreate(&caller_stream) == cudaSuccess, "cudaStreamCreate failed"); + caller_stream_guard.emplace(caller_stream); + } for (int64_t pos : positions) { for (size_t i = 1; i < num_inputs; ++i) { std::fill(data[i].begin(), data[i].end(), pos); @@ -157,6 +184,20 @@ static std::vector run_decode(Program& program, const char* method_name, } ET_CHECK_MSG(method->execute() == Error::Ok, "execute() failed at pos %" PRId64, pos); } + caller_stream_guard.reset(); + if (use_caller_stream) { + // The last execute() may have returned with the enqueue still running, so a + // runner owning the stream owes it this before reading an output. It is what + // the contract asks for rather than what makes the read correct here: + // cudaStreamCreate returns a blocking stream, and the synchronous cudaMemcpy + // below runs on the legacy default stream, which waits for every blocking + // stream in the context. So this exercises the branch but cannot detect a + // missing synchronization -- deleting it leaves the check passing. + ET_CHECK_MSG(cudaStreamSynchronize(caller_stream) == cudaSuccess, "cudaStreamSynchronize failed"); + if (cudaStreamDestroy(caller_stream) != cudaSuccess) { + ET_LOG(Error, "cudaStreamDestroy failed"); + } + } EValue out; ET_CHECK_MSG(method->get_outputs(&out, 1) == Error::Ok, "get_outputs failed"); @@ -164,7 +205,9 @@ static std::vector run_decode(Program& program, const char* method_name, exec_aten::Tensor t = out.toTensor(); ET_CHECK_MSG(t.scalar_type() == exec_aten::ScalarType::Float, "expected float logits output"); // The output may be device-resident; cudaMemcpyDefault copies from host or - // device. execute() synchronized (no caller stream) so the result is ready. + // device. The work is finished either way: with no caller stream the backend + // synchronized at the end of execute(), and with one the stream was + // synchronized above. std::vector result(static_cast(t.numel())); ET_CHECK_MSG( cudaMemcpy(result.data(), t.const_data_ptr(), result.size() * sizeof(float), cudaMemcpyDefault) == cudaSuccess, @@ -191,37 +234,59 @@ int main(int argc, char** argv) { const char* method_name = *name; ET_LOG(Info, "Loaded '%s' method '%s'", model_path, method_name); - // A: pos=1 from a zeroed cache. B: pos=0 then pos=1 (second step sees pos 0). - std::vector a = run_decode(*program, method_name, {1}); - std::vector b = run_decode(*program, method_name, {0, 1}); - - ET_CHECK_MSG(a.size() == b.size() && !a.empty(), "output size mismatch (%zu vs %zu)", a.size(), b.size()); - double max_abs_diff = 0.0; - bool saw_nan = false; - for (size_t i = 0; i < a.size(); ++i) { - // std::max(0.0, fabs(NaN)) is 0.0 (NaN compares false), so a NaN logit would - // otherwise leave max_abs_diff at 0.0 and be misreported as "identical". - // Track NaNs explicitly and fail on them below. - if (std::isnan(a[i]) || std::isnan(b[i])) { - saw_nan = true; + // Both modes have to hold. Without a caller stream the backend synchronizes at + // the end of every execute(); with one it may skip that, which is the branch + // zero-copy KV depends on and the one a wrong result shows up in only + // sometimes. + struct Mode { + const char* label; + bool use_caller_stream; + }; + for (const Mode& mode : {Mode{"none", false}, Mode{"own", true}}) { + // A: pos=1 from a zeroed cache. B: pos=0 then pos=1 (second step sees pos 0). + std::vector a = run_decode(*program, method_name, {1}, mode.use_caller_stream); + std::vector b = run_decode(*program, method_name, {0, 1}, mode.use_caller_stream); + + ET_CHECK_MSG(a.size() == b.size() && !a.empty(), "output size mismatch (%zu vs %zu)", a.size(), b.size()); + double max_abs_diff = 0.0; + bool saw_nan = false; + for (size_t i = 0; i < a.size(); ++i) { + // std::max(0.0, fabs(NaN)) is 0.0 (NaN compares false), so a NaN logit would + // otherwise leave max_abs_diff at 0.0 and be misreported as "identical". + // Track NaNs explicitly and fail on them below. + if (std::isnan(a[i]) || std::isnan(b[i])) { + saw_nan = true; + } + max_abs_diff = std::max(max_abs_diff, std::fabs(static_cast(a[i]) - static_cast(b[i]))); } - max_abs_diff = std::max(max_abs_diff, std::fabs(static_cast(a[i]) - static_cast(b[i]))); - } - fprintf( - stderr, - "[kv-check] logits numel=%zu max|A(no-history) - B(with-history)| = %.6g (tol=%.3g)\n", - a.size(), - max_abs_diff, - tol); - if (saw_nan) { - fprintf(stderr, "[kv-check] FAIL: logits contain NaN -> the decode produced invalid output.\n"); - return 1; - } - if (max_abs_diff > tol) { - fprintf(stderr, "[kv-check] PASS: decode at pos=1 observed the KV written at pos=0 across execute() calls.\n"); - return 0; + fprintf( + stderr, + "[kv-check] caller stream: %s logits numel=%zu max|A(no-history) - B(with-history)| = %.6g (tol=%.3g)\n", + mode.label, + a.size(), + max_abs_diff, + tol); + if (saw_nan) { + fprintf( + stderr, + "[kv-check] FAIL: logits contain NaN -> the decode produced invalid output (caller stream: %s).\n", + mode.label); + return 1; + } + if (max_abs_diff <= tol) { + fprintf( + stderr, + "[kv-check] FAIL: outputs are identical -> the KV write did not persist across execute() calls " + "(caller stream: %s).\n", + mode.label); + return 1; + } + fprintf( + stderr, + "[kv-check] PASS: decode at pos=1 observed the KV written at pos=0 across execute() calls " + "(caller stream: %s).\n", + mode.label); } - fprintf(stderr, "[kv-check] FAIL: outputs are identical -> the KV write did not persist across execute() calls.\n"); - return 1; + return 0; } From a48164adddb626a789e5c67351fa10036f788226 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Fri, 4 Sep 2026 16:51:01 -0700 Subject: [PATCH 12/22] fix(executorch): key the un-staging on what the program ends up looking like Two defects in `_unstage_aliased_buffers`, both from deciding by the edit it made rather than by the shape it has to leave behind. **It keyed success on having deleted a staging copy.** The property zero-copy needs is that the marked buffer is a direct argument of a TensorRT delegate; removing an `_h2d_copy` is only the usual route there. `ExecutorchBackendConfig(enable_non_cpu_memory_planning=False)` is a supported configuration -- the field defaults to True and `zero_copy_backend_config` preserves whatever the caller set -- and under it `PropagateDevicePass` inserts no staging copies at all. Measured against that pass directly: with the flag True the delegate's argument is an `_h2d_copy` and the placeholder stays on the host; with it False there is no copy node, the placeholder *is* the argument, and its spec is already `cuda:0`. That program is in the shape zero-copy wants and the pass rejected it, naming a caller two causes -- the buffer never reached a delegate, or the pass was installed twice -- neither of which had happened, and telling them to install it once. `check_zero_copy_kv` accepted the same graph, so the two halves disagreed about one program. The pass now counts a marked buffer as satisfied when a TensorRT delegate takes it, whether or not this pass is what put it there, and refuses only when no TensorRT delegate takes it at all -- directly or through a staging copy the pass can remove. That is the condition `check_zero_copy_kv` reads, so the two accept and refuse the same graphs. A consequence is that installing the pass twice is now a no-op rather than an error: the second run finds the first run's work in place. The docstring, the `save()` comment and the user guide each promised that finalization raises there, and now describe what it does. **It allowed a second staging copy to the same GPU to survive the move.** A marked buffer feeding two same-GPU `_h2d_copy` nodes, one to a TensorRT delegate and one to another backend, had the TensorRT staging erased and its own spec flipped host->CUDA, leaving the other copy reading a source that is now device memory. `_h2d_copy_out` requires a host source and fails `InvalidArgument` on a device one (portable kernel; the ATen branch has no such check). The rule is now that any staging copy this pass does not itself remove blocks the move, same GPU or not, and a copy is removed exactly when every user of it is a TensorRT delegate whose argument gets rewired. That reverses a decision earlier in this stack. `test_unstage_leaves_another_backends_same_gpu_staging_in_place` was rewritten to pin the allowance -- that the buffer still moves and the other backend's copy stays -- and now pins the refusal, under a name that says so. Two neighbouring tests describe which comparison catches which shape and are corrected with it: the different-GPU-other-backend shape is no longer refused by the device index alone, and the two-TensorRT-GPUs shape, where every staging does feed a TensorRT delegate, now asserts that nothing moved rather than only that something raised -- without that, dropping the index comparison would still leave it green. --- .../runtime_performance/saving_models.rst | 7 +- py/torch_tensorrt/_compile.py | 6 +- py/torch_tensorrt/executorch/_zero_copy.py | 167 ++++++++++-------- .../py/dynamo/executorch/test_zero_copy_kv.py | 159 +++++++++++++---- 4 files changed, 227 insertions(+), 112 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 34b56af8543..18397c4bfb9 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -416,9 +416,10 @@ one silently would break a runner built before this feature. zero_copy_kv=True, ) -It installs ``zero_copy_backend_config`` for you, so do not hand it one as -``backend_config`` as well: the pass would be installed twice and finalization -raises. The two entry points are alternatives, not a pair. +It installs ``zero_copy_backend_config`` for you, so there is no need to hand it +one as ``backend_config`` as well: that installs the pass twice, which is +redundant rather than an error -- the second run finds the buffers already +un-staged. The two entry points are alternatives, not a pair. Two further responsibilities are the caller's, and neither raises: diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index 93461b0e64b..270f33c1976 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -1463,9 +1463,9 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None # Unlike the direct export()+to_executorch() path -- where the two steps # belong to different owners and pairing them is the caller's job -- save() # owns both, so it installs the finalization pass itself. Wrapping preserves - # every field of the caller's config. save() installs the pass once; a - # backend_config that already carries it is wrapped again here, and - # finalization then raises. + # every field of the caller's config. A backend_config that already carries + # the pass is wrapped again here; the second run finds the buffers already + # un-staged and changes nothing. backend_config = kwargs.get("backend_config") if zero_copy_kv: backend_config = zero_copy_backend_config(backend_config) diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index f96cb028e48..6bee3ceb332 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -335,9 +335,9 @@ def _delegate_declares_zero_copy( several TensorRT delegates therefore marks only the KV one, never the plain compute engines beside it -- which is what keeps this cross-check from demanding an aliased buffer from a delegate that never had one. A delegate - that declares it but un-staged nothing is a lost KV update -- the mark that - would have driven the un-staging did not survive to this pass -- and is caught - in :func:`_unstage_aliased_buffers`. + that declares it but ends up taking no marked buffer is a lost KV update -- + the mark that would have driven the un-staging did not survive to this pass + -- and is caught in :func:`_unstage_aliased_buffers`. """ from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY @@ -352,22 +352,29 @@ def _delegate_declares_zero_copy( def _device_move_is_safe( - source: Node, h2d_copy: Any, target_device: Any, target_device_index: Any + graph_module: torch.fx.GraphModule, + source: Node, + h2d_copy: Any, + target_device: Any, + target_device_index: Any, ) -> bool: """True when moving ``source``'s spec device disturbs no other consumer. A placeholder's device is shared by every user, so it can only be retargeted - to the delegate's device when nothing *reads* it on another one. ExecuTorch + to the delegate's device when nothing else *reads* it afterwards. ExecuTorch guards the same hazard, more strictly and only under its opt-in ``skip_h2d_for_method_inputs``: it demands the placeholder have exactly one - user. The rule here is looser because two kinds of user impose no such - constraint and are allowed: the graph ``output`` node -- the buffer is its + user. The rule here is looser because two kinds of user survive the move + unaffected and are allowed: the graph ``output`` node -- the buffer is its own BUFFER_MUTATION result, which is exactly what zero-copy sets up and - which carries no device of its own -- and another ``_h2d_copy`` staging to - the same GPU, superseded by the un-staging when it feeds a TensorRT - delegate, and otherwise left in place, still reading a buffer that now lives - on the GPU it was copying to. Any other reader (a compute op, or a staging - to a different device) makes the move unsafe. + which carries no device of its own -- and an ``_h2d_copy`` to the same GPU + that this pass removes, because every one of its users is a TensorRT + delegate whose argument the pass rewires to the buffer. + + A staging copy that outlives the pass is *not* allowed, even on the same + GPU. It would go on reading the buffer as its source after the move has put + the buffer in device memory, and ``_h2d_copy_out`` requires a host source: + the portable kernel checks it and fails ``InvalidArgument``. The index is compared as well as the type, because ``spec.device`` is only ``CUDA``/``CPU``: two engines resolved to ``cuda:0`` and ``cuda:1`` stage the @@ -389,13 +396,17 @@ def _device_move_is_safe( return False if spec.device_index != target_device_index: return False + if not user.users or not all( + _is_tensorrt_delegate(graph_module, copy_user) for copy_user in user.users + ): + return False return True def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: """Route TensorRT delegate inputs from their staging copy back to the buffer. - A delegate input qualifies only when it is an ``_h2d_copy`` of a placeholder + An input is un-staged only when it is an ``_h2d_copy`` of a placeholder carrying the mark left by :func:`rewire_aliased_mutations_to_buffers`. Every other input keeps its staging, including a mutable buffer the engine does not write in place. @@ -404,32 +415,42 @@ def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: planning puts the buffer in the delegate's device arena rather than a host one. That is what makes handing the buffer straight to the engine valid at all: a host-arena pointer is not something the engine can write. That move is - refused when the buffer has another consumer (see - :func:`_device_move_is_safe`), which would otherwise have its device silently - changed too. - - A failure here is a lost KV update -- unless the program has already been - through this pass, the one case where nothing is lost -- so it is raised - rather than logged: export has already removed the copy-back, so a marked - buffer left staged has the engine write per-call scratch that is then - discarded and the buffer never updates. It raises when the staging copy has - no spec or is not on CUDA, when the device move is unsafe, and -- so a - discovery miss cannot pass silently -- after the loop when any marked buffer - was never un-staged, cross-checked against each delegate's own - ``zero_copy_kv`` spec: a TensorRT delegate that declares zero-copy but - un-staged nothing is broken and names the buffer. - - Returns the number of delegate inputs un-staged. + refused when the buffer has a consumer this pass leaves behind (see + :func:`_device_move_is_safe`), which would otherwise be reading the buffer + from somewhere other than where it was put. + + What the pass has to establish is the *post-condition*: every marked buffer + is a direct argument of a TensorRT delegate. Removing a staging copy is only + the usual way of getting there, not the goal, and a marked buffer that + already is such an argument is left alone and counts as satisfied. Two things + produce that shape. Setting ``enable_non_cpu_memory_planning=False`` on the + ``ExecutorchBackendConfig`` makes ``PropagateDevicePass`` tag the placeholder + with the delegate's device instead of inserting a staging copy, so there is + never one to remove; and running this pass a second time over a program it + has already un-staged finds its own work in place. + + A failure here is a lost KV update, so it is raised rather than logged: + export has already removed the copy-back, so a marked buffer left staged has + the engine write per-call scratch that is then discarded and the buffer never + updates. It raises when the staging copy has no spec or is not on CUDA, when + the device move is unsafe, and -- so a discovery miss cannot pass silently -- + after the loop when the post-condition does not hold for some marked buffer, + cross-checked against each delegate's own ``zero_copy_kv`` spec: a TensorRT + delegate that declares zero-copy but takes no marked buffer is broken, and + that refusal names the delegate and lists what it does take. + + Returns the number of delegate inputs un-staged, which is zero for a program + that already satisfied the post-condition. """ from executorch.exir.schema import DeviceType from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY h2d_copy = torch.ops.et_copy._h2d_copy.default unstaged = 0 - unstaged_placeholders: Set[Node] = set() + satisfied_placeholders: Set[Node] = set() orphaned_stagings: List[Node] = [] zero_copy_delegates: List[Node] = [] - unstaged_per_delegate: Dict[Node, int] = {} + satisfied_per_delegate: Dict[Node, int] = {} for node in list(graph_module.graph.nodes): if not _is_tensorrt_delegate(graph_module, node): @@ -437,10 +458,18 @@ def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: declares_zero_copy = _delegate_declares_zero_copy(graph_module, node) if declares_zero_copy: zero_copy_delegates.append(node) - unstaged_per_delegate[node] = 0 + satisfied_per_delegate[node] = 0 new_args = list(node.args) for i, arg in enumerate(node.args[1:], start=1): - if not isinstance(arg, Node) or arg.target is not h2d_copy: + if not isinstance(arg, Node): + continue + if arg.op == "placeholder": + if arg.meta.get("_torch_tensorrt_aliased_buffer"): + satisfied_placeholders.add(arg) + if declares_zero_copy: + satisfied_per_delegate[node] += 1 + continue + if arg.target is not h2d_copy: continue source = arg.args[0] if not isinstance(source, Node) or source.op != "placeholder": @@ -475,66 +504,60 @@ def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: staged_spec.device_index, ) if not already_placed and not _device_move_is_safe( - source, h2d_copy, staged_spec.device, staged_spec.device_index + graph_module, + source, + h2d_copy, + staged_spec.device, + staged_spec.device_index, ): raise RuntimeError( "TensorRT zero-copy KV: buffer " - f"'{source.name}' is read by a consumer the move would " - "disturb -- an op outside the delegate, or a staging copy " - "bound for a different GPU -- so placing it on this engine's " - "device would silently change that consumer's device too. " - "Export this method without zero_copy_kv, or stop sharing the " - "aliased buffer." + f"'{source.name}' is read by a consumer this pass leaves in " + "place, so placing the buffer on this engine's device would " + "change the device that consumer reads it from. Export this " + "method without zero_copy_kv, or stop sharing the aliased " + "buffer." ) source_spec.device = staged_spec.device source_spec.device_index = staged_spec.device_index new_args[i] = source unstaged += 1 - unstaged_placeholders.add(source) + satisfied_placeholders.add(source) orphaned_stagings.append(arg) if declares_zero_copy: - unstaged_per_delegate[node] += 1 + satisfied_per_delegate[node] += 1 node.args = tuple(new_args) - marked_but_unstaged = [ + marked_but_unsatisfied = [ node for node in graph_module.graph.nodes if node.op == "placeholder" and node.meta.get("_torch_tensorrt_aliased_buffer") - and node not in unstaged_placeholders + and node not in satisfied_placeholders ] - if marked_but_unstaged: - names = ", ".join(repr(node.name) for node in marked_but_unstaged) + if marked_but_unsatisfied: + names = ", ".join(repr(node.name) for node in marked_but_unsatisfied) raise RuntimeError( "TensorRT zero-copy KV: buffer(s) " f"{names} were marked for in-place update but no TensorRT delegate " - "staging was found to un-stage. Either they never reached a " - "TensorRT delegate, which is a broken zero-copy program -- export " - "removed their copy-back, so leaving them staged has the engine " - "write per-call scratch that is discarded and the buffer never " - "updates -- or this pass has already run over the program and they " - "are wired straight to the engine already, which happens whenever " - "it is installed twice: nesting zero_copy_backend_config, " - "finalizing the same program twice, or passing " - "save(zero_copy_kv=True) a config that already carries the pass. " - "Install it once." + "takes them, either directly or through a staging copy this pass " + "could remove. Export removed their copy-back on the promise that a " + "TensorRT engine writes them in place, so as this program stands " + "nothing updates them. Export this method without zero_copy_kv, or " + "keep the aliased buffer on a TensorRT delegate." ) for delegate in zero_copy_delegates: - if unstaged_per_delegate[delegate] == 0: - staged_inputs = [ - arg.args[0].name - for arg in delegate.args[1:] - if isinstance(arg, Node) - and arg.target is h2d_copy - and isinstance(arg.args[0], Node) + if satisfied_per_delegate[delegate] == 0: + delegate_inputs = [ + arg.name for arg in delegate.args[1:] if isinstance(arg, Node) ] raise RuntimeError( "TensorRT zero-copy KV: delegate " f"'{delegate.name}' declares zero-copy KV " - f"(compile spec '{ZERO_COPY_KV_COMPILE_SPEC_KEY}') but no aliased " - f"buffer was un-staged for it (staged inputs: {staged_inputs}). Export " - "elided its aliased outputs, so the engine now writes per-call " - "scratch that is discarded and the cache never updates." + f"(compile spec '{ZERO_COPY_KV_COMPILE_SPEC_KEY}') but takes no " + f"buffer marked for in-place update (inputs: {delegate_inputs}). " + "Export elided its aliased outputs, so the engine now writes " + "per-call scratch that is discarded and the cache never updates." ) if unstaged: @@ -688,10 +711,12 @@ def zero_copy_backend_config( becomes an error; ``torch_tensorrt.save(..., zero_copy_kv=True)`` runs the check for you. - The opposite mistake does raise. ``save(..., zero_copy_kv=True)`` - installs this pass itself, so handing it the result of this function as - ``backend_config`` applies the pass twice and finalization fails. The - two entry points are mutually exclusive: use one or the other. + ``save(..., zero_copy_kv=True)`` installs this pass itself, so handing + it the result of this function as ``backend_config`` applies the pass + twice. That is redundant rather than an error -- the second run finds + the buffers already wired straight to their delegates and changes + nothing -- but the two entry points are alternatives: use one or the + other. """ from dataclasses import replace diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index f97e9b7dbbc..ebc3561354b 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -430,6 +430,73 @@ def test_unstage_keeps_staging_for_an_unmarked_buffer(): assert k_buffer.meta["spec"].device == DeviceType.CPU +def _direct_delegate_graph(*, compile_specs=None): + """A lowered graph with no staging at all: delegate(lowered, k_buffer). + + This is what ``PropagateDevicePass`` produces under + ``ExecutorchBackendConfig(enable_non_cpu_memory_planning=False)``: instead of + inserting an ``_h2d_copy`` before each delegate input it tags the argument's + own spec with the delegate's device, so the placeholder stays the delegate's + argument and there is no staging copy to remove. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + lowered = graph.get_attr("lowered_module_0") + delegate = graph.call_function(executorch_call_delegate, (lowered, k_buffer)) + graph.output((k_buffer, delegate)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=compile_specs + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + return graph_module, k_buffer, delegate + + +@pytest.mark.unit +@pytest.mark.parametrize( + "compile_specs", + [None, [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")]], + ids=["plain", "declares-zero-copy"], +) +def test_unstage_accepts_a_buffer_that_never_had_a_staging_copy(compile_specs): + """A marked buffer already handed straight to its delegate needs no work. + + What the pass has to leave behind is a marked buffer that is a delegate + argument; removing a staging copy is only the usual route there. Keying + success on having removed one instead rejects this program, which is a + supported ExecuTorch configuration and is already in the shape zero-copy + wants. The delegate's own zero-copy declaration is cross-checked against the + same count, so it is parametrized here too. + """ + graph_module, k_buffer, delegate = _direct_delegate_graph( + compile_specs=compile_specs + ) + + assert Z._unstage_aliased_buffers(graph_module) == 0 + + assert delegate.args[1] is k_buffer + assert k_buffer.meta["spec"].device == DeviceType.CUDA + + +@pytest.mark.unit +def test_unstage_runs_a_second_time_without_raising(): + """Installing the pass twice is redundant rather than an error. + + ``save(zero_copy_kv=True)`` installs it, so a caller who also passes + ``zero_copy_backend_config()`` as ``backend_config`` gets two of them. The + second run finds the buffer already wired to the delegate and returns + without un-staging anything. + """ + graph_module, k_buffer, _, delegate = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + assert Z._unstage_aliased_buffers(graph_module) == 1 + + assert Z._unstage_aliased_buffers(graph_module) == 0 + assert delegate.args[1] is k_buffer + + @pytest.mark.unit def test_unstage_raises_for_a_marked_buffer_on_another_backends_delegate(): """Only a TensorRT engine promises the in-place write, so a marked buffer @@ -441,7 +508,7 @@ def test_unstage_raises_for_a_marked_buffer_on_another_backends_delegate(): ) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - with pytest.raises(RuntimeError, match="no TensorRT delegate staging"): + with pytest.raises(RuntimeError, match="no TensorRT delegate takes them"): Z._unstage_aliased_buffers(graph_module) @@ -457,7 +524,7 @@ def test_unstage_raises_when_a_marked_buffer_is_never_unstaged(): k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - with pytest.raises(RuntimeError, match="no TensorRT delegate staging"): + with pytest.raises(RuntimeError, match="no TensorRT delegate takes them"): Z._unstage_aliased_buffers(graph_module) @@ -476,17 +543,15 @@ def test_unstage_raises_when_a_zero_copy_delegate_unstaged_nothing(): @pytest.mark.unit -def test_unstage_leaves_another_backends_same_gpu_staging_in_place(): - """A marked buffer read by a second backend on the *same* GPU is still moved. - - This is the accept side of the ``_h2d_copy`` allowance in - ``_device_move_is_safe``: the other backend's delegate keeps its staging copy, - which after the move reads a buffer already resident on the GPU it was copying - to, so nothing it sees changes. The function's other accept branch, the graph - ``output`` node, is pinned by - ``test_unstage_allows_a_buffer_that_is_also_its_mutation_output``; every other - unit test that gives the buffer a second ``_h2d_copy`` pins a refusal, so a - rule that allowed no second staging at all would still pass all of those. +def test_unstage_refuses_a_buffer_another_backend_stages_on_the_same_gpu(): + """A second backend's staging copy on the *same* GPU is refused, not kept. + + The other backend's ``_h2d_copy`` outlives this pass and goes on reading the + buffer as its source, but the move has just put the buffer in device memory. + ``_h2d_copy_out`` requires a host source and fails ``InvalidArgument`` on a + device one, so leaving the copy in place produces a program that does not + run. Same GPU is what makes this shape distinct: the device and index both + match, so neither of ``_device_move_is_safe``'s spec comparisons rejects it. """ graph = torch.fx.Graph() k_buffer = graph.placeholder("b_k_0") @@ -515,13 +580,14 @@ def test_unstage_leaves_another_backends_same_gpu_staging_in_place(): staged_other.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - assert Z._unstage_aliased_buffers(graph_module) == 1 + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): + Z._unstage_aliased_buffers(graph_module) - assert delegate_trt.args[1] is k_buffer + # Refused before anything moved: the buffer is where the other backend's + # staging copy expects to read it. + assert k_buffer.meta["spec"].device == DeviceType.CPU + assert delegate_trt.args[1] is staged_trt assert delegate_other.args[1] is staged_other - assert staged_other in graph_module.graph.nodes - assert k_buffer.meta["spec"].device == DeviceType.CUDA - assert k_buffer.meta["spec"].device_index == 0 @pytest.mark.unit @@ -595,7 +661,7 @@ def test_unstage_refuses_to_move_a_shared_buffer(): staged_k.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=3) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - with pytest.raises(RuntimeError, match="consumer the move would disturb"): + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): Z._unstage_aliased_buffers(graph_module) @@ -605,13 +671,16 @@ def test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus(): un-staged for either: a spec carries one device index, so whichever engine lost the race would be handed an address on the other's GPU. ``spec.device`` is only CUDA/CPU, so it is the device-index comparison in - ``_device_move_is_safe`` that refuses the first delegate here. That is what - separates this from the supported two-delegates-one-GPU shape, but it is not - what this test pins: two branches of that function refuse the shape in - sequence, and were the index comparison gone, delegate 0 would be un-staged - and the direct-consumer branch would refuse delegate 1 instead, leaving this - test green. The index comparison itself is pinned by - ``test_unstage_refuses_a_buffer_another_backend_stages_to_a_different_gpu``. + ``_device_move_is_safe`` that refuses the first delegate here -- both + stagings feed a TensorRT delegate, which is what separates this from the + two-backends shapes and leaves the index the only comparison that can catch + it. + + Which is why the refusal alone would not pin it. Were the index comparison + gone, delegate 0 would be un-staged and the direct-consumer branch would then + refuse delegate 1, and this test would still see a RuntimeError. What it + checks is that nothing moved: the buffer is still on the host and delegate 0 + still reads its staging copy. """ graph = torch.fx.Graph() k_buffer = graph.placeholder("b_k_0") @@ -636,9 +705,12 @@ def test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus(): staged_1.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=1) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - with pytest.raises(RuntimeError, match="consumer the move would disturb"): + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): Z._unstage_aliased_buffers(graph_module) + assert k_buffer.meta["spec"].device == DeviceType.CPU + assert delegate_0.args[1] is staged_0 + @pytest.mark.unit def test_unstage_refuses_a_buffer_another_backend_stages_to_a_different_gpu(): @@ -648,10 +720,12 @@ def test_unstage_refuses_a_buffer_another_backend_stages_to_a_different_gpu(): Un-staging skips the other backend's delegate, so its staging copy keeps reading the buffer while staging it to cuda:1, and re-homing the buffer onto the TensorRT engine's cuda:0 would move the source of that read to the wrong - GPU. The device-index comparison is the only thing that refuses this shape: - the other staging is an ``_h2d_copy`` like every supported one, and because - it is never un-staged its delegate never becomes the kind of direct consumer - the shared-buffer rule catches. + GPU. Two of ``_device_move_is_safe``'s comparisons refuse this shape -- the + device index, which runs first, and the surviving copy's non-TensorRT user + -- so this test does not discriminate between them. The index comparison on + its own is pinned by + ``test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus``, where every + staging does feed a TensorRT delegate. """ graph = torch.fx.Graph() k_buffer = graph.placeholder("b_k_0") @@ -680,7 +754,7 @@ def test_unstage_refuses_a_buffer_another_backend_stages_to_a_different_gpu(): staged_other.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=1) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - with pytest.raises(RuntimeError, match="consumer the move would disturb"): + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): Z._unstage_aliased_buffers(graph_module) @@ -718,7 +792,7 @@ def test_unstage_refuses_to_rehome_a_buffer_already_on_another_gpu(): staged_1.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=1) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - with pytest.raises(RuntimeError, match="consumer the move would disturb"): + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): Z._unstage_aliased_buffers(graph_module) assert k_buffer.meta["spec"].device_index == 0 @@ -797,6 +871,21 @@ def test_check_zero_copy_kv_rejects_a_still_staged_buffer(): Z.check_zero_copy_kv(_finalized_program(graph_module)) +@pytest.mark.unit +def test_check_zero_copy_kv_accepts_a_buffer_that_never_had_a_staging_copy(): + """The checker and the un-staging pass read the same post-condition. + + A program finalized with ``enable_non_cpu_memory_planning=False`` hands the + buffer straight to the delegate without a staging copy ever existing. The + pass accepts it (``test_unstage_accepts_a_buffer_that_never_had_a_staging_copy``) + and so must this: the two disagreeing is what would let one path refuse a + program the other calls correct. + """ + graph_module, _, _ = _direct_delegate_graph() + + Z.check_zero_copy_kv(_finalized_program(graph_module)) + + @pytest.mark.unit def test_check_zero_copy_kv_rejects_a_buffer_only_another_backend_takes(): """Another backend's delegate taking the buffer directly is not zero-copy. @@ -1056,8 +1145,8 @@ def test_multi_delegate_zero_copy_lowers_without_false_raise(monkeypatch): The KV buffer is un-staged and the plain delegate is left alone. A plain delegate stamped zero-copy would instead make _unstage_aliased_buffers raise - "declares zero-copy KV ... but no aliased buffer was un-staged for it" over a - program that is correct. + "declares zero-copy KV ... but takes no buffer marked for in-place update" + over a program that is correct. """ program, engine_a, engine_b = _two_engine_program() result = _partition_two_engines(program, engine_a, engine_b, monkeypatch) From f41338f7d2d714cf65fb638e700b9af2a0eb230f Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sat, 5 Sep 2026 19:18:06 -0700 Subject: [PATCH 13/22] docs(executorch): correct four zero-copy claims a reader would act on Four pieces of prose that a reader can check against the code and find disagreeing. All four sit in commits already pushed, so they are collected here rather than folded back into them. `_aliased_inputs_by_output_index` said the skipped case is "an output binding absent from the delegate's inputs". What the two `continue`s skip is an `aliased_io` entry whose *input* does not resolve -- a name that is not one of the engine's input bindings, or an index past the delegate's argument list. `_declare_aliased_kv_mutations_on_ep` warns on both, which is the sentence's point and is why neither is reported twice. `rewire_aliased_mutations_to_buffers` said "an engine mixing the two is caught". An engine carrying a rewired aliased output beside an un-rewired user alias is not caught and must not be: it exports cleanly, because the un-rewired output is still a delegate output and that is what the binding check wants. What is caught is a delegate that dropped the un-rewired one as well. The user guide's zero-copy example passed `CudaPartitioner([])` for each of two methods. `export()`'s own docstring says a partitioner whose specs name no method leaves a backend that reads its method name from them -- the CUDA backend -- unable to find it, so the snippet as written raises during lowering. Nothing in the example needs a second backend; drop the argument. `test_validate_output_binding_order_still_accepts_aliased_outputs_threaded` was described as covering "the pre-existing shape", which tells a later reader nothing: there is no before. It covers that naming a binding elidable permits the drop without requiring it. Also two smaller ones: a comment in `test_edge_cases.py` explained the `zero_copy_kv=False` assertion in terms of a KV buffer that model does not have; and the `**ExecuTorch lowering options**` heading had no blank line above it, so it ran into the paragraph before. That last one is on `main` and predates this stack, but the paragraph it runs into is one this stack rewrote. --- .../runtime_performance/saving_models.rst | 2 +- py/torch_tensorrt/executorch/_zero_copy.py | 13 ++++++++----- tests/py/dynamo/executorch/test_backend.py | 4 ++-- tests/py/dynamo/executorch/test_edge_cases.py | 4 ++-- 4 files changed, 13 insertions(+), 10 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 18397c4bfb9..33d82355a9c 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -377,7 +377,6 @@ calls, one at each end of the Edge boundary: edge = export( {"prefill": prefill_program, "decode": decode_program}, - partitioners={"prefill": [CudaPartitioner([])], "decode": [CudaPartitioner([])]}, zero_copy_kv=True, ) @@ -502,6 +501,7 @@ for the engine before it returns. Under :ref:`zero-copy KV ` those outputs are elided, so there is nothing to reflect and the delegate returns with the engine still running -- the synchronization is then the only thing making a host read see the new values. + **ExecuTorch lowering options** ``torch_tensorrt.save`` takes these extra keyword arguments. They are only diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index 6bee3ceb332..43a0bc74a22 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -63,9 +63,11 @@ def _aliased_inputs_by_output_index( graph. The graph cannot tell the difference: an aliased KV mutation and a copy-back mutation are both a ``getitem`` off the engine node whose buffer is also an engine input, and rewiring a copy-back would silently drop a real - update. An output binding absent from the delegate's inputs is skipped - rather than reported -- ``_declare_aliased_kv_mutations_on_ep`` has already - warned about that same engine, and there is no mutation to rewire either way. + update. An entry whose aliased *input* does not resolve -- the name is not one + of the engine's input bindings, or its index is past the delegate's argument + list -- is skipped rather than reported: + ``_declare_aliased_kv_mutations_on_ep`` warns on both of those for the same + engine, and neither leaves a mutation to rewire. """ from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ( ALIASED_IO_IDX, @@ -222,8 +224,9 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> List[str]: one per rewired mutation. Only these names may later be exempted from the backend's output-binding check -- every *other* aliased output (a user alias on a plain, non-buffer input, which export never rewired) must still be a - delegate output, so an engine mixing the two is caught rather than silently - dropping the un-rewired one's update into scratch. + delegate output. So on an engine mixing the two, a delegate that dropped the + un-rewired one as well is caught rather than silently writing that update + into scratch. """ from torch.export.graph_signature import ( ExportGraphSignature, diff --git a/tests/py/dynamo/executorch/test_backend.py b/tests/py/dynamo/executorch/test_backend.py index c4eed938cc9..fe25896a12c 100644 --- a/tests/py/dynamo/executorch/test_backend.py +++ b/tests/py/dynamo/executorch/test_backend.py @@ -475,8 +475,8 @@ def test_validate_output_binding_order_rejects_elision_that_was_not_requested(): @pytest.mark.unit def test_validate_output_binding_order_still_accepts_aliased_outputs_threaded(): - """Not eliding is the pre-existing shape and stays legal: without zero-copy - each aliased output is a delegate output like any other.""" + """Naming a binding elidable permits the drop, it does not require it: a + delegate that still carries every aliased output is accepted too.""" from torch_tensorrt.executorch.backend import _validate_output_binding_order ep, engine = _engine_partition([0, 1, 2]) diff --git a/tests/py/dynamo/executorch/test_edge_cases.py b/tests/py/dynamo/executorch/test_edge_cases.py index f5cf9b417c9..a6223d67b9f 100644 --- a/tests/py/dynamo/executorch/test_edge_cases.py +++ b/tests/py/dynamo/executorch/test_edge_cases.py @@ -71,8 +71,8 @@ def test_save_as_executorch_uses_public_lowering_and_persists_data( weight_streaming_budget_per_engine=None, zero_copy_kv=False, ) - # zero_copy_kv defaults off, so the backend_config reaches to_executorch() - # unwrapped and the KV buffer keeps its staging and its copy-back. + # With zero_copy_kv off the caller's backend_config reaches to_executorch() + # exactly as given: save() wraps it only to install the un-staging pass. edge.to_executorch.assert_called_once_with(config=backend_config) program.write_to_file.assert_called_once() program.write_tensor_data_to_file.assert_called_once_with(str(tmp_path)) From f9fea6b62ee733dcb8f9dcf22b3967428c5be3a7 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sun, 6 Sep 2026 01:00:57 -0700 Subject: [PATCH 14/22] fix(executorch): require the aliased buffer to be planned in device memory `_unstage_aliased_buffers` counted a marked buffer as satisfied the moment a TensorRT delegate took it directly, on the strength of the mark and the delegate edge alone. Being a direct argument is half of what zero-copy needs. The other half is that memory planning puts the buffer in a device arena, and nothing checked it. `ExecutorchBackendConfig(enable_non_cpu_memory_planning=False)` produces exactly that shape and does not satisfy it. `PropagateDevicePass` inserts no staging copy under that flag and writes the delegate's device straight onto the placeholder's own spec, so the buffer looks placed; memory planning with the flag off then ignores every spec device and puts the whole program in one host arena. Measured on a real export, the marked placeholders carry `DeviceType.CUDA` either way, and land in a CUDA arena with the flag on and in the single host arena with it off. The `.pte` that came out exported clean, passed `check_zero_copy_kv`, and failed its first `execute()` on the runtime's alias-target guard -- "aliased input 'buf_k_cache' must be device-resident", `Error::InvalidArgument`. That was run on device rather than inferred. The spec's device is therefore not on its own the answer either. Two things decide where the buffer is planned and only one of them is in the graph, so `zero_copy_backend_config` reads `enable_non_cpu_memory_planning` off the configuration it is building from and hands it to the pass. The pass refuses a direct argument that is either off-CUDA or host-planned, naming the buffer, since that buffer is the one whose update is lost. The docstring named that configuration as one of the two supported ways to reach the direct-argument shape. It is not one, and the sentence now says what being a direct argument is and is not enough for; the pass's own second run over its output is the route that remains. Both tests pinning the shape set the spec to CUDA by hand, which is what the real configuration does too, so neither could have caught this. The new ones cover each conjunct on its own and the wiring that carries the planning mode in. --- .../runtime_performance/saving_models.rst | 6 + py/torch_tensorrt/_compile.py | 4 +- py/torch_tensorrt/executorch/_zero_copy.py | 124 ++++++++++++++---- .../py/dynamo/executorch/test_zero_copy_kv.py | 99 +++++++++++--- 4 files changed, 188 insertions(+), 45 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 33d82355a9c..d59c8aefdcd 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -384,6 +384,12 @@ calls, one at each end of the Edge boundary: # (memory planning, passes) is preserved. program = edge.to_executorch(zero_copy_backend_config(backend_config)) +One field of that config is also *read*. The engine writes the cache wherever +memory planning put it, so ``enable_non_cpu_memory_planning=False`` -- which +plans every tensor into a single host arena whatever device its ``TensorSpec`` +asks for -- cannot be combined with zero-copy KV: ``to_executorch`` raises +instead of writing a ``.pte`` whose every ``execute()`` fails on a host pointer. + It is opt-in rather than automatic because the resulting ``.pte`` needs a runtime that understands a delegate whose aliased outputs are elided. Producing one silently would break a runner built before this feature. diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index 270f33c1976..0a66b04edbb 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -1463,7 +1463,9 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None # Unlike the direct export()+to_executorch() path -- where the two steps # belong to different owners and pairing them is the caller's job -- save() # owns both, so it installs the finalization pass itself. Wrapping preserves - # every field of the caller's config. A backend_config that already carries + # every field of the caller's config and reads one of them: a config that + # turns non-CPU memory planning off cannot place the caches where the engine + # writes, and the pass refuses it. A backend_config that already carries # the pass is wrapped again here; the second run finds the buffers already # un-staged and changes nothing. backend_config = kwargs.get("backend_config") diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index 43a0bc74a22..ce47c25a5e2 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -406,7 +406,9 @@ def _device_move_is_safe( return True -def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: +def _unstage_aliased_buffers( + graph_module: torch.fx.GraphModule, *, device_memory_planning: bool = True +) -> int: """Route TensorRT delegate inputs from their staging copy back to the buffer. An input is un-staged only when it is an ``_h2d_copy`` of a placeholder @@ -414,33 +416,44 @@ def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: other input keeps its staging, including a mutable buffer the engine does not write in place. - The placeholder's spec takes over the staging copy's device, so memory - planning puts the buffer in the delegate's device arena rather than a host - one. That is what makes handing the buffer straight to the engine valid at - all: a host-arena pointer is not something the engine can write. That move is - refused when the buffer has a consumer this pass leaves behind (see + The placeholder's spec takes over the staging copy's device, which is what + asks memory planning for the delegate's device arena rather than a host one + (asks, not settles -- see ``device_memory_planning`` below). That is what + makes handing the buffer straight to the engine valid at all: a host-arena + pointer is not something the engine can write. That move is refused when the + buffer has a consumer this pass leaves behind (see :func:`_device_move_is_safe`), which would otherwise be reading the buffer from somewhere other than where it was put. What the pass has to establish is the *post-condition*: every marked buffer - is a direct argument of a TensorRT delegate. Removing a staging copy is only - the usual way of getting there, not the goal, and a marked buffer that - already is such an argument is left alone and counts as satisfied. Two things - produce that shape. Setting ``enable_non_cpu_memory_planning=False`` on the - ``ExecutorchBackendConfig`` makes ``PropagateDevicePass`` tag the placeholder - with the delegate's device instead of inserting a staging copy, so there is - never one to remove; and running this pass a second time over a program it - has already un-staged finds its own work in place. + is a direct argument of a TensorRT delegate *and* ends up planned in device + memory. Removing a staging copy is only the usual way of getting there, not + the goal, and a marked buffer that already satisfies both is left alone and + counts as satisfied -- which is what running this pass a second time over a + program it has already un-staged finds. + + Being a direct argument is not on its own enough, so it is not on its own + accepted. Two things decide where the buffer is planned, and the second is + not visible in the graph: the spec's own device, and whether memory planning + reads spec devices at all. ``enable_non_cpu_memory_planning=False`` on the + ``ExecutorchBackendConfig`` produces exactly the shape above -- no staging + copy is inserted and ``PropagateDevicePass`` writes the delegate's device + straight onto the placeholder's spec -- and then plans every tensor into the + one host arena regardless, so the engine is handed a host pointer for a + buffer it must write on the device. ``device_memory_planning`` carries that + configuration in; :func:`zero_copy_backend_config` takes it off the + ``ExecutorchBackendConfig`` it is building from. A failure here is a lost KV update, so it is raised rather than logged: export has already removed the copy-back, so a marked buffer left staged has the engine write per-call scratch that is then discarded and the buffer never updates. It raises when the staging copy has no spec or is not on CUDA, when - the device move is unsafe, and -- so a discovery miss cannot pass silently -- - after the loop when the post-condition does not hold for some marked buffer, - cross-checked against each delegate's own ``zero_copy_kv`` spec: a TensorRT - delegate that declares zero-copy but takes no marked buffer is broken, and - that refusal names the delegate and lists what it does take. + a direct argument would not be planned in device memory, when the device move + is unsafe, and -- so a discovery miss cannot pass silently -- after the loop + when the post-condition does not hold for some marked buffer, cross-checked + against each delegate's own ``zero_copy_kv`` spec: a TensorRT delegate that + declares zero-copy but takes no marked buffer is broken, and that refusal + names the delegate and lists what it does take. Returns the number of delegate inputs un-staged, which is zero for a program that already satisfied the post-condition. @@ -468,6 +481,46 @@ def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: continue if arg.op == "placeholder": if arg.meta.get("_torch_tensorrt_aliased_buffer"): + direct_spec = arg.meta.get("spec") + placement = "" + remedy = "" + if direct_spec is None: + placement = "it carries no TensorSpec" + remedy = ( + "The specs exist only while this runs as the " + "ExecutorchBackendConfig to_out_var_pass, which is " + "where torch_tensorrt.executorch." + "zero_copy_backend_config() installs it." + ) + elif direct_spec.device != DeviceType.CUDA: + placement = f"its TensorSpec asks for {direct_spec.device!r}" + remedy = ( + "Give the buffer to a TensorRT delegate on CUDA, or " + "export this method without zero_copy_kv." + ) + elif not device_memory_planning: + placement = ( + "memory planning is configured with " + "enable_non_cpu_memory_planning=False, which puts " + "every tensor in the one host arena whatever its " + "TensorSpec says" + ) + remedy = ( + "Finalize over a configuration that leaves " + "enable_non_cpu_memory_planning on, so a CUDA " + "TensorSpec is given a CUDA arena." + ) + if placement: + raise RuntimeError( + "TensorRT zero-copy KV: buffer " + f"'{arg.name}' reaches a TensorRT delegate directly, " + f"with no staging copy to remove, but {placement}, so " + "it is not planned in device memory and the engine is " + "handed a host pointer it cannot write. The engine " + "writes this buffer in place and its copy-back has " + "already been removed, so the update would be lost. " + f"{remedy}" + ) satisfied_placeholders.add(arg) if declares_zero_copy: satisfied_per_delegate[node] += 1 @@ -574,7 +627,9 @@ def _unstage_aliased_buffers(graph_module: torch.fx.GraphModule) -> int: return unstaged -def unstage_aliased_buffers_pass(inner_pass: Optional[Any] = None) -> Any: +def unstage_aliased_buffers_pass( + inner_pass: Optional[Any] = None, *, device_memory_planning: bool = True +) -> Any: """Build a ``to_out_var_pass`` that un-stages aliased buffers, then delegates. ``to_out_var_pass`` is the last hook that runs after ``PropagateDevicePass`` @@ -584,6 +639,11 @@ def unstage_aliased_buffers_pass(inner_pass: Optional[Any] = None) -> Any: ``inner_pass`` is the ``to_out_var_pass`` that would otherwise have run; it runs after the un-staging. Omit it for ExecuTorch's default. + + ``device_memory_planning`` is the ``enable_non_cpu_memory_planning`` the + program will be finalized with. Nothing in the graph records it, and it is + half of what decides whether a marked buffer ends up somewhere the engine can + write, so the pass has to be told: see :func:`_unstage_aliased_buffers`. """ from executorch.exir import ExecutorchBackendConfig from executorch.exir.pass_base import PassBase @@ -596,7 +656,9 @@ def unstage_aliased_buffers_pass(inner_pass: Optional[Any] = None) -> Any: class _UnstageThenToOutVar(PassBase): # type: ignore[misc] def call(self, graph_module: torch.fx.GraphModule) -> Any: - unstaged = _unstage_aliased_buffers(graph_module) + unstaged = _unstage_aliased_buffers( + graph_module, device_memory_planning=device_memory_planning + ) _LOGGER.debug("un-staged %d aliased delegate buffer(s)", unstaged) return inner(graph_module) @@ -637,8 +699,12 @@ def check_zero_copy_kv(program: Any) -> None: not an error, so a model that rewires only its decode step is accepted. This reads the graph, so it says what the program does rather than what the - passes recorded. It says nothing about whether the engine's write is correct, - only that the buffer it writes is the caller's. + passes recorded. What it establishes is the wiring: the caller's buffer is + what each TensorRT delegate is handed. It says nothing about whether the + engine's write is correct, nor about where the buffer is planned -- the pass + :func:`zero_copy_backend_config` installs owns that half, and a program whose + buffer reaches its delegate directly out of a host arena passes here and then + fails every ``execute()`` on the runtime's alias-target guard. """ method_names = sorted(program.methods) staged_by_method: Dict[str, List[str]] = {} @@ -703,7 +769,11 @@ def zero_copy_backend_config( ``config`` is your own configuration -- every field is preserved, and a ``to_out_var_pass`` you already set runs after the un-staging. Omit it to - start from ExecuTorch's defaults. + start from ExecuTorch's defaults. One field is not merely carried but read: + zero-copy needs the caches planned in device memory, so a config with + ``enable_non_cpu_memory_planning=False`` -- which plans every tensor into the + one host arena -- has the pass refuse each cache it finds rather than write a + ``.pte`` whose every ``execute()`` fails. .. warning:: Finalizing a ``zero_copy_kv=True`` program *without* this config does @@ -727,5 +797,9 @@ def zero_copy_backend_config( base = config if config is not None else ExecutorchBackendConfig() return replace( - base, to_out_var_pass=unstage_aliased_buffers_pass(base.to_out_var_pass) + base, + to_out_var_pass=unstage_aliased_buffers_pass( + base.to_out_var_pass, + device_memory_planning=base.enable_non_cpu_memory_planning, + ), ) diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index ebc3561354b..c6b325f7979 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -430,14 +430,17 @@ def test_unstage_keeps_staging_for_an_unmarked_buffer(): assert k_buffer.meta["spec"].device == DeviceType.CPU -def _direct_delegate_graph(*, compile_specs=None): +def _direct_delegate_graph(*, compile_specs=None, device=DeviceType.CUDA): """A lowered graph with no staging at all: delegate(lowered, k_buffer). - This is what ``PropagateDevicePass`` produces under - ``ExecutorchBackendConfig(enable_non_cpu_memory_planning=False)``: instead of - inserting an ``_h2d_copy`` before each delegate input it tags the argument's - own spec with the delegate's device, so the placeholder stays the delegate's - argument and there is no staging copy to remove. + The shape the un-staging pass itself leaves behind, and so the shape its own + second run is handed: the marked buffer is the delegate's argument outright + and there is no ``_h2d_copy`` to remove. + + ``device`` is what the buffer's own spec asks for, which is only half of + where memory planning ends up putting it. The other half is the + ``enable_non_cpu_memory_planning`` the program is finalized with; the graph + does not record it and the pass is told separately. """ graph = torch.fx.Graph() k_buffer = graph.placeholder("b_k_0") @@ -449,7 +452,7 @@ def _direct_delegate_graph(*, compile_specs=None): backend_id="TensorRTBackend", compile_specs=compile_specs ) graph_module = torch.fx.GraphModule(root, graph) - k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["spec"] = SimpleNamespace(device=device, device_index=0) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True return graph_module, k_buffer, delegate @@ -464,11 +467,11 @@ def test_unstage_accepts_a_buffer_that_never_had_a_staging_copy(compile_specs): """A marked buffer already handed straight to its delegate needs no work. What the pass has to leave behind is a marked buffer that is a delegate - argument; removing a staging copy is only the usual route there. Keying - success on having removed one instead rejects this program, which is a - supported ExecuTorch configuration and is already in the shape zero-copy - wants. The delegate's own zero-copy declaration is cross-checked against the - same count, so it is parametrized here too. + argument planned in device memory; removing a staging copy is only the usual + route there. Keying success on having removed one instead rejects this + program, which is already in the shape zero-copy wants -- and it is the shape + the pass's own second run sees. The delegate's own zero-copy declaration is + cross-checked against the same count, so it is parametrized here too. """ graph_module, k_buffer, delegate = _direct_delegate_graph( compile_specs=compile_specs @@ -480,6 +483,63 @@ def test_unstage_accepts_a_buffer_that_never_had_a_staging_copy(compile_specs): assert k_buffer.meta["spec"].device == DeviceType.CUDA +@pytest.mark.unit +def test_unstage_refuses_a_direct_buffer_whose_spec_stays_on_the_host(): + """Reaching the delegate directly is not enough; the buffer has to be there. + + A marked buffer whose own spec asks for the host is planned in a host arena, + and the engine cannot write a host pointer in place. Accepting it on the + strength of the mark and the delegate edge alone writes a ``.pte`` whose + every ``execute()`` fails on the alias-target guard. + """ + graph_module, _, _ = _direct_delegate_graph(device=DeviceType.CPU) + + with pytest.raises(RuntimeError, match="not planned in device memory"): + Z._unstage_aliased_buffers(graph_module) + + +@pytest.mark.unit +def test_unstage_refuses_a_direct_buffer_under_host_only_memory_planning(): + """A CUDA spec does not mean CUDA memory when planning ignores spec devices. + + ``enable_non_cpu_memory_planning=False`` leaves exactly the graph above -- + ``PropagateDevicePass`` writes the delegate's device onto the placeholder's + spec and inserts no staging copy -- and then plans every tensor into the one + host arena regardless. Nothing in the graph distinguishes it from the + already-un-staged shape, which is why the spec device is asserted here to be + CUDA: the refusal can only come from the planning mode the pass was told. + """ + graph_module, k_buffer, _ = _direct_delegate_graph() + assert k_buffer.meta["spec"].device == DeviceType.CUDA, ( + "this test only discriminates while the buffer's spec is CUDA; on a host " + "spec the refusal below could come from the spec-device check instead of " + "from the planning mode" + ) + + with pytest.raises(RuntimeError, match="not planned in device memory"): + Z._unstage_aliased_buffers(graph_module, device_memory_planning=False) + + +@pytest.mark.unit +def test_zero_copy_backend_config_carries_the_planning_mode_into_the_pass(): + """The refusal is only reachable if the config's flag actually gets there. + + ``zero_copy_backend_config`` is the only place that sees the + ``ExecutorchBackendConfig``, so building the pass without reading + ``enable_non_cpu_memory_planning`` off it leaves the check above unreachable + from any real finalization. + """ + from executorch.exir import ExecutorchBackendConfig + + config = Z.zero_copy_backend_config( + ExecutorchBackendConfig(enable_non_cpu_memory_planning=False) + ) + graph_module, _, _ = _direct_delegate_graph() + + with pytest.raises(RuntimeError, match="not planned in device memory"): + config.to_out_var_pass(graph_module) + + @pytest.mark.unit def test_unstage_runs_a_second_time_without_raising(): """Installing the pass twice is redundant rather than an error. @@ -873,13 +933,14 @@ def test_check_zero_copy_kv_rejects_a_still_staged_buffer(): @pytest.mark.unit def test_check_zero_copy_kv_accepts_a_buffer_that_never_had_a_staging_copy(): - """The checker and the un-staging pass read the same post-condition. - - A program finalized with ``enable_non_cpu_memory_planning=False`` hands the - buffer straight to the delegate without a staging copy ever existing. The - pass accepts it (``test_unstage_accepts_a_buffer_that_never_had_a_staging_copy``) - and so must this: the two disagreeing is what would let one path refuse a - program the other calls correct. + """The checker and the un-staging pass read the same graph post-condition. + + A program the pass has already un-staged hands the buffer straight to the + delegate with no staging copy left. The pass accepts that shape + (``test_unstage_accepts_a_buffer_that_never_had_a_staging_copy``) and so must + this: the two disagreeing is what would let one path refuse a program the + other calls correct. The checker reads only the graph, so the placement half + of the pass's post-condition is not its to enforce. """ graph_module, _, _ = _direct_delegate_graph() From 8adba5a52937d4c00d1e941fa47ba38d9f3a455d Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Sun, 6 Sep 2026 01:01:07 -0700 Subject: [PATCH 15/22] fix(executorch): refuse an engine whose aliased outputs are only partly elided Export decides elision per binding name; the runtime decides it with one subtraction. The two agree only while every aliased output of an engine is a rewired buffer mutation. An engine that also aliases onto a plain input -- one buffer KV cache beside an argument cache whose updated value is returned -- has the narrower set elided, passes the output-binding check, passes `check_zero_copy_kv`, and writes a `.pte` that cannot run: the runtime reads elision by subtracting the engine's whole aliased-output count from the argument count it was handed, so a delegate short of only some of them reads as not elided at all. Loaded through the ExecuTorch runtime with this backend, such a file opens, binds both aliases at init, and then fails every `execute()` on the argument count -- "expected at least 7 args, got 6" for the two-alias engine measured -- with `Error::InvalidArgument`. That was run on device rather than inferred from the arity code. `preprocess` already holds the elidable names and the engine's whole `aliased_io` ten lines apart. It now compares them and refuses, so the failure lands where the export can still be re-run instead of on the device, and it says plainly that partial elision is not expressible by the runtime. Teaching the runtime to read the elided names from the compile spec it is already handed is the other way out of this, and is deliberately not taken: refusing is much smaller, and going from refused to supported later is a compatible progression. Deriving the narrower set is not what is wrong here, so the test that pins that derivation keeps its assertion and gains the refusal beside it. --- py/torch_tensorrt/executorch/_zero_copy.py | 8 ++-- py/torch_tensorrt/executorch/backend.py | 32 +++++++++++--- py/torch_tensorrt/executorch/partitioner.py | 5 +++ .../py/dynamo/executorch/test_zero_copy_kv.py | 43 ++++++++++++++++--- 4 files changed, 74 insertions(+), 14 deletions(-) diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index ce47c25a5e2..a19a1489f1c 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -224,9 +224,11 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> List[str]: one per rewired mutation. Only these names may later be exempted from the backend's output-binding check -- every *other* aliased output (a user alias on a plain, non-buffer input, which export never rewired) must still be a - delegate output. So on an engine mixing the two, a delegate that dropped the - un-rewired one as well is caught rather than silently writing that update - into scratch. + delegate output, so a delegate that dropped one of those as well is caught + rather than silently writing that update into scratch. An engine that mixes + the two kinds does not lower at all: ``TensorRTBackend.preprocess`` refuses + it, because the runtime reads elision off a single argument count and so + cannot express eliding only part of one engine's aliased outputs. """ from torch.export.graph_signature import ( ExportGraphSignature, diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index c38091934d0..2034bc1ea49 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -306,10 +306,13 @@ def _validate_output_binding_order( ``elidable_output_names`` names the bindings the delegate is *allowed* to have dropped, which zero-copy KV sets to exactly the aliased outputs export - rewired to write in place (never the whole aliased_io): the engine's in-place - write through the aliased input already is the buffer update, so no argument - is passed for them. Pass ``None`` (the default) when elision was not asked - for, and the delegate must carry every binding. + rewired to write in place: the engine's in-place write through the aliased + input already is the buffer update, so no argument is passed for them. Pass + ``None`` (the default) when elision was not asked for, and the delegate must + carry every binding. This check does not require that set to cover the + engine's whole ``aliased_io``; ``preprocess`` requires a non-empty one to, + because the runtime can only take all of an engine's aliased outputs as + elided or none of them. A delegate that dropped its aliased outputs because nothing declared them as mutations looks exactly like a zero-copy one, and the runtime reads elision off @@ -464,11 +467,12 @@ def preprocess( output_names = _split_binding_names( _get_str(engine_info, OUTPUT_BINDING_NAMES_IDX) ) + elidable_output_names = _elided_output_names(compile_specs) _validate_output_binding_order( edge_program, engine_node, output_names, - _elided_output_names(compile_specs), + elidable_output_names, ) io_bindings = [ TensorRTIOBinding(name=name, is_input=True) for name in input_names @@ -478,6 +482,24 @@ def preprocess( # C++ backend binds each aliased output to its aliased input's tensor # (in-place) and reflects the update back into the delegate output. aliased_io = deserialize_aliased_io(_get_str(engine_info, ALIASED_IO_IDX)) + if elidable_output_names and set(elidable_output_names) != set(aliased_io): + raise ValueError( + "TensorRT ExecuTorch backend: engine " + f"'{engine_node.name}' aliases the outputs {sorted(aliased_io)}, " + f"but only {sorted(elidable_output_names)} of them are elided -- " + "dropped from the delegate's arguments, because the engine's " + "in-place write through the aliased input already is that " + "output. Partial elision is not expressible by the runtime: it " + "takes the aliased outputs as elided only when the argument " + "count is short by the engine's whole aliased-output count, so a " + "delegate short of only some of them reads as not elided at all " + "and every execute() fails with an argument-count error. An " + "output is elided when export rewired a buffer mutation to write " + "it in place, so this engine mixes such a buffer with an aliased " + "input that is not one -- a plain input, say. Export this method " + "without zero_copy_kv, or keep every aliased input of this engine " + "a mutated buffer." + ) metadata = TensorRTBlobMetadata( io_bindings=io_bindings, diff --git a/py/torch_tensorrt/executorch/partitioner.py b/py/torch_tensorrt/executorch/partitioner.py index cc9dc0c16f8..7c0fd47e955 100644 --- a/py/torch_tensorrt/executorch/partitioner.py +++ b/py/torch_tensorrt/executorch/partitioner.py @@ -217,6 +217,11 @@ def _partition_elided_output_names( Any extraction failure returns an empty set: the delegate then carries every binding and a genuinely missing aliased output stays an error in the backend's ``_validate_output_binding_order``. + + A set that is neither empty nor the engine's whole ``aliased_io`` is the + right answer and still not a lowerable one, since the runtime reads + elision off a single argument count: ``TensorRTBackend.preprocess`` + refuses that engine. """ from torch_tensorrt.executorch._zero_copy import _aliased_inputs_by_output_index diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index c6b325f7979..5d9f38d89cf 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -1074,13 +1074,19 @@ def inner(gm): # -------------------------------------------------------------------------- -def _no_op_engine_node(graph, input_nodes, *, aliased_io, input_names, output_names): +def _no_op_engine_node( + graph, input_nodes, *, aliased_io, input_names, output_names, engine="" +): """A no_op_placeholder_for_execute_engine node with inlined engine info. Mirrors what replace_execute_engine() produces before partitioning: args are ``(input_list, *engine_info)`` with the binding names and aliased_io in their serialized wire form, so the partitioner's real per-engine resolution (_resolve_engine_info / _aliased_inputs_by_output_index) runs unmocked. + + ``engine`` is the serialized-plan slot, which the partitioner never reads. + Give it bytes to hand the same node to ``TensorRTBackend.preprocess``, which + does. """ from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ( ALIASED_IO_IDX, @@ -1094,7 +1100,7 @@ def _no_op_engine_node(graph, input_nodes, *, aliased_io, input_names, output_na from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import serialize_aliased_io info = [""] * SERIALIZATION_LEN - info[ENGINE_IDX] = "" # not read by the partitioner's elision resolution + info[ENGINE_IDX] = engine info[DEVICE_IDX] = "0" info[INPUT_BINDING_NAMES_IDX] = SERIALIZED_ENGINE_BINDING_DELIM.join(input_names) info[OUTPUT_BINDING_NAMES_IDX] = SERIALIZED_ENGINE_BINDING_DELIM.join(output_names) @@ -1302,9 +1308,9 @@ def test_unstage_raises_when_the_plain_delegate_is_wrongly_stamped(): @pytest.mark.unit -def test_partition_elides_only_the_outputs_aliased_onto_a_marked_buffer(): - """One engine, two aliased outputs, one marked buffer: only the marked one is - named elidable. +def test_a_mixed_alias_engine_derives_the_narrower_set_and_is_then_refused(): + """One engine, two aliased outputs, one marked buffer: the narrower elidable + set is right, and an engine that needs it cannot be lowered. An aliased output whose input is not a buffer export rewired -- a user alias, which nothing rewires and whose placeholder therefore carries no @@ -1312,8 +1318,18 @@ def test_partition_elides_only_the_outputs_aliased_onto_a_marked_buffer(): elidable set from the engine's aliased_io alone would exempt it too, and the backend would then accept a delegate that dropped a mutation nothing writes back. + + The derivation is right and the engine is still unusable, because the runtime + reads elision off one argument count and subtracts the engine's *whole* + aliased-output count. A .pte written for this shape loads and then fails + every ``execute()`` with an argument-count error, so ``preprocess`` refuses + it where the export can still be re-run. """ from torch_tensorrt.executorch._zero_copy import _aliased_inputs_by_output_index + from torch_tensorrt.executorch.backend import ( + TensorRTBackend, + _serialize_elided_output_names, + ) from torch_tensorrt.executorch.partitioner import TensorRTPartitioner graph = torch.fx.Graph() @@ -1328,7 +1344,10 @@ def test_partition_elides_only_the_outputs_aliased_onto_a_marked_buffer(): }, input_names=["k_in", "u_in"], output_names=["out_k", "out_u"], + engine=b"engine-bytes", ) + # out_k has already left the partition with the mutation that was rewired + # onto the buffer, so the delegate returns the one surviving binding. graph.output((engine,)) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) @@ -1348,7 +1367,19 @@ def test_partition_elides_only_the_outputs_aliased_onto_a_marked_buffer(): partitioner = TensorRTPartitioner( compile_specs=[CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b'["out_u"]')] ) - assert partitioner._partition_elided_output_names(program, partition) == {"out_k"} + elided = partitioner._partition_elided_output_names(program, partition) + assert elided == {"out_k"} + + with pytest.raises(ValueError, match="Partial elision is not expressible"): + TensorRTBackend.preprocess( + program, + [ + CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(elided), + ) + ], + ) # -------------------------------------------------------------------------- From 36b83e1cdd576e6ad6ccc8c7081df126b4895f99 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Mon, 7 Sep 2026 00:26:37 -0700 Subject: [PATCH 16/22] fix(executorch): close the four measured holes in the zero-copy guards The un-staging pass froze `enable_non_cpu_memory_planning` into itself at config-build time and read it on only one of its two branches. `ExecutorchBackendConfig` is a plain mutable dataclass, so a caller who turns the flag off on the config `zero_copy_backend_config()` returned -- the config `to_executorch` then uses -- got a program the pass accepted and the finalizer planned into the host arena. The pass now reads the flag off that config when it runs, and refuses a host-only planning mode once for the whole graph rather than on one branch, so the staged shape is covered too. `check_zero_copy_kv`, which the docs and `save()` tell people to rely on, unioned the arguments of every TensorRT delegate and never looked at placement. Measured at head: it accepts a program finalized with `ExecutorchBackendConfig(enable_non_cpu_memory_planning=False)` whose caches are direct delegate arguments planned in the single host arena -- the `.pte` whose first `execute()` fails on the alias-target guard. It now narrows that union to the delegates carrying the zero-copy spec, so an unrelated engine reading the same buffer cannot stand in for the one whose write was elided, and requires the buffer's `mem_id` to name one of the CUDA arenas memory planning recorded in `non_const_buffer_device`. In the same function, the completeness check accepted partial marker loss: one zero-copy delegate taking two staged buffers with only one still marked returned 1 and raised nothing, leaving the second cache wired through its `_h2d_copy`. It now compares against the number of aliased outputs the delegate's compile spec says it elided. A method mixing zero-copy caches with a copy-back buffer finalized with a crossed mutation map -- reproduced on the real pipeline as `{'copy__default': 'k_cache', 'b_k_cache': 'v_cache', 'b_v_cache': 'conv_state'}` against the unrewired control's correct pairing. Upstream's write-back pass inserts a copy only for a mutation whose value is not already its buffer placeholder, moves those copies to the front of the output tuple, and reassigns each mutation spec's argument by position; rewiring a cache to its own placeholder takes it out of that leading run. `order_copyback_mutations_first` puts the mutations that still get a copy first. It runs on the Edge program, not beside the rewiring: `to_edge_transform_and_lower` re-derives the signature, so an order set earlier is discarded (measured -- the pre-edge program comes out `conv_state, k_cache, v_cache` and the Edge program `k_cache, v_cache, conv_state`). The `.pte` was written correctly either way; what was wrong is the finalized signature and ExecuTorch's eager call path, which walks `buffers_to_mutate` writing the graph's leading results into the state dict in that order. The blob parser refused a repeated `aliased_io` entry but accepted a repeated binding name in `io_bindings`, which is the same hazard one step earlier: a TensorRT engine has one name space, so both slots resolve to one tensor, `init` records the alias on the first, and `execute` re-binds the same name for the later one. Confirmed on TensorRT 11.1.0.106 that a second `setTensorAddress` for one name replaces the first and the engine's write lands entirely in the second buffer, so the address replaced is the caller's cache. The parser now refuses the repeat. Tests: the duplicate-`aliased_io` test repeated the entry exactly, so a parser keyed on the output/input pair -- the regression that reopens the inflated-count bug -- passed the whole file; the same-output-different-input case is added alongside three for duplicate binding names. On the Python side the mixed test now asserts the finalized signature, which it did not before, and each new refusal has a test that fails when its clause is mutated away. One hazard the aliased-output elision looks like it introduces, and does not: eliding the aliased reflects does drop the term that used to force the end-of-execute sync, but the engine's other outputs are ordered by the caller-stream contract instead. Every host consumer ExecuTorch inserts for a device delegate output -- including the write-back of a copy-back buffer beside the caches -- is an `et_copy::_d2h_copy`, whose kernel issues its copy on `getCallerStream()`, the same stream the engine enqueued on, and then synchronizes it. With no caller stream that kernel is a blocking `cudaMemcpy` and `must_sync` is true anyway. Written down at the `must_sync` computation. --- .../executorch/TensorRTBackend.cpp | 14 +- .../executorch/TensorRTBlobHeader.cpp | 12 + .../runtime_performance/saving_models.rst | 4 +- py/torch_tensorrt/_compile.py | 10 +- py/torch_tensorrt/executorch/_export.py | 10 +- py/torch_tensorrt/executorch/_zero_copy.py | 402 +++++++++++++----- .../test_executorch_blob_header.cpp | 50 +++ tests/py/dynamo/executorch/test_export.py | 18 +- .../py/dynamo/executorch/test_zero_copy_kv.py | 277 ++++++++++-- 9 files changed, 654 insertions(+), 143 deletions(-) diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 5c4b0085a9a..9513c00f44c 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -442,7 +442,10 @@ Result TensorRTBackend::init( // Non-aliased models have an empty header.aliased_io -> all -1, unchanged path. // The parser has already refused a blob claiming one output twice, so the count // built below is one per distinct output binding, which is what execute() - // subtracts on. + // subtracts on. It has also refused a blob repeating a binding name, so the + // first-match scans below reach the only slot carrying that name -- otherwise + // the alias would be recorded on the first slot while execute() re-bound the + // same TensorRT name for the later one, replacing the caller's buffer address. handle->output_aliased_input_idx.assign(handle->num_outputs, -1); handle->input_is_alias_target.assign(handle->num_inputs, false); for (const auto& ab : header.aliased_io) { @@ -990,6 +993,15 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // contract every coalesced .pte already depends on. A host reader that // inspected it immediately after execute() returns would see stale data unless // it synchronized `stream` itself; ExecuTorch's KV path never does such a read. + // Eliding the reflects drops aliased_reflect_pending, so an execute() with a + // caller stream and no host staging no longer syncs at all. What orders the + // engine's *other* outputs is then that same contract: the host consumer + // ExecuTorch inserts for a device delegate output is an et_copy::_d2h_copy, + // whose kernel issues its copy on getCallerStream() -- this stream -- and then + // synchronizes it; a device consumer of that output is another delegate, + // enqueued on the same stream. That covers the write-back of a copy-back buffer + // sitting beside the zero-copy caches. With no caller stream set that kernel + // falls back to a blocking cudaMemcpy, and must_sync is true here anyway. const bool aliased_reflect_pending = !aliased_reflects.empty(); const bool must_sync = output_staged_to_host || input_staged_from_host || aliased_reflect_pending || !caller_stream_set; diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp index 1eaeae996f0..73d4d0bd81e 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp @@ -153,6 +153,15 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { return false; } + // A TensorRT engine has one name space for its tensors, so a name repeated + // across io_bindings -- in either list -- cannot be two bindings. It is + // accepted by every name lookup, which stops at the first match, and then + // contradicted by every address bind, which is keyed on the name and so + // overwrites whatever the earlier slot bound. For an aliased output the + // address overwritten is the caller's buffer, and the engine's in-place + // update lands somewhere else. Refuse the blob here, where the repeat is + // visible from the bytes alone. + std::unordered_set claimed_bindings; std::size_t pos = arr_start + 1; while (true) { pos = skip_ws(json, pos); @@ -226,6 +235,9 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { } if (saw_name && !name.empty()) { + if (!claimed_bindings.insert(name).second) { + return false; + } if (is_input) { out.input_binding_names.push_back(name); } else { diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index d59c8aefdcd..6eb3c45515e 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -401,8 +401,8 @@ one silently would break a runner built before this feature. the engine writes a per-call staging copy that is discarded, and the cache never updates. For a KV cache that is wrong output, not a crash. Pass the finalized program to ``torch_tensorrt.executorch.check_zero_copy_kv``, which - reads the graph back and refuses one whose caches are still staged, before - writing the ``.pte``:: + reads it back and refuses one whose caches are still staged, or planned + somewhere the engine cannot write them, before writing the ``.pte``:: program = edge.to_executorch(zero_copy_backend_config(backend_config)) torch_tensorrt.executorch.check_zero_copy_kv(program) diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index 0a66b04edbb..d62ca480970 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -1473,11 +1473,11 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None backend_config = zero_copy_backend_config(backend_config) executorch_program = edge_program.to_executorch(config=backend_config) if zero_copy_kv: - # save() holds the finalized program here, which is the only place the - # graph shows whether the caches actually reach the engine un-staged. - # Both halves of zero-copy no-op quietly when they find nothing to do, so - # without this a save() that asked for zero-copy could still write an - # ordinary staged .pte. + # save() holds the finalized program here, which is the only place that + # shows whether the caches actually reach the engine un-staged and in an + # arena it can write. Both halves of zero-copy no-op quietly when they + # find nothing to do, so without this a save() that asked for zero-copy + # could still write an ordinary staged .pte. check_zero_copy_kv(executorch_program) with open(file_path, "wb") as f: executorch_program.write_to_file(f) diff --git a/py/torch_tensorrt/executorch/_export.py b/py/torch_tensorrt/executorch/_export.py index 04c4a64e703..b3fa0ca0e19 100644 --- a/py/torch_tensorrt/executorch/_export.py +++ b/py/torch_tensorrt/executorch/_export.py @@ -562,6 +562,7 @@ def export( stage_exported_program, validate_engine_program, ) + from torch_tensorrt.executorch._zero_copy import order_copyback_mutations_first from torch_tensorrt.executorch.backend import ( ZERO_COPY_KV_COMPILE_SPEC_KEY, _serialize_elided_output_names, @@ -727,7 +728,7 @@ def export( edge_programs = rewritten["forward"] partitioner_pipeline = method_partitioners["forward"] - return to_edge_transform_and_lower( + edge_manager = to_edge_transform_and_lower( edge_programs, transform_passes=transform_passes, partitioner=partitioner_pipeline, @@ -737,3 +738,10 @@ def export( ), generate_etrecord=generate_etrecord, ) + if zero_copy_methods: + # After lowering, not beside the rewiring: to_edge re-derives the graph + # signature, so an order set earlier does not reach the finalizer. See + # order_copyback_mutations_first. + for name in edge_manager.methods: + order_copyback_mutations_first(edge_manager.exported_program(name)) + return edge_manager diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index a19a1489f1c..9a5b72f6743 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -21,7 +21,8 @@ * :func:`rewire_aliased_mutations_to_buffers`, on the exported program before partitioning, drops the copy-back by declaring that the buffer *is* the mutation's result. The aliased output then has no user and disappears from the - partition. + partition. :func:`order_copyback_mutations_first` then repairs, on the Edge + program, the one thing that declaration disturbs downstream. * :func:`unstage_aliased_buffers_pass`, as a ``to_out_var_pass``, drops the staging so the engine writes the caller's buffer rather than a copy. @@ -30,10 +31,11 @@ pass is public on its own: the rewiring is reached only through ``export(..., zero_copy_kv=True)``, and the un-staging only through :func:`zero_copy_backend_config`. That, plus :func:`check_zero_copy_kv` -- which -reads a finalized program back and refuses one where the pairing did not happen --- is what this module exports. +reads a finalized program back and refuses one in which the engine does not end +up writing the caller's buffer -- is what this module exports. """ +import json import logging import operator from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Set @@ -310,6 +312,81 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> List[str]: return elided_output_names +def order_copyback_mutations_first(exported_program: Any) -> int: + """Reorder one Edge method's mutations so ExecuTorch pairs them up correctly. + + ExecuTorch finalizes a mutation by inserting a ``copy_`` for it, but only + when its value is not already the buffer placeholder. It moves those copies + to the front of the output tuple, leaves everything else behind them in + order, and then walks the mutation specs reassigning each one's argument + *by position* over the result (``insert_write_back_for_buffers_pass``). + :func:`rewire_aliased_mutations_to_buffers` is what makes a mutation's value + its own placeholder, so a rewired cache gets no copy and drops out of that + leading run -- and in a method that also has a copy-back buffer, every + mutation spec from the first rewired one on then comes out of finalization + naming a different buffer's value. The specs ahead of it are unaffected, + since the copies keep their order among themselves. The ``.pte`` is written + correctly either way, because the emitter and the memory planner read only + which buffers are mutated and not what by. What the pairing decides is the + finalized signature, which anyone inspecting the program reads, and + ExecuTorch's eager call path, which walks + ``buffers_to_mutate`` writing the graph's leading results into the state dict + in that order and so updates each buffer from another one's value. + + Putting the mutations that still get a copy first restores the + correspondence. Only the mutations this module rewired are moved: a mutation + already bound to its own placeholder before zero-copy ran would need moving + too, but that shape is not one this feature creates. + + This runs on the *Edge* program rather than beside the rewiring, because + ``to_edge_transform_and_lower`` re-derives the whole graph signature -- the + order it hands back is the buffers' own order, whatever order it was given. + Nothing between here and the write-back pass re-derives it again. + + Returns zero when the order already holds, and otherwise the number of + rewired mutations, all of which the reorder leaves behind the copy-back + ones. + """ + from torch.export.graph_signature import ExportGraphSignature, OutputKind + + signature = exported_program.graph_signature + specs = list(signature.output_specs) + output_node = exported_program.graph_module.graph.output_node() + args = list(output_node.args[0]) + slots = [ + index + for index, spec in enumerate(specs) + if spec.kind in (OutputKind.BUFFER_MUTATION, OutputKind.USER_INPUT_MUTATION) + and index < len(args) + ] + rewired = { + index + for index in slots + if isinstance(args[index], Node) + and args[index].op == "placeholder" + and args[index].meta.get("_torch_tensorrt_aliased_buffer") + } + source = [index for index in slots if index not in rewired] + [ + index for index in slots if index in rewired + ] + if source == slots: + return 0 + + new_args, new_specs = list(args), list(specs) + for slot, index in zip(slots, source): + new_args[slot] = args[index] + new_specs[slot] = specs[index] + output_node.args = (tuple(new_args),) + exported_program.graph_module.recompile() + exported_program._graph_signature = ExportGraphSignature( + input_specs=list(signature.input_specs), output_specs=new_specs + ) + _LOGGER.debug( + "moved %d rewired mutation(s) behind the copy-back ones", len(rewired) + ) + return len(rewired) + + def _is_tensorrt_delegate(graph_module: torch.fx.GraphModule, node: Node) -> bool: """True when ``node`` is a call_delegate dispatching to the TensorRT backend. @@ -328,6 +405,46 @@ def _is_tensorrt_delegate(graph_module: torch.fx.GraphModule, node: Node) -> boo return bool(getattr(module, "backend_id", None) == TensorRTBackend.__name__) +def _zero_copy_compile_spec(graph_module: torch.fx.GraphModule, node: Node) -> Any: + """One delegate's zero-copy KV compile spec, or ``None`` if it carries none.""" + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + lowered = node.args[0] if node.args else None + if not isinstance(lowered, Node) or lowered.op != "get_attr": + return None + module = getattr(graph_module, lowered.target, None) + for spec in getattr(module, "compile_specs", None) or []: + if getattr(spec, "key", None) == ZERO_COPY_KV_COMPILE_SPEC_KEY: + return spec + return None + + +def _delegate_elided_output_names( + graph_module: torch.fx.GraphModule, node: Node +) -> Set[str]: + """The aliased output binding names one delegate's zero-copy spec claims. + + ``TensorRTPartitioner`` writes one name per aliased output it elided on that + delegate's own engine, and an aliased output is elided exactly when a buffer + mutation writes it in place, so the size of this set is how many marked + buffers the delegate has to take. Empty when the delegate carries no such + spec, and also when it carries one whose value does not decode into a list of + names -- a spec built by hand rather than by the partitioner, which always + writes the JSON list. Callers read empty as "cannot tell", not as "none". + """ + spec = _zero_copy_compile_spec(graph_module, node) + if spec is None: + return set() + value = getattr(spec, "value", None) + if not isinstance(value, (str, bytes, bytearray)): + return set() + try: + names = json.loads(value) + except ValueError: + return set() + return {str(name) for name in names} if isinstance(names, list) else set() + + def _delegate_declares_zero_copy( graph_module: torch.fx.GraphModule, node: Node ) -> bool: @@ -340,20 +457,12 @@ def _delegate_declares_zero_copy( several TensorRT delegates therefore marks only the KV one, never the plain compute engines beside it -- which is what keeps this cross-check from demanding an aliased buffer from a delegate that never had one. A delegate - that declares it but ends up taking no marked buffer is a lost KV update -- - the mark that would have driven the un-staging did not survive to this pass - -- and is caught in :func:`_unstage_aliased_buffers`. + that declares it but ends up taking fewer marked buffers than it elided + aliased outputs has lost a KV update -- the mark that would have driven the + un-staging did not survive to this pass -- and is caught in + :func:`_unstage_aliased_buffers`. """ - from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY - - lowered = node.args[0] if node.args else None - if not isinstance(lowered, Node) or lowered.op != "get_attr": - return False - module = getattr(graph_module, lowered.target, None) - return any( - getattr(spec, "key", None) == ZERO_COPY_KV_COMPILE_SPEC_KEY - for spec in (getattr(module, "compile_specs", None) or []) - ) + return _zero_copy_compile_spec(graph_module, node) is not None def _device_move_is_safe( @@ -438,24 +547,27 @@ def _unstage_aliased_buffers( accepted. Two things decide where the buffer is planned, and the second is not visible in the graph: the spec's own device, and whether memory planning reads spec devices at all. ``enable_non_cpu_memory_planning=False`` on the - ``ExecutorchBackendConfig`` produces exactly the shape above -- no staging - copy is inserted and ``PropagateDevicePass`` writes the delegate's device - straight onto the placeholder's spec -- and then plans every tensor into the - one host arena regardless, so the engine is handed a host pointer for a - buffer it must write on the device. ``device_memory_planning`` carries that - configuration in; :func:`zero_copy_backend_config` takes it off the - ``ExecutorchBackendConfig`` it is building from. + ``ExecutorchBackendConfig`` plans every tensor into the one host arena + whatever its spec says, so the engine is handed a host pointer for a buffer + it must write on the device -- and it does that to every marked buffer, the + ones this pass un-stages as much as the ones that already reach their + delegate directly. That is why it is checked once for the whole graph, up + front, rather than on one of the two branches below. + ``device_memory_planning`` carries the configuration in; + :func:`unstage_aliased_buffers_pass` reads it off the finalization config + when it runs. A failure here is a lost KV update, so it is raised rather than logged: export has already removed the copy-back, so a marked buffer left staged has the engine write per-call scratch that is then discarded and the buffer never - updates. It raises when the staging copy has no spec or is not on CUDA, when - a direct argument would not be planned in device memory, when the device move - is unsafe, and -- so a discovery miss cannot pass silently -- after the loop - when the post-condition does not hold for some marked buffer, cross-checked - against each delegate's own ``zero_copy_kv`` spec: a TensorRT delegate that - declares zero-copy but takes no marked buffer is broken, and that refusal - names the delegate and lists what it does take. + updates. It raises when memory planning is host-only, when the staging copy + has no spec or is not on CUDA, when a direct argument has no spec of its own + or that spec is not on CUDA, when the device move is unsafe, and -- so a + discovery miss cannot pass silently -- after the loop when the post-condition + does not hold for some marked buffer, cross-checked against each delegate's + own ``zero_copy_kv`` spec: a TensorRT delegate that takes fewer marked + buffers than the aliased outputs that spec says it elided is broken, and that + refusal names the delegate and lists what it does take. Returns the number of delegate inputs un-staged, which is zero for a program that already satisfied the post-condition. @@ -463,6 +575,25 @@ def _unstage_aliased_buffers( from executorch.exir.schema import DeviceType from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + marked_placeholders = [ + node + for node in graph_module.graph.nodes + if node.op == "placeholder" and node.meta.get("_torch_tensorrt_aliased_buffer") + ] + if marked_placeholders and not device_memory_planning: + names = ", ".join(repr(node.name) for node in marked_placeholders) + raise RuntimeError( + "TensorRT zero-copy KV: buffer(s) " + f"{names} are marked for in-place update, but memory planning is " + "configured with enable_non_cpu_memory_planning=False, which puts " + "every tensor in the one host arena whatever its TensorSpec says. " + "The engine would be handed a host pointer it cannot write, and " + "export has already removed the copy-back, so the update would be " + "lost. Finalize over a configuration that leaves " + "enable_non_cpu_memory_planning on, so a CUDA TensorSpec is given a " + "CUDA arena." + ) + h2d_copy = torch.ops.et_copy._h2d_copy.default unstaged = 0 satisfied_placeholders: Set[Node] = set() @@ -500,18 +631,6 @@ def _unstage_aliased_buffers( "Give the buffer to a TensorRT delegate on CUDA, or " "export this method without zero_copy_kv." ) - elif not device_memory_planning: - placement = ( - "memory planning is configured with " - "enable_non_cpu_memory_planning=False, which puts " - "every tensor in the one host arena whatever its " - "TensorSpec says" - ) - remedy = ( - "Finalize over a configuration that leaves " - "enable_non_cpu_memory_planning on, so a CUDA " - "TensorSpec is given a CUDA arena." - ) if placement: raise RuntimeError( "TensorRT zero-copy KV: buffer " @@ -587,11 +706,7 @@ def _unstage_aliased_buffers( node.args = tuple(new_args) marked_but_unsatisfied = [ - node - for node in graph_module.graph.nodes - if node.op == "placeholder" - and node.meta.get("_torch_tensorrt_aliased_buffer") - and node not in satisfied_placeholders + node for node in marked_placeholders if node not in satisfied_placeholders ] if marked_but_unsatisfied: names = ", ".join(repr(node.name) for node in marked_but_unsatisfied) @@ -605,18 +720,34 @@ def _unstage_aliased_buffers( "keep the aliased buffer on a TensorRT delegate." ) for delegate in zero_copy_delegates: - if satisfied_per_delegate[delegate] == 0: - delegate_inputs = [ - arg.name for arg in delegate.args[1:] if isinstance(arg, Node) - ] - raise RuntimeError( - "TensorRT zero-copy KV: delegate " - f"'{delegate.name}' declares zero-copy KV " - f"(compile spec '{ZERO_COPY_KV_COMPILE_SPEC_KEY}') but takes no " - f"buffer marked for in-place update (inputs: {delegate_inputs}). " - "Export elided its aliased outputs, so the engine now writes " - "per-call scratch that is discarded and the cache never updates." + # One marked buffer per aliased output the spec says this delegate + # elided. Demanding only one would accept a delegate that lost all but + # one of its marks, whose remaining caches are still wired through their + # staging copies. A spec whose value does not decode names cannot say how + # many to expect, so it falls back to demanding at least one. + elided = _delegate_elided_output_names(graph_module, delegate) + satisfied = satisfied_per_delegate[delegate] + if satisfied >= max(len(elided), 1): + continue + delegate_inputs = [ + arg.name for arg in delegate.args[1:] if isinstance(arg, Node) + ] + expectation = ( + "takes no buffer marked for in-place update" + if not elided + else ( + f"elided the aliased output(s) {sorted(elided)} but takes only " + f"{satisfied} of the {len(elided)} buffers that implies" ) + ) + raise RuntimeError( + "TensorRT zero-copy KV: delegate " + f"'{delegate.name}' declares zero-copy KV " + f"(compile spec '{ZERO_COPY_KV_COMPILE_SPEC_KEY}') but " + f"{expectation} (inputs: {delegate_inputs}). " + "Export elided its aliased outputs, so the engine now writes " + "per-call scratch that is discarded and those caches never update." + ) if unstaged: # Erase only the stagings we orphaned. A graph-wide eliminate_dead_code() @@ -646,6 +777,15 @@ def unstage_aliased_buffers_pass( program will be finalized with. Nothing in the graph records it, and it is half of what decides whether a marked buffer ends up somewhere the engine can write, so the pass has to be told: see :func:`_unstage_aliased_buffers`. + It is only the fallback. Set ``finalization_config`` on the returned pass to + the ``ExecutorchBackendConfig`` the program will be finalized with and the + flag is read off that config when the pass runs, which is the value memory + planning will use a few passes later. + ``ExecutorchBackendConfig`` is a plain mutable dataclass, so the field can be + set again on the very config being finalized after the pass has captured its + value; a pass still reading the captured copy would then accept a program the + finalizer plans into the host arena. :func:`zero_copy_backend_config` sets + the attribute for that reason. """ from executorch.exir import ExecutorchBackendConfig from executorch.exir.pass_base import PassBase @@ -657,9 +797,17 @@ def unstage_aliased_buffers_pass( ) class _UnstageThenToOutVar(PassBase): # type: ignore[misc] + finalization_config: Optional["ExecutorchBackendConfig"] = None + def call(self, graph_module: torch.fx.GraphModule) -> Any: + config = self.finalization_config + planning = ( + device_memory_planning + if config is None + else config.enable_non_cpu_memory_planning + ) unstaged = _unstage_aliased_buffers( - graph_module, device_memory_planning=device_memory_planning + graph_module, device_memory_planning=planning ) _LOGGER.debug("un-staged %d aliased delegate buffer(s)", unstaged) return inner(graph_module) @@ -667,6 +815,32 @@ def call(self, graph_module: torch.fx.GraphModule) -> Any: return _UnstageThenToOutVar() +def _device_planned_arenas(graph_module: torch.fx.GraphModule) -> Set[int]: + """The ``mem_id``s of the finalized program's CUDA arenas. + + Memory planning partitions the specs by device, gives each device its own + arena, and records the non-CPU ones on the graph module as + ``non_const_buffer_device``. The key is absent from a program planned + entirely on the host, which is what an ``enable_non_cpu_memory_planning=False`` + finalization produces however the specs are marked. + """ + from executorch.exir.schema import DeviceType + + return { + entry.buffer_idx + for entry in (graph_module.meta.get("non_const_buffer_device") or []) + if getattr(entry, "device_type", None) == DeviceType.CUDA + } + + +def _name_detail(names_by_method: Dict[str, List[str]]) -> str: + return ", ".join( + f"'{name}' in method '{method}'" + for method, names in names_by_method.items() + for name in names + ) + + def check_zero_copy_kv(program: Any) -> None: """Raise unless a finalized program really updates its KV buffers in place. @@ -678,38 +852,51 @@ def check_zero_copy_kv(program: Any) -> None: and stages its cache like any other, which for a KV cache is wrong output rather than a crash. - Two shapes are refused: a marked buffer that is not a direct argument of any - *TensorRT* delegate -- it still reaches one through an ``_h2d_copy`` staging, - or reaches none -- and a program with no marked buffer in any method. The - first is what finalizing without :func:`zero_copy_backend_config` leaves - behind; when that config *is* installed, ``_unstage_aliased_buffers`` has - already raised on the same condition. - - Only a TensorRT delegate counts, matching what that pass un-stages. The mark - is put on a buffer because a TensorRT engine writes it in place, so another - backend's delegate taking the buffer directly says nothing about whether the - engine did: it can be handed the buffer while the engine beside it still - reads a staging copy whose contents are discarded. + Three shapes are refused: a marked buffer that is not a direct argument of a + TensorRT delegate carrying the zero-copy compile spec, a marked buffer that + reaches such a delegate directly but is not planned in device memory, and a + program with no marked buffer in any method. The first two are what + finalizing without :func:`zero_copy_backend_config` leaves behind. That + config's pass refuses both earlier, off the configuration and the graph; this + reads the placement memory planning went on to choose. + + The spec is what narrows the delegates that count. The mark is put on a + buffer because one TensorRT engine writes it in place, and only a delegate + whose own engine elided an aliased output is stamped, so another backend's + delegate taking the buffer says nothing about whether the engine did, and + neither does an unrelated TensorRT engine that happens to read it. Either + would stand in for the delegate whose write was elided while that one still + reads a staging copy whose contents are discarded. Stamped delegates are read + as one set rather than matched to the buffer each of them elided, so in a + method holding more than one they can still stand in for each other here; + :func:`_unstage_aliased_buffers`, which counts each delegate's own buffers + against its own spec, is what separates them. + + Placement is read from where memory planning actually put the buffer: its + ``TensorSpec``'s ``mem_id`` has to name one of the arenas the finalized + program records as CUDA in ``non_const_buffer_device``. The spec's own device + does not settle it -- ``PropagateDevicePass`` writes CUDA onto the spec of a + buffer that reaches a CUDA delegate directly, and + ``enable_non_cpu_memory_planning=False`` then plans that same buffer into the + one host arena, which is the shape whose every ``execute()`` fails on the + runtime's alias-target guard. Every method is read, not only ``forward``. ``export()`` rewires each method on its own, so a check that stopped at ``forward`` would pass a program whose decode had degenerated to staged -- and on the prefill/decode pair the user guide's zero-copy example exports it would not get that far, since a - multi-method program need not have a ``forward`` at all. The second refusal is + multi-method program need not have a ``forward`` at all. The last refusal is about the program rather than about one method, matching the warning ``export()`` emits: a method with no aliased buffer mutation of its own is not an error, so a model that rewires only its decode step is accepted. - This reads the graph, so it says what the program does rather than what the - passes recorded. What it establishes is the wiring: the caller's buffer is - what each TensorRT delegate is handed. It says nothing about whether the - engine's write is correct, nor about where the buffer is planned -- the pass - :func:`zero_copy_backend_config` installs owns that half, and a program whose - buffer reaches its delegate directly out of a host arena passes here and then - fails every ``execute()`` on the runtime's alias-target guard. + This reads the graph and the finalized specs, so it says what the program + does rather than what the passes recorded. It still says nothing about + whether the engine's write itself is correct. """ method_names = sorted(program.methods) staged_by_method: Dict[str, List[str]] = {} + host_planned_by_method: Dict[str, List[str]] = {} marked_anywhere = False for method_name in method_names: graph_module = program.exported_program(method_name).graph_module @@ -722,15 +909,25 @@ def check_zero_copy_kv(program: Any) -> None: if not marked: continue marked_anywhere = True - delegate_args = { + zero_copy_delegate_args = { arg for node in graph_module.graph.nodes if _is_tensorrt_delegate(graph_module, node) + and _delegate_declares_zero_copy(graph_module, node) for arg in node.args[1:] } - staged = [node.name for node in marked if node not in delegate_args] + staged = [node.name for node in marked if node not in zero_copy_delegate_args] if staged: staged_by_method[method_name] = staged + device_arenas = _device_planned_arenas(graph_module) + host_planned = [ + node.name + for node in marked + if node in zero_copy_delegate_args + and getattr(node.meta.get("spec"), "mem_id", None) not in device_arenas + ] + if host_planned: + host_planned_by_method[method_name] = host_planned if not marked_anywhere: raise RuntimeError( "TensorRT zero-copy KV: no buffer in this program is marked for " @@ -740,17 +937,23 @@ def check_zero_copy_kv(program: Any) -> None: "mutation was found -- export logs a warning for that case." ) if staged_by_method: - detail = ", ".join( - f"'{name}' in method '{method}'" - for method, names in staged_by_method.items() - for name in names + raise RuntimeError( + f"TensorRT zero-copy KV: buffer(s) {_name_detail(staged_by_method)} " + "are marked for in-place update but do not reach the TensorRT " + "delegate that elided them directly, so the engine writes a staging " + "copy that is discarded and the cache never updates. Export removed " + "their copy-back, so nothing else would restore it. Finalize with " + "torch_tensorrt.executorch.zero_copy_backend_config()." ) + if host_planned_by_method: raise RuntimeError( - f"TensorRT zero-copy KV: buffer(s) {detail} are marked for in-place " - "update but do not reach a TensorRT delegate directly, so the engine " - "writes a staging copy that is discarded and the cache never updates. Export " - "removed their copy-back, so nothing else would restore it. Finalize " - "with torch_tensorrt.executorch.zero_copy_backend_config()." + f"TensorRT zero-copy KV: buffer(s) " + f"{_name_detail(host_planned_by_method)} reach their TensorRT " + "delegate directly but are not planned in any of the program's CUDA " + "arenas, so the engine is handed a host pointer it cannot write and " + "every execute() fails on the runtime's alias-target guard. Finalize " + "with torch_tensorrt.executorch.zero_copy_backend_config() over a " + "configuration that leaves enable_non_cpu_memory_planning on." ) @@ -772,10 +975,13 @@ def zero_copy_backend_config( ``config`` is your own configuration -- every field is preserved, and a ``to_out_var_pass`` you already set runs after the un-staging. Omit it to start from ExecuTorch's defaults. One field is not merely carried but read: - zero-copy needs the caches planned in device memory, so a config with + zero-copy needs the caches planned in device memory, so ``enable_non_cpu_memory_planning=False`` -- which plans every tensor into the one host arena -- has the pass refuse each cache it finds rather than write a - ``.pte`` whose every ``execute()`` fails. + ``.pte`` whose every ``execute()`` fails. It is read off the config returned + here, at the moment the pass runs, so setting the field on that config + afterwards is honoured: the pass and the finalizer then cannot disagree about + it. .. warning:: Finalizing a ``zero_copy_kv=True`` program *without* this config does @@ -798,10 +1004,10 @@ def zero_copy_backend_config( from executorch.exir import ExecutorchBackendConfig base = config if config is not None else ExecutorchBackendConfig() - return replace( - base, - to_out_var_pass=unstage_aliased_buffers_pass( - base.to_out_var_pass, - device_memory_planning=base.enable_non_cpu_memory_planning, - ), + unstage = unstage_aliased_buffers_pass( + base.to_out_var_pass, + device_memory_planning=base.enable_non_cpu_memory_planning, ) + wrapped = replace(base, to_out_var_pass=unstage) + unstage.finalization_config = wrapped + return wrapped diff --git a/tests/cpp/executorch/test_executorch_blob_header.cpp b/tests/cpp/executorch/test_executorch_blob_header.cpp index 90782ff4a3a..1d7935e40ec 100644 --- a/tests/cpp/executorch/test_executorch_blob_header.cpp +++ b/tests/cpp/executorch/test_executorch_blob_header.cpp @@ -192,6 +192,56 @@ TEST(ExecuTorchTensorRTBlobHeader, RejectsRepeatedAliasedIoOutput) { EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); } +TEST(ExecuTorchTensorRTBlobHeader, RejectsRepeatedAliasedIoOutputWithADifferentInput) { + // The same repeat with the second entry naming a different input. This is the + // shape the check above exists for: keying the refusal on the output/input + // pair instead of the output alone would accept it, and out_k's aliased-output + // count would be two for one output binding, which is the arity the backend + // subtracts on. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"out_v","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"},)" + R"({"output":"out_k","input":"in_v","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsRepeatedOutputBindingName) { + // One name, two output slots. Every name lookup stops at the first slot, so + // init would record the alias there, and execute() would then bind the second + // slot's ExecuTorch storage to the same TensorRT name -- replacing the address + // of the caller's buffer that the alias exists to write. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsRepeatedInputBindingName) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsOneNameUsedAsBothAnInputAndAnOutput) { + // TensorRT has one name space for its tensors, so this is the same collision + // as the two above rather than a distinct input and output that happen to + // share a spelling. + const std::string metadata = R"({"io_bindings":[{"name":"kv","is_input":true},{"name":"kv","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + TEST(ExecuTorchTensorRTBlobHeader, ParsesAliasedIoEntriesForDistinctOutputs) { // The minimal pair for the test above: the same blob with the second entry // claiming its own output. A second entry is not itself the defect. diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py index f10c12ea243..219e435f5bb 100644 --- a/tests/py/dynamo/executorch/test_export.py +++ b/tests/py/dynamo/executorch/test_export.py @@ -63,6 +63,22 @@ def __init__(self): self.graph_signature = SimpleNamespace(inputs_to_buffers={}, output_specs=[]) +class FakeEdgeProgramManager: + """The ``EdgeProgramManager`` stand-in ``to_edge_transform_and_lower`` returns. + + ``export()`` reads back the methods of the manager it is handed, to put each + zero-copy method's mutations in the order ExecuTorch finalizes them in. These + programs declare no mutation, so that call finds nothing to reorder. + """ + + def __init__(self): + self._programs = {"forward": FakeExportedProgram()} + self.methods = set(self._programs) + + def exported_program(self, method_name="forward"): + return self._programs[method_name] + + class FakeTensorRTPartitioner: def __init__(self, compile_specs): self.compile_specs = compile_specs @@ -241,7 +257,7 @@ def _patch_lowering(monkeypatch, engine_counts=None): ) export_module = importlib.import_module("torch_tensorrt.executorch._export") engine_counts = engine_counts or {} - lower = MagicMock(return_value=object()) + lower = MagicMock(return_value=FakeEdgeProgramManager()) monkeypatch.setattr(executorch.exir, "to_edge_transform_and_lower", lower) monkeypatch.setattr(executorch_api, "TensorRTPartitioner", FakeTensorRTPartitioner) monkeypatch.setattr(executorch_api, "get_edge_compile_config", lambda: "default") diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index 5d9f38d89cf..e84c750621c 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -14,6 +14,7 @@ is never un-staged is caught rather than dropped. """ +import json import operator from types import SimpleNamespace @@ -367,6 +368,20 @@ def test_rewire_rejects_an_engine_whose_only_other_output_is_dead( Z.rewire_aliased_mutations_to_buffers(program) +def _zero_copy_specs(*names): + """What ``TensorRTPartitioner`` stamps on the delegate whose engine elided. + + The value is the list of aliased output binding names, one per buffer the + engine writes in place, which is what the count checks read. + """ + return [ + CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + json.dumps(list(names or ("out_k",))).encode(), + ) + ] + + def _staged_delegate_graph( *, backend_id="TensorRTBackend", device=DeviceType.CUDA, compile_specs=None ): @@ -499,24 +514,30 @@ def test_unstage_refuses_a_direct_buffer_whose_spec_stays_on_the_host(): @pytest.mark.unit -def test_unstage_refuses_a_direct_buffer_under_host_only_memory_planning(): +@pytest.mark.parametrize("shape", ["direct", "staged"]) +def test_unstage_refuses_a_marked_buffer_under_host_only_memory_planning(shape): """A CUDA spec does not mean CUDA memory when planning ignores spec devices. - ``enable_non_cpu_memory_planning=False`` leaves exactly the graph above -- - ``PropagateDevicePass`` writes the delegate's device onto the placeholder's - spec and inserts no staging copy -- and then plans every tensor into the one - host arena regardless. Nothing in the graph distinguishes it from the - already-un-staged shape, which is why the spec device is asserted here to be - CUDA: the refusal can only come from the planning mode the pass was told. + ``enable_non_cpu_memory_planning=False`` plans every tensor into the one host + arena whatever its ``TensorSpec`` says, so no marked buffer can be written in + place under it. Both graph shapes are covered because the pass reaches them + by different branches -- the buffer that already is a delegate argument, and + the one whose staging copy the pass removes -- and the refusal belongs to + neither, so it is checked once for the whole graph. """ - graph_module, k_buffer, _ = _direct_delegate_graph() - assert k_buffer.meta["spec"].device == DeviceType.CUDA, ( - "this test only discriminates while the buffer's spec is CUDA; on a host " - "spec the refusal below could come from the spec-device check instead of " - "from the planning mode" - ) + if shape == "direct": + graph_module, k_buffer, _ = _direct_delegate_graph() + assert k_buffer.meta["spec"].device == DeviceType.CUDA, ( + "this shape only stands for the planning-mode hazard while the " + "buffer's spec is CUDA: what the refusal below exists for is a CUDA " + "spec that still lands in the host arena, and a host spec is refused " + "under any planning mode" + ) + else: + graph_module, k_buffer, _, _ = _staged_delegate_graph() + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - with pytest.raises(RuntimeError, match="not planned in device memory"): + with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): Z._unstage_aliased_buffers(graph_module, device_memory_planning=False) @@ -536,10 +557,37 @@ def test_zero_copy_backend_config_carries_the_planning_mode_into_the_pass(): ) graph_module, _, _ = _direct_delegate_graph() - with pytest.raises(RuntimeError, match="not planned in device memory"): + with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): config.to_out_var_pass(graph_module) +@pytest.mark.unit +def test_zero_copy_backend_config_reads_the_planning_mode_when_the_pass_runs(): + """The flag has to be read where it is in effect, not captured when it is set. + + ``ExecutorchBackendConfig`` is a plain mutable dataclass and the config + returned here is the one the caller hands to ``to_executorch``, so the value + memory planning uses is whatever the field holds by then. A pass that froze + the field when the config was built disagrees with the finalizer in both + directions, and the first of those writes a ``.pte`` whose caches are planned + in the host arena and whose every ``execute()`` fails. + """ + from executorch.exir import ExecutorchBackendConfig + + turned_off = Z.zero_copy_backend_config() + turned_off.enable_non_cpu_memory_planning = False + graph_module, _, _ = _direct_delegate_graph() + with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): + turned_off.to_out_var_pass(graph_module) + + turned_on = Z.zero_copy_backend_config( + ExecutorchBackendConfig(enable_non_cpu_memory_planning=False) + ) + turned_on.enable_non_cpu_memory_planning = True + graph_module, _, _ = _direct_delegate_graph() + turned_on.to_out_var_pass(graph_module) + + @pytest.mark.unit def test_unstage_runs_a_second_time_without_raising(): """Installing the pass twice is redundant rather than an error. @@ -602,6 +650,46 @@ def test_unstage_raises_when_a_zero_copy_delegate_unstaged_nothing(): Z._unstage_aliased_buffers(graph_module) +@pytest.mark.unit +def test_unstage_raises_when_a_zero_copy_delegate_unstaged_only_some(): + """One surviving mark must not stand in for the ones that were lost. + + The delegate's spec names both aliased outputs it elided, so it has to take + two buffers written in place. Only one still carries the mark, and demanding + merely one would un-stage that one, raise nothing, and leave the second cache + wired through the staging copy whose contents are discarded -- with its + copy-back already gone, a silently frozen cache. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + v_buffer = graph.placeholder("b_v_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + staged_v = graph.call_function(h2d, (v_buffer,)) + lowered = graph.get_attr("lowered_module_0") + delegate = graph.call_function( + executorch_call_delegate, (lowered, staged_k, staged_v) + ) + graph.output((k_buffer, v_buffer, delegate)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs("out_k", "out_v") + ) + graph_module = torch.fx.GraphModule(root, graph) + for node, dev in ( + (k_buffer, DeviceType.CPU), + (v_buffer, DeviceType.CPU), + (staged_k, DeviceType.CUDA), + (staged_v, DeviceType.CUDA), + ): + node.meta["spec"] = SimpleNamespace(device=dev, device_index=0) + # v_buffer's mark did not survive lowering; k_buffer's did. + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="takes only 1 of the 2 buffers"): + Z._unstage_aliased_buffers(graph_module) + + @pytest.mark.unit def test_unstage_refuses_a_buffer_another_backend_stages_on_the_same_gpu(): """A second backend's staging copy on the *same* GPU is refused, not kept. @@ -902,21 +990,44 @@ def _finalized_program(forward=None, **methods): ) +CUDA_ARENA = 2 + + +def _planned(graph_module, *, arena=CUDA_ARENA, on_device=True): + """Add what memory planning leaves behind, which the checker reads. + + The graphs above are built for the passes that run before planning, so they + carry no ``mem_id`` and the module records no arena devices. Planning assigns + both, and ``on_device=False`` is what it leaves when it puts everything in + the one host arena -- it records no device arenas at all then. + """ + from executorch.exir.schema import NonConstBufferDevice + + for node in graph_module.graph.nodes: + if node.op == "placeholder" and node.meta.get("spec") is not None: + node.meta["spec"].mem_id = arena + if on_device: + graph_module.meta["non_const_buffer_device"] = [ + NonConstBufferDevice( + buffer_idx=arena, device_type=DeviceType.CUDA, device_index=0 + ) + ] + return graph_module + + def _unstaged_graph(): """A graph whose marked buffer already reaches its delegate directly.""" - graph_module, k_buffer, _, _ = _staged_delegate_graph() + graph_module, k_buffer, _, _ = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True Z._unstage_aliased_buffers(graph_module) - return graph_module + return _planned(graph_module) @pytest.mark.unit def test_check_zero_copy_kv_accepts_an_unstaged_buffer(): - graph_module, k_buffer, _, delegate = _staged_delegate_graph() - k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - Z._unstage_aliased_buffers(graph_module) - - Z.check_zero_copy_kv(_finalized_program(graph_module)) + Z.check_zero_copy_kv(_finalized_program(_unstaged_graph())) @pytest.mark.unit @@ -924,27 +1035,84 @@ def test_check_zero_copy_kv_rejects_a_still_staged_buffer(): """The shape a program finalized without zero_copy_backend_config has: the buffer is marked, so export dropped its copy-back, but it still reaches the delegate through a staging copy the engine's write is thrown away with.""" - graph_module, k_buffer, _, _ = _staged_delegate_graph() + graph_module, k_buffer, _, _ = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - with pytest.raises(RuntimeError, match="do not reach a TensorRT delegate"): - Z.check_zero_copy_kv(_finalized_program(graph_module)) + with pytest.raises(RuntimeError, match="do not reach the TensorRT delegate"): + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module))) @pytest.mark.unit def test_check_zero_copy_kv_accepts_a_buffer_that_never_had_a_staging_copy(): - """The checker and the un-staging pass read the same graph post-condition. + """The checker and the un-staging pass read the same post-condition. A program the pass has already un-staged hands the buffer straight to the delegate with no staging copy left. The pass accepts that shape (``test_unstage_accepts_a_buffer_that_never_had_a_staging_copy``) and so must this: the two disagreeing is what would let one path refuse a program the - other calls correct. The checker reads only the graph, so the placement half - of the pass's post-condition is not its to enforce. + other calls correct. """ - graph_module, _, _ = _direct_delegate_graph() + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module))) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_direct_buffer_planned_on_the_host(): + """Reaching the delegate directly is only half of it; placement is the rest. + + Finalizing with ``enable_non_cpu_memory_planning=False`` and no zero-copy + config gives exactly this: no staging copy is inserted, so the wiring is what + zero-copy wants, and every tensor still lands in the one host arena. Nothing + else refuses it -- the un-staging pass never ran -- and the engine is handed a + host pointer it cannot write, so the ``.pte`` fails its first ``execute()``. + The graph is the accepted one above with only the arena changed, so the + refusal can come from nothing else. + """ + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + + with pytest.raises(RuntimeError, match="not planned in any of the program's CUDA"): + Z.check_zero_copy_kv( + _finalized_program(_planned(graph_module, arena=1, on_device=False)) + ) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_buffer_staged_at_its_own_delegate(): + """A second TensorRT delegate must not stand in for the one that elided. + + The engine whose aliased output was elided -- the one carrying the zero-copy + spec -- still reads a staging copy, so its write is discarded and the cache + never updates. Another TensorRT engine happens to read the same buffer + directly, which says nothing about that write. Taking the union over every + TensorRT delegate accepts this program. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_k = graph.call_function(h2d, (k_buffer,)) + lowered_kv = graph.get_attr("lowered_module_0") + lowered_plain = graph.get_attr("lowered_module_1") + delegate_kv = graph.call_function(executorch_call_delegate, (lowered_kv, staged_k)) + delegate_plain = graph.call_function( + executorch_call_delegate, (lowered_plain, k_buffer) + ) + graph.output((k_buffer, delegate_kv, delegate_plain)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs() + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - Z.check_zero_copy_kv(_finalized_program(graph_module)) + with pytest.raises(RuntimeError, match="do not reach the TensorRT delegate"): + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module))) @pytest.mark.unit @@ -953,7 +1121,8 @@ def test_check_zero_copy_kv_rejects_a_buffer_only_another_backend_takes(): The mark is on this buffer because a TensorRT engine writes it in place, and that engine here is still reading a staging copy whose contents are thrown - away. Counting any backend's delegate would pass this program. + away. Both delegates carry the zero-copy spec, so the backend is the only + thing separating them: counting any backend's delegate passes this program. """ graph = torch.fx.Graph() k_buffer = graph.placeholder("b_k_0") @@ -970,16 +1139,17 @@ def test_check_zero_copy_kv_rejects_a_buffer_only_another_backend_takes(): graph.output((k_buffer, delegate_trt, delegate_other)) root = torch.nn.Module() root.lowered_module_0 = SimpleNamespace( - backend_id="TensorRTBackend", compile_specs=None + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs() ) root.lowered_module_1 = SimpleNamespace( - backend_id="CudaBackend", compile_specs=None + backend_id="CudaBackend", compile_specs=_zero_copy_specs() ) graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - with pytest.raises(RuntimeError, match="do not reach a TensorRT delegate"): - Z.check_zero_copy_kv(_finalized_program(graph_module)) + with pytest.raises(RuntimeError, match="do not reach the TensorRT delegate"): + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module))) @pytest.mark.unit @@ -1580,6 +1750,42 @@ def _assert_marked_buffers_reach_the_engine_unstaged(program): ) +def _assert_each_mutation_names_its_own_value(program): + """The finalized signature pairs every mutated buffer with its own new value. + + ExecuTorch finalizes the mutations by inserting a copy for each one whose + value is not already its buffer placeholder, moving those copies to the front + of the output tuple, and then reassigning the mutation specs' arguments by + position. Rewiring a cache to its own placeholder takes it out of that + leading run, so a method that mixes the two kinds comes out of finalization + with each buffer named against another buffer's value unless the mutations + were declared in the order the pass assumes. Nothing inside the finalizer + reads the pairing, so the ``.pte`` is written either way -- what reads it is + anyone inspecting the program, and the eager call path that copies mutated + values back into the state dict in this order. + """ + signature = program.exported_program().graph_signature + placeholder_of = {fqn: name for name, fqn in signature.inputs_to_buffers.items()} + mutated = { + spec.target: spec.arg.name + for spec in signature.output_specs + if spec.kind == OutputKind.BUFFER_MUTATION + } + assert set(mutated) == {"k_cache", "v_cache", "conv_state"} + for cache in ("k_cache", "v_cache"): + assert mutated[cache] == placeholder_of[cache], ( + f"the finalized signature gives {cache} the value " + f"{mutated[cache]!r}, but zero-copy left that cache as its own " + "mutation result, so its value is its own placeholder " + f"{placeholder_of[cache]!r}" + ) + assert mutated["conv_state"] not in placeholder_of.values(), ( + "the finalized signature gives conv_state a buffer placeholder as its " + f"value ({mutated['conv_state']!r}); it is copied back, so its value is " + "the copy ExecuTorch inserted" + ) + + @pytest.mark.skipif( not torch.cuda.is_available(), reason="requires CUDA + TensorRT for a real engine" ) @@ -1654,6 +1860,7 @@ def test_zero_copy_kv_keeps_a_copyback_buffer_in_the_same_method(retrace): config=torch_tensorrt.executorch.zero_copy_backend_config() ) _assert_marked_buffers_reach_the_engine_unstaged(program) + _assert_each_mutation_names_its_own_value(program) class _SplitRolesDecodeStep(_MixedDecodeStep): From 69e2daab6b8a22a79c3f5fefc628cb07f23e24b9 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Mon, 7 Sep 2026 14:01:32 -0700 Subject: [PATCH 17/22] fix(executorch): correct four zero-copy guards, and bound the blob extents Three of the fixes below are in the machinery that threads the planning flag into the un-staging pass and reorders the copy-back mutations; the other two are in the partitioner's alias-resolution fallback and the blob parser. `zero_copy_backend_config` preserved `skip_h2d_for_method_inputs`, and handed back a config that cannot finalize at all. Measured on the real mixed-KV model: `to_executorch` raises `skip_h2d_for_method_inputs=True requires placeholder 'b_k_cache' to have exactly one user, but it has 2 users`. That option is ExecuTorch's own un-staging of method inputs and it demands a single user, while a rewired cache always has two -- the delegate, and the graph output it is its own mutation result for -- so every zero-copy graph has the shape it refuses. This is the path the guide tells people to use, so the wrapper now refuses the option itself, naming it: the bool form, and any method a per-method dict asks for. `PropagateDevicePass` only tests that field for truth rather than resolving it per method, so a dict whose entries are all `False` is carried through here and still read as on a layer down; refusing that one too is left open. `order_copyback_mutations_first` did not do what it claimed. It split the mutations on this feature's own mark instead of on the predicate upstream keys on, and `run_reinplace_pass` and `reinplace_extra_ops` are supported `ExecutorchBackendConfig` fields whose pass runs just before the write-back and produces the same copy-free shape from an ordinary mutation. Measured against ExecuTorch's real `insert_write_back_for_buffers_pass`: an in-place-lineage mutation ahead of a copy-back one comes out with each buffer named against the other's value, and the reorder reports zero moved. It now asks upstream's own `_inplace_lineage`, imported rather than reimplemented, and returns the number of slots whose value changed. `check_zero_copy_kv` read an absent `non_const_buffer_device` as proof of a host placement. `apply_algo` is the only thing in ExecuTorch that writes that key and `to_executorch` takes any callable as `memory_planning_pass` -- which the guide tells people to bring for a cache shared between prefill and decode -- so a correct program was refused and `save` wrote nothing. Absence is no longer read as proof of a host placement; it is refused under a message of its own, because `MethodMeta::memory_planned_buffer_device` answers CPU for an arena the `.pte` records nothing for whoever planned it. Where the program's host tensors share the cache's arena -- which is what host-only planning produces and no device-aware planner would do -- the more specific host-arena refusal is what fires. All three shapes are measured on the real model -- the good one, the hole, and a caller-supplied planner that records no devices. A resolver failure inside `_partition_elided_output_names` did not degrade safely. Injecting one, the export died with "the aliased-buffer mark did not survive lowering", which is not what happened: the aliased outputs were removed before partitioning, so returning an empty set stamps nothing and guarantees a downstream failure that names the wrong cause. It re-raises now when any buffer in the method carries the mark, and keeps the fallback for the case where the claim is true. Two in the blob parser, both memory safety. The engine extent check added a 64-bit attacker-controlled `engine_size` to the offset before comparing against the blob length, so the sum wraps: an 8 KiB blob declaring 2^64-4086 parses clean and that length reaches `deserializeCudaEngine`. Every extent is now checked by subtracting the offset from the bound instead; the two metadata extents cannot wrap on a 64-bit `size_t` but are written the same way so the form, not the width of each field, is what makes them safe. Separately, the repeated-entry refusal covered only the output side of `aliased_io`, so two entries naming different outputs and one input both parsed and both resolved to one caller pointer. Confirmed by driving TensorRT 11.1 directly: it accepts two output bindings at one address, runs without error, and only the second write survives. The `kv_cache_update` kind is ruled out by the engine cross-check, but the `user` kind is compared only on shape. An empty or absent binding name is now refused in the same place rather than silently skipped, which used to shorten the recorded list while the delegate's argument list kept its length. Also: the two public helpers are in the API reference; the user guide's snippet no longer passes an unbound `backend_config` and no longer scopes the caller-stream synchronization duty to a coalesced `.pte`; the two readers of the zero-copy compile spec take the same values, so a hand-built one raises naming the key instead of decoding a JSON string into its own characters; a delegate with no outputs at all is refused after lowering, which naming every binding elidable used to satisfy; `save()`'s options message is built from the option table rather than hand-maintained; and the spec-missing refusal names whichever node is bare. Ten tests were added or widened for clauses that had none: the reorder against the real upstream pass, the skip-H2D refusal, both halves of the placement inference, the CUDA filter on the recorded arenas, the no-spec arms on both sides of the device move, the device-type half of the surviving-consumer check, the two unresolvable-alias skips, the empty delegate output list, and the compile-spec value shapes. Four C++ parser tests fail against the old parser. The multi-delegate CUDA-neighbour test is parametrized over the exporter, the mark-survival test asserts both cache names rather than a non-empty list, and the declaration-before-rewiring test records both passes in one tagged log so reversing them fails. --- .../executorch/TensorRTBlobHeader.cpp | 52 ++- docsrc/py_api/executorch.rst | 2 + .../runtime_performance/saving_models.rst | 46 +- examples/executorch_reference_runner/BUILD | 4 + py/torch_tensorrt/_compile.py | 55 ++- py/torch_tensorrt/executorch/_export.py | 11 +- py/torch_tensorrt/executorch/_zero_copy.py | 310 ++++++++++--- py/torch_tensorrt/executorch/backend.py | 65 ++- py/torch_tensorrt/executorch/partitioner.py | 29 +- .../test_executorch_blob_header.cpp | 72 +++ tests/py/dynamo/executorch/test_backend.py | 84 ++++ tests/py/dynamo/executorch/test_export.py | 48 +- .../test_weight_streaming_budget.py | 11 +- .../py/dynamo/executorch/test_zero_copy_kv.py | 428 ++++++++++++++++-- 14 files changed, 1037 insertions(+), 180 deletions(-) diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp index 73d4d0bd81e..a37050167c2 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp @@ -234,15 +234,22 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { } } - if (saw_name && !name.empty()) { - if (!claimed_bindings.insert(name).second) { - return false; - } - if (is_input) { - out.input_binding_names.push_back(name); - } else { - out.output_binding_names.push_back(name); - } + // A nameless entry cannot be refused earlier because the keys may arrive in + // any order, so it is refused here, beside the repeat. Skipping it instead + // would shorten the recorded list while the delegate's argument list keeps + // its full length, and the two are only inferred from the engine when both + // are empty -- so one real name beside a blank leaves a short list that no + // longer lines up with the engine's bindings. + if (!saw_name || name.empty()) { + return false; + } + if (!claimed_bindings.insert(name).second) { + return false; + } + if (is_input) { + out.input_binding_names.push_back(name); + } else { + out.output_binding_names.push_back(name); } } @@ -260,6 +267,7 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { } ++apos; std::unordered_set claimed_outputs; + std::unordered_set claimed_inputs; while (true) { apos = skip_ws(json, apos); if (apos >= json.size()) { @@ -325,6 +333,18 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { if (!claimed_outputs.insert(ab.output).second) { return false; } + // An input may likewise be claimed by at most one entry. Two entries + // naming different outputs and one input record the same input index + // for both, and execute() binds each aliased output to that index's + // caller pointer -- one address with two writers, so whichever the + // engine writes second wins and the other update disappears with no + // error. TensorRT's own aliasing rules out the kv_cache_update kind + // (init cross-checks it against getAliasedInputTensor), but the user + // kind is only compared on shape, so two same-shaped outputs onto one + // input would pass. + if (!claimed_inputs.insert(ab.input).second) { + return false; + } // The current Python serializer always writes "kind" (serialization.py), // and older blobs carry no aliased_io array at all, so this default is // defensive: it only fires for a blob that has an aliased_io entry but @@ -382,13 +402,21 @@ bool TensorRTBlobHeader::parse(const void* data, std::size_t size, TensorRTBlobH if (out.engine_offset % ENGINE_ALIGNMENT != 0) { return false; } - if (static_cast(out.metadata_offset) + out.metadata_size > size) { + // Every extent below is checked by subtracting the offset from the bound + // rather than by adding the size to the offset. engine_size is a 64-bit field + // read straight from the file, so the sum form wraps: a blob claiming a size + // just under 2^64 produces a small total, passes, and hands + // deserializeCudaEngine a pointer plus a length far past the end of the file. + // The two metadata extents cannot wrap on a 64-bit size_t -- both operands are + // 32-bit fields, so their sum is at most 2^33 -- but they are written the same + // way so that the form, not the width of each field, is what makes them safe. + if (out.metadata_offset > size || out.metadata_size > size - out.metadata_offset) { return false; } - if (static_cast(out.engine_offset) + out.engine_size > size) { + if (out.engine_offset > size || out.engine_size > size - out.engine_offset) { return false; } - if (static_cast(out.metadata_offset) + out.metadata_size > out.engine_offset) { + if (out.metadata_offset > out.engine_offset || out.metadata_size > out.engine_offset - out.metadata_offset) { return false; } diff --git a/docsrc/py_api/executorch.rst b/docsrc/py_api/executorch.rst index aefb26756cb..13dfbd965a2 100644 --- a/docsrc/py_api/executorch.rst +++ b/docsrc/py_api/executorch.rst @@ -37,6 +37,8 @@ Functions .. autofunction:: export .. autofunction:: get_edge_compile_config +.. autofunction:: zero_copy_backend_config +.. autofunction:: check_zero_copy_kv Classes -------- diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 6eb3c45515e..e8bb910494f 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -380,15 +380,21 @@ calls, one at each end of the Edge boundary: zero_copy_kv=True, ) - # zero_copy_backend_config composes onto your own config; every other field - # (memory planning, passes) is preserved. - program = edge.to_executorch(zero_copy_backend_config(backend_config)) - -One field of that config is also *read*. The engine writes the cache wherever -memory planning put it, so ``enable_non_cpu_memory_planning=False`` -- which -plans every tensor into a single host arena whatever device its ``TensorSpec`` -asks for -- cannot be combined with zero-copy KV: ``to_executorch`` raises -instead of writing a ``.pte`` whose every ``execute()`` fails on a host pointer. + # The argument is optional. Pass your own ExecutorchBackendConfig and + # zero_copy_backend_config composes onto it; every other field (memory + # planning, passes) is preserved. + program = edge.to_executorch(zero_copy_backend_config()) + +Two fields of that config are not merely carried. The engine writes the cache +wherever memory planning put it, so ``enable_non_cpu_memory_planning=False`` -- +which plans every tensor into a single host arena whatever device its +``TensorSpec`` asks for -- cannot be combined with zero-copy KV: +``to_executorch`` raises instead of writing a ``.pte`` whose every ``execute()`` +fails on a host pointer. And ``propagate_device_config.skip_h2d_for_method_inputs`` +is refused outright: ``PropagateDevicePass`` refuses to un-stage a method input +whose placeholder does not have exactly one user, and a zero-copy cache always +has two, so preserving that option would hand back a configuration that cannot +finalize at all. It is opt-in rather than automatic because the resulting ``.pte`` needs a runtime that understands a delegate whose aliased outputs are elided. Producing @@ -404,7 +410,7 @@ one silently would break a runner built before this feature. reads it back and refuses one whose caches are still staged, or planned somewhere the engine cannot write them, before writing the ``.pte``:: - program = edge.to_executorch(zero_copy_backend_config(backend_config)) + program = edge.to_executorch(zero_copy_backend_config()) torch_tensorrt.executorch.check_zero_copy_kv(program) ``torch_tensorrt.save`` owns both ends and runs that check itself, so there @@ -426,14 +432,24 @@ one as ``backend_config`` as well: that installs the pass twice, which is redundant rather than an error -- the second run finds the buffers already un-staged. The two entry points are alternatives, not a pair. -Two further responsibilities are the caller's, and neither raises: +Three further responsibilities are the caller's, and none raises: -* **One CUDA stream for every delegate**, if the ``.pte`` is coalesced -- and the - synchronization it calls for, which zero-copy makes load-bearing. Getting this - wrong is a race, not a deterministic error: it is intermittent and can surface - as wrong results *or* as an illegal memory access. See +* **One CUDA stream for every delegate**, if the ``.pte`` is coalesced. Getting + this wrong is a race, not a deterministic error: it is intermittent and can + surface as wrong results *or* as an illegal memory access. See :ref:`Running a coalesced .pte `. +* **Synchronizing that stream before reading a cache on the host**, for any + zero-copy ``.pte`` a runner drives on a caller stream, coalesced or not. A + delegate whose aliased outputs are threaded through it reflects each one into + its delegate output and so waits for the engine before returning; zero-copy + elides those outputs, so there is nothing to reflect and ``execute()`` returns + with the engine still running. A single-delegate decode loop owes the + synchronization as much as a coalesced program does. It is only the caller + stream that brings the duty: with none installed ``execute()`` synchronizes + before it returns, as it does whenever a delegate input or output is staged + through the host. + * **Sharing one cache between methods.** Zero-copy is per method: it makes each method's engine write that method's buffer. Giving a prefill and a decode method *the same* cache is a memory-planning question -- their mutable buffers diff --git a/examples/executorch_reference_runner/BUILD b/examples/executorch_reference_runner/BUILD index 5d85db9f5be..e6d3bc2d1b5 100644 --- a/examples/executorch_reference_runner/BUILD +++ b/examples/executorch_reference_runner/BUILD @@ -43,5 +43,9 @@ cc_binary( "@cuda//:cudart", "@executorch//:executorch_core", "@executorch//:executorch_file_data_loader", + # Included and called directly (CallerStreamGuard), so declared directly + # rather than reached through a transitive header, matching the sibling + # runner above. + "@executorch//:extension_cuda", ], ) diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index d62ca480970..d9967363a90 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -685,6 +685,23 @@ def load( ) +# The keyword arguments save() consumes only for output_format="executorch", +# each with the default it is popped with. One table, because the unexpected-keyword +# error spells the supported set out for the caller: an option added to the pops +# alone would leave that message telling someone their flag is unsupported. +_EXECUTORCH_SAVE_OPTIONS: Dict[str, Any] = { + "partitioners": None, + "compile_specs": None, + "backend_config": None, + "constant_methods": None, + "transform_passes": None, + "compile_config": None, + "generate_etrecord": False, + "weight_streaming_budget_per_engine": None, + "zero_copy_kv": False, +} + + def save( module: Any, file_path: str = "", @@ -781,10 +798,10 @@ def save( parameter takes precedence. kwargs: Additional format-specific kwargs. ``partitioners=``, ``compile_specs=``, ``backend_config=``, ``constant_methods=``, - ``transform_passes=``, ``compile_config=``, ``generate_etrecord=`` - and ``weight_streaming_budget_per_engine=`` are only used with - ``output_format="executorch"``; otherwise they are ignored with a - warning. Pass ``compile_specs=[CompileSpec("target_device", + ``transform_passes=``, ``compile_config=``, ``generate_etrecord=``, + ``weight_streaming_budget_per_engine=`` and ``zero_copy_kv=`` are + only used with ``output_format="executorch"``; otherwise they are + ignored with a warning. Pass ``compile_specs=[CompileSpec("target_device", b"cuda:")]`` to override the default target device (``cuda:0``). ``backend_config=`` takes an ``Optional[ExecutorchBackendConfig]`` and is forwarded to ``to_executorch(config=...)`` to customize @@ -847,17 +864,21 @@ def save( if kwarg_inputs and any(value is None for value in kwarg_inputs.values()): raise ValueError("kwargs should not include None.") - executorch_partitioners = kwargs.pop("partitioners", None) - executorch_compile_specs = kwargs.pop("compile_specs", None) - executorch_backend_config = kwargs.pop("backend_config", None) - executorch_constant_methods = kwargs.pop("constant_methods", None) - executorch_transform_passes = kwargs.pop("transform_passes", None) - executorch_compile_config = kwargs.pop("compile_config", None) - executorch_generate_etrecord = kwargs.pop("generate_etrecord", False) - executorch_weight_streaming_budget_per_engine = kwargs.pop( - "weight_streaming_budget_per_engine", None - ) - executorch_zero_copy_kv = kwargs.pop("zero_copy_kv", False) + executorch_options = { + name: kwargs.pop(name, default) + for name, default in _EXECUTORCH_SAVE_OPTIONS.items() + } + executorch_partitioners = executorch_options["partitioners"] + executorch_compile_specs = executorch_options["compile_specs"] + executorch_backend_config = executorch_options["backend_config"] + executorch_constant_methods = executorch_options["constant_methods"] + executorch_transform_passes = executorch_options["transform_passes"] + executorch_compile_config = executorch_options["compile_config"] + executorch_generate_etrecord = executorch_options["generate_etrecord"] + executorch_weight_streaming_budget_per_engine = executorch_options[ + "weight_streaming_budget_per_engine" + ] + executorch_zero_copy_kv = executorch_options["zero_copy_kv"] if output_format not in accepted_formats: raise ValueError( @@ -873,11 +894,11 @@ def save( # Every executorch option is popped above, so a leftover kwarg is a typo. Fail # here rather than silently ignoring it, since nothing downstream reads kwargs. if kwargs: + supported = ", ".join(repr(name) for name in _EXECUTORCH_SAVE_OPTIONS) raise TypeError( "save() received unexpected keyword argument(s) for " f"output_format='executorch': {sorted(kwargs)}. Supported executorch " - "options are 'partitioners', 'compile_specs', 'backend_config', and " - "'weight_streaming_budget_per_engine'." + f"options are {supported}." ) # Validate the budget before the input and model-shape checks below, so a wrong # type is not reported as an unrelated failure. diff --git a/py/torch_tensorrt/executorch/_export.py b/py/torch_tensorrt/executorch/_export.py index b3fa0ca0e19..e4ea70d1c56 100644 --- a/py/torch_tensorrt/executorch/_export.py +++ b/py/torch_tensorrt/executorch/_export.py @@ -699,15 +699,18 @@ def export( trt_compile_specs = list(method_compile_specs[name]) if name in zero_copy_methods: # Signal to TensorRTPartitioner that this method elided aliased - # outputs, and carry the method-wide binding names it elided. The - # partitioner does not apply this spec to every partition: it + # outputs. Only the presence of the key is read: the partitioner + # drops this spec from the list it applies to every partition and # re-derives, per engine, exactly which of a delegate's aliased - # outputs were elided and stamps only those onto only that delegate - # (see TensorRTPartitioner._partition_elided_output_names), so a + # outputs were elided, stamping only those onto only that delegate + # (see TensorRTPartitioner._partition_elided_output_names). So a # method that lowers to several TensorRT delegates marks only the KV # one and a plain-compute delegate beside it carries no zero-copy # spec. Without any spec the backend rejects a delegate short of its # bindings, which keeps an accidentally dropped output an error. + # The method-wide names go in the value so this spec reads like the + # per-partition one the partitioner writes, but nothing decodes this + # one: changing which names are listed here changes no delegate. trt_compile_specs.append( _CompileSpec( ZERO_COPY_KV_COMPILE_SPEC_KEY, diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index 9a5b72f6743..c83b7e4c9b8 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -22,7 +22,7 @@ partitioning, drops the copy-back by declaring that the buffer *is* the mutation's result. The aliased output then has no user and disappears from the partition. :func:`order_copyback_mutations_first` then repairs, on the Edge - program, the one thing that declaration disturbs downstream. + program, the mutation-spec pairing that declaration disturbs downstream. * :func:`unstage_aliased_buffers_pass`, as a ``to_out_var_pass``, drops the staging so the engine writes the caller's buffer rather than a copy. @@ -312,41 +312,76 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> List[str]: return elided_output_names +def _mutation_targets_a_lifted_input(signature: Any) -> Set[str]: + """The mutation targets ``insert_write_back_for_buffers_pass`` can resolve. + + Mirrors the ``lifted_inputs`` map that pass builds: a buffer, constant, + parameter or custom object contributes its ``target``, a user input its + argument name. A mutation whose target is not in here gets no copy either, + so it belongs with the copy-free ones. + """ + from torch.export.graph_signature import InputKind, TensorArgument + + lifted: Set[str] = set() + for spec in signature.input_specs: + if spec.kind in ( + InputKind.BUFFER, + InputKind.CONSTANT_TENSOR, + InputKind.PARAMETER, + InputKind.CUSTOM_OBJ, + ): + if spec.target is not None: + lifted.add(spec.target) + elif spec.kind is InputKind.USER_INPUT and isinstance(spec.arg, TensorArgument): + lifted.add(spec.arg.name) + return lifted + + def order_copyback_mutations_first(exported_program: Any) -> int: """Reorder one Edge method's mutations so ExecuTorch pairs them up correctly. ExecuTorch finalizes a mutation by inserting a ``copy_`` for it, but only - when its value is not already the buffer placeholder. It moves those copies - to the front of the output tuple, leaves everything else behind them in - order, and then walks the mutation specs reassigning each one's argument - *by position* over the result (``insert_write_back_for_buffers_pass``). - :func:`rewire_aliased_mutations_to_buffers` is what makes a mutation's value - its own placeholder, so a rewired cache gets no copy and drops out of that - leading run -- and in a method that also has a copy-back buffer, every - mutation spec from the first rewired one on then comes out of finalization - naming a different buffer's value. The specs ahead of it are unaffected, - since the copies keep their order among themselves. The ``.pte`` is written - correctly either way, because the emitter and the memory planner read only - which buffers are mutated and not what by. What the pairing decides is the + when its target is one of the lifted inputs and its value is not already + reached, through in-place ops, from a placeholder of the mutation's own kind + -- any buffer placeholder for a buffer mutation, not necessarily the one it + targets. It moves those copies to the front of the output tuple, leaves + everything else behind them in order, and then walks the mutation specs + reassigning each one's argument *by position* over the result + (``insert_write_back_for_buffers_pass``). + :func:`rewire_aliased_mutations_to_buffers` makes a mutation's value its own + placeholder, so a rewired cache gets no copy and drops out of that leading + run -- and in a method that also has a copy-back buffer, every mutation spec + from the first copy-free one on then comes out of finalization naming a + different buffer's value. The specs ahead of it are unaffected, since the + copies keep their order among themselves. The ``.pte`` is written correctly + either way, because the emitter and the memory planner read only which + buffers are mutated and not what by. What the pairing decides is the finalized signature, which anyone inspecting the program reads, and - ExecuTorch's eager call path, which walks - ``buffers_to_mutate`` writing the graph's leading results into the state dict - in that order and so updates each buffer from another one's value. + ExecuTorch's eager call path, which walks ``buffers_to_mutate`` writing the + graph's leading results into the state dict in that order and so updates + each buffer from another one's value. Putting the mutations that still get a copy first restores the - correspondence. Only the mutations this module rewired are moved: a mutation - already bound to its own placeholder before zero-copy ran would need moving - too, but that shape is not one this feature creates. + correspondence. Which ones those are is decided by asking upstream's own + predicate (``_inplace_lineage``, imported rather than reimplemented) rather + than by asking which mutations this module rewired. The two answers differ: + ``run_reinplace_pass`` and ``reinplace_extra_ops`` are supported + ``ExecutorchBackendConfig`` fields whose pass runs just before the write-back + and turns ordinary mutations into in-place ones, and such a mutation gets no + copy either. Keyed on this module's own mark, the reorder would leave that + pair crossed and report nothing moved. This runs on the *Edge* program rather than beside the rewiring, because ``to_edge_transform_and_lower`` re-derives the whole graph signature -- the order it hands back is the buffers' own order, whatever order it was given. Nothing between here and the write-back pass re-derives it again. - Returns zero when the order already holds, and otherwise the number of - rewired mutations, all of which the reorder leaves behind the copy-back - ones. + Returns the number of mutation slots whose value changed, which is zero when + the order already holds. """ + from executorch.exir.passes.insert_write_back_for_buffers_pass import ( + _inplace_lineage, + ) from torch.export.graph_signature import ExportGraphSignature, OutputKind signature = exported_program.graph_signature @@ -359,15 +394,24 @@ def order_copyback_mutations_first(exported_program: Any) -> int: if spec.kind in (OutputKind.BUFFER_MUTATION, OutputKind.USER_INPUT_MUTATION) and index < len(args) ] - rewired = { - index - for index in slots - if isinstance(args[index], Node) - and args[index].op == "placeholder" - and args[index].meta.get("_torch_tensorrt_aliased_buffer") - } - source = [index for index in slots if index not in rewired] + [ - index for index in slots if index in rewired + if not slots: + return 0 + lifted = _mutation_targets_a_lifted_input(signature) + + def gets_a_copy(index: int) -> bool: + value = args[index] + if not isinstance(value, Node): + # Upstream reads a non-Node value as needing a copy, and then raises + # walking it. Grouping it with the copies keeps this reorder from + # being what raises first. + return True + if specs[index].target not in lifted: + return False + return not _inplace_lineage(value, signature, specs[index].kind) + + copied = {index for index in slots if gets_a_copy(index)} + source = [index for index in slots if index in copied] + [ + index for index in slots if index not in copied ] if source == slots: return 0 @@ -381,10 +425,9 @@ def order_copyback_mutations_first(exported_program: Any) -> int: exported_program._graph_signature = ExportGraphSignature( input_specs=list(signature.input_specs), output_specs=new_specs ) - _LOGGER.debug( - "moved %d rewired mutation(s) behind the copy-back ones", len(rewired) - ) - return len(rewired) + moved = sum(1 for slot, index in zip(slots, source) if slot != index) + _LOGGER.debug("moved %d mutation(s) so the copy-back ones come first", moved) + return moved def _is_tensorrt_delegate(graph_module: torch.fx.GraphModule, node: Node) -> bool: @@ -431,6 +474,15 @@ def _delegate_elided_output_names( spec, and also when it carries one whose value does not decode into a list of names -- a spec built by hand rather than by the partitioner, which always writes the JSON list. Callers read empty as "cannot tell", not as "none". + + ``backend._elided_output_names`` reads the same key on the same spec and + takes the same shapes of value. It differs in what it does with the rest: it + raises naming the key, because an undecodable spec leaves it unable to say + which outputs the delegate was allowed to drop, while here it only weakens a + cross-check that then falls back to demanding at least one buffer. The one + value the two read differently is bytes that are not valid UTF-8: it + replaces the bad units and reads a name out of them, ``json.loads`` refuses + them here, and here that is another "cannot tell". """ spec = _zero_copy_compile_spec(graph_module, node) if spec is None: @@ -560,8 +612,9 @@ def _unstage_aliased_buffers( A failure here is a lost KV update, so it is raised rather than logged: export has already removed the copy-back, so a marked buffer left staged has the engine write per-call scratch that is then discarded and the buffer never - updates. It raises when memory planning is host-only, when the staging copy - has no spec or is not on CUDA, when a direct argument has no spec of its own + updates. It raises when memory planning is host-only, when either end of the + move -- the staging copy or the buffer -- has no spec, when the staging copy + is not on CUDA, when a direct argument has no spec of its own or that spec is not on CUDA, when the device move is unsafe, and -- so a discovery miss cannot pass silently -- after the loop when the post-condition does not hold for some marked buffer, cross-checked against each delegate's @@ -656,15 +709,20 @@ def _unstage_aliased_buffers( staged_spec = arg.meta.get("spec") source_spec = source.meta.get("spec") if staged_spec is None or source_spec is None: + missing = ( + f"the staging copy '{arg.name}'" + if staged_spec is None + else f"the buffer placeholder '{source.name}'" + ) raise RuntimeError( - "TensorRT zero-copy KV: no TensorSpec on the staging copy of " - f"buffer '{source.name}', so it cannot be moved to the " - "delegate's device. The TensorRT engine writes this buffer in " - "place and its copy-back has already been removed, so the " - "update would be lost. This pass has to run as the " - "ExecutorchBackendConfig to_out_var_pass, which is where the " - "specs exist; torch_tensorrt.executorch.zero_copy_backend_config " - "installs it there." + f"TensorRT zero-copy KV: no TensorSpec on {missing}, so buffer " + f"'{source.name}' cannot be moved to the delegate's device. The " + "TensorRT engine writes this buffer in place and its copy-back " + "has already been removed, so the update would be lost. This " + "pass has to run as the ExecutorchBackendConfig to_out_var_pass, " + "which is where the specs exist; " + "torch_tensorrt.executorch.zero_copy_backend_config installs it " + "there." ) # spec.device is an exir schema DeviceType, not a torch.device. if staged_spec.device != DeviceType.CUDA: @@ -815,24 +873,76 @@ def call(self, graph_module: torch.fx.GraphModule) -> Any: return _UnstageThenToOutVar() -def _device_planned_arenas(graph_module: torch.fx.GraphModule) -> Set[int]: - """The ``mem_id``s of the finalized program's CUDA arenas. +def _device_planned_arenas(graph_module: torch.fx.GraphModule) -> Optional[Set[int]]: + """The ``mem_id``s of the finalized program's CUDA arenas, or ``None``. Memory planning partitions the specs by device, gives each device its own arena, and records the non-CPU ones on the graph module as - ``non_const_buffer_device``. The key is absent from a program planned - entirely on the host, which is what an ``enable_non_cpu_memory_planning=False`` - finalization produces however the specs are marked. + ``non_const_buffer_device``. ``None`` means the program records no arena + devices at all, which does *not* mean the host: ``apply_algo`` is the only + thing in ExecuTorch that writes the key, and ``to_executorch`` accepts any + callable as ``memory_planning_pass``, so a caller-supplied planner -- which + the user guide tells people to bring for a cache shared between prefill and + decode -- can plan onto a device and still leave the key unwritten. """ from executorch.exir.schema import DeviceType + entries = graph_module.meta.get("non_const_buffer_device") + if not entries: + return None return { entry.buffer_idx - for entry in (graph_module.meta.get("non_const_buffer_device") or []) + for entry in entries if getattr(entry, "device_type", None) == DeviceType.CUDA } +def _host_planned_arenas(graph_module: torch.fx.GraphModule) -> Set[int]: + """The ``mem_id``s of arenas that hold at least one host tensor. + + Planning gives each device its own arena, so an arena holding a tensor whose + spec is CPU is a host arena -- an argument from the graph rather than from + the planner's records, which is what makes it usable when those records are + absent. It is what separates the two shapes that both record no arena + devices: ``enable_non_cpu_memory_planning=False`` puts every tensor in one + bucket, so a CUDA-spec cache ends up sharing an arena with the host tensors, + while a caller-supplied device-aware planner keeps them apart. + + A method with no CPU tensor at all is the residual: nothing here can then + tell a device arena from a host one, and if the program records no arena + devices either, the buffer is accepted. + """ + from executorch.exir.schema import DeviceType + + host: Set[int] = set() + for node in graph_module.graph.nodes: + specs = node.meta.get("spec") + for spec in specs if isinstance(specs, (list, tuple)) else [specs]: + if ( + spec is not None + and getattr(spec, "device", None) == DeviceType.CPU + and getattr(spec, "mem_id", None) is not None + ): + host.add(spec.mem_id) + return host + + +def _is_host_planned( + node: Node, device_arenas: Optional[Set[int]], host_arenas: Set[int] +) -> bool: + """True when memory planning put ``node`` somewhere the engine cannot write. + + Two independent grounds, because either record may be the only one there. + Sharing an arena with a host tensor settles it whatever the program records; + otherwise, when the program does record its arena devices, an arena missing + from that record is not a CUDA one. + """ + mem_id = getattr(node.meta.get("spec"), "mem_id", None) + if mem_id in host_arenas: + return True + return device_arenas is not None and mem_id not in device_arenas + + def _name_detail(names_by_method: Dict[str, List[str]]) -> str: return ", ".join( f"'{name}' in method '{method}'" @@ -848,9 +958,12 @@ def check_zero_copy_kv(program: Any) -> None: nothing quietly when they find nothing to do: ``zero_copy_kv=True`` warns and carries on when the model holds no aliased buffer mutation, and :func:`unstage_aliased_buffers_pass`, handed a program with nothing marked, - un-stages nothing and returns. Either way the ``.pte`` that comes out runs - and stages its cache like any other, which for a KV cache is wrong output - rather than a crash. + un-stages nothing and returns. Neither of those is wrong output -- nothing + removed a copy-back, so the ``.pte`` stages its cache and updates it like any + other -- but the optimization the caller asked for is silently not there, and + a caller who reads the successful ``save`` as proof it is gets neither an + error nor the speedup. The wrong-output case is the one below: a rewiring + that did happen and then lost its mark. Three shapes are refused: a marked buffer that is not a direct argument of a TensorRT delegate carrying the zero-copy compile spec, a marked buffer that @@ -872,14 +985,21 @@ def check_zero_copy_kv(program: Any) -> None: :func:`_unstage_aliased_buffers`, which counts each delegate's own buffers against its own spec, is what separates them. - Placement is read from where memory planning actually put the buffer: its - ``TensorSpec``'s ``mem_id`` has to name one of the arenas the finalized - program records as CUDA in ``non_const_buffer_device``. The spec's own device - does not settle it -- ``PropagateDevicePass`` writes CUDA onto the spec of a - buffer that reaches a CUDA delegate directly, and + Placement is read from where memory planning actually put the buffer -- the + ``mem_id`` on its ``TensorSpec`` -- rather than from the spec's own device, + which does not settle it: ``PropagateDevicePass`` writes CUDA onto the spec + of a buffer that reaches a CUDA delegate directly, and ``enable_non_cpu_memory_planning=False`` then plans that same buffer into the one host arena, which is the shape whose every ``execute()`` fails on the - runtime's alias-target guard. + runtime's alias-target guard. Two independent grounds answer whether that + ``mem_id`` names a host arena. An arena that also holds one of the program's + host tensors is one whatever the program records (see + :func:`_host_planned_arenas`); failing that, an arena missing from the CUDA + ones the program *does* record in ``non_const_buffer_device`` is one too. + Absence of that record is not itself read as the host: only ``apply_algo`` + writes it, so a caller-supplied ``memory_planning_pass`` that does not go + through it records nothing whatever it planned, and refusing on that would + block the one thing the guide says to bring your own planner for. Every method is read, not only ``forward``. ``export()`` rewires each method on its own, so a check that stopped at ``forward`` would pass a program whose @@ -920,11 +1040,12 @@ def check_zero_copy_kv(program: Any) -> None: if staged: staged_by_method[method_name] = staged device_arenas = _device_planned_arenas(graph_module) + host_arenas = _host_planned_arenas(graph_module) host_planned = [ node.name for node in marked if node in zero_copy_delegate_args - and getattr(node.meta.get("spec"), "mem_id", None) not in device_arenas + and _is_host_planned(node, device_arenas, host_arenas) ] if host_planned: host_planned_by_method[method_name] = host_planned @@ -949,14 +1070,49 @@ def check_zero_copy_kv(program: Any) -> None: raise RuntimeError( f"TensorRT zero-copy KV: buffer(s) " f"{_name_detail(host_planned_by_method)} reach their TensorRT " - "delegate directly but are not planned in any of the program's CUDA " - "arenas, so the engine is handed a host pointer it cannot write and " + "delegate directly but memory planning put them in an arena that " + "also holds the program's host tensors, or in one it does not record " + "as CUDA, so the engine is handed a host pointer it cannot write and " "every execute() fails on the runtime's alias-target guard. Finalize " "with torch_tensorrt.executorch.zero_copy_backend_config() over a " - "configuration that leaves enable_non_cpu_memory_planning on." + "configuration that leaves enable_non_cpu_memory_planning on. If it " + "is already on, the memory_planning_pass in use is what put this " + "buffer among the host tensors, and it has to give the delegate's " + "device an arena of its own." ) +def _refuse_skip_h2d(config: "ExecutorchBackendConfig") -> None: + """Raise if the config asks ExecuTorch to un-stage method inputs as well. + + The field is typed as a bool or a per-method dict, and every truthy entry is + refused, naming the method that asked. That is narrower than the pass it + guards: ``PropagateDevicePass`` takes the field as one value and only tests + it for truth, so it reads *any* non-empty dict as on -- including one whose + entries are all ``False``, which this carries through and which then fails + to finalize a layer down. + """ + propagate = getattr(config, "propagate_device_config", None) + skip = getattr(propagate, "skip_h2d_for_method_inputs", False) + if isinstance(skip, dict): + asked_for = [name for name, value in skip.items() if value] + else: + asked_for = ["every method"] if skip else [] + if not asked_for: + return + raise ValueError( + "TensorRT zero-copy KV: this configuration sets " + f"skip_h2d_for_method_inputs for {', '.join(asked_for)}, which cannot be " + "combined with zero-copy KV. PropagateDevicePass refuses to un-stage a " + "method input whose placeholder does not have exactly one user, and a " + "buffer zero-copy rewired has two -- the TensorRT delegate, and the graph " + "output it is its own mutation result for -- so finalization raises " + "there. Zero-copy already un-stages the aliased buffers; leave " + "skip_h2d_for_method_inputs off and pass the method's own inputs on the " + "host." + ) + + def zero_copy_backend_config( config: Optional["ExecutorchBackendConfig"] = None, ) -> "ExecutorchBackendConfig": @@ -974,14 +1130,23 @@ def zero_copy_backend_config( ``config`` is your own configuration -- every field is preserved, and a ``to_out_var_pass`` you already set runs after the un-staging. Omit it to - start from ExecuTorch's defaults. One field is not merely carried but read: - zero-copy needs the caches planned in device memory, so - ``enable_non_cpu_memory_planning=False`` -- which plans every tensor into the - one host arena -- has the pass refuse each cache it finds rather than write a - ``.pte`` whose every ``execute()`` fails. It is read off the config returned - here, at the moment the pass runs, so setting the field on that config - afterwards is honoured: the pass and the finalizer then cannot disagree about - it. + start from ExecuTorch's defaults. Two fields are not merely carried: + + * ``enable_non_cpu_memory_planning`` is *read*. Zero-copy needs the caches + planned in device memory, so ``False`` -- which plans every tensor into + the one host arena -- has the pass refuse each cache it finds rather than + write a ``.pte`` whose every ``execute()`` fails. It is read off the + config returned here, at the moment the pass runs, so setting the field on + that config afterwards is honoured: the pass and the finalizer then cannot + disagree about it. + * ``propagate_device_config.skip_h2d_for_method_inputs`` is *refused*, for + any method. It is ExecuTorch's own un-staging of method inputs, and it + requires each placeholder it un-stages to have exactly one user. A rewired + cache always has two -- the delegate, and the graph output it is its own + mutation result for -- so ``PropagateDevicePass`` raises on every + zero-copy graph. Returning the option unchanged would hand back a config + that cannot finalize at all; this says so here instead, where the caller + can act on it. .. warning:: Finalizing a ``zero_copy_kv=True`` program *without* this config does @@ -1004,6 +1169,7 @@ def zero_copy_backend_config( from executorch.exir import ExecutorchBackendConfig base = config if config is not None else ExecutorchBackendConfig() + _refuse_skip_h2d(base) unstage = unstage_aliased_buffers_pass( base.to_out_var_pass, device_memory_planning=base.enable_non_cpu_memory_planning, diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index 2034bc1ea49..975d934a323 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -329,6 +329,28 @@ def _validate_output_binding_order( node for node in edge_program.graph_module.graph.nodes if node.op == "output" ) out_args = list(output_node.args[0]) + if not out_args: + # Naming every binding elidable empties unaliased_indices too, so the + # comparison below would match. That is the zero-output delegate + # rewire_aliased_mutations_to_buffers raises about -- nothing reads it, + # so a later graph-wide dead-code elimination can erase the computation. + # The rewiring's guard runs before partitioning; this is the only one + # left after lowering. The check is unconditional because a delegate with + # no outputs is wrong however it got that way; only the remedy below is + # about zero-copy. + raise ValueError( + "TensorRT ExecuTorch backend: the delegate has no outputs at all, " + f"but the engine declares {len(output_names)} output binding(s). A " + "delegate nothing reads is a pure node a later dead-code elimination " + "can erase, taking the engine with it." + + ( + " Every one of this engine's outputs was declared elidable, so " + "nothing was left to thread out; export this method without " + "zero_copy_kv." + if elidable_output_names is not None and not unaliased_indices + else "" + ) + ) # A single-output engine is returned directly rather than through a getitem, # and one binding has no order to get wrong. The same holds under elision # when exactly one binding is left unaliased. @@ -405,14 +427,53 @@ def _elided_output_names(compile_specs: List[CompileSpec]) -> Optional[Set[str]] ``None`` when no zero-copy spec is present, which keeps a missing output an error: only a caller who asked for zero-copy may drop the aliased outputs, and then only exactly the ones export rewired to write in place. + + A spec built by hand rather than by ``TensorRTPartitioner`` may carry + anything, so the value is type-checked and the decode is caught, and both + raise naming this key. Without that a bare ``b"1"`` comes out as an + unattributed ``TypeError``, and -- the quieter one -- a JSON *string* + decodes into a set of its own characters, exempting every one-character + binding name and not the real one. + + ``_zero_copy._delegate_elided_output_names`` reads the same key on the same + spec and takes the same shapes of value. It differs in what it does with the + rest: where this raises, it returns the empty set, because for it an + undecodable spec merely weakens a cross-check while here it leaves the + backend unable to say which outputs may be missing. The one value the two + read differently is bytes that are not valid UTF-8, replaced here and + refused there by ``json.loads``. """ for spec in compile_specs: if getattr(spec, "key", None) != ZERO_COPY_KV_COMPILE_SPEC_KEY: continue value = spec.value + if not isinstance(value, (str, bytes, bytearray)): + raise ValueError( + "TensorRT ExecuTorch backend: compile spec " + f"'{ZERO_COPY_KV_COMPILE_SPEC_KEY}' must hold a JSON list of " + f"engine output binding names, not {type(value).__name__}. It is " + "written by TensorRTPartitioner; a hand-built spec has to match." + ) if isinstance(value, (bytes, bytearray)): - value = bytes(value).decode("utf-8") - return set(json.loads(value)) + value = bytes(value).decode("utf-8", "replace") + try: + names = json.loads(value) + except ValueError as e: + raise ValueError( + "TensorRT ExecuTorch backend: compile spec " + f"'{ZERO_COPY_KV_COMPILE_SPEC_KEY}' does not decode as JSON " + f"({e}). It holds a JSON list of engine output binding names, " + "written by TensorRTPartitioner." + ) from e + if not isinstance(names, list): + raise ValueError( + "TensorRT ExecuTorch backend: compile spec " + f"'{ZERO_COPY_KV_COMPILE_SPEC_KEY}' decoded to " + f"{type(names).__name__}, not a list of engine output binding " + "names. A JSON string would decode into its own characters and " + "exempt every one-character binding name." + ) + return {str(name) for name in names} return None diff --git a/py/torch_tensorrt/executorch/partitioner.py b/py/torch_tensorrt/executorch/partitioner.py index 7c0fd47e955..3cdde4d9d05 100644 --- a/py/torch_tensorrt/executorch/partitioner.py +++ b/py/torch_tensorrt/executorch/partitioner.py @@ -214,9 +214,13 @@ def _partition_elided_output_names( binding. That is what lets a real lost output on the plain delegate still raise while the KV delegate's genuine elision is exempted. - Any extraction failure returns an empty set: the delegate then carries every - binding and a genuinely missing aliased output stays an error in the - backend's ``_validate_output_binding_order``. + An extraction failure is survivable only for a method in which nothing was + rewired: the delegate then really does carry every binding, and a genuinely + missing aliased output stays an error in the backend's + ``_validate_output_binding_order``. Where a buffer *was* rewired the aliased + outputs are already gone from the graph -- that happened before partitioning + -- so an empty set stamps nothing, and the export dies downstream blaming a + lost aliased-buffer mark, which is not what went wrong. That case re-raises. A set that is neither empty nor the engine's whole ``aliased_io`` is the right answer and still not a lowerable one, since the runtime reads @@ -251,13 +255,22 @@ def _partition_elided_output_names( elided.add(output_names[output_index]) return elided except Exception as e: - # Broad by design, mirroring _resolve_target_device_for_partition: any - # extraction failure must not abort the export. It degrades safely -- - # the delegate keeps every binding, so a truly-elided output is caught - # downstream rather than dropped. + # Broad in the same shape as _resolve_target_device_for_partition, but + # not with the same licence: that fallback picks a default device and + # is genuinely harmless, while this one can only be harmless where no + # output was elided. The question is asked of the graph, not of the + # engine metadata that just failed to read, so it cannot fail the same + # way. + if any( + node.op == "placeholder" + and node.meta.get("_torch_tensorrt_aliased_buffer") + for node in exported_program.graph_module.graph.nodes + ): + raise logger.warning( "zero-copy KV: could not resolve elided outputs for partition %s " - "(%s); the delegate will carry every binding.", + "(%s); no buffer in this method was rewired, so the delegate " + "carries every binding.", getattr(partition, "id", "?"), e, ) diff --git a/tests/cpp/executorch/test_executorch_blob_header.cpp b/tests/cpp/executorch/test_executorch_blob_header.cpp index 1d7935e40ec..5f2fbe693b8 100644 --- a/tests/cpp/executorch/test_executorch_blob_header.cpp +++ b/tests/cpp/executorch/test_executorch_blob_header.cpp @@ -258,6 +258,78 @@ TEST(ExecuTorchTensorRTBlobHeader, ParsesAliasedIoEntriesForDistinctOutputs) { EXPECT_EQ(header.aliased_io[1].output, "out_v"); } +TEST(ExecuTorchTensorRTBlobHeader, RejectsRepeatedAliasedIoInputForDifferentOutputs) { + // Two entries naming different outputs and one input. Both resolve to the same + // input index, so execute() binds both output bindings to that input's caller + // pointer -- one address with two writers, the second of which erases the + // first with no error. The kind here is "user", which init validates only by + // comparing shapes, so two same-shaped outputs get past it. + const std::string metadata = R"({"io_bindings":[{"name":"in_0","is_input":true},)" + R"({"name":"out_0","is_input":false},{"name":"out_1","is_input":false}],)" + R"("aliased_io":[{"output":"out_0","input":"in_0","kind":"user"},)" + R"({"output":"out_1","input":"in_0","kind":"user"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsEmptyBindingName) { + // Skipping a blank name instead of refusing it shortens the recorded output + // list while the delegate's argument list keeps its length: here one aliased + // output would be recorded and one real output dropped, so the aliased + // binding would consume the argument belonging to the dropped one. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsBindingEntryWithNoNameKey) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsWrappingEngineExtent) { + // engine_size is a 64-bit field read straight from the file, so adding it to + // engine_offset before comparing against the blob length wraps: 4096 plus + // 2^64-4086 is 10, which is comfortably inside an 8 KiB blob. The pointer and + // that length are what TensorRTBackend hands deserializeCudaEngine. + const std::string metadata = R"({"io_bindings":[{"name":"x","is_input":true}]})"; + constexpr std::size_t kBlobSize = 8192; + constexpr uint32_t kMetadataOffset = HEADER_SIZE; + constexpr uint32_t kEngineOffset = 4096; + + std::vector blob(kBlobSize, 0); + std::memcpy(blob.data(), TENSORRT_MAGIC, sizeof(TENSORRT_MAGIC)); + write_field(blob, METADATA_OFFSET_FIELD_OFFSET, kMetadataOffset); + write_field(blob, METADATA_SIZE_FIELD_OFFSET, static_cast(metadata.size())); + write_field(blob, ENGINE_OFFSET_FIELD_OFFSET, kEngineOffset); + write_field(blob, ENGINE_SIZE_FIELD_OFFSET, ~uint64_t{0} - (kEngineOffset - 11)); + std::memcpy(blob.data() + kMetadataOffset, metadata.data(), metadata.size()); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsEngineOffsetPastEndOfBlob) { + // The offset alone is out of range. The subtraction form has to refuse that + // before it evaluates size - engine_offset, which would itself wrap. + auto blob = make_blob(R"({"io_bindings":[{"name":"x","is_input":true}]})"); + const auto past_end = static_cast(align_up(blob.size() + ENGINE_ALIGNMENT, ENGINE_ALIGNMENT)); + write_field(blob, ENGINE_OFFSET_FIELD_OFFSET, past_end); + write_field(blob, ENGINE_SIZE_FIELD_OFFSET, uint64_t{0}); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + TEST(ExecuTorchTensorRTBlobHeader, RejectsUnknownFutureMagic) { constexpr char kFutureMagic[4] = {'T', 'R', '0', '3'}; const std::string metadata = R"({"io_bindings":[{"name":"x","is_input":true}]})"; diff --git a/tests/py/dynamo/executorch/test_backend.py b/tests/py/dynamo/executorch/test_backend.py index fe25896a12c..b20c09cc1ba 100644 --- a/tests/py/dynamo/executorch/test_backend.py +++ b/tests/py/dynamo/executorch/test_backend.py @@ -611,3 +611,87 @@ def test_preprocess_rejects_a_non_buffer_alias_elided_alongside_a_buffer_alias() ) with pytest.raises(ValueError, match="engine output indices"): TensorRTBackend.preprocess(edge_program, [spec]) + + +@pytest.mark.unit +def test_preprocess_rejects_a_delegate_with_no_outputs_at_all(): + """Naming every binding elidable leaves an empty expected index list, which + an empty delegate output list matches. That is the zero-output delegate the + rewiring pass raises about -- nothing reads it, so a later graph-wide + dead-code elimination can erase the engine with it -- and the rewiring's + guard runs before partitioning, so this is the only check left after + lowering.""" + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + TensorRTBackend, + _serialize_elided_output_names, + ) + + edge_program = _aliased_edge_program( + present_indices=[], + out_names=["out_k", "out_v"], + aliased_io_map={ + "out_k": ("tokens", "kv_cache_update"), + "out_v": ("tokens", "kv_cache_update"), + }, + ) + spec = CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(["out_k", "out_v"]), + ) + with pytest.raises(ValueError, match="no outputs at all"): + TensorRTBackend.preprocess(edge_program, [spec]) + + +@pytest.mark.unit +@pytest.mark.parametrize( + "value", + [1, True, None, b"", b"kv0", b'"kv0"', b"[1, 2", b"\xff\xfe"], + ids=[ + "int", + "bool", + "none", + "empty-bytes", + "bare-name", + "json-string", + "truncated-json", + "invalid-utf8", + ], +) +def test_elided_output_names_refuses_a_value_it_cannot_read(value): + """A hand-built spec may carry anything, and every shape has to name the key. + + Two of these are the interesting ones. ``1`` is the obvious value for a key + that reads like a flag, and without a type check it surfaces as an + unattributed ``TypeError``. ``'"kv0"'`` is the quiet one: a JSON *string* + decodes into a set of its own characters, so nothing raises, the real binding + name is not exempted and every one-character binding name is. + """ + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _elided_output_names, + ) + + specs = [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, value)] + + with pytest.raises(ValueError, match=ZERO_COPY_KV_COMPILE_SPEC_KEY): + _elided_output_names(specs) + + +@pytest.mark.unit +def test_elided_output_names_reads_the_list_the_partitioner_writes(): + """The control for the refusals above, so they cannot pass vacuously.""" + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _elided_output_names, + _serialize_elided_output_names, + ) + + specs = [ + CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, _serialize_elided_output_names(["kv0"]) + ) + ] + + assert _elided_output_names(specs) == {"kv0"} + assert _elided_output_names([]) is None diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py index 219e435f5bb..8c42437dfa8 100644 --- a/tests/py/dynamo/executorch/test_export.py +++ b/tests/py/dynamo/executorch/test_export.py @@ -516,12 +516,16 @@ def test_export_returns_edge_and_forwards_all_options(monkeypatch): assert compile_specs == [compile_spec] -def _patch_declare(monkeypatch): +def _patch_declare(monkeypatch, log=None): """Record the programs the declaration pass sees, and hand each one back. export() imports the symbol inside its own body, so the patch has to land on the module that owns it. The stub takes ``**kw`` because export() passes ``copyback_buffers=``. + + ``log`` is a shared, tagged call log. Two separate recorders would each be + satisfied by their own calls whichever order the two passes ran in, and the + order is the thing under test. """ import torch_tensorrt.dynamo._exporter as dynamo_exporter @@ -529,6 +533,8 @@ def _patch_declare(monkeypatch): def _declare(program, **kw): seen.append(program) + if log is not None: + log.append(("declare", program)) return program monkeypatch.setattr( @@ -537,13 +543,15 @@ def _declare(program, **kw): return seen -def _patch_rewire(monkeypatch, elided_names=("kv",)): +def _patch_rewire(monkeypatch, elided_names=("kv",), log=None): import torch_tensorrt.executorch._zero_copy as zero_copy seen = [] def _rewire(program): seen.append(program) + if log is not None: + log.append(("rewire", program)) return list(elided_names) monkeypatch.setattr(zero_copy, "rewire_aliased_mutations_to_buffers", _rewire) @@ -558,8 +566,9 @@ def test_export_zero_copy_kv_rewires_every_method(monkeypatch): mutations that declaration produced. """ export_module, lower = _patch_lowering(monkeypatch) - declared = _patch_declare(monkeypatch) - rewired = _patch_rewire(monkeypatch) + calls = [] + declared = _patch_declare(monkeypatch, log=calls) + rewired = _patch_rewire(monkeypatch, log=calls) prefill = FakeExportedProgram() decode = FakeExportedProgram() @@ -571,20 +580,33 @@ def test_export_zero_copy_kv_rewires_every_method(monkeypatch): assert declared == [prefill, decode] assert rewired == [prefill, decode] + # One log, so the order between the two passes is pinned and not just the + # order within each. Rewiring works from the mutations declaration produced, + # so running it first would find nothing to rewire. + assert calls == [ + ("declare", prefill), + ("declare", decode), + ("rewire", prefill), + ("rewire", decode), + ] # The backend rejects a delegate missing its aliased outputs unless it is - # told the omission was deliberate, and this spec is the only channel. - from torch_tensorrt.executorch.backend import ( - ZERO_COPY_KV_COMPILE_SPEC_KEY, - _elided_output_names, - ) + # told the omission was deliberate, and the presence of this key on the + # partitioner is the only channel that says so. The partitioner drops the + # spec itself and re-derives the names per engine, so its value is not read. + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + from torch_tensorrt.executorch.partitioner import TensorRTPartitioner for pipeline in lower.call_args.kwargs["partitioner"].values(): - assert any( + specs = pipeline[0].compile_specs + assert any(spec.key == ZERO_COPY_KV_COMPILE_SPEC_KEY for spec in specs) + # The partitioner here is a stub, so what the real one makes of these + # specs is asserted against the real one. + real = TensorRTPartitioner(compile_specs=specs) + assert real._zero_copy_requested + assert not any( spec.key == ZERO_COPY_KV_COMPILE_SPEC_KEY - for spec in pipeline[0].compile_specs + for spec in real._base_compile_specs ) - # The spec carries the elided binding names, not a bare flag. - assert _elided_output_names(pipeline[0].compile_specs) == {"kv"} @pytest.mark.unit diff --git a/tests/py/dynamo/executorch/test_weight_streaming_budget.py b/tests/py/dynamo/executorch/test_weight_streaming_budget.py index cce7194f3c5..42b10ea97d8 100644 --- a/tests/py/dynamo/executorch/test_weight_streaming_budget.py +++ b/tests/py/dynamo/executorch/test_weight_streaming_budget.py @@ -277,7 +277,7 @@ def test_save_rejects_negative_budget(tmp_path): @pytest.mark.unit def test_save_rejects_unknown_executorch_kwarg(tmp_path): - with pytest.raises(TypeError, match="unexpected keyword argument"): + with pytest.raises(TypeError, match="unexpected keyword argument") as excinfo: save( torch.nn.Linear(1, 1), str(tmp_path / "model.pte"), @@ -285,6 +285,15 @@ def test_save_rejects_unknown_executorch_kwarg(tmp_path): weight_streaming_budget_per_enginet=4096, ) + # The message spells the supported set out, so someone who mistyped an option + # is told what to type instead. A hand-written list drifts the moment an + # option is added, and then tells them the flag they wanted is unsupported. + from torch_tensorrt._compile import _EXECUTORCH_SAVE_OPTIONS + + assert len(_EXECUTORCH_SAVE_OPTIONS) > 1 + for name in _EXECUTORCH_SAVE_OPTIONS: + assert repr(name) in str(excinfo.value) + @pytest.mark.unit def test_save_warns_when_budget_used_with_non_executorch_format(tmp_path, caplog): diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index e84c750621c..b7aefbeba2f 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -16,6 +16,7 @@ import json import operator +import re from types import SimpleNamespace import pytest @@ -175,6 +176,35 @@ def test_rewire_leaves_mutations_the_engine_does_not_alias(monkeypatch, mutation assert "_torch_tensorrt_aliased_buffer" not in k_buffer.meta +@pytest.mark.unit +@pytest.mark.parametrize("unresolvable", ["unknown-input", "index-past-the-args"]) +def test_rewire_skips_an_alias_whose_input_does_not_resolve(monkeypatch, unresolvable): + """An aliased_io entry naming an input this delegate does not take is skipped. + + Two ways it can fail to resolve: the name is not one of the engine's input + bindings at all, or it is but its index is past the end of the delegate's + argument list. Neither leaves a mutation that could be rewired, and + ``_declare_aliased_kv_mutations_on_ep`` has already warned about both for the + same engine, so both are skipped rather than reported here. + """ + if unresolvable == "unknown-input": + aliased_input, input_names = "not_an_input", ["k_in", "tokens"] + else: + # A real binding name, but the third one, while the engine node takes two + # arguments -- so the index is past the end of the argument list. + aliased_input, input_names = "spare", ["k_in", "tokens", "spare"] + program, k_buffer, _ = _kv_program() + _patch_engine_metadata( + monkeypatch, + aliased_io={"out_k": (aliased_input, "kv_cache_update")}, + input_names=input_names, + output_names=["logits", "out_k"], + ) + + assert Z.rewire_aliased_mutations_to_buffers(program) == [] + assert "_torch_tensorrt_aliased_buffer" not in k_buffer.meta + + def _mixed_program(): """One engine, one method, both kinds of mutation at once. @@ -499,17 +529,27 @@ def test_unstage_accepts_a_buffer_that_never_had_a_staging_copy(compile_specs): @pytest.mark.unit -def test_unstage_refuses_a_direct_buffer_whose_spec_stays_on_the_host(): +@pytest.mark.parametrize("spec", ["absent", "host"]) +def test_unstage_refuses_a_direct_buffer_that_is_not_on_the_device(spec): """Reaching the delegate directly is not enough; the buffer has to be there. A marked buffer whose own spec asks for the host is planned in a host arena, and the engine cannot write a host pointer in place. Accepting it on the strength of the mark and the delegate edge alone writes a ``.pte`` whose - every ``execute()`` fails on the alias-target guard. + every ``execute()`` fails on the alias-target guard. A buffer with no spec at + all is the other half of the same refusal, and means the pass is running + somewhere the specs do not exist yet. """ - graph_module, _, _ = _direct_delegate_graph(device=DeviceType.CPU) + graph_module, k_buffer, _ = _direct_delegate_graph( + device=DeviceType.CPU if spec == "host" else DeviceType.CUDA + ) + if spec == "absent": + del k_buffer.meta["spec"] + expected = ( + "carries no TensorSpec" if spec == "absent" else "its TensorSpec asks for" + ) - with pytest.raises(RuntimeError, match="not planned in device memory"): + with pytest.raises(RuntimeError, match=expected): Z._unstage_aliased_buffers(graph_module) @@ -750,12 +790,27 @@ def test_unstage_raises_when_the_staging_copy_is_not_on_cuda(): @pytest.mark.unit -def test_unstage_raises_when_the_staging_copy_has_no_spec(): +@pytest.mark.parametrize("missing", ["staging-copy", "buffer"]) +def test_unstage_raises_when_either_side_of_the_move_has_no_spec(missing): + """The move reads a spec on both nodes, and the message names the bare one. + + One condition covers both, so a message that always blamed the staging copy + would send someone whose copy has a spec to look at the wrong node. + """ graph_module, k_buffer, staged_k, _ = _staged_delegate_graph() k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - del staged_k.meta["spec"] + bare, kind = ( + (staged_k, "staging copy") + if missing == "staging-copy" + else (k_buffer, "buffer placeholder") + ) + del bare.meta["spec"] - with pytest.raises(RuntimeError, match="no TensorSpec"): + # Anchored on the node kind as well as the name: the buffer's name appears + # again later in the message, so a loose pattern would match either wording. + with pytest.raises( + RuntimeError, match=re.escape(f"no TensorSpec on the {kind} '{bare.name}'") + ): Z._unstage_aliased_buffers(graph_module) @@ -860,6 +915,50 @@ def test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus(): assert delegate_0.args[1] is staged_0 +@pytest.mark.unit +def test_unstage_refuses_to_move_a_buffer_a_second_consumer_stages_to_the_host(): + """The device-*type* half of ``_device_move_is_safe``'s spec comparison. + + Its sibling, the device index, is pinned by + ``test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus``. Here the second + consumer is another ``_h2d_copy`` that stays on the host, so the indices + agree and only the type comparison separates the two: deleting it accepts + this graph, leaving a copy that reads a buffer the move has put on the GPU. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_trt = graph.call_function(h2d, (k_buffer,)) + staged_host = graph.call_function(h2d, (k_buffer,)) + lowered_trt = graph.get_attr("lowered_module_0") + lowered_other = graph.get_attr("lowered_module_1") + delegate_trt = graph.call_function( + executorch_call_delegate, (lowered_trt, staged_trt) + ) + delegate_other = graph.call_function( + executorch_call_delegate, (lowered_other, staged_host) + ) + graph.output((k_buffer, delegate_trt, delegate_other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + staged_trt.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_host.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): + Z._unstage_aliased_buffers(graph_module) + + assert k_buffer.meta["spec"].device == DeviceType.CPU + assert delegate_trt.args[1] is staged_trt + + @pytest.mark.unit def test_unstage_refuses_a_buffer_another_backend_stages_to_a_different_gpu(): """A marked buffer staged to a TensorRT delegate on cuda:0 and to a @@ -993,23 +1092,51 @@ def _finalized_program(forward=None, **methods): CUDA_ARENA = 2 -def _planned(graph_module, *, arena=CUDA_ARENA, on_device=True): +HOST_ARENA = 1 + + +def _planned( + graph_module, + *, + arena=CUDA_ARENA, + on_device=True, + host_tensor_arena=None, + device_type=DeviceType.CUDA, +): """Add what memory planning leaves behind, which the checker reads. The graphs above are built for the passes that run before planning, so they carry no ``mem_id`` and the module records no arena devices. Planning assigns - both, and ``on_device=False`` is what it leaves when it puts everything in - the one host arena -- it records no device arenas at all then. + both. ``on_device=False`` drops the arena-device record, which is what + ``enable_non_cpu_memory_planning=False`` leaves -- and also what a + caller-supplied planner that does not go through ``apply_algo`` leaves, that + being the only thing that writes it. What separates those two is where the + program's *host* tensors ended up: host-only planning puts every tensor in + one bucket, so they share the buffer's arena, while a device-aware planner + keeps them apart. + ``host_tensor_arena`` is where the CPU-spec tensors go, and it is also put on + the output node's spec, which is where a real finalized program carries one. + Leaving it unset gives the two shapes above: a separate arena when the + program records devices, and the buffer's own when it does not. """ from executorch.exir.schema import NonConstBufferDevice + if host_tensor_arena is None: + host_tensor_arena = HOST_ARENA if on_device else arena for node in graph_module.graph.nodes: - if node.op == "placeholder" and node.meta.get("spec") is not None: - node.meta["spec"].mem_id = arena + spec = node.meta.get("spec") + if node.op == "placeholder" and spec is not None: + spec.mem_id = arena if spec.device == DeviceType.CUDA else host_tensor_arena + if node.op == "output": + node.meta["spec"] = [ + SimpleNamespace( + device=DeviceType.CPU, device_index=0, mem_id=host_tensor_arena + ) + ] if on_device: graph_module.meta["non_const_buffer_device"] = [ NonConstBufferDevice( - buffer_idx=arena, device_type=DeviceType.CUDA, device_index=0 + buffer_idx=arena, device_type=device_type, device_index=0 ) ] return graph_module @@ -1069,13 +1196,66 @@ def test_check_zero_copy_kv_rejects_a_direct_buffer_planned_on_the_host(): else refuses it -- the un-staging pass never ran -- and the engine is handed a host pointer it cannot write, so the ``.pte`` fails its first ``execute()``. The graph is the accepted one above with only the arena changed, so the - refusal can come from nothing else. + refusal can come from nothing else. The program records no arena devices -- + host-only planning writes none -- so what identifies the arena as the host's + is that the program's host tensors are in it too. + """ + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + + with pytest.raises(RuntimeError, match="also holds the program's host tensors"): + Z.check_zero_copy_kv( + _finalized_program( + _planned(graph_module, arena=HOST_ARENA, on_device=False) + ) + ) + + +@pytest.mark.unit +def test_check_zero_copy_kv_accepts_a_device_arena_a_custom_planner_did_not_record(): + """An absent arena-device record is "cannot tell", not "on the host". + + ``apply_algo`` is the only thing in ExecuTorch that writes + ``non_const_buffer_device``, and ``to_executorch`` takes any callable as + ``memory_planning_pass`` -- which the user guide tells people to supply for a + cache shared between prefill and decode. Reading the absent record as a host + placement refuses that program, and ``save`` runs this check with no way to + opt out, so nothing is written at all. This is the previous test's graph with + the host tensors moved out of the buffer's arena, which is the whole + difference between the two. + """ + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + + Z.check_zero_copy_kv( + _finalized_program( + _planned( + graph_module, + arena=CUDA_ARENA, + on_device=False, + host_tensor_arena=HOST_ARENA, + ) + ) + ) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_an_arena_the_program_records_as_non_cuda(): + """When the record *is* present it is read, and only its CUDA entries count. + + Nothing in ExecuTorch emits a CPU entry today -- the builder filters them -- + so this pins the filter against a program that carries one, hand-built or + from a future planner, rather than against the stock one. """ graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) - with pytest.raises(RuntimeError, match="not planned in any of the program's CUDA"): + with pytest.raises(RuntimeError, match="does not record as CUDA"): Z.check_zero_copy_kv( - _finalized_program(_planned(graph_module, arena=1, on_device=False)) + _finalized_program( + _planned( + graph_module, + device_type=DeviceType.CPU, + host_tensor_arena=HOST_ARENA, + ) + ) ) @@ -1215,6 +1395,55 @@ def test_zero_copy_backend_config_defaults_to_executorch_defaults(): assert config.emit_stacktrace == defaults.emit_stacktrace +@pytest.mark.unit +@pytest.mark.parametrize( + "skip", + [True, {"decode": True}, {"prefill": False, "decode": True}], + ids=["bool", "one-method", "one-of-two-methods"], +) +def test_zero_copy_backend_config_refuses_skip_h2d_for_method_inputs(skip): + """The one option that cannot be carried through, refused where it is set. + + ``skip_h2d_for_method_inputs`` is ExecuTorch's own un-staging of method + inputs and it demands each placeholder it un-stages have exactly one user. A + rewired cache has two -- the delegate, and the graph output it is its own + mutation result for -- so ``PropagateDevicePass`` raises on every zero-copy + graph. Preserving the option hands back a config that cannot finalize at + all, which is a failure a long way from the line that caused it. + """ + from executorch.exir import ExecutorchBackendConfig + from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig + + base = ExecutorchBackendConfig( + propagate_device_config=PropagateDeviceConfig(skip_h2d_for_method_inputs=skip) + ) + + with pytest.raises(ValueError, match="skip_h2d_for_method_inputs"): + Z.zero_copy_backend_config(base) + + +@pytest.mark.unit +@pytest.mark.parametrize("skip", [False, {"decode": False}], ids=["bool", "one-method"]) +def test_zero_copy_backend_config_carries_skip_h2d_left_off(skip): + """A field set and left off is carried through, bool or dict alike. + + The dict half is narrower than ``PropagateDevicePass``, which tests the + field for truth rather than resolving it per method and so reads any + non-empty dict as on. Such a config is carried here and still fails to + finalize. + """ + from executorch.exir import ExecutorchBackendConfig + from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig + + base = ExecutorchBackendConfig( + propagate_device_config=PropagateDeviceConfig(skip_h2d_for_method_inputs=skip) + ) + + config = Z.zero_copy_backend_config(base) + + assert config.propagate_device_config.skip_h2d_for_method_inputs == skip + + @pytest.mark.unit def test_unstage_pass_runs_the_inner_pass_after_unstaging(): """A caller's own to_out_var_pass has to survive being composed with.""" @@ -1638,12 +1867,15 @@ def test_aliased_buffer_mark_survives_real_lowering(generate_etrecord): ) ep = edge.exported_program() - marked = [ - node + marked = { + ep.graph_signature.inputs_to_buffers[node.name] for node in ep.graph_module.graph.nodes if node.op == "placeholder" and node.meta.get("_torch_tensorrt_aliased_buffer") - ] - assert marked, "no placeholder kept _torch_tensorrt_aliased_buffer through lowering" + } + # Both, by name. The model registers two caches, so asserting the list is + # non-empty would pass on partial marker loss -- which is the very shape the + # per-delegate count check elsewhere in this feature exists for. + assert marked == {"k_cache", "v_cache"} class _MixedDecodeStep(torch.nn.Module): @@ -1719,12 +1951,16 @@ def _assert_marked_buffers_reach_the_engine_unstaged(program): whether the rewiring happened, and getting it wrong is silent -- the engine writes per-call scratch that is discarded and the cache never updates. - The library's own check runs first, on the whole program. It is weaker than - what follows -- it accepts a marked buffer reaching any backend's delegate -- - but it is the only place the container API it reads, ``methods`` and - ``exported_program(name)``, meets a real ``ExecutorchProgramManager``: every - other test of it builds the program itself, so an upstream rename would - leave those green and break ``save(zero_copy_kv=True)`` for every caller. + The library's own check runs first, on the whole program. It is the stronger + of the two -- it also requires the delegate to carry the zero-copy compile + spec and the buffer to be planned somewhere the engine can write it, neither + of which the assertions below read -- and it is the only place the container + API it reads, ``methods`` and ``exported_program(name)``, meets a real + ``ExecutorchProgramManager``: every other test of it builds the program + itself, so an upstream rename would leave those green and break + ``save(zero_copy_kv=True)`` for every caller. What follows is kept because it + reads the graph rather than the marks, which is what says the rewiring + actually happened. """ torch_tensorrt.executorch.check_zero_copy_kv(program) graph_module = program.exported_program().graph_module @@ -1750,17 +1986,128 @@ def _assert_marked_buffers_reach_the_engine_unstaged(program): ) +class _StubProgram: + """The three attributes the reorder and upstream's write-back pass read. + + ``graph_signature`` is a property over ``_graph_signature`` because that is + how ``ExportedProgram`` exposes it, and the reorder replaces the signature by + assigning the private name. + """ + + def __init__(self, graph_module, signature): + self.graph_module = graph_module + self._graph_signature = signature + + @property + def graph_signature(self): + return self._graph_signature + + @property + def graph(self): + return self.graph_module.graph + + +def _two_mutation_program(first_value_is_inplace): + """A two-mutation program in the shape the write-back pass sees. + + Slot 0 mutates ``first``; its value is either an in-place op on its own + buffer -- which upstream reads as needing no copy -- or an ordinary + functional result, which does. Slot 1 is always an ordinary copy-back on + ``cb``. Nothing here carries the zero-copy mark, so a reorder keyed on that + mark cannot see slot 0 at all. + """ + from torch.export.exported_program import ExportGraphSignature + from torch.export.graph_signature import InputKind, InputSpec + + graph = torch.fx.Graph() + b_first = graph.placeholder("b_first") + b_cb = graph.placeholder("b_cb") + x = graph.placeholder("x") + first_value = ( + graph.call_function(torch.ops.aten.add_.Tensor, (b_first, x)) + if first_value_is_inplace + else graph.call_function(torch.ops.aten.add.Tensor, (b_first, x)) + ) + cb_value = graph.call_function(torch.ops.aten.add.Tensor, (b_cb, x)) + user = graph.call_function(torch.ops.aten.mul.Tensor, (x, x)) + graph.output((first_value, cb_value, user)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + signature = ExportGraphSignature( + input_specs=[ + InputSpec(InputKind.BUFFER, TensorArgument("b_first"), "first", False), + InputSpec(InputKind.BUFFER, TensorArgument("b_cb"), "cb", False), + InputSpec(InputKind.USER_INPUT, TensorArgument("x"), None), + ], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(first_value.name), "first" + ), + OutputSpec(OutputKind.BUFFER_MUTATION, TensorArgument(cb_value.name), "cb"), + OutputSpec(OutputKind.USER_OUTPUT, TensorArgument(user.name), None), + ], + ) + return _StubProgram(graph_module, signature) + + +@pytest.mark.unit +@pytest.mark.parametrize("reorder", [False, True], ids=["without", "with"]) +def test_reorder_moves_an_inplace_mutation_this_feature_did_not_create(reorder): + """The reorder keys on upstream's predicate, not on this feature's own mark. + + ``run_reinplace_pass`` and ``reinplace_extra_ops`` are ``ExecutorchBackendConfig`` + fields whose pass runs just before the write-back and rewrites an ordinary + mutation into an in-place one. Upstream then inserts no copy for it, exactly + as for a rewired cache -- and a reorder that asks "did zero-copy rewire + this?" instead of "will upstream copy this?" leaves the resulting pair + crossed while reporting that it moved nothing. + """ + from executorch.exir.passes.insert_write_back_for_buffers_pass import ( + insert_write_back_for_buffers_pass, + ) + + program = _two_mutation_program(first_value_is_inplace=True) + moved = Z.order_copyback_mutations_first(program) if reorder else 0 + _, signature = insert_write_back_for_buffers_pass(program) + value_of = {buffer: value for value, buffer in signature.buffers_to_mutate.items()} + + if not reorder: + # Pin the defect too, so the assertions below cannot pass vacuously. + assert value_of["first"].startswith("copy_") + return + assert moved == 2 + assert not value_of["first"].startswith("copy_"), ( + "'first' is mutated in place, so upstream inserts no copy for it and its " + f"finalized value must not be one; got {value_of['first']!r}" + ) + assert value_of["cb"].startswith("copy_"), ( + "'cb' is copied back, so its finalized value is the copy upstream " + f"inserted; got {value_of['cb']!r}" + ) + + +@pytest.mark.unit +def test_reorder_leaves_an_already_correct_order_alone(): + """Two copy-back mutations need no move, and the function says so.""" + program = _two_mutation_program(first_value_is_inplace=False) + before = [spec.arg.name for spec in program.graph_signature.output_specs] + + assert Z.order_copyback_mutations_first(program) == 0 + assert [spec.arg.name for spec in program.graph_signature.output_specs] == before + + def _assert_each_mutation_names_its_own_value(program): """The finalized signature pairs every mutated buffer with its own new value. ExecuTorch finalizes the mutations by inserting a copy for each one whose - value is not already its buffer placeholder, moving those copies to the front - of the output tuple, and then reassigning the mutation specs' arguments by - position. Rewiring a cache to its own placeholder takes it out of that - leading run, so a method that mixes the two kinds comes out of finalization - with each buffer named against another buffer's value unless the mutations - were declared in the order the pass assumes. Nothing inside the finalizer - reads the pairing, so the ``.pte`` is written either way -- what reads it is + value is not already reached from a buffer placeholder through in-place ops, + moving those copies to the front of the output tuple, and then reassigning + the mutation specs' arguments by position. Rewiring a cache to its own + placeholder takes it out of that leading run, and so does any other + mutation upstream reads as in-place, so a method that mixes the two kinds + comes out of finalization with each buffer named against another buffer's + value unless the mutations were declared in the order the pass assumes. + Nothing inside the finalizer reads the pairing, so the ``.pte`` is written + either way -- what reads it is anyone inspecting the program, and the eager call path that copies mutated values back into the state dict in this order. """ @@ -1856,8 +2203,16 @@ def test_zero_copy_kv_keeps_a_copyback_buffer_in_the_same_method(retrace): # Everything above is the export half. The staging the other half removes does # not exist until PropagateDevicePass runs inside to_executorch, so this is the # earliest point at which the caches can be seen reaching the engine directly. + # Composed onto a caller's own config, the optional form the user guide + # describes. It is the only shape under which a preserved field can make the + # returned config unfinalizable; a no-argument zero_copy_backend_config() + # starts from the defaults and so has nothing to preserve. + from executorch.exir import ExecutorchBackendConfig + program = edge.to_executorch( - config=torch_tensorrt.executorch.zero_copy_backend_config() + config=torch_tensorrt.executorch.zero_copy_backend_config( + ExecutorchBackendConfig(extract_delegate_segments=False) + ) ) _assert_marked_buffers_reach_the_engine_unstaged(program) _assert_each_mutation_names_its_own_value(program) @@ -2009,7 +2364,8 @@ def test_zero_copy_kv_with_the_copyback_on_a_second_delegate(retrace): @pytest.mark.skipif( not torch.cuda.is_available(), reason="requires CUDA + TensorRT for a real engine" ) -def test_zero_copy_kv_beside_an_executorch_cuda_delegate(): +@pytest.mark.parametrize("retrace", [False, True], ids=["legacy", "retrace"]) +def test_zero_copy_kv_beside_an_executorch_cuda_delegate(retrace): """An aliased KV cache in a method that also holds an ExecuTorch CUDA delegate. ``erfinv`` has no TensorRT converter, so with a ``CudaPartitioner`` catch-all @@ -2053,7 +2409,7 @@ def forward(self, tokens, input_pos): edge = torch_tensorrt.executorch.export( trt_gm, arg_inputs=(tokens, input_pos), - retrace=False, + retrace=retrace, zero_copy_kv=True, partitioners=[ cuda_partitioner.CudaPartitioner( From 36e286af914ce5ab866c90f666575dcf4ef30286 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Mon, 7 Sep 2026 15:15:43 -0700 Subject: [PATCH 18/22] fix(executorch): close three holes the zero-copy guards still had `zero_copy_backend_config` refused `skip_h2d_for_method_inputs` only where its value was true, and only in a single `PropagateDeviceConfig`. Both are narrower than the pass they guard. `PropagateDevicePass` is handed the field whole and only tests it for truth (in `_insert_h2d_copies`), so it reads any non-empty dict as on for every method; and `propagate_device_config` is itself typed as one `PropagateDeviceConfig` or a dict of them keyed by method, with ExecuTorch handing the pass the entry for the method being finalized. Measured on the real mixed-KV model, both `PropagateDeviceConfig(skip_h2d_for_method_inputs={"decode": False})` and `{"forward": PropagateDeviceConfig(skip_h2d_for_method_inputs=True)}` were accepted here and then died inside `to_executorch` with `requires placeholder 'b_k_cache' to have exactly one user, but it has 2 users`, which is the failure this refusal exists to prevent. It now refuses the option wherever it is written, and on every value that pass reads as on rather than only on `True`. `False` and `{}` are still carried, because those are exactly what that pass reads as off. `check_zero_copy_kv` read the stamped delegates as one set. A method that lowers to two of them can hand both caches to one and leave the other reading a staging copy of the cache whose copy-back export had already removed: every marked buffer does reach a stamped delegate, so nothing refused it, and that engine's write was discarded. Each stamped delegate is now counted against its own compile spec -- one marked buffer per aliased output the spec says it elided, falling back to demanding at least one when the spec names none -- which is the cross-check `_unstage_aliased_buffers` already made before planning, now made again after it. This is the function `save()` and the user guide tell people to rely on, so it is where the last stand-in had to be closed. What neither count can separate is an exact swap, since the mark is a bare flag naming no engine; that is stated in the docstring and under Known gaps. `zero_copy_backend_config` passed `unstage_aliased_buffers_pass` a `device_memory_planning` that the next line then overrode for every call, by binding `finalization_config` to the config it returns. Proved dead by poisoning it to `False` and re-running the suite unchanged. The argument is gone from that call site; the parameter stays on the builder, where a pass constructed by hand and left unbound does read it. What that binding does not cover is documented rather than fixed. The planning flag is a bool copied by value and the pass an object copied by reference, so a config derived from the returned one with `dataclasses.replace` carries a pass still reading the original, and turning planning off on the derived config alone is not refused there. Making the pass follow the derivation was considered and rejected: nothing on the config points back at the pass, and nothing on the pass can see which config `to_executorch` is finalizing, so the only complete fixes are returning a dynamically built subclass of the caller's config class -- which changes the returned type and breaks pickling -- or moving the guard behind a wrapped `memory_planning_pass`, which relocates a tested refusal to a later point. Both are out of proportion to a defect whose measured outcome is already a later refusal: calling `zero_copy_backend_config` again on the derived config is the remedy -- the pass it builds is bound to that one -- and `check_zero_copy_kv` refuses the program the mistake produces for any method holding a host tensor to give the shared arena away. Measured, documented, and pinned by a test. Five tests: the all-`False` dict, the per-method config dict, a crossed pair of zero-copy delegates against the matched pair that is accepted, a stamped delegate whose spec names no aliased output, and the derived-config remedy. Each new clause was mutated away and kills only the tests belonging to it. --- .../runtime_performance/saving_models.rst | 7 +- py/torch_tensorrt/executorch/_zero_copy.py | 211 ++++++++++++------ .../py/dynamo/executorch/test_zero_copy_kv.py | 172 ++++++++++++-- 3 files changed, 313 insertions(+), 77 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index e8bb910494f..e34055c2f8f 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -394,7 +394,12 @@ fails on a host pointer. And ``propagate_device_config.skip_h2d_for_method_input is refused outright: ``PropagateDevicePass`` refuses to un-stage a method input whose placeholder does not have exactly one user, and a zero-copy cache always has two, so preserving that option would hand back a configuration that cannot -finalize at all. +finalize at all. It is refused wherever it is written -- in one +``PropagateDeviceConfig`` or in a per-method dict of them -- and on every value +that pass reads as on rather than only on ``True``, because it tests the field +for truth without ever resolving it per method, so even a dict of ``False`` is +on for every method. ``False`` and the empty dict are what it reads as off, and +those are carried unchanged. It is opt-in rather than automatic because the resulting ``.pte`` needs a runtime that understands a delegate whose aliased outputs are elided. Producing diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index c83b7e4c9b8..f4924e26fef 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -471,18 +471,19 @@ def _delegate_elided_output_names( delegate's own engine, and an aliased output is elided exactly when a buffer mutation writes it in place, so the size of this set is how many marked buffers the delegate has to take. Empty when the delegate carries no such - spec, and also when it carries one whose value does not decode into a list of - names -- a spec built by hand rather than by the partitioner, which always - writes the JSON list. Callers read empty as "cannot tell", not as "none". + spec, when it carries one whose value does not decode into a list of names, + and when that value decodes into an empty list -- the last two are shapes a + spec built by hand produces, since the partitioner writes one JSON name per + aliased output it elided. Callers read empty as "cannot tell", not as "none". ``backend._elided_output_names`` reads the same key on the same spec and takes the same shapes of value. It differs in what it does with the rest: it raises naming the key, because an undecodable spec leaves it unable to say - which outputs the delegate was allowed to drop, while here it only weakens a - cross-check that then falls back to demanding at least one buffer. The one - value the two read differently is bytes that are not valid UTF-8: it - replaces the bad units and reads a name out of them, ``json.loads`` refuses - them here, and here that is another "cannot tell". + which outputs the delegate was allowed to drop, while here it only weakens + the two cross-checks that read it, both of which then fall back to demanding + at least one buffer. The one value the two read differently is bytes that + are not valid UTF-8: it replaces the bad units and reads a name out of them, + ``json.loads`` refuses them here, and here that is another "cannot tell". """ spec = _zero_copy_compile_spec(graph_module, node) if spec is None: @@ -511,8 +512,9 @@ def _delegate_declares_zero_copy( demanding an aliased buffer from a delegate that never had one. A delegate that declares it but ends up taking fewer marked buffers than it elided aliased outputs has lost a KV update -- the mark that would have driven the - un-staging did not survive to this pass -- and is caught in - :func:`_unstage_aliased_buffers`. + un-staging did not survive, or another delegate took the buffer -- and is + caught in :func:`_unstage_aliased_buffers` before planning and again in + :func:`check_zero_copy_kv` after it. """ return _zero_copy_compile_spec(graph_module, node) is not None @@ -781,8 +783,9 @@ def _unstage_aliased_buffers( # One marked buffer per aliased output the spec says this delegate # elided. Demanding only one would accept a delegate that lost all but # one of its marks, whose remaining caches are still wired through their - # staging copies. A spec whose value does not decode names cannot say how - # many to expect, so it falls back to demanding at least one. + # staging copies. A spec that names none -- listing none, or not + # decoding -- cannot say how many to expect, so it falls back to + # demanding at least one. elided = _delegate_elided_output_names(graph_module, delegate) satisfied = satisfied_per_delegate[delegate] if satisfied >= max(len(elided), 1): @@ -835,15 +838,17 @@ def unstage_aliased_buffers_pass( program will be finalized with. Nothing in the graph records it, and it is half of what decides whether a marked buffer ends up somewhere the engine can write, so the pass has to be told: see :func:`_unstage_aliased_buffers`. - It is only the fallback. Set ``finalization_config`` on the returned pass to - the ``ExecutorchBackendConfig`` the program will be finalized with and the - flag is read off that config when the pass runs, which is the value memory + It is what a pass built here and left unbound uses. Set + ``finalization_config`` on the returned pass to the + ``ExecutorchBackendConfig`` the program will be finalized with and the flag + is read off that config instead, on every call, which is the value memory planning will use a few passes later. ``ExecutorchBackendConfig`` is a plain mutable dataclass, so the field can be - set again on the very config being finalized after the pass has captured its - value; a pass still reading the captured copy would then accept a program the - finalizer plans into the host arena. :func:`zero_copy_backend_config` sets - the attribute for that reason. + set again on the very config being finalized after the pass was built; a pass + reading a value captured here would then accept a program the finalizer plans + into the host arena. :func:`zero_copy_backend_config` binds the attribute for + that reason, and passes no ``device_memory_planning`` at all, since the + binding would override it on every call. """ from executorch.exir import ExecutorchBackendConfig from executorch.exir.pass_base import PassBase @@ -965,13 +970,17 @@ def check_zero_copy_kv(program: Any) -> None: error nor the speedup. The wrong-output case is the one below: a rewiring that did happen and then lost its mark. - Three shapes are refused: a marked buffer that is not a direct argument of a - TensorRT delegate carrying the zero-copy compile spec, a marked buffer that - reaches such a delegate directly but is not planned in device memory, and a - program with no marked buffer in any method. The first two are what - finalizing without :func:`zero_copy_backend_config` leaves behind. That - config's pass refuses both earlier, off the configuration and the graph; this - reads the placement memory planning went on to choose. + Four shapes are refused: a marked buffer that is not a direct argument of a + TensorRT delegate carrying the zero-copy compile spec, a stamped delegate + that takes fewer marked buffers than its own spec says it elided aliased + outputs, a marked buffer that reaches such a delegate directly but is not + planned in device memory, and a program with no marked buffer in any method. + The first three are what finalizing without :func:`zero_copy_backend_config` + leaves behind. That config's pass gets to all three earlier, off the + configuration and the graph: it removes the staging copy the first two come + from, and refuses outright when there is none to remove or when the + configuration plans nothing onto a device. What it cannot see is the arena + planning then chose, which is what this reads. The spec is what narrows the delegates that count. The mark is put on a buffer because one TensorRT engine writes it in place, and only a delegate @@ -979,11 +988,19 @@ def check_zero_copy_kv(program: Any) -> None: delegate taking the buffer says nothing about whether the engine did, and neither does an unrelated TensorRT engine that happens to read it. Either would stand in for the delegate whose write was elided while that one still - reads a staging copy whose contents are discarded. Stamped delegates are read - as one set rather than matched to the buffer each of them elided, so in a - method holding more than one they can still stand in for each other here; - :func:`_unstage_aliased_buffers`, which counts each delegate's own buffers - against its own spec, is what separates them. + reads a staging copy whose contents are discarded. Being stamped is not + enough on its own either: in a method that lowers to two of them, one + stamped delegate holding both caches leaves every marked buffer reaching + *some* stamped delegate while the other's write is still thrown away. So + each stamped delegate is also counted against its own spec, the same + cross-check :func:`_unstage_aliased_buffers` makes, against the same spec and + with the same fallback: a spec that names none -- listing none, or not + decoding -- cannot say how many to expect, so it demands at least one. What + that count cannot separate is an exact swap -- two stamped delegates each + taking one marked buffer, each the other's. The mark records only that some + engine writes the buffer in place, never which one, and the spec lists + engine output binding names rather than buffers, so there is nothing left to + match on; that pass has the same blind spot for the same reason. Placement is read from where memory planning actually put the buffer -- the ``mem_id`` on its ``TensorSpec`` -- rather than from the spec's own device, @@ -1016,6 +1033,7 @@ def check_zero_copy_kv(program: Any) -> None: """ method_names = sorted(program.methods) staged_by_method: Dict[str, List[str]] = {} + short_by_method: Dict[str, List[str]] = {} host_planned_by_method: Dict[str, List[str]] = {} marked_anywhere = False for method_name in method_names: @@ -1029,16 +1047,42 @@ def check_zero_copy_kv(program: Any) -> None: if not marked: continue marked_anywhere = True - zero_copy_delegate_args = { - arg + zero_copy_delegates = [ + node for node in graph_module.graph.nodes if _is_tensorrt_delegate(graph_module, node) and _delegate_declares_zero_copy(graph_module, node) - for arg in node.args[1:] + ] + zero_copy_delegate_args = { + arg for node in zero_copy_delegates for arg in node.args[1:] } staged = [node.name for node in marked if node not in zero_copy_delegate_args] if staged: staged_by_method[method_name] = staged + marked_nodes = set(marked) + short = [] + for delegate in zero_copy_delegates: + elided = _delegate_elided_output_names(graph_module, delegate) + taken = sum( + 1 + for arg in delegate.args[1:] + if isinstance(arg, Node) and arg in marked_nodes + ) + if taken >= max(len(elided), 1): + continue + if elided: + short.append( + f"'{delegate.name}' takes {taken} marked buffer(s), not the " + f"{len(elided)} its elided aliased output(s) {sorted(elided)} " + "imply" + ) + else: + short.append( + f"'{delegate.name}' names no aliased output this can count, so " + "it must take at least one marked buffer and takes none" + ) + if short: + short_by_method[method_name] = short device_arenas = _device_planned_arenas(graph_module) host_arenas = _host_planned_arenas(graph_module) host_planned = [ @@ -1066,6 +1110,21 @@ def check_zero_copy_kv(program: Any) -> None: "their copy-back, so nothing else would restore it. Finalize with " "torch_tensorrt.executorch.zero_copy_backend_config()." ) + if short_by_method: + raise RuntimeError( + "TensorRT zero-copy KV: delegate(s) " + + "; ".join( + f"{detail} in method '{method}'" + for method, details in short_by_method.items() + for detail in details + ) + + ". Every marked buffer does reach a delegate declaring zero-copy " + "KV, so either another one in the same method is holding this " + "engine's cache or a mark was lost; either way this engine still " + "reads a staging copy that is discarded, and export removed the " + "copy-back. Export this method without zero_copy_kv, or keep each " + "aliased buffer on the delegate whose engine elided it." + ) if host_planned_by_method: raise RuntimeError( f"TensorRT zero-copy KV: buffer(s) " @@ -1085,19 +1144,31 @@ def check_zero_copy_kv(program: Any) -> None: def _refuse_skip_h2d(config: "ExecutorchBackendConfig") -> None: """Raise if the config asks ExecuTorch to un-stage method inputs as well. - The field is typed as a bool or a per-method dict, and every truthy entry is - refused, naming the method that asked. That is narrower than the pass it - guards: ``PropagateDevicePass`` takes the field as one value and only tests - it for truth, so it reads *any* non-empty dict as on -- including one whose - entries are all ``False``, which this carries through and which then fails - to finalize a layer down. + Both places the option can be written are read. ``propagate_device_config`` + is one ``PropagateDeviceConfig`` or a dict of them keyed by method, and + within either, ``skip_h2d_for_method_inputs`` is a bool or a second + per-method dict. + + What is refused is every value that pass reads as on, which is every truthy + one rather than only ``True``. ``PropagateDevicePass`` is handed the field + whole and only tests it for truth (``propagate_device_pass.py:216``), never + resolving it per method, so it reads any non-empty dict as on for every + method -- one whose entries are all ``False`` included. Refusing only the + true entries would carry such a dict through and hand back a config that + raises a layer down, which is the failure this exists to prevent. ``False`` + and the empty dict are what that pass reads as off, and both are carried. """ propagate = getattr(config, "propagate_device_config", None) - skip = getattr(propagate, "skip_h2d_for_method_inputs", False) - if isinstance(skip, dict): - asked_for = [name for name, value in skip.items() if value] - else: - asked_for = ["every method"] if skip else [] + per_method = ( + sorted(propagate.items()) + if isinstance(propagate, dict) + else [("every method", propagate)] + ) + asked_for = [ + method + for method, entry in per_method + if getattr(entry, "skip_h2d_for_method_inputs", False) + ] if not asked_for: return raise ValueError( @@ -1107,9 +1178,12 @@ def _refuse_skip_h2d(config: "ExecutorchBackendConfig") -> None: "method input whose placeholder does not have exactly one user, and a " "buffer zero-copy rewired has two -- the TensorRT delegate, and the graph " "output it is its own mutation result for -- so finalization raises " - "there. Zero-copy already un-stages the aliased buffers; leave " - "skip_h2d_for_method_inputs off and pass the method's own inputs on the " - "host." + "there. Every value that pass reads as on is refused, not only True: it " + "tests skip_h2d_for_method_inputs for truth without ever resolving it " + "per method, so setting it to a dict whose entries are all False still " + "turns it on for every method. Zero-copy already un-stages the aliased " + "buffers; leave skip_h2d_for_method_inputs at False or unset -- both of " + "which are carried -- and pass the method's own inputs on the host." ) @@ -1136,17 +1210,29 @@ def zero_copy_backend_config( planned in device memory, so ``False`` -- which plans every tensor into the one host arena -- has the pass refuse each cache it finds rather than write a ``.pte`` whose every ``execute()`` fails. It is read off the - config returned here, at the moment the pass runs, so setting the field on - that config afterwards is honoured: the pass and the finalizer then cannot - disagree about it. - * ``propagate_device_config.skip_h2d_for_method_inputs`` is *refused*, for - any method. It is ExecuTorch's own un-staging of method inputs, and it - requires each placeholder it un-stages to have exactly one user. A rewired - cache always has two -- the delegate, and the graph output it is its own - mutation result for -- so ``PropagateDevicePass`` raises on every - zero-copy graph. Returning the option unchanged would hand back a config - that cannot finalize at all; this says so here instead, where the caller - can act on it. + config returned here, at the moment the pass runs, so setting the field + *on that object* afterwards is honoured: the pass and the finalizer then + cannot disagree about it. Building a *new* config out of this one with + ``dataclasses.replace`` is the case that does not carry -- the field is a + bool, copied by value, while the pass is copied by reference and goes on + reading the config returned here -- so call this function again on the + derived config and the pass it builds reads that one. Not noticing is + mostly not silent either: :func:`check_zero_copy_kv` reads the arena + memory planning actually chose and refuses the program that mistake + produces, except in a method holding no host tensor to give the shared + arena away. + * ``propagate_device_config.skip_h2d_for_method_inputs`` is *refused*, + wherever it is written -- in the single ``PropagateDeviceConfig`` or in a + per-method dict of them -- and on every value ``PropagateDevicePass`` + reads as on rather than only on ``True``: it tests the field for truth + without ever resolving it per method, so any non-empty dict is on for + every method, one of ``False`` included. It is ExecuTorch's own un-staging + of method inputs, and it requires each placeholder it un-stages to have + exactly one user. A rewired cache always has two -- the delegate, and the + graph output it is its own mutation result for -- so + ``PropagateDevicePass`` raises on every zero-copy graph. Returning the + option unchanged would hand back a config that cannot finalize at all; + this says so here instead, where the caller can act on it. .. warning:: Finalizing a ``zero_copy_kv=True`` program *without* this config does @@ -1170,10 +1256,9 @@ def zero_copy_backend_config( base = config if config is not None else ExecutorchBackendConfig() _refuse_skip_h2d(base) - unstage = unstage_aliased_buffers_pass( - base.to_out_var_pass, - device_memory_planning=base.enable_non_cpu_memory_planning, - ) + # No device_memory_planning here: setting finalization_config below overrides + # it for every call, so a value passed would never be read. + unstage = unstage_aliased_buffers_pass(base.to_out_var_pass) wrapped = replace(base, to_out_var_pass=unstage) unstage.finalization_config = wrapped return wrapped diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index b7aefbeba2f..a647648e2ad 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -628,6 +628,34 @@ def test_zero_copy_backend_config_reads_the_planning_mode_when_the_pass_runs(): turned_on.to_out_var_pass(graph_module) +@pytest.mark.unit +def test_zero_copy_backend_config_rebuilt_over_a_derived_config_reads_it(): + """Deriving a config with ``dataclasses.replace`` is the case that needs it. + + The flag is a bool, copied by value; the pass is copied by reference. So a + config derived from the one this returns carries a pass still reading the + original, and the first half below is that known gap, pinned: turning + planning off on the derived config alone is not refused by the inherited + pass. It is mostly not silent either -- ``check_zero_copy_kv`` reads the + arena memory planning chose and refuses the program it produces, for any + method holding a host tensor to give the shared arena away -- and the remedy + the docstring gives is the second half: call the function again on the + derived config and the pass it builds is bound to that one. + """ + import dataclasses + + derived = dataclasses.replace( + Z.zero_copy_backend_config(), enable_non_cpu_memory_planning=False + ) + graph_module, _, _ = _direct_delegate_graph() + derived.to_out_var_pass(graph_module) + + rebuilt = Z.zero_copy_backend_config(derived) + graph_module, _, _ = _direct_delegate_graph() + with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): + rebuilt.to_out_var_pass(graph_module) + + @pytest.mark.unit def test_unstage_runs_a_second_time_without_raising(): """Installing the pass twice is redundant rather than an error. @@ -1261,13 +1289,15 @@ def test_check_zero_copy_kv_rejects_an_arena_the_program_records_as_non_cuda(): @pytest.mark.unit def test_check_zero_copy_kv_rejects_a_buffer_staged_at_its_own_delegate(): - """A second TensorRT delegate must not stand in for the one that elided. + """An unstamped TensorRT delegate must not stand in for the one that elided. The engine whose aliased output was elided -- the one carrying the zero-copy spec -- still reads a staging copy, so its write is discarded and the cache never updates. Another TensorRT engine happens to read the same buffer directly, which says nothing about that write. Taking the union over every - TensorRT delegate accepts this program. + TensorRT delegate accepts this program; narrowing to the stamped ones is what + refuses it here, and ``..._a_second_zero_copy_delegate_standing_in`` is the + case where both are stamped and that narrowing is not enough. """ graph = torch.fx.Graph() k_buffer = graph.placeholder("b_k_0") @@ -1295,6 +1325,86 @@ def test_check_zero_copy_kv_rejects_a_buffer_staged_at_its_own_delegate(): Z.check_zero_copy_kv(_finalized_program(_planned(graph_module))) +def _two_zero_copy_delegate_graph(*, crossed, first_specs=None): + """Two stamped TensorRT delegates in one method, one elided output each. + + ``crossed`` gives the second delegate both caches and leaves the first + reading a staging copy of the one it elided, which is the shape a check that + reads the stamped delegates as one set cannot see: every marked buffer does + reach a stamped delegate, just not the one whose write was removed. + ``first_specs`` overrides what the first delegate's spec claims it elided. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + v_buffer = graph.placeholder("b_v_0") + lowered_k = graph.get_attr("lowered_module_0") + lowered_v = graph.get_attr("lowered_module_1") + if crossed: + staged_k = graph.call_function(torch.ops.et_copy._h2d_copy.default, (k_buffer,)) + k_args, v_args = (lowered_k, staged_k), (lowered_v, k_buffer, v_buffer) + else: + k_args, v_args = (lowered_k, k_buffer), (lowered_v, v_buffer) + delegate_k = graph.call_function(executorch_call_delegate, k_args) + delegate_v = graph.call_function(executorch_call_delegate, v_args) + graph.output((k_buffer, v_buffer, delegate_k, delegate_v)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", + compile_specs=( + _zero_copy_specs("out_k") if first_specs is None else first_specs + ), + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs("out_v") + ) + graph_module = torch.fx.GraphModule(root, graph) + for buffer in (k_buffer, v_buffer): + buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + buffer.meta["_torch_tensorrt_aliased_buffer"] = True + return _planned(graph_module) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_second_zero_copy_delegate_standing_in(): + """A stamped delegate is counted against its own spec, not pooled with the rest. + + Both delegates here declare zero-copy KV, so narrowing by the spec does not + separate them, and both marked buffers reach one of the two directly -- so + the still-staged refusal has nothing to say. What is wrong is which delegate + took which: ``lowered_module_0`` elided ``out_k`` and reads a staging copy of + the cache that write was removed for, so that cache never updates. The + matched half is the same graph with each delegate holding the cache it + elided, and it is accepted, so the refusal can come only from the crossing. + """ + Z.check_zero_copy_kv( + _finalized_program(_two_zero_copy_delegate_graph(crossed=False)) + ) + + with pytest.raises(RuntimeError, match="takes 0 marked buffer"): + Z.check_zero_copy_kv( + _finalized_program(_two_zero_copy_delegate_graph(crossed=True)) + ) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_stamped_delegate_with_no_names_to_count(): + """A spec listing no name still says its engine elided an aliased output. + + Only a delegate whose own engine had one is stamped, so the count it cannot + read off the spec falls back to at least one -- the same fallback + ``_unstage_aliased_buffers`` makes. Reading "no names" as "no buffers owed" + would accept the crossing above whenever the partitioner's list is empty or + unreadable. + """ + graph_module = _two_zero_copy_delegate_graph( + crossed=True, + first_specs=[CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")], + ) + + with pytest.raises(RuntimeError, match="must take at least one marked buffer"): + Z.check_zero_copy_kv(_finalized_program(graph_module)) + + @pytest.mark.unit def test_check_zero_copy_kv_rejects_a_buffer_only_another_backend_takes(): """Another backend's delegate taking the buffer directly is not zero-copy. @@ -1398,11 +1508,11 @@ def test_zero_copy_backend_config_defaults_to_executorch_defaults(): @pytest.mark.unit @pytest.mark.parametrize( "skip", - [True, {"decode": True}, {"prefill": False, "decode": True}], - ids=["bool", "one-method", "one-of-two-methods"], + [True, {"decode": True}, {"prefill": False, "decode": True}, {"decode": False}], + ids=["bool", "dict-one-true", "dict-one-of-two-true", "dict-all-false"], ) def test_zero_copy_backend_config_refuses_skip_h2d_for_method_inputs(skip): - """The one option that cannot be carried through, refused where it is set. + """The one option that cannot be carried through, refused wherever it is on. ``skip_h2d_for_method_inputs`` is ExecuTorch's own un-staging of method inputs and it demands each placeholder it un-stages have exactly one user. A @@ -1410,6 +1520,17 @@ def test_zero_copy_backend_config_refuses_skip_h2d_for_method_inputs(skip): mutation result for -- so ``PropagateDevicePass`` raises on every zero-copy graph. Preserving the option hands back a config that cannot finalize at all, which is a failure a long way from the line that caused it. + + ``dict-all-false`` is the case that makes the refusal key on truthiness + rather than on ``True``: that pass is handed the field whole and only tests + it for truth, so it reads a dict of ``False`` as on for every method, and + finalizing such a config raises with ``placeholder 'b_k_cache' to have + exactly one user``. ``False`` and ``{}`` are read as off there and are + carried, in ``..._carries_skip_h2d_left_falsy``. + + The ids name what each value *is*, not who it applies to: that pass never + resolves this field per method, so every dict here is on for every method + whatever key it carries. """ from executorch.exir import ExecutorchBackendConfig from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig @@ -1423,14 +1544,39 @@ def test_zero_copy_backend_config_refuses_skip_h2d_for_method_inputs(skip): @pytest.mark.unit -@pytest.mark.parametrize("skip", [False, {"decode": False}], ids=["bool", "one-method"]) -def test_zero_copy_backend_config_carries_skip_h2d_left_off(skip): - """A field set and left off is carried through, bool or dict alike. - - The dict half is narrower than ``PropagateDevicePass``, which tests the - field for truth rather than resolving it per method and so reads any - non-empty dict as on. Such a config is carried here and still fails to - finalize. +def test_zero_copy_backend_config_refuses_skip_h2d_in_a_per_method_config(): + """``propagate_device_config`` is itself one config or a dict of them. + + ExecuTorch resolves a dict of ``PropagateDeviceConfig`` by method name + (``_program.py``, ``edge_to_executorch_passes``) and hands the chosen one's + ``skip_h2d_for_method_inputs`` to the pass, so the option reaches + ``PropagateDevicePass`` from here exactly as it does from the single-config + form. Reading only the single form leaves that route open: measured, such a + config finalizes into the same ``exactly one user`` failure. + """ + from executorch.exir import ExecutorchBackendConfig + from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig + + base = ExecutorchBackendConfig( + propagate_device_config={ + "prefill": PropagateDeviceConfig(), + "decode": PropagateDeviceConfig(skip_h2d_for_method_inputs=True), + } + ) + + with pytest.raises(ValueError, match="skip_h2d_for_method_inputs for decode"): + Z.zero_copy_backend_config(base) + + +@pytest.mark.unit +@pytest.mark.parametrize("skip", [False, {}], ids=["bool", "empty-dict"]) +def test_zero_copy_backend_config_carries_skip_h2d_left_falsy(skip): + """A field left falsy *to PropagateDevicePass* is carried through. + + That pass only tests the field for truth, so what it reads as off is exactly + ``False`` and the empty dict -- and those two are what this carries. A + non-empty dict of ``False`` is not one of them; it is refused, in + ``..._refuses_skip_h2d_for_method_inputs[dict-all-false]``. """ from executorch.exir import ExecutorchBackendConfig from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig From 204432ddd164e94873529989d90ccf6454ddd960 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Mon, 7 Sep 2026 21:50:54 -0700 Subject: [PATCH 19/22] fix(executorch): refuse a name the parser and TensorRT read differently The blob metadata is copied out with an explicit length, so a binding name can hold a NUL byte, while every consumer resolves the name through `c_str()` and TensorRT stops at the first one. The parser has three repeat refusals and one emptiness check, and a NUL walks past every one of them. Past the repeat refusals in four ways: two names differing only after a NUL are two entries here and one tensor to the engine -- the exact collision those refusals exist to stop -- and the set covering `io_bindings` spans both lists, so an input colliding with an output defeats it as well as two outputs do. Past the emptiness check because that check draws its line where `size()` does rather than where `c_str()` does, so a name that is only a NUL is non-empty here and empty to the engine. One predicate beside the emptiness check closes all five routes, since every name the engine is asked for has to be an `io_bindings` name first. Measured against a standalone build of the parser: two output bindings `"x\0y"`/`"x\0z"`, an input and an output colliding across the two lists, both `aliased_io` shapes and the NUL-only name all parsed clean and reported the same `c_str()`; all five are refused now and a plain blob is unaffected. The same predicate on the two `aliased_io` names stops no collapse of its own -- `init()` compares those whole, so a NUL there matches no binding and `init()` refuses the blob -- but it keeps the invariant that every recorded name is one the engine can be asked for, and moves that failure to parse. No writer emits a NUL -- `json.dumps` escapes it -- so this refuses only a blob assembled some other way. An `aliased_io` entry naming only one of its two bindings was silently skipped and now shares that refusal. `init()` counts the entries it accepts and `execute()` subtracts that count from the delegate's argument list, so a dropped entry surfaces as an argument-count error on every call, in a message that never mentions aliasing, instead of at parse where the blob-header tests reach it without a GPU. The writer always emits both keys (`serialization.py`) and an older blob carries no `aliased_io` array at all. `check_zero_copy_kv` passed over any method with no marked buffer before it looked at the delegates, so a method still carrying a stamped zero-copy delegate whose mark was lost was invisible to it while `_unstage_aliased_buffers` refuses that same graph -- measured on a two-method program, the pass raising and the check returning clean. That skip is one side of a two-sided question: the function exists to catch the marks and the stamped delegates disagreeing, and driving the walk from the marks alone leaves anything recorded only on the delegate side outside it. Both records are now enumerated in every method and a method is passed over only when it carries neither; the per-delegate count already in place then refuses the lost-mark case. The program-wide "nothing is marked" refusal moved below the two that fire on such a disagreement, because on a single-method program in that state it answered "probably not exported with zero_copy_kv=True" while the compile spec said it was. `order_copyback_mutations_first`'s docstring claimed that asking upstream's predicate rather than this module's mark is what covers a mutation `run_reinplace_pass` turns in place. It does not: `reinplace_pass` runs inside `to_executorch` (`_program.py:1706`), after the reorder, so such a mutation is ordinary when the predicate is asked and in-place when the write-back reads it. Measured on a two-buffer model with no TensorRT in it, the pair finalizes crossed with and without the reorder and `moved` is 0 either way. The docstring and the test that cited the case now claim only what they cover -- any mutation already in place at the Edge boundary -- and name the reinplace ordering as upstream's. Coverage the mutation runs showed missing. Dropping the clause that compares the metadata extent against `engine_offset` left every blob-header test green, so a blob whose metadata sits past the engine parsed; one case pins it. The partial-elision refusal in `preprocess` was pinned only from the zero-copy file, so replacing it with a constant false left `test_backend.py` green; a case now sits beside its neighbours. The two new public names were in `__all__` and on the API page but not in the test that pins the package's surface. The four tests that build a real engine gated CUDA with a decorator `skipif`, which resolves during collection -- off the GPU host on a remote-GPU runner, so the skip freezes in before a GPU is attached and the only real-engine coverage this feature has disappears with the lane still green. They use the runtime gate the rest of this directory documents three times over. A note on the bazel target for the KV decode check records that nothing depends on it and the CMake build beside it is what CI compiles, so its dependency list is kept in step by hand. Blob-header gtests 22 -> 28, Python suite 325 -> 327 collected, one pre-existing wheel-pin failure unchanged. Every behavioural clause was mutated away and the test that bites it named. Two mutants of the metadata extent survive and cannot be killed: rewriting it in addition form, and deleting it altogether -- 0 verdict differences over 1,946,720 header-field combinations, against 3,876 and 904 for deleting either of the other two extent checks. The comment there says why it is kept anyway. --- .../executorch/TensorRTBlobHeader.cpp | 102 +++++++++++------- examples/executorch_reference_runner/BUILD | 8 ++ py/torch_tensorrt/executorch/_zero_copy.py | 73 ++++++++----- .../test_executorch_blob_header.cpp | 101 +++++++++++++++++ tests/py/dynamo/executorch/test_api.py | 2 + tests/py/dynamo/executorch/test_backend.py | 33 ++++++ .../py/dynamo/executorch/test_zero_copy_kv.py | 80 ++++++++++---- 7 files changed, 317 insertions(+), 82 deletions(-) diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp index a37050167c2..35a1d5cdd87 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp @@ -137,6 +137,23 @@ bool parse_int_after_key(const std::string& json, std::size_t search_from, const return true; } +// Every recorded name is resolved against the engine through c_str(), which +// stops at the first NUL, while the refusals below compare whole std::string +// values. The metadata is copied out of the blob with an explicit length, so a +// name can hold a NUL: two names differing only after one are distinct here and +// are the same tensor to TensorRT, which is exactly what those refusals exist to +// stop, and a name that is only a NUL is non-empty here and empty to TensorRT. +// Refusing a NUL outright makes what the parser compares be what the engine +// will compare. No writer emits one -- json.dumps escapes it -- so this refuses +// only a blob assembled some other way. +// Like every refusal in this parser it is silent: parse() returns false and +// the caller reports its own generic parse failure, so a NUL is not +// distinguishable at load time from a bad magic or a wrapped extent. Refusing +// uniformly is the deliberate choice here, not a missing diagnostic. +bool usable_as_binding_name(const std::string& name) { + return !name.empty() && name.find('\0') == std::string::npos; +} + bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { out.input_binding_names.clear(); out.output_binding_names.clear(); @@ -234,13 +251,13 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { } } - // A nameless entry cannot be refused earlier because the keys may arrive in - // any order, so it is refused here, beside the repeat. Skipping it instead - // would shorten the recorded list while the delegate's argument list keeps - // its full length, and the two are only inferred from the engine when both - // are empty -- so one real name beside a blank leaves a short list that no - // longer lines up with the engine's bindings. - if (!saw_name || name.empty()) { + // An entry with no usable name cannot be refused earlier because the keys + // may arrive in any order, so it is refused here, beside the repeat. + // Skipping it instead would shorten the recorded list while the delegate's + // argument list keeps its full length, and the two are only inferred from + // the engine when both are empty -- so one real name beside a blank leaves + // a short list that no longer lines up with the engine's bindings. + if (!saw_name || !usable_as_binding_name(name)) { return false; } if (!claimed_bindings.insert(name).second) { @@ -323,39 +340,46 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { return false; } } - if (!ab.output.empty() && !ab.input.empty()) { - // An output binding may be claimed by at most one entry. A second entry - // for the same output names the same binding, so nothing that resolves - // the names can tell the two apart -- but a reader that counts aliased - // outputs per entry, as TensorRTBackend does to size the delegate - // argument list, counts one output twice. Refuse the blob here, where - // the repeat is visible from the bytes alone. - if (!claimed_outputs.insert(ab.output).second) { - return false; - } - // An input may likewise be claimed by at most one entry. Two entries - // naming different outputs and one input record the same input index - // for both, and execute() binds each aliased output to that index's - // caller pointer -- one address with two writers, so whichever the - // engine writes second wins and the other update disappears with no - // error. TensorRT's own aliasing rules out the kv_cache_update kind - // (init cross-checks it against getAliasedInputTensor), but the user - // kind is only compared on shape, so two same-shaped outputs onto one - // input would pass. - if (!claimed_inputs.insert(ab.input).second) { - return false; - } - // The current Python serializer always writes "kind" (serialization.py), - // and older blobs carry no aliased_io array at all, so this default is - // defensive: it only fires for a blob that has an aliased_io entry but - // omits "kind". Default to the TRT-enforced kind so init()'s kind - // validation treats an absent key the same as the Python runtime rather - // than rejecting it as unknown. - if (ab.kind.empty()) { - ab.kind = "kv_cache_update"; - } - out.aliased_io.push_back(std::move(ab)); + // An entry missing either name is refused rather than skipped, for the + // reason the binding walk above gives and one more: init() counts the + // entries it accepts and execute() subtracts that count from the delegate + // argument list, so a dropped entry surfaces as an argument-count error at + // every execute, which never mentions aliasing, instead of at parse, which + // the blob-header tests reach without a GPU. + if (!usable_as_binding_name(ab.output) || !usable_as_binding_name(ab.input)) { + return false; + } + // An output binding may be claimed by at most one entry. A second entry + // for the same output names the same binding, so nothing that resolves + // the names can tell the two apart -- but a reader that counts aliased + // outputs per entry, as TensorRTBackend does to size the delegate + // argument list, counts one output twice. Refuse the blob here, where + // the repeat is visible from the bytes alone. + if (!claimed_outputs.insert(ab.output).second) { + return false; + } + // An input may likewise be claimed by at most one entry. Two entries + // naming different outputs and one input record the same input index + // for both, and execute() binds each aliased output to that index's + // caller pointer -- one address with two writers, so whichever the + // engine writes second wins and the other update disappears with no + // error. TensorRT's own aliasing rules out the kv_cache_update kind + // (init cross-checks it against getAliasedInputTensor), but the user + // kind is only compared on shape, so two same-shaped outputs onto one + // input would pass. + if (!claimed_inputs.insert(ab.input).second) { + return false; + } + // The current Python serializer always writes "kind" (serialization.py), + // and older blobs carry no aliased_io array at all, so this default is + // defensive: it only fires for a blob that has an aliased_io entry but + // omits "kind". Default to the TRT-enforced kind so init()'s kind + // validation treats an absent key the same as the Python runtime rather + // than rejecting it as unknown. + if (ab.kind.empty()) { + ab.kind = "kv_cache_update"; } + out.aliased_io.push_back(std::move(ab)); } } diff --git a/examples/executorch_reference_runner/BUILD b/examples/executorch_reference_runner/BUILD index e6d3bc2d1b5..30025c19463 100644 --- a/examples/executorch_reference_runner/BUILD +++ b/examples/executorch_reference_runner/BUILD @@ -33,6 +33,14 @@ cc_binary( ], ) +# A convenience target for building this check out of a bazel checkout. Unlike +# the runner above, which //:bin packages into the release tar wherever this +# backend ships and so compiles on those packaging builds, nothing depends on +# this one: it ships as source in //:executorch_source_package and is compiled +# from the CMakeLists beside it, which is the path +# verify-executorch-reference-runner.sh drives. So this dependency list has to +# be kept in step with that one by hand -- an error here surfaces only when +# someone builds this target. cc_binary( name = "kv_cache_decode_check", srcs = ["kv_cache_decode_check.cpp"], diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index f4924e26fef..d3091d893e1 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -364,12 +364,18 @@ def order_copyback_mutations_first(exported_program: Any) -> int: Putting the mutations that still get a copy first restores the correspondence. Which ones those are is decided by asking upstream's own predicate (``_inplace_lineage``, imported rather than reimplemented) rather - than by asking which mutations this module rewired. The two answers differ: - ``run_reinplace_pass`` and ``reinplace_extra_ops`` are supported - ``ExecutorchBackendConfig`` fields whose pass runs just before the write-back - and turns ordinary mutations into in-place ones, and such a mutation gets no - copy either. Keyed on this module's own mark, the reorder would leave that - pair crossed and report nothing moved. + than by asking which mutations this module rewired, so the answer cannot + drift from the one the write-back pass will give, and any mutation the graph + already presents as in-place is covered whatever put it there. + + Only what the graph presents *here* is covered, though. ``run_reinplace_pass`` + and ``reinplace_extra_ops`` are supported ``ExecutorchBackendConfig`` fields + whose pass runs inside ``to_executorch``, after this and immediately before + the write-back: a mutation it rewrites is ordinary when this reads it and + in-place when the write-back does, so that pair comes out crossed anyway and + this reports nothing moved. Reordering cannot reach it from here. The same + crossing reproduces on a model using no zero-copy at all, so it is a pass + ordering upstream owns rather than one this creates. This runs on the *Edge* program rather than beside the rewiring, because ``to_edge_transform_and_lower`` re-derives the whole graph signature -- the @@ -973,10 +979,13 @@ def check_zero_copy_kv(program: Any) -> None: Four shapes are refused: a marked buffer that is not a direct argument of a TensorRT delegate carrying the zero-copy compile spec, a stamped delegate that takes fewer marked buffers than its own spec says it elided aliased - outputs, a marked buffer that reaches such a delegate directly but is not - planned in device memory, and a program with no marked buffer in any method. - The first three are what finalizing without :func:`zero_copy_backend_config` - leaves behind. That config's pass gets to all three earlier, off the + outputs -- including one in a method with no marked buffer at all, which is + the lost-mark case -- a marked buffer that reaches such a delegate directly + but is not planned in device memory, and a program carrying neither a marked + buffer nor a stamped delegate in any of its methods. The first three are + what finalizing without :func:`zero_copy_backend_config` leaves behind -- + all but the lost-mark case folded into the second, which no finalization + choice produces. That config's pass gets to all three earlier, off the configuration and the graph: it removes the staging copy the first two come from, and refuses outright when there is none to remove or when the configuration plans nothing onto a device. What it cannot see is the arena @@ -1022,10 +1031,15 @@ def check_zero_copy_kv(program: Any) -> None: on its own, so a check that stopped at ``forward`` would pass a program whose decode had degenerated to staged -- and on the prefill/decode pair the user guide's zero-copy example exports it would not get that far, since a - multi-method program need not have a ``forward`` at all. The last refusal is - about the program rather than about one method, matching the warning - ``export()`` emits: a method with no aliased buffer mutation of its own is - not an error, so a model that rewires only its decode step is accepted. + multi-method program need not have a ``forward`` at all. Within a method the + marks and the stamped delegates are enumerated independently and a method is + passed over only when it carries neither, because the disagreement between + those two records is the whole subject: starting from the marks and looking + the delegates up leaves anything recorded only on the delegate side outside + the walk. The last refusal is about the program rather than about one method, + matching the warning ``export()`` emits: a method with no aliased buffer + mutation of its own is not an error, so a model that rewires only its decode + step is accepted. This reads the graph and the finalized specs, so it says what the program does rather than what the passes recorded. It still says nothing about @@ -1044,15 +1058,20 @@ def check_zero_copy_kv(program: Any) -> None: if node.op == "placeholder" and node.meta.get("_torch_tensorrt_aliased_buffer") ] - if not marked: - continue - marked_anywhere = True zero_copy_delegates = [ node for node in graph_module.graph.nodes if _is_tensorrt_delegate(graph_module, node) and _delegate_declares_zero_copy(graph_module, node) ] + # Both records have to be read before a method can be passed over. A + # method carrying neither is one zero-copy never touched; a method + # carrying a stamped delegate and no mark is the disagreement this + # exists to catch, which skipping on the marks alone would hide. + if not marked and not zero_copy_delegates: + continue + if marked: + marked_anywhere = True zero_copy_delegate_args = { arg for node in zero_copy_delegates for arg in node.args[1:] } @@ -1093,14 +1112,6 @@ def check_zero_copy_kv(program: Any) -> None: ] if host_planned: host_planned_by_method[method_name] = host_planned - if not marked_anywhere: - raise RuntimeError( - "TensorRT zero-copy KV: no buffer in this program is marked for " - f"in-place update, in any of its methods ({', '.join(method_names)}), " - "so it stages its caches like any other .pte. Either it was not " - "exported with zero_copy_kv=True, or it was and no aliased buffer " - "mutation was found -- export logs a warning for that case." - ) if staged_by_method: raise RuntimeError( f"TensorRT zero-copy KV: buffer(s) {_name_detail(staged_by_method)} " @@ -1125,6 +1136,18 @@ def check_zero_copy_kv(program: Any) -> None: "copy-back. Export this method without zero_copy_kv, or keep each " "aliased buffer on the delegate whose engine elided it." ) + # Ordered after the still-staged and short-count refusals because both of + # those fire on a disagreement between the marks and the stamped delegates, + # and this one would answer such a program with "it was probably not exported + # with zero_copy_kv=True" while its own compile specs say it was. + if not marked_anywhere: + raise RuntimeError( + "TensorRT zero-copy KV: no buffer in this program is marked for " + f"in-place update, in any of its methods ({', '.join(method_names)}), " + "so it stages its caches like any other .pte. Either it was not " + "exported with zero_copy_kv=True, or it was and no aliased buffer " + "mutation was found -- export logs a warning for that case." + ) if host_planned_by_method: raise RuntimeError( f"TensorRT zero-copy KV: buffer(s) " diff --git a/tests/cpp/executorch/test_executorch_blob_header.cpp b/tests/cpp/executorch/test_executorch_blob_header.cpp index 5f2fbe693b8..eee75049ad9 100644 --- a/tests/cpp/executorch/test_executorch_blob_header.cpp +++ b/tests/cpp/executorch/test_executorch_blob_header.cpp @@ -288,6 +288,83 @@ TEST(ExecuTorchTensorRTBlobHeader, RejectsEmptyBindingName) { EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); } +// The metadata is copied out of the blob with an explicit length, so a name can +// hold a NUL byte. The refusals above compare whole std::string values; every +// consumer hands name.c_str() to TensorRT, which stops at the first NUL. The +// three cases below that name a NUL build their metadata by std::string +// concatenation, because a raw string literal cannot carry an embedded one. +const std::string kNul(1, '\0'); + +TEST(ExecuTorchTensorRTBlobHeader, RejectsBindingNamesDifferingOnlyAfterANul) { + // Two distinct std::strings, one TensorRT tensor. The repeated-name refusal + // above compares the whole value and lets this through, and then the collision + // it exists to stop happens anyway: execute() binds an address for that one + // TensorRT name twice, the second replacing the first, so the engine's write + // lands entirely in the later output and the earlier one is never written. + // This blob declares no alias and needs none for that; with an aliased_io + // entry beside it the address replaced is the caller's cache, which is the + // case init() says it relies on this refusal for. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"out_k)" + kNul + + R"(a","is_input":false},{"name":"out_k)" + kNul + R"(b","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAnAliasedIoNameCarryingANul) { + // The binding names here are clean, so the walk above accepts them and only the + // alias walk's own check can refuse this. It is not the collision case -- init + // compares alias names to binding names whole, so this pair matches no binding + // and init would refuse it. Refusing at parse keeps the invariant that every + // name the header records is one the engine can be asked for, and puts the + // failure where the blob-header tests reach it without a GPU. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"output":"out_k)" + + kNul + R"(a","input":"in_k)" + kNul + R"(a","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsABindingNameThatIsOnlyANul) { + // Non-empty to the refusal above, empty to TensorRT -- so the emptiness check + // has to draw its line where c_str() draws it, not where size() does. + const std::string metadata = + R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":")" + kNul + R"(","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAliasedIoEntryWithABlankInput) { + // Skipping this entry rather than refusing it records one alias for the two + // the engine has, and execute() subtracts the recorded count from the delegate + // argument list, so the .pte fails its arity check on every call with a + // message that never mentions aliasing. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v","is_input":true},)" + R"({"name":"out_k","is_input":false},{"name":"out_v","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"},)" + R"({"output":"out_v","input":"","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAliasedIoEntryWithNoOutputKey) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + TEST(ExecuTorchTensorRTBlobHeader, RejectsBindingEntryWithNoNameKey) { const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"is_input":false}]})"; const auto blob = make_blob(metadata); @@ -330,6 +407,30 @@ TEST(ExecuTorchTensorRTBlobHeader, RejectsEngineOffsetPastEndOfBlob) { EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); } +TEST(ExecuTorchTensorRTBlobHeader, RejectsMetadataThatStartsPastTheEngine) { + // Both metadata fields are in range of the blob, so only the clause comparing + // them against engine_offset refuses this. That clause is written in the same + // subtraction form as the engine extent, and the subtraction is what needs the + // ordering test in front of it: engine_offset - metadata_offset is unsigned, + // so with the metadata past the engine it wraps to nearly 2^32 and any + // metadata_size fits under it. + const std::string metadata = R"({"io_bindings":[]})"; + constexpr std::size_t kBlobSize = 8192; + constexpr uint32_t kMetadataOffset = 4096; + constexpr uint32_t kEngineOffset = 48; + + std::vector blob(kBlobSize, 0); + std::memcpy(blob.data(), TENSORRT_MAGIC, sizeof(TENSORRT_MAGIC)); + write_field(blob, METADATA_OFFSET_FIELD_OFFSET, kMetadataOffset); + write_field(blob, METADATA_SIZE_FIELD_OFFSET, static_cast(metadata.size())); + write_field(blob, ENGINE_OFFSET_FIELD_OFFSET, kEngineOffset); + write_field(blob, ENGINE_SIZE_FIELD_OFFSET, uint64_t{4}); + std::memcpy(blob.data() + kMetadataOffset, metadata.data(), metadata.size()); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + TEST(ExecuTorchTensorRTBlobHeader, RejectsUnknownFutureMagic) { constexpr char kFutureMagic[4] = {'T', 'R', '0', '3'}; const std::string metadata = R"({"io_bindings":[{"name":"x","is_input":true}]})"; diff --git a/tests/py/dynamo/executorch/test_api.py b/tests/py/dynamo/executorch/test_api.py index 7de4e563fef..2e0b849aeb8 100644 --- a/tests/py/dynamo/executorch/test_api.py +++ b/tests/py/dynamo/executorch/test_api.py @@ -99,6 +99,8 @@ def test_public_api_symbols_present(): assert "TensorRTPartitioner" in module.__all__ assert "TensorRTBackend" in module.__all__ assert "export" in module.__all__ + assert "zero_copy_backend_config" in module.__all__ + assert "check_zero_copy_kv" in module.__all__ assert "Program" not in module.__all__ assert "load" not in module.__all__ assert "to_executorch" not in module.__all__ diff --git a/tests/py/dynamo/executorch/test_backend.py b/tests/py/dynamo/executorch/test_backend.py index b20c09cc1ba..791221ea6c1 100644 --- a/tests/py/dynamo/executorch/test_backend.py +++ b/tests/py/dynamo/executorch/test_backend.py @@ -613,6 +613,39 @@ def test_preprocess_rejects_a_non_buffer_alias_elided_alongside_a_buffer_alias() TensorRTBackend.preprocess(edge_program, [spec]) +@pytest.mark.unit +def test_preprocess_rejects_an_engine_whose_aliased_outputs_are_partly_elided(): + """Two aliased outputs, one elided: the arity the runtime reads cannot say so. + + ``out_v`` is still a delegate output, so the binding-order check above is + satisfied and only the whole-set comparison refuses this. The runtime takes + the aliased outputs as elided only when the argument count is short by the + engine's whole aliased-output count, so a delegate short of one of two reads + as not elided at all and every ``execute()`` fails on the argument count. The + same clause is reached from the zero-copy suite through the partitioner's own + derivation; this pins it beside the other ``preprocess`` refusals. + """ + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + TensorRTBackend, + _serialize_elided_output_names, + ) + + edge_program = _aliased_edge_program( + present_indices=[0, 2], # logits and out_v kept, only out_k dropped + out_names=["logits", "out_k", "out_v"], + aliased_io_map={ + "out_k": ("tokens", "kv_cache_update"), + "out_v": ("tokens", "kv_cache_update"), + }, + ) + spec = CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, _serialize_elided_output_names(["out_k"]) + ) + with pytest.raises(ValueError, match="Partial elision is not expressible"): + TensorRTBackend.preprocess(edge_program, [spec]) + + @pytest.mark.unit def test_preprocess_rejects_a_delegate_with_no_outputs_at_all(): """Naming every binding elidable leaves an empty expected index list, which diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index a647648e2ad..0b1f41d4d51 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -46,6 +46,19 @@ ) +def _require_real_engine(): + """Gate the tests that build a real engine, at run time rather than collection. + + A decorator ``skipif`` resolves while pytest collects, and on a remote-GPU + runner collection happens off the GPU host, so the skip is frozen in before + any GPU is attached. These are the only tests that put this feature on a real + engine, so where that happens the lane stays green with that coverage gone. + The other CUDA gates in this directory are runtime gates for the same reason. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA + TensorRT for a real engine") + + def _patch_engine_metadata(monkeypatch, *, aliased_io, input_names, output_names): """Make every engine node report one fixed set of bindings and aliases.""" import torch_tensorrt.dynamo.runtime._serialized_engine_layout as layout @@ -1488,6 +1501,44 @@ def test_check_zero_copy_kv_rejects_a_multi_method_program_with_nothing_marked() Z.check_zero_copy_kv(_finalized_program(prefill=first, decode=second)) +def _stamped_delegate_with_no_mark(): + """A stamped zero-copy delegate whose buffer lost its mark. + + Nothing but the compile spec is left recording that this engine's aliased + output was elided, the mark being the only other witness to it. + """ + graph_module, k_buffer, _ = _direct_delegate_graph( + compile_specs=_zero_copy_specs("out_k") + ) + del k_buffer.meta["_torch_tensorrt_aliased_buffer"] + return _planned(graph_module) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_a_stamped_delegate_whose_method_lost_its_mark(): + """A method is read for both records, not passed over on the marks alone. + + ``_unstage_aliased_buffers`` refuses this graph -- the spec says an aliased + output was elided and no marked buffer arrives to hold it -- so a check that + passed it would call correct a method the pass calls broken. Both halves + matter: the two-method program pins that a marked method does not vouch for + an unmarked one, and the single-method program pins the message, which would + otherwise be the program-wide "probably not exported with zero_copy_kv=True" + on a program whose own compile spec says it was. + """ + with pytest.raises(RuntimeError, match=r"takes 0 marked buffer"): + Z.check_zero_copy_kv( + _finalized_program( + prefill=_unstaged_graph(), decode=_stamped_delegate_with_no_mark() + ) + ) + + with pytest.raises(RuntimeError, match=r"takes 0 marked buffer"): + Z.check_zero_copy_kv( + _finalized_program(decode=_stamped_delegate_with_no_mark()) + ) + + @pytest.mark.unit def test_zero_copy_backend_config_defaults_to_executorch_defaults(): """Called with no config it starts from ExecuTorch's defaults, and the one @@ -1977,9 +2028,6 @@ def split_heads(proj: torch.Tensor) -> torch.Tensor: return self.lm(self.o(out)) -@pytest.mark.skipif( - not torch.cuda.is_available(), reason="requires CUDA + TensorRT for a real engine" -) @pytest.mark.parametrize("generate_etrecord", [False, True], ids=["plain", "etrecord"]) def test_aliased_buffer_mark_survives_real_lowering(generate_etrecord): """After a real export(..., zero_copy_kv=True), the KV buffer placeholder in @@ -1991,6 +2039,7 @@ def test_aliased_buffer_mark_survives_real_lowering(generate_etrecord): raise here: it surfaces later as the un-staging pass finding a marked buffer it never un-staged, by which point the connection to this option is gone. """ + _require_real_engine() with torch.no_grad(): torch.manual_seed(0) model = _KVDecodeStep().eval().cuda() @@ -2200,12 +2249,13 @@ def _two_mutation_program(first_value_is_inplace): def test_reorder_moves_an_inplace_mutation_this_feature_did_not_create(reorder): """The reorder keys on upstream's predicate, not on this feature's own mark. - ``run_reinplace_pass`` and ``reinplace_extra_ops`` are ``ExecutorchBackendConfig`` - fields whose pass runs just before the write-back and rewrites an ordinary - mutation into an in-place one. Upstream then inserts no copy for it, exactly - as for a rewired cache -- and a reorder that asks "did zero-copy rewire - this?" instead of "will upstream copy this?" leaves the resulting pair - crossed while reporting that it moved nothing. + Slot 0 is in-place in the graph and carries no zero-copy mark, so upstream + inserts no copy for it exactly as for a rewired cache, and a reorder that + asked "did zero-copy rewire this?" instead of "will upstream copy this?" + would leave the pair crossed while reporting that it moved nothing. A + mutation ``reinplace_pass`` rewrites is *not* this case: that pass runs after + the reorder, so such a mutation is still ordinary when the predicate is asked + and no reorder here can pre-empt it -- see ``order_copyback_mutations_first``. """ from executorch.exir.passes.insert_write_back_for_buffers_pass import ( insert_write_back_for_buffers_pass, @@ -2279,9 +2329,6 @@ def _assert_each_mutation_names_its_own_value(program): ) -@pytest.mark.skipif( - not torch.cuda.is_available(), reason="requires CUDA + TensorRT for a real engine" -) @pytest.mark.parametrize("retrace", [False, True], ids=["legacy", "retrace"]) def test_zero_copy_kv_keeps_a_copyback_buffer_in_the_same_method(retrace): """A real method holding both kinds of mutable buffer exports and keeps both. @@ -2302,6 +2349,7 @@ def test_zero_copy_kv_keeps_a_copyback_buffer_in_the_same_method(retrace): the fx boundary and leaves the copy-back value as a plain return -- so that pass is what separates the two kinds, by reading each engine's ``aliased_io``. """ + _require_real_engine() with torch.no_grad(): torch.manual_seed(0) model = _MixedDecodeStep().eval().cuda() @@ -2402,9 +2450,6 @@ def split_heads(proj: torch.Tensor) -> torch.Tensor: return self.lm(h + self.conv_state.sum(dim=2).reshape(1, 1, DIM)) -@pytest.mark.skipif( - not torch.cuda.is_available(), reason="requires CUDA + TensorRT for a real engine" -) @pytest.mark.parametrize("retrace", [False, True], ids=["legacy", "retrace"]) def test_zero_copy_kv_with_the_copyback_on_a_second_delegate(retrace): """Two TensorRT delegates, one with the aliased caches and one with the copy-back. @@ -2422,6 +2467,7 @@ def test_zero_copy_kv_with_the_copyback_on_a_second_delegate(retrace): different engine than the caches it declares. The legacy exporter declares all of that while inlining, so that scan runs only on this parameter. """ + _require_real_engine() with torch.no_grad(): torch.manual_seed(0) model = _SplitRolesDecodeStep().eval().cuda() @@ -2507,9 +2553,6 @@ def test_zero_copy_kv_with_the_copyback_on_a_second_delegate(retrace): _assert_marked_buffers_reach_the_engine_unstaged(program) -@pytest.mark.skipif( - not torch.cuda.is_available(), reason="requires CUDA + TensorRT for a real engine" -) @pytest.mark.parametrize("retrace", [False, True], ids=["legacy", "retrace"]) def test_zero_copy_kv_beside_an_executorch_cuda_delegate(retrace): """An aliased KV cache in a method that also holds an ExecuTorch CUDA delegate. @@ -2522,6 +2565,7 @@ def test_zero_copy_kv_beside_an_executorch_cuda_delegate(retrace): That the gate itself refuses a marked buffer on another backend is pinned by ``test_unstage_raises_for_a_marked_buffer_on_another_backends_delegate``. """ + _require_real_engine() cuda_backend = pytest.importorskip("executorch.backends.cuda.cuda_backend") cuda_partitioner = pytest.importorskip("executorch.backends.cuda.cuda_partitioner") From 8fc298e3e19c89cd1553a4d113b3a4fe0ca530ae Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Tue, 8 Sep 2026 03:46:47 -0700 Subject: [PATCH 20/22] fix(executorch): let the documented zero-copy path refuse its own bad program `export(zero_copy_kv=True)` removes the copy-back before it returns, so from that point every way of finalizing that does not also un-stage the caches produces a program whose caches never update -- for a KV cache, wrong output rather than a crash. `to_executorch()` with ExecuTorch's defaults is one of those, and it is the call the two-step API documents: measured on a real engine, it succeeded and wrote a 113 KB `.pte` with both caches still staged. What had been standing behind that gap was `check_zero_copy_kv`, which does refuse such a program -- but only for a caller who runs it, and it is opt-in. The manager `export` returns now runs it on whatever its own `to_executorch` produces, which is what `save(..., zero_copy_kv=True)` already did at the same point. A correct program comes back untouched; the bad one raises, with the refusal that already names `zero_copy_backend_config`. The hook is bound to the instance, because `EdgeProgramManager` is ExecuTorch's, so it reaches that manager and not one `transform()` or `to_backend()` derives from it -- the docs and docstrings say so and say to call the check by hand there. `save` keeps its own call because it finalizes the program itself; `export` installs the hook whenever zero-copy was asked for, so the two entry points refuse the same models -- including a `zero_copy_kv=True` that found nothing to do. A buffer whose spec already named the staging copy's device skipped the surviving-consumer check entirely, and the pass then rewired it while `_device_move_is_safe` said no -- another backend's `_h2d_copy` left reading device memory as a host source, which fails `InvalidArgument` on every call. Nothing moves in that shape, but where the buffer ends up is the same either way, so the question is now asked either way and the helper is named for the placement rather than the move. Two users it must allow, or asking it always would refuse programs this pass itself builds: a TensorRT delegate already taking the buffer, and a staging copy this run has already detached. The second is passed in rather than inferred from having no users, because a *foreign* dead copy looks the same and nothing erases that one. Three refusals read a value they are not finalized with, or read it too late: * The host-planning refusal reads `enable_non_cpu_memory_planning` off the config, but `to_executorch` does not pass that flag to the memory planner -- it assigns it, and only onto a planner that already has an attribute of that name. A caller who brings their own planner, which the user guide tells people to do for a cache shared between prefill and decode, never receives it, so the field decides nothing there in either direction. It is now resolved the way ExecuTorch resolves it and is a ground to refuse on only where it reaches the planner; `check_zero_copy_kv` reads the arena that planner actually chose, which is the answer this cannot give. * The placement check compared the arena's device *type* and dropped its index, so a cache asking for `cuda:0` passed while its arena was recorded for `cuda:1` -- an address on a GPU the engine is not running on, which fails as a host pointer does. Both are compared now, with either index left unrecorded still saying nothing. * `save`'s one refusal that reads nothing but the config fired after export had partitioned the graph and built every engine. It now fires beside the weight-streaming budget, before the compile, for the reason that one gives. The per-delegate satisfaction count, in the pass and in the finalized-program check both, summed over argument slots rather than distinct buffers, so one buffer in two slots satisfied a spec naming two elided outputs -- the exact shape that count exists to catch. Both count distinct buffers now, together, so the two cannot start disagreeing about how many one graph satisfies. Four more holes in the blob parser, of the same kind as the ones already refused there: * An entry with no `is_input`, or with the key misspelled by one byte, kept the initializer and landed in the output list, while `TensorRTIOBinding` in `serialization.py` defaults the same field to an input -- the two readers of one blob disagreeing, and not about one slot: a binding changing list shifts every index after it, which a static-shape engine runs without complaint on the wrong tensors. Refused beside the nameless entry. * `parse_string` drops a backslash and keeps what follows, so `"a\nb"` is recorded as `anb` and a `\u` escape as its literal digits -- names the engine does not have, which is the property the NUL refusal beside it establishes. A name carrying an escape this parser reads differently from a JSON reader is refused by the same predicate. An escaped quote or backslash is not: those two come out exactly as `json.dumps` wrote them, and refusing them would refuse names the merge-base parser loads. * TR02 declares that the metadata carries `aliased_io` and nothing compared that against what was parsed, so a blob whose array the walk cannot find -- absent, truncated, one byte wrong in the key, or written before `io_bindings` -- loaded as alias-free. The threaded shape does not fail afterwards: every aliased output gets storage of its own and the caller's cache quietly stops updating. The converse is left alone deliberately: a TR01 blob carrying an array is read as aliased and runs correctly here. * `parse_int_after_key` accumulated file-controlled digits into a plain `int`. That is undefined behaviour on overflow and in an ordinary build it wraps: `4294967297` came back as device 1, a GPU that exists, so `cudaSetDevice` succeeds and the engine deserializes where nobody asked. Accumulated in 64 bits and bounded on every digit now. Two siblings found beside them. The `hardware_compatible` and `device_id` scans started past `io_bindings` but before the alias array that is written between, so an aliased binding named like either key was matched as the key, the scan walked to the next colon, met the kind string, and failed a good blob; they start past the array now. And a scan matched any occurrence of its key, value or not, so a string value that reads like the key -- which an older blob with no `device_id` of its own can carry -- was read as the field. Only a match in key position, its own colon next, is taken. The rest is prose the code no longer matched. The paragraph justifying the aliased-count guard in the runtime said dropping it would accept any argument count and that both subtractions could wrap in one call; enumerated, it lets through exactly one count and they never both wrap, which is what it says now. The argument-count error could only ever print the threaded expectation, so a zero-copy program arriving short was told to supply the count that puts it in the other branch; it names both, and the aliasing that is the reason there are two. The `EngineHandle` contract described only the threaded layout. And `_engine_info_str` was a third copy of `backend._get_str`, exact over every slot value the two see. One gap is documented rather than closed, because closing it is a decision this does not own: a config derived with `dataclasses.replace` still carries a pass bound to the original. It says so where the caller reads, and the check the export path just gained catches what the mistake produces -- it reads the arena planning chose, which is what a lost flag changes. The neighbouring case is refused rather than documented: a `memory_planning_pass` that does not record its arenas leaves the `.pte` reporting every planned buffer as CPU, and a runner that honours that -- the reference runner in this repo does -- backs the cache with host memory, so an absent record is refused under a message of its own. --- .../executorch/TensorRTBackend.h | 14 +- .../executorch/TensorRTBackend.cpp | 49 ++- .../executorch/TensorRTBlobHeader.cpp | 155 +++++-- .../runtime_performance/saving_models.rst | 54 ++- py/torch_tensorrt/_compile.py | 14 +- py/torch_tensorrt/executorch/_export.py | 25 +- py/torch_tensorrt/executorch/_zero_copy.py | 385 ++++++++++++++---- py/torch_tensorrt/executorch/partitioner.py | 8 +- .../test_executorch_blob_header.cpp | 146 +++++++ tests/py/dynamo/executorch/test_export.py | 7 + .../py/dynamo/executorch/test_zero_copy_kv.py | 307 +++++++++++++- 11 files changed, 992 insertions(+), 172 deletions(-) diff --git a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h index b33d712d403..dc168a03367 100644 --- a/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h +++ b/cpp/include/torch_tensorrt/executorch/TensorRTBackend.h @@ -61,11 +61,15 @@ struct EngineHandle { size_t num_outputs = 0; // Per output binding [0..num_outputs): index into input_binding_names of the // input it aliases (in-place KV-cache / user alias), or -1 for a normal output. - // Built at init from the blob's aliased_io. The KV buffers are threaded by - // ExecuTorch as caller-owned mutable-buffer delegate args (input AND aliased - // output): execute() binds each aliased TRT output binding to its aliased - // input's caller-provided pointer (in-place) and reflects the result into the - // delegate output EValue, which ExecuTorch's write-back copy_ then reads. + // Built at init from the blob's aliased_io. Either way execute() binds the + // aliased TRT output binding to its aliased input's caller-provided pointer, + // so the engine's write lands in the caller's buffer; what differs is how the + // .pte carries the buffer. Threaded: the buffer is both a delegate input arg + // and a delegate output arg (the caller-owned mutable buffer's mutation slot), + // and execute() reflects the result into that output EValue for ExecuTorch's + // write-back copy_ to read. Elided -- zero-copy KV -- the buffer is an input + // arg only, the delegate has no output for it, and execute() skips the + // reflect: the in-place write already is the update. std::vector output_aliased_input_idx; // Per input binding [0..num_inputs): true if any output aliases this input, so // its in-place (KV/user) update must land in the caller-owned storage. Built at diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 9513c00f44c..73646867918 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -582,11 +582,22 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // The blob parser refuses a second aliased_io entry for an output already // claimed, and init refuses an entry whose output is not one of the recorded // output bindings, so this holds for any header that reached here. It is - // checked anyway because both subtractions below are unsigned: on a header - // built some other way they wrap and the length check then accepts any - // argument count. That wrap is all this catches. A duplicate entry inflates - // the count while staying within num_outputs, passes both checks, and still - // indexes one past the end of args; the parser's refusal is what stops that. + // checked anyway because the subtractions below are unsigned: on a header + // built some other way one of them wraps, and the length check then accepts + // an argument count it should not. Exactly one such count, not any: with the + // aliased count above the outputs, the only count that can set the elided + // flag is inputs plus outputs minus aliased, and once it is set + // num_delegate_outputs is the wrapped outputs-minus-aliased, which adding the + // inputs back wraps round to that same count -- so the length check passes + // and the call goes on to index past the end of args. Every other count below + // inputs plus outputs is still rejected. The two subtractions never wrap in + // the same call either: with the aliased count above the outputs but not + // above inputs plus outputs only the second wraps, and above both the first + // wraps to a value no argument count can match, so the flag is never set and + // the second never runs. A duplicate entry is the case this does not catch: + // it inflates the count while staying within num_outputs, passes both checks, + // and still indexes one past the end of args; the parser's refusal is what + // stops that. if (num_aliased_outputs > num_outputs) { ET_LOG( Error, @@ -599,11 +610,29 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* num_aliased_outputs > 0 && args.size() == num_delegate_inputs + num_outputs - num_aliased_outputs; const size_t num_delegate_outputs = num_outputs - (aliased_outputs_elided ? num_aliased_outputs : 0); if (args.size() < num_delegate_inputs + num_delegate_outputs) { - ET_LOG( - Error, - "TensorRTBackend::execute: expected at least %zu args, got %zu", - num_delegate_inputs + num_delegate_outputs, - args.size()); + // With aliased outputs there are two right answers and only one of them can + // ever be num_delegate_outputs here: the elided flag is set only by an + // argument count that already fits, so a program that arrives short is + // always measured against the threaded count. Reporting that alone tells a + // zero-copy .pte to supply the longer list, which is the shape that consumes + // its real outputs as mutation slots -- so both counts are named, and the + // aliasing that is the reason for the two. + if (num_aliased_outputs > 0) { + ET_LOG( + Error, + "TensorRTBackend::execute: expected %zu args with the engine's %zu aliased output(s) threaded, " + "or %zu with them elided (zero-copy KV), got %zu", + num_delegate_inputs + num_outputs, + num_aliased_outputs, + num_delegate_inputs + num_outputs - num_aliased_outputs, + args.size()); + } else { + ET_LOG( + Error, + "TensorRTBackend::execute: expected at least %zu args, got %zu", + num_delegate_inputs + num_delegate_outputs, + args.size()); + } return Error::InvalidArgument; } diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp index 35a1d5cdd87..a113e0cfa0d 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -10,8 +11,9 @@ namespace torch_tensorrt { namespace executorch_backend { namespace { -// TR02 marks a blob whose metadata carries aliased_io; TR01 is one without. This -// parser handles aliased_io, so it accepts either. +// TR02 marks a blob whose metadata carries aliased_io; TR01 is one without. +// This parser handles aliased_io, so it accepts either, and holds TR02 to its +// promise: see the cross-check at the end of parse_metadata_json. constexpr char TENSORRT_MAGIC[4] = {'T', 'R', '0', '1'}; constexpr char TENSORRT_MAGIC_ALIASED_IO[4] = {'T', 'R', '0', '2'}; constexpr uint32_t METADATA_OFFSET_FIELD_OFFSET = 4; @@ -28,7 +30,19 @@ std::size_t skip_ws(const std::string& s, std::size_t pos) { return pos; } -std::size_t parse_string(const std::string& s, std::size_t pos, std::string& out) { +// A backslash is dropped and the character after it kept, which is not what a +// JSON escape means: "\n" comes out as the letter n, and a \u escape as the +// letter u and its four digits. The writer is json.dumps, which escapes every +// control character and everything non-ASCII, so those are exactly the strings +// the two ends would read differently. Rather than decode them, the two callers +// that record a name refuse one that arrived escaped (see +// usable_as_binding_name); saw_escape is how they are told. Every other caller +// -- the keys, and the values skip_value walks past -- only has to find the end +// of the string, which this does correctly either way. +std::size_t parse_string(const std::string& s, std::size_t pos, std::string& out, bool* saw_escape = nullptr) { + if (saw_escape != nullptr) { + *saw_escape = false; + } if (pos >= s.size() || s[pos] != '"') { return std::string::npos; } @@ -36,6 +50,9 @@ std::size_t parse_string(const std::string& s, std::size_t pos, std::string& out out.clear(); while (pos < s.size() && s[pos] != '"') { if (s[pos] == '\\' && pos + 1 < s.size()) { + if (saw_escape != nullptr) { + *saw_escape = true; + } ++pos; } out += s[pos++]; @@ -87,16 +104,36 @@ std::size_t skip_value(const std::string& s, std::size_t pos) { return pos; } +// The two scalar fields are found by searching the metadata text for the key, +// quotes included, so anything else quoted the same way would be matched +// instead: a binding named device_id sitting in the alias array the search runs +// over, or a string value that is exactly the key name. Requiring the match to +// be in key position -- its own colon next, whitespace aside -- is what tells +// the two apart, since a value is followed by a comma or a closing brace. A +// value that merely *contains* the key text needs no rule: json.dumps escapes +// the quotes it carries, so the closing quote of the search never lines up and +// there is no match to reject. An occurrence that is not in key position is +// passed over rather than refused: the key may still be ahead of it, and if it +// is not, the field is absent and keeps its default, which is what a blob +// written before the field existed wants. +std::size_t value_pos_after_key(const std::string& json, std::size_t search_from, const char* key) { + const std::size_t key_len = std::strlen(key); + std::size_t pos = search_from; + while ((pos = json.find(key, pos)) != std::string::npos) { + const std::size_t colon = skip_ws(json, pos + key_len); + if (colon < json.size() && json[colon] == ':') { + return skip_ws(json, colon + 1); + } + pos += key_len; + } + return std::string::npos; +} + bool parse_bool_after_key(const std::string& json, std::size_t search_from, const char* key, bool& value) { - const std::size_t key_pos = json.find(key, search_from); - if (key_pos == std::string::npos) { + const std::size_t val = value_pos_after_key(json, search_from, key); + if (val == std::string::npos) { return true; } - const std::size_t colon = json.find(':', key_pos); - if (colon == std::string::npos) { - return false; - } - const std::size_t val = skip_ws(json, colon + 1); if (json.compare(val, 4, "true") == 0) { value = true; return true; @@ -109,31 +146,36 @@ bool parse_bool_after_key(const std::string& json, std::size_t search_from, cons } bool parse_int_after_key(const std::string& json, std::size_t search_from, const char* key, int& value) { - const std::size_t key_pos = json.find(key, search_from); - if (key_pos == std::string::npos) { + std::size_t pos = value_pos_after_key(json, search_from, key); + if (pos == std::string::npos) { return true; } - const std::size_t colon = json.find(':', key_pos); - if (colon == std::string::npos) { - return false; - } - std::size_t pos = skip_ws(json, colon + 1); bool neg = false; if (pos < json.size() && json[pos] == '-') { neg = true; ++pos; } - int parsed = 0; + // Accumulated in 64 bits and bounded on every digit, because the digit count + // is the blob's to choose. Overflowing an int here is undefined behaviour -- + // a trapping build aborts on it -- and in an ordinary one it wraps, which is + // the bad case: the only field parsed this way is device_id, and a value just + // over four billion wraps onto a device that exists, so cudaSetDevice then + // succeeds and the engine deserializes on a GPU nobody asked for. Anything + // outside an int is refused instead. + int64_t parsed = 0; bool saw_digit = false; while (pos < json.size() && json[pos] >= '0' && json[pos] <= '9') { saw_digit = true; parsed = parsed * 10 + (json[pos] - '0'); + if (parsed > std::numeric_limits::max()) { + return false; + } ++pos; } if (!saw_digit) { return false; } - value = neg ? -parsed : parsed; + value = static_cast(neg ? -parsed : parsed); return true; } @@ -143,18 +185,23 @@ bool parse_int_after_key(const std::string& json, std::size_t search_from, const // name can hold a NUL: two names differing only after one are distinct here and // are the same tensor to TensorRT, which is exactly what those refusals exist to // stop, and a name that is only a NUL is non-empty here and empty to TensorRT. -// Refusing a NUL outright makes what the parser compares be what the engine -// will compare. No writer emits one -- json.dumps escapes it -- so this refuses -// only a blob assembled some other way. +// A name that arrived escaped is the same problem one step earlier: parse_string +// does not decode escapes, so the name recorded here is not the one the writer +// wrote and so not one the engine has either, and two names differing only in an +// escape collapse into one and are refused as a repeat. Refusing both outright +// makes what the parser compares be what the engine will compare. No writer +// emits a NUL -- json.dumps escapes it -- and json.dumps escapes only what a +// binding name has no business holding, so this refuses only a blob assembled +// some other way. // Like every refusal in this parser it is silent: parse() returns false and // the caller reports its own generic parse failure, so a NUL is not // distinguishable at load time from a bad magic or a wrapped extent. Refusing // uniformly is the deliberate choice here, not a missing diagnostic. -bool usable_as_binding_name(const std::string& name) { - return !name.empty() && name.find('\0') == std::string::npos; +bool usable_as_binding_name(const std::string& name, bool escaped) { + return !escaped && !name.empty() && name.find('\0') == std::string::npos; } -bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { +bool parse_metadata_json(const std::string& json, bool expects_aliased_io, TensorRTBlobHeader& out) { out.input_binding_names.clear(); out.output_binding_names.clear(); out.aliased_io.clear(); @@ -201,6 +248,8 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { std::string name; bool is_input = false; bool saw_name = false; + bool saw_is_input = false; + bool name_escaped = false; while (true) { pos = skip_ws(json, pos); @@ -228,12 +277,13 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { pos = skip_ws(json, pos + 1); if (key == "name") { - pos = parse_string(json, pos, name); + pos = parse_string(json, pos, name, &name_escaped); saw_name = pos != std::string::npos; if (!saw_name) { return false; } } else if (key == "is_input") { + saw_is_input = true; if (json.compare(pos, 4, "true") == 0) { is_input = true; pos += 4; @@ -257,7 +307,19 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { // argument list keeps its full length, and the two are only inferred from // the engine when both are empty -- so one real name beside a blank leaves // a short list that no longer lines up with the engine's bindings. - if (!saw_name || !usable_as_binding_name(name)) { + if (!saw_name || !usable_as_binding_name(name, name_escaped)) { + return false; + } + // is_input has no safe default, so an entry without it is refused beside + // the nameless one. The initializer here reads the binding as an output + // and TensorRTIOBinding.is_input in serialization.py defaults to an input, + // so the two readers of these bytes would disagree -- and one binding + // changing list shifts every index after it, which on a static-shape + // engine leaves nothing for shape inference to object to: it runs on the + // wrong tensors. A key misspelled by one byte is the same case, since it + // falls through to skip_value and leaves the initializer standing. The + // writer always emits the key. + if (!saw_is_input) { return false; } if (!claimed_bindings.insert(name).second) { @@ -276,6 +338,7 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { // // Search from pos (past the io_bindings array) so a model input literally // named "aliased_io" isn't matched as the array key. + std::size_t scalars_from = pos; const std::size_t alias_key = json.find("\"aliased_io\"", pos); if (alias_key != std::string::npos) { std::size_t apos = json.find('[', alias_key); @@ -304,6 +367,8 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { ++apos; AliasedBinding ab; + bool output_escaped = false; + bool input_escaped = false; while (true) { apos = skip_ws(json, apos); if (apos >= json.size()) { @@ -328,9 +393,9 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { } apos = skip_ws(json, apos + 1); if (key == "output") { - apos = parse_string(json, apos, ab.output); + apos = parse_string(json, apos, ab.output, &output_escaped); } else if (key == "input") { - apos = parse_string(json, apos, ab.input); + apos = parse_string(json, apos, ab.input, &input_escaped); } else if (key == "kind") { apos = parse_string(json, apos, ab.kind); } else { @@ -346,7 +411,7 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { // argument list, so a dropped entry surfaces as an argument-count error at // every execute, which never mentions aliasing, instead of at parse, which // the blob-header tests reach without a GPU. - if (!usable_as_binding_name(ab.output) || !usable_as_binding_name(ab.input)) { + if (!usable_as_binding_name(ab.output, output_escaped) || !usable_as_binding_name(ab.input, input_escaped)) { return false; } // An output binding may be claimed by at most one entry. A second entry @@ -381,10 +446,30 @@ bool parse_metadata_json(const std::string& json, TensorRTBlobHeader& out) { } out.aliased_io.push_back(std::move(ab)); } + // Past the alias array, not merely past io_bindings: the array is written + // between the two and this walk advances a position of its own, so the + // scalar scans below would otherwise search across the alias entries and + // match an aliased binding named like one of their keys. + scalars_from = apos; + } + + // The magic says which of the two shapes this blob is, and until here nothing + // compared that against what was parsed. TR02 means the metadata carries + // aliased_io, so an empty list is a blob whose alias array this walk did not + // find -- absent, truncated away, one byte wrong in the key, or written + // before io_bindings, since the search starts past that array. The threaded + // shape does not fail afterwards: every aliased output gets storage of its + // own and the caller's cache quietly stops being updated, which is the + // outcome the aliasing code exists to make impossible. The converse is not + // refused: a TR01 blob carrying an alias array is read as aliased and run + // correctly here, and refusing it would only turn a blob this runtime handles + // into one it does not. + if (expects_aliased_io && out.aliased_io.empty()) { + return false; } - return parse_bool_after_key(json, pos, "\"hardware_compatible\"", out.hardware_compatible) && - parse_int_after_key(json, pos, "\"device_id\"", out.device_id); + return parse_bool_after_key(json, scalars_from, "\"hardware_compatible\"", out.hardware_compatible) && + parse_int_after_key(json, scalars_from, "\"device_id\"", out.device_id); } } // namespace @@ -399,8 +484,8 @@ bool TensorRTBlobHeader::parse(const void* data, std::size_t size, TensorRTBlobH } const auto* bytes = static_cast(data); - if (std::memcmp(bytes, TENSORRT_MAGIC, sizeof(TENSORRT_MAGIC)) != 0 && - std::memcmp(bytes, TENSORRT_MAGIC_ALIASED_IO, sizeof(TENSORRT_MAGIC_ALIASED_IO)) != 0) { + const bool aliased_io_magic = std::memcmp(bytes, TENSORRT_MAGIC_ALIASED_IO, sizeof(TENSORRT_MAGIC_ALIASED_IO)) == 0; + if (!aliased_io_magic && std::memcmp(bytes, TENSORRT_MAGIC, sizeof(TENSORRT_MAGIC)) != 0) { return false; } @@ -445,7 +530,7 @@ bool TensorRTBlobHeader::parse(const void* data, std::size_t size, TensorRTBlobH } std::string json(reinterpret_cast(bytes + out.metadata_offset), out.metadata_size); - return parse_metadata_json(json, out); + return parse_metadata_json(json, aliased_io_magic, out); } } // namespace executorch_backend diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index e34055c2f8f..1d9b86b8bed 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -390,8 +390,21 @@ wherever memory planning put it, so ``enable_non_cpu_memory_planning=False`` -- which plans every tensor into a single host arena whatever device its ``TensorSpec`` asks for -- cannot be combined with zero-copy KV: ``to_executorch`` raises instead of writing a ``.pte`` whose every ``execute()`` -fails on a host pointer. And ``propagate_device_config.skip_h2d_for_method_inputs`` -is refused outright: ``PropagateDevicePass`` refuses to un-stage a method input +fails on a host pointer. That refusal reads the config object +``zero_copy_backend_config`` returned, so two ways of setting the field are +outside it. Building a new config out of the returned one with +``dataclasses.replace`` copies the field by value and the pass by reference, so +the pass goes on reading the config it was built for: call +``zero_copy_backend_config`` again on the derived config. And a +``memory_planning_pass`` of your own that does not already have an attribute of +that name never receives the field at all -- ``to_executorch`` assigns it onto +the planner rather than passing it -- so where the caches land is that planner's +own business. Neither is left to the runtime to discover: ``check_zero_copy_kv`` +reads the arena planning actually chose, and the manager ``export`` returns runs +it for you -- so long as that planner records its arenas, which is its own +responsibility and is covered below. And +``propagate_device_config.skip_h2d_for_method_inputs`` is refused outright: +``PropagateDevicePass`` refuses to un-stage a method input whose placeholder does not have exactly one user, and a zero-copy cache always has two, so preserving that option would hand back a configuration that cannot finalize at all. It is refused wherever it is written -- in one @@ -407,19 +420,24 @@ one silently would break a runner built before this feature. .. warning:: - **Both calls are required.** Exporting with ``zero_copy_kv=True`` and then - finalizing without ``zero_copy_backend_config`` does not raise on its own: - the engine writes a per-call staging copy that is discarded, and the cache - never updates. For a KV cache that is wrong output, not a crash. Pass the - finalized program to ``torch_tensorrt.executorch.check_zero_copy_kv``, which - reads it back and refuses one whose caches are still staged, or planned - somewhere the engine cannot write them, before writing the ``.pte``:: - - program = edge.to_executorch(zero_copy_backend_config()) + **Both calls are required.** Exporting with ``zero_copy_kv=True`` removes + the copy-back; finalizing without ``zero_copy_backend_config`` leaves the + engine writing a per-call staging copy that is discarded, so the cache never + updates. For a KV cache that is wrong output, not a crash. The manager + ``export`` returns is what stops that: its ``to_executorch`` reads its own + finalized program through + ``torch_tensorrt.executorch.check_zero_copy_kv``, which refuses one whose + caches are still staged, or planned somewhere the engine cannot write them, + before there is a ``.pte`` to write. + + ``transform()`` and ``to_backend()`` return a *new* manager, which that + check does not travel to, so a program finalized off one of those needs it + by hand:: + + program = edge.transform(passes).to_executorch(zero_copy_backend_config()) torch_tensorrt.executorch.check_zero_copy_kv(program) - ``torch_tensorrt.save`` owns both ends and runs that check itself, so there - is nothing to pair on that path. + ``torch_tensorrt.save`` owns both ends and runs the same check itself. ``torch_tensorrt.save`` finalizes the program itself, so a single ``zero_copy_kv=True`` covers both steps: @@ -461,6 +479,16 @@ Three further responsibilities are the caller's, and none raises: have to land at the same arena offsets -- which ExecuTorch's memory planner owns and which is deployment-specific. Supply your own ``memory_planning_pass`` for it; ``zero_copy_backend_config`` preserves it. + Such a planner has to record the device of each arena it places, which means + running ExecuTorch's ``apply_algo`` -- the only thing that writes + ``non_const_buffer_device`` -- and running it with + ``enable_non_cpu_memory_planning=True``, since that parameter defaults to + ``False`` and with it off ``apply_algo`` plans every spec into one CPU bucket + and records nothing either. Without that record the ``.pte`` reports every + planned buffer as CPU, and a runner that honours it backs the cache with host + memory the engine cannot write. A missing record is not something + ``check_zero_copy_kv`` refuses on: it is also what a device-aware planner of + your own leaves, so it is read as "cannot tell" rather than as the host. **Coalesced TensorRT + CUDA .pte** diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index d9967363a90..fcf3408472d 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -909,6 +909,15 @@ def save( normalize_weight_streaming_budget_per_engine( executorch_weight_streaming_budget_per_engine ) + # For the same reason, the one refusal zero-copy makes on the config + # alone. _save_as_executorch reaches it only through + # zero_copy_backend_config, which it calls after export() has partitioned + # the graph and built every engine -- a whole compile before the caller + # is told the combination is not allowed. + if executorch_zero_copy_kv and executorch_backend_config is not None: + from torch_tensorrt.executorch._zero_copy import _refuse_skip_h2d + + _refuse_skip_h2d(executorch_backend_config) def _all_are_input_objects(obj: Any) -> bool: """Recursively check if all elements in nested collections are Input objects.""" @@ -1498,7 +1507,10 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None # shows whether the caches actually reach the engine un-staged and in an # arena it can write. Both halves of zero-copy no-op quietly when they # find nothing to do, so without this a save() that asked for zero-copy - # could still write an ordinary staged .pte. + # could still write an ordinary staged .pte. export() installs the same + # check on the manager it returns, but only for a method it rewired + # something in: the case this call adds is the one where it rewired + # nothing, which is a zero_copy_kv=True that bought the caller nothing. check_zero_copy_kv(executorch_program) with open(file_path, "wb") as f: executorch_program.write_to_file(f) diff --git a/py/torch_tensorrt/executorch/_export.py b/py/torch_tensorrt/executorch/_export.py index e4ea70d1c56..093ec0eb37b 100644 --- a/py/torch_tensorrt/executorch/_export.py +++ b/py/torch_tensorrt/executorch/_export.py @@ -454,8 +454,12 @@ def export( producing one silently would break an older runner; and it is only half the change. Finalize such a program with ``to_executorch(torch_tensorrt.executorch.zero_copy_backend_config(config))`` - -- without it the buffer is still staged and its updates are discarded, with - no error. + -- without it the buffer is still staged and its updates would be discarded. + The returned manager reads its own finalized program back through + :func:`torch_tensorrt.executorch.check_zero_copy_kv` and raises on that + rather than handing back a program whose caches never update. That reaches + the manager returned here and no other: ``transform()`` and ``to_backend()`` + build a new one, so a program finalized off either owes that call by hand. Only a buffer the engine declares aliased is affected. A method may hold both kinds at once: a mutable buffer with no aliasing available -- a convolution @@ -513,9 +517,10 @@ def export( zero_copy_kv (bool): Let a TensorRT engine write its aliased mutable buffer -- a KV cache -- in place, instead of receiving a staging copy that ExecuTorch copies back. Requires finalizing the returned program with - :func:`torch_tensorrt.executorch.zero_copy_backend_config`; without that - second call the buffer is still staged and every update is discarded, with - no error. ``torch_tensorrt.save(output_format="executorch", + :func:`torch_tensorrt.executorch.zero_copy_backend_config`; the returned + manager's ``to_executorch`` refuses a program finalized without it, since + the buffer would still be staged and every update discarded. + ``torch_tensorrt.save(output_format="executorch", zero_copy_kv=True)`` owns both steps and needs only the one flag. generate_etrecord (bool): Ask ExecuTorch for an ETRecord for later debugging. This copies the whole program, engines included. @@ -562,7 +567,10 @@ def export( stage_exported_program, validate_engine_program, ) - from torch_tensorrt.executorch._zero_copy import order_copyback_mutations_first + from torch_tensorrt.executorch._zero_copy import ( + _check_zero_copy_kv_when_finalized, + order_copyback_mutations_first, + ) from torch_tensorrt.executorch.backend import ( ZERO_COPY_KV_COMPILE_SPEC_KEY, _serialize_elided_output_names, @@ -747,4 +755,9 @@ def export( # order_copyback_mutations_first. for name in edge_manager.methods: order_copyback_mutations_first(edge_manager.exported_program(name)) + # The copy-back is gone from here on, so finalizing without the matching + # un-staging writes a .pte whose caches never update. This makes the + # manager's own to_executorch read the finalized program back and refuse + # that, which the two-call API otherwise leaves to the caller. + _check_zero_copy_kv_when_finalized(edge_manager) return edge_manager diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index d3091d893e1..f6e902c820e 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -35,6 +35,7 @@ up writing the caller's buffer -- is what this module exports. """ +import functools import json import logging import operator @@ -49,13 +50,6 @@ _LOGGER = logging.getLogger(__name__) -def _engine_info_str(engine_info: List[Any], index: int) -> str: - if index < 0 or index >= len(engine_info) or engine_info[index] is None: - return "" - value = engine_info[index] - return value.decode("utf-8", "replace") if isinstance(value, bytes) else str(value) - - def _aliased_inputs_by_output_index( exported_program: Any, engine_node: Node ) -> Dict[int, Node]: @@ -81,19 +75,20 @@ def _aliased_inputs_by_output_index( deserialize_aliased_io, ) from torch_tensorrt.executorch._export_utils import _resolve_engine_info + from torch_tensorrt.executorch.backend import _get_str # Only aliased_io and the binding names are read, never the engine itself. engine_info = _resolve_engine_info( exported_program, engine_node, metadata_only=True ) - aliased_io = deserialize_aliased_io(_engine_info_str(engine_info, ALIASED_IO_IDX)) + aliased_io = deserialize_aliased_io(_get_str(engine_info, ALIASED_IO_IDX)) if not aliased_io: return {} input_names = deserialize_binding_names( - _engine_info_str(engine_info, INPUT_BINDING_NAMES_IDX) + _get_str(engine_info, INPUT_BINDING_NAMES_IDX) ) output_names = deserialize_binding_names( - _engine_info_str(engine_info, OUTPUT_BINDING_NAMES_IDX) + _get_str(engine_info, OUTPUT_BINDING_NAMES_IDX) ) input_nodes = list(engine_node.args[0]) @@ -125,12 +120,13 @@ def _engine_output_binding_names(exported_program: Any, engine_node: Node) -> Li deserialize_binding_names, ) from torch_tensorrt.executorch._export_utils import _resolve_engine_info + from torch_tensorrt.executorch.backend import _get_str engine_info = _resolve_engine_info( exported_program, engine_node, metadata_only=True ) names: List[str] = deserialize_binding_names( - _engine_info_str(engine_info, OUTPUT_BINDING_NAMES_IDX) + _get_str(engine_info, OUTPUT_BINDING_NAMES_IDX) ) return names @@ -525,45 +521,62 @@ def _delegate_declares_zero_copy( return _zero_copy_compile_spec(graph_module, node) is not None -def _device_move_is_safe( +def _device_placement_is_safe( graph_module: torch.fx.GraphModule, source: Node, h2d_copy: Any, target_device: Any, target_device_index: Any, + already_removed: Any = (), ) -> bool: - """True when moving ``source``'s spec device disturbs no other consumer. + """True when every other consumer of ``source`` survives its device placement. - A placeholder's device is shared by every user, so it can only be retargeted - to the delegate's device when nothing else *reads* it afterwards. ExecuTorch - guards the same hazard, more strictly and only under its opt-in - ``skip_h2d_for_method_inputs``: it demands the placeholder have exactly one - user. The rule here is looser because two kinds of user survive the move + A placeholder's device is shared by every user, so it can only be planned in + the delegate's device memory when nothing else *reads* it from somewhere + else. ExecuTorch guards the same hazard, more strictly and only under its + opt-in ``skip_h2d_for_method_inputs``: it demands the placeholder have + exactly one user. The rule here is looser because some users survive unaffected and are allowed: the graph ``output`` node -- the buffer is its own BUFFER_MUTATION result, which is exactly what zero-copy sets up and - which carries no device of its own -- and an ``_h2d_copy`` to the same GPU - that this pass removes, because every one of its users is a TensorRT - delegate whose argument the pass rewires to the buffer. + which carries no device of its own -- an ``_h2d_copy`` to the same GPU that + this pass removes, because every one of its users is a TensorRT delegate + whose argument the pass rewires to the buffer, and a TensorRT delegate + already taking the buffer itself, which is the shape this pass leaves behind + and so what a second marked delegate, or a second run, finds. + + ``already_removed`` is the last of them: the staging copies this run has + detached, which are still in the graph because they are erased only once the + walk has succeeded. A detached one has no users left, which is the shape a + *foreign* dead copy has too, and that one is refused -- nothing erases it, + so the emitter keeps it and it reads device memory as a host source. Ours + are named here rather than inferred, so the two cannot be confused. A staging copy that outlives the pass is *not* allowed, even on the same - GPU. It would go on reading the buffer as its source after the move has put - the buffer in device memory, and ``_h2d_copy_out`` requires a host source: - the portable kernel checks it and fails ``InvalidArgument``. + GPU. It would go on reading the buffer as its source once the buffer is in + device memory, and ``_h2d_copy_out`` requires a host source: the portable + kernel checks it and fails ``InvalidArgument``. The index is compared as well as the type, because ``spec.device`` is only ``CUDA``/``CPU``: two engines resolved to ``cuda:0`` and ``cuda:1`` stage the same buffer to different GPUs, and un-staging both would leave whichever ran last owning the buffer while the other engine writes an address on the wrong device. + + This is asked whether or not the buffer's spec already names the target + device. Where it does, the pass changes no device and the surviving copy is + already reading device memory as a host source -- a program that was broken + before this pass touched it -- but the pass is about to hand that buffer to + an engine on the strength of the same post-condition, so it is refused here + rather than left to fail on first execution. """ for user in source.users: if user.op == "output": continue - if not ( - isinstance(user, Node) - and user.op == "call_function" - and user.target is h2d_copy - ): + if not isinstance(user, Node) or user.op != "call_function": + return False + if user in already_removed or _is_tensorrt_delegate(graph_module, user): + continue + if user.target is not h2d_copy: return False spec = user.meta.get("spec") if spec is None or spec.device != target_device: @@ -591,10 +604,11 @@ def _unstage_aliased_buffers( asks memory planning for the delegate's device arena rather than a host one (asks, not settles -- see ``device_memory_planning`` below). That is what makes handing the buffer straight to the engine valid at all: a host-arena - pointer is not something the engine can write. That move is refused when the - buffer has a consumer this pass leaves behind (see - :func:`_device_move_is_safe`), which would otherwise be reading the buffer - from somewhere other than where it was put. + pointer is not something the engine can write. It is refused when the buffer + has a consumer this pass leaves behind that does not survive being handed + the buffer from there (see :func:`_device_placement_is_safe`) -- asked + whether or not the spec already names that device, since it is the placement + and not the change of device that such a consumer does not survive. What the pass has to establish is the *post-condition*: every marked buffer is a direct argument of a TensorRT delegate *and* ends up planned in device @@ -613,9 +627,11 @@ def _unstage_aliased_buffers( ones this pass un-stages as much as the ones that already reach their delegate directly. That is why it is checked once for the whole graph, up front, rather than on one of the two branches below. - ``device_memory_planning`` carries the configuration in; - :func:`unstage_aliased_buffers_pass` reads it off the finalization config - when it runs. + ``device_memory_planning`` carries the configuration in; the pass + :func:`unstage_aliased_buffers_pass` builds reads it off the finalization + config when it runs, resolved as :func:`_config_plans_on_devices` describes, + so a configuration whose planner never receives the flag is not refused on + it. A failure here is a lost KV update, so it is raised rather than logged: export has already removed the copy-back, so a marked buffer left staged has @@ -623,7 +639,7 @@ def _unstage_aliased_buffers( updates. It raises when memory planning is host-only, when either end of the move -- the staging copy or the buffer -- has no spec, when the staging copy is not on CUDA, when a direct argument has no spec of its own - or that spec is not on CUDA, when the device move is unsafe, and -- so a + or that spec is not on CUDA, when the placement is unsafe, and -- so a discovery miss cannot pass silently -- after the loop when the post-condition does not hold for some marked buffer, cross-checked against each delegate's own ``zero_copy_kv`` spec: a TensorRT delegate that takes fewer marked @@ -658,9 +674,15 @@ def _unstage_aliased_buffers( h2d_copy = torch.ops.et_copy._h2d_copy.default unstaged = 0 satisfied_placeholders: Set[Node] = set() - orphaned_stagings: List[Node] = [] + # A dict rather than a list: the walk asks whether a staging copy has + # already been detached, and the erase below wants each one once. + orphaned_stagings: Dict[Node, None] = {} zero_copy_delegates: List[Node] = [] - satisfied_per_delegate: Dict[Node, int] = {} + # Distinct buffers, not argument slots: one buffer occupying two of a + # delegate's slots is one cache written in place, and counting the slots + # would let it satisfy a spec naming two elided outputs. The finalized-program + # check counts the same way, so the two cannot disagree about one graph. + satisfied_per_delegate: Dict[Node, Set[Node]] = {} for node in list(graph_module.graph.nodes): if not _is_tensorrt_delegate(graph_module, node): @@ -668,7 +690,7 @@ def _unstage_aliased_buffers( declares_zero_copy = _delegate_declares_zero_copy(graph_module, node) if declares_zero_copy: zero_copy_delegates.append(node) - satisfied_per_delegate[node] = 0 + satisfied_per_delegate[node] = set() new_args = list(node.args) for i, arg in enumerate(node.args[1:], start=1): if not isinstance(arg, Node): @@ -705,7 +727,7 @@ def _unstage_aliased_buffers( ) satisfied_placeholders.add(arg) if declares_zero_copy: - satisfied_per_delegate[node] += 1 + satisfied_per_delegate[node].add(arg) continue if arg.target is not h2d_copy: continue @@ -742,22 +764,20 @@ def _unstage_aliased_buffers( "buffer in place and its copy-back has already been removed, " "so the update would be lost." ) - already_placed = (source_spec.device, source_spec.device_index) == ( - staged_spec.device, - staged_spec.device_index, - ) - if not already_placed and not _device_move_is_safe( + if not _device_placement_is_safe( graph_module, source, h2d_copy, staged_spec.device, staged_spec.device_index, + orphaned_stagings, ): raise RuntimeError( "TensorRT zero-copy KV: buffer " f"'{source.name}' is read by a consumer this pass leaves in " - "place, so placing the buffer on this engine's device would " - "change the device that consumer reads it from. Export this " + "place that does not survive the buffer being planned in this " + "engine's device memory -- a staging copy left behind reads it " + "as a host source and fails InvalidArgument. Export this " "method without zero_copy_kv, or stop sharing the aliased " "buffer." ) @@ -766,9 +786,9 @@ def _unstage_aliased_buffers( new_args[i] = source unstaged += 1 satisfied_placeholders.add(source) - orphaned_stagings.append(arg) + orphaned_stagings[arg] = None if declares_zero_copy: - satisfied_per_delegate[node] += 1 + satisfied_per_delegate[node].add(source) node.args = tuple(new_args) marked_but_unsatisfied = [ @@ -793,7 +813,7 @@ def _unstage_aliased_buffers( # decoding -- cannot say how many to expect, so it falls back to # demanding at least one. elided = _delegate_elided_output_names(graph_module, delegate) - satisfied = satisfied_per_delegate[delegate] + satisfied = len(satisfied_per_delegate[delegate]) if satisfied >= max(len(elided), 1): continue delegate_inputs = [ @@ -819,7 +839,7 @@ def _unstage_aliased_buffers( if unstaged: # Erase only the stagings we orphaned. A graph-wide eliminate_dead_code() # in a to_out_var_pass could delete another backend's unused delegate. - for staging in dict.fromkeys(orphaned_stagings): + for staging in orphaned_stagings: if not staging.users: graph_module.graph.erase_node(staging) graph_module.graph.lint() @@ -827,6 +847,33 @@ def _unstage_aliased_buffers( return unstaged +def _config_plans_on_devices(config: "ExecutorchBackendConfig") -> bool: + """Whether ``enable_non_cpu_memory_planning`` is anything this config decides by. + + ``to_executorch`` does not hand the flag to the memory planner it is given. + It *assigns* it, and only onto a planner that already has an attribute of + that name (``exir/program/_program.py``), which in practice means + ``MemoryPlanningPass`` or a subclass. A caller who brings their own planner + -- which the user guide tells people to do for a cache shared between + prefill and decode -- never receives the flag, and what that planner does + with the specs is its own business: ``False`` does not put the caches in the + host arena, and ``True`` does not keep them out of it. So the flag answers + nothing for such a config and ``True`` is returned, meaning only "not a + ground to refuse on"; :func:`check_zero_copy_kv` reads the arena that + planner actually chose, which is the answer this cannot give. + + ``memory_planning_pass`` may also be a per-method dict, and this pass sees + one graph module without its method name, so a dict is read as deciding + nothing unless every planner in it takes the flag. A method the dict omits + gets ExecuTorch's default planner, which does. + """ + planner = config.memory_planning_pass + planners = list(planner.values()) if isinstance(planner, dict) else [planner] + if not all(hasattr(p, "enable_non_cpu_memory_planning") for p in planners): + return True + return bool(config.enable_non_cpu_memory_planning) + + def unstage_aliased_buffers_pass( inner_pass: Optional[Any] = None, *, device_memory_planning: bool = True ) -> Any: @@ -846,9 +893,11 @@ def unstage_aliased_buffers_pass( write, so the pass has to be told: see :func:`_unstage_aliased_buffers`. It is what a pass built here and left unbound uses. Set ``finalization_config`` on the returned pass to the - ``ExecutorchBackendConfig`` the program will be finalized with and the flag - is read off that config instead, on every call, which is the value memory - planning will use a few passes later. + ``ExecutorchBackendConfig`` the program will be finalized with and the pass + reads the flag off that config instead, on every call, as + :func:`_config_plans_on_devices` resolves it -- which is what memory planning + will do with it a few passes later, and is nothing at all when the config + carries a planner that does not take the flag. ``ExecutorchBackendConfig`` is a plain mutable dataclass, so the field can be set again on the very config being finalized after the pass was built; a pass reading a value captured here would then accept a program the finalizer plans @@ -873,7 +922,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> Any: planning = ( device_memory_planning if config is None - else config.enable_non_cpu_memory_planning + else _config_plans_on_devices(config) ) unstaged = _unstage_aliased_buffers( graph_module, device_memory_planning=planning @@ -884,8 +933,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> Any: return _UnstageThenToOutVar() -def _device_planned_arenas(graph_module: torch.fx.GraphModule) -> Optional[Set[int]]: - """The ``mem_id``s of the finalized program's CUDA arenas, or ``None``. +def _device_planned_arenas( + graph_module: torch.fx.GraphModule, +) -> Optional[Dict[int, Any]]: + """The finalized program's CUDA arenas, ``mem_id -> device index``, or ``None``. Memory planning partitions the specs by device, gives each device its own arena, and records the non-CPU ones on the graph module as @@ -895,6 +946,12 @@ def _device_planned_arenas(graph_module: torch.fx.GraphModule) -> Optional[Set[i callable as ``memory_planning_pass``, so a caller-supplied planner -- which the user guide tells people to bring for a cache shared between prefill and decode -- can plan onto a device and still leave the key unwritten. + + The index is carried, not only the type, because the record is what the + runtime allocates from: an arena recorded for ``cuda:1`` holding a cache the + engine writes on ``cuda:0`` is a pointer on the wrong GPU, which the device + *type* alone cannot tell from a correct program. ``apply_algo`` writes one + entry per arena, so a ``mem_id`` names at most one device. """ from executorch.exir.schema import DeviceType @@ -902,7 +959,7 @@ def _device_planned_arenas(graph_module: torch.fx.GraphModule) -> Optional[Set[i if not entries: return None return { - entry.buffer_idx + entry.buffer_idx: getattr(entry, "device_index", None) for entry in entries if getattr(entry, "device_type", None) == DeviceType.CUDA } @@ -939,7 +996,7 @@ def _host_planned_arenas(graph_module: torch.fx.GraphModule) -> Set[int]: def _is_host_planned( - node: Node, device_arenas: Optional[Set[int]], host_arenas: Set[int] + node: Node, device_arenas: Optional[Dict[int, Any]], host_arenas: Set[int] ) -> bool: """True when memory planning put ``node`` somewhere the engine cannot write. @@ -954,6 +1011,32 @@ def _is_host_planned( return device_arenas is not None and mem_id not in device_arenas +def _planned_on_another_gpu( + node: Node, device_arenas: Optional[Dict[int, Any]] +) -> Optional[str]: + """How ``node``'s arena and its own spec disagree about which GPU, if they do. + + Being in a CUDA arena is not enough: the engine writes the cache through the + pointer the runtime allocates out of that arena, so an arena the program + records for another GPU is an address the engine cannot write, exactly as a + host one is. Only a disagreement between two recorded indices is read as + one -- either side left unrecorded says nothing, and the arena's own device + type has already been established by the caller. + """ + if device_arenas is None: + return None + spec = node.meta.get("spec") + mem_id = getattr(spec, "mem_id", None) + arena_index = None if mem_id is None else device_arenas.get(mem_id) + spec_index = getattr(spec, "device_index", None) + if arena_index is None or spec_index is None or arena_index == spec_index: + return None + return ( + f"'{node.name}' asks for cuda:{spec_index} and was planned in an arena " + f"the program records as cuda:{arena_index}" + ) + + def _name_detail(names_by_method: Dict[str, List[str]]) -> str: return ", ".join( f"'{name}' in method '{method}'" @@ -967,21 +1050,22 @@ def check_zero_copy_kv(program: Any) -> None: ``program`` is what ``to_executorch()`` returns. Both halves of zero-copy do nothing quietly when they find nothing to do: ``zero_copy_kv=True`` warns and - carries on when the model holds no aliased buffer mutation, and - :func:`unstage_aliased_buffers_pass`, handed a program with nothing marked, - un-stages nothing and returns. Neither of those is wrong output -- nothing + carries on when the model holds no aliased buffer mutation, and the pass + :func:`unstage_aliased_buffers_pass` builds, handed a program with nothing + marked, un-stages nothing and returns. Neither of those is wrong output -- nothing removed a copy-back, so the ``.pte`` stages its cache and updates it like any other -- but the optimization the caller asked for is silently not there, and a caller who reads the successful ``save`` as proof it is gets neither an error nor the speedup. The wrong-output case is the one below: a rewiring that did happen and then lost its mark. - Four shapes are refused: a marked buffer that is not a direct argument of a + Five shapes are refused: a marked buffer that is not a direct argument of a TensorRT delegate carrying the zero-copy compile spec, a stamped delegate that takes fewer marked buffers than its own spec says it elided aliased outputs -- including one in a method with no marked buffer at all, which is the lost-mark case -- a marked buffer that reaches such a delegate directly - but is not planned in device memory, and a program carrying neither a marked + but is not planned in device memory, one planned in a device arena the + program records for another GPU, and a program carrying neither a marked buffer nor a stamped delegate in any of its methods. The first three are what finalizing without :func:`zero_copy_backend_config` leaves behind -- all but the lost-mark case folded into the second, which no finalization @@ -1025,7 +1109,22 @@ def check_zero_copy_kv(program: Any) -> None: Absence of that record is not itself read as the host: only ``apply_algo`` writes it, so a caller-supplied ``memory_planning_pass`` that does not go through it records nothing whatever it planned, and refusing on that would - block the one thing the guide says to bring your own planner for. + block the one thing the guide says to bring your own planner for. That + acceptance is about what this function can *tell*, not about what the + runtime does with such a program: ``MethodMeta::memory_planned_buffer_device`` + answers ``CPU`` for a buffer the ``.pte`` records nothing for, so a runner + that honours it -- ``examples/executorch_reference_runner`` does -- backs + that arena with host memory and the engine then fails the alias-target guard + on every call. A planner for a zero-copy cache has to leave the record + behind, which means calling ``apply_algo`` with + ``enable_non_cpu_memory_planning=True``: that parameter defaults to + ``False``, and with it off ``apply_algo`` plans every spec into one CPU + bucket and writes no record either. Where the + record does name a GPU it is read as one: an arena recorded for ``cuda:1`` + holding a cache whose own spec asks for ``cuda:0`` is an address on the + wrong device, which fails the same way a host one does, so the index is + compared and not only the type. Either index left unrecorded says nothing + and is accepted, as an unrecorded arena is. Every method is read, not only ``forward``. ``export()`` rewires each method on its own, so a check that stopped at ``forward`` would pass a program whose @@ -1044,11 +1143,30 @@ def check_zero_copy_kv(program: Any) -> None: This reads the graph and the finalized specs, so it says what the program does rather than what the passes recorded. It still says nothing about whether the engine's write itself is correct. + + Both entry points run it for you: ``save(..., zero_copy_kv=True)`` before it + writes the file, and the ``EdgeProgramManager`` that + ``export(..., zero_copy_kv=True)`` returns on whatever its ``to_executorch`` + produces. Call it by hand for a program that reached finalization some other + way -- one derived from that manager by ``transform()`` or ``to_backend()``, + which are new managers without the hook. + + Arguments: + program (executorch.exir.ExecutorchProgramManager): The finalized program + ``to_executorch()`` returned. Every method it holds is read. + + Returns: + None: a program that passes is left exactly as it was. + + Raises: + RuntimeError: If the program does not update its KV buffers in place, + naming the buffers or delegates and the method each is in. """ method_names = sorted(program.methods) staged_by_method: Dict[str, List[str]] = {} short_by_method: Dict[str, List[str]] = {} host_planned_by_method: Dict[str, List[str]] = {} + wrong_gpu_by_method: Dict[str, List[str]] = {} marked_anywhere = False for method_name in method_names: graph_module = program.exported_program(method_name).graph_module @@ -1082,10 +1200,14 @@ def check_zero_copy_kv(program: Any) -> None: short = [] for delegate in zero_copy_delegates: elided = _delegate_elided_output_names(graph_module, delegate) - taken = sum( - 1 - for arg in delegate.args[1:] - if isinstance(arg, Node) and arg in marked_nodes + # Distinct buffers rather than argument slots, as the un-staging + # pass counts them: one buffer in two slots is one cache. + taken = len( + { + arg + for arg in delegate.args[1:] + if isinstance(arg, Node) and arg in marked_nodes + } ) if taken >= max(len(elided), 1): continue @@ -1104,14 +1226,25 @@ def check_zero_copy_kv(program: Any) -> None: short_by_method[method_name] = short device_arenas = _device_planned_arenas(graph_module) host_arenas = _host_planned_arenas(graph_module) + reaching = [node for node in marked if node in zero_copy_delegate_args] host_planned = [ node.name - for node in marked - if node in zero_copy_delegate_args - and _is_host_planned(node, device_arenas, host_arenas) + for node in reaching + if _is_host_planned(node, device_arenas, host_arenas) ] if host_planned: host_planned_by_method[method_name] = host_planned + wrong_gpu = [ + detail + for detail in ( + _planned_on_another_gpu(node, device_arenas) + for node in reaching + if not _is_host_planned(node, device_arenas, host_arenas) + ) + if detail is not None + ] + if wrong_gpu: + wrong_gpu_by_method[method_name] = wrong_gpu if staged_by_method: raise RuntimeError( f"TensorRT zero-copy KV: buffer(s) {_name_detail(staged_by_method)} " @@ -1162,6 +1295,56 @@ def check_zero_copy_kv(program: Any) -> None: "buffer among the host tensors, and it has to give the delegate's " "device an arena of its own." ) + if wrong_gpu_by_method: + raise RuntimeError( + "TensorRT zero-copy KV: buffer(s) " + + "; ".join( + f"{detail} in method '{method}'" + for method, details in wrong_gpu_by_method.items() + for detail in details + ) + + ". The engine writes the cache through the pointer the runtime " + "allocates out of that arena, so it would write the wrong GPU. The " + "memory_planning_pass in use has to give each delegate's own device " + "an arena, and put every buffer its engine writes in place in that " + "device's one." + ) + + +def _check_zero_copy_kv_when_finalized(edge_manager: Any) -> None: + """Make one manager's own ``to_executorch`` run :func:`check_zero_copy_kv`. + + Export removes the copy-back of the rewired caches before it returns, so + from that point on every way of finalizing the program that does not also + un-stage them produces a ``.pte`` whose caches never update -- silently, and + for a KV cache that is wrong output rather than a crash. ``to_executorch()`` + with ExecuTorch's defaults is one such way, and it is the call the two-step + API documents. Reading the finalized program back is what tells the two + apart, and this is the only place holding it that knows the export asked for + zero-copy, so the check runs here rather than being left for the caller to + remember. It is the same check ``save(..., zero_copy_kv=True)`` runs at the + same point for the same reason. + + Nothing about the program changes: a correct one is returned untouched, and + a broken one raises instead of being handed back. The refusals name what to + do -- finalize through :func:`zero_copy_backend_config` -- so this needs no + message of its own. + + Bound to the instance rather than to a type, because ``EdgeProgramManager`` + is ExecuTorch's. That reaches the manager ``export()`` hands back and no + other: ``transform()`` and ``to_backend()`` build a *new* manager, which + this hook does not travel to, so a caller who takes either detour owes + :func:`check_zero_copy_kv` by hand. + """ + finalize = edge_manager.to_executorch + + @functools.wraps(finalize) + def to_executorch(*args: Any, **kwargs: Any) -> Any: + program = finalize(*args, **kwargs) + check_zero_copy_kv(program) + return program + + edge_manager.to_executorch = to_executorch def _refuse_skip_h2d(config: "ExecutorchBackendConfig") -> None: @@ -1235,15 +1418,23 @@ def zero_copy_backend_config( write a ``.pte`` whose every ``execute()`` fails. It is read off the config returned here, at the moment the pass runs, so setting the field *on that object* afterwards is honoured: the pass and the finalizer then - cannot disagree about it. Building a *new* config out of this one with - ``dataclasses.replace`` is the case that does not carry -- the field is a + cannot disagree about it. There are two cases the field decides nothing + in, and that the pass therefore refuses nothing on. One is building a + *new* config out of this one with ``dataclasses.replace``: the field is a bool, copied by value, while the pass is copied by reference and goes on - reading the config returned here -- so call this function again on the - derived config and the pass it builds reads that one. Not noticing is - mostly not silent either: :func:`check_zero_copy_kv` reads the arena - memory planning actually chose and refuses the program that mistake - produces, except in a method holding no host tensor to give the shared - arena away. + reading the config returned here, so call this function again on the + derived config and the pass it builds reads that one. The other is a + ``memory_planning_pass`` of your own that does not already carry an + attribute of that name -- ``to_executorch`` assigns the flag onto the + planner rather than passing it, and only onto a planner that has it, so + for any other the field reaches nothing and where the caches land is that + planner's own business. Neither is left to the runtime to discover: + :func:`check_zero_copy_kv` reads the arena memory planning actually chose + and refuses the program either mistake produces, and the manager + ``export(..., zero_copy_kv=True)`` returns runs that check itself. Two + placements it cannot settle are a method holding no host tensor to give + the shared arena away, and a planner that leaves no arena-device record + for it to read; :func:`check_zero_copy_kv` describes both. * ``propagate_device_config.skip_h2d_for_method_inputs`` is *refused*, wherever it is written -- in the single ``PropagateDeviceConfig`` or in a per-method dict of them -- and on every value ``PropagateDevicePass`` @@ -1258,13 +1449,16 @@ def zero_copy_backend_config( this says so here instead, where the caller can act on it. .. warning:: - Finalizing a ``zero_copy_kv=True`` program *without* this config does - not raise on its own. The engine writes a per-call staging copy that is - then discarded and the buffer never updates, which for a KV cache is - wrong output rather than a crash. Hand the finalized program to - :func:`check_zero_copy_kv` before writing the ``.pte`` and that mistake - becomes an error; ``torch_tensorrt.save(..., zero_copy_kv=True)`` runs - the check for you. + Finalizing a ``zero_copy_kv=True`` program *without* this config leaves + the engine writing a per-call staging copy that is then discarded, so + the buffer never updates -- for a KV cache, wrong output rather than a + crash. The manager ``export()`` hands back refuses that itself, reading + its own finalized program through :func:`check_zero_copy_kv`, and + ``torch_tensorrt.save(..., zero_copy_kv=True)`` runs the same check + before it writes the file. A program finalized off some *other* + manager -- one ``transform()`` or ``to_backend()`` derived from that + one -- is not covered by it, so hand that program to + :func:`check_zero_copy_kv` yourself before writing the ``.pte``. ``save(..., zero_copy_kv=True)`` installs this pass itself, so handing it the result of this function as ``backend_config`` applies the pass @@ -1272,6 +1466,21 @@ def zero_copy_backend_config( the buffers already wired straight to their delegates and changes nothing -- but the two entry points are alternatives: use one or the other. + + Arguments: + config (Optional[executorch.exir.ExecutorchBackendConfig]): The + configuration to compose onto. Omit it to start from ExecuTorch's + defaults. + + Returns: + executorch.exir.ExecutorchBackendConfig: A new config, every field of + the given one preserved, whose ``to_out_var_pass`` un-stages the aliased + buffers before running the ``to_out_var_pass`` that was there. + + Raises: + ValueError: If the configuration sets + ``propagate_device_config.skip_h2d_for_method_inputs``, which cannot + be combined with zero-copy KV. """ from dataclasses import replace diff --git a/py/torch_tensorrt/executorch/partitioner.py b/py/torch_tensorrt/executorch/partitioner.py index 3cdde4d9d05..f6cd1b892b6 100644 --- a/py/torch_tensorrt/executorch/partitioner.py +++ b/py/torch_tensorrt/executorch/partitioner.py @@ -23,6 +23,7 @@ TensorRTBackend, _get_engine_info_for_node, _get_engine_nodes_in, + _get_str, _parse_device_id, _serialize_elided_output_names, ) @@ -241,10 +242,9 @@ def _partition_elided_output_names( engine_info = _get_engine_info_for_node( exported_program, engine, metadata_only=True ) - raw = engine_info[OUTPUT_BINDING_NAMES_IDX] - if isinstance(raw, bytes): - raw = raw.decode("utf-8", "replace") - output_names = deserialize_binding_names(str(raw or "")) + output_names = deserialize_binding_names( + _get_str(engine_info, OUTPUT_BINDING_NAMES_IDX) + ) elided: Set[str] = set() for output_index, input_node in aliased.items(): if ( diff --git a/tests/cpp/executorch/test_executorch_blob_header.cpp b/tests/cpp/executorch/test_executorch_blob_header.cpp index eee75049ad9..1eaa5abaf42 100644 --- a/tests/cpp/executorch/test_executorch_blob_header.cpp +++ b/tests/cpp/executorch/test_executorch_blob_header.cpp @@ -431,6 +431,152 @@ TEST(ExecuTorchTensorRTBlobHeader, RejectsMetadataThatStartsPastTheEngine) { EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); } +TEST(ExecuTorchTensorRTBlobHeader, RejectsBindingEntryWithNoIsInputKey) { + // Without the key the initializer in the parser reads the binding as an output + // while serialization.py's TensorRTIOBinding reads it as an input, so the two + // readers of these bytes disagree -- and the disagreement is not one slot: it + // moves in_v out of the input list, which shifts every index after it. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v"},)" + R"({"name":"out_k","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsBindingEntryWithAMisspelledIsInputKey) { + // One byte wrong is the same case: the key falls through to skip_value, which + // consumes the value and leaves the initializer standing. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},{"name":"in_v","is_inout":true},)" + R"({"name":"out_k","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsABindingNameThatArrivedEscaped) { + // parse_string drops the backslash and keeps what follows, so this name is + // recorded as the three characters anb and TensorRT has no such tensor. The + // writer is json.dumps, which escapes every control character and everything + // non-ASCII, so an escape is exactly where the two ends read a name + // differently. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"a\nb","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAnAliasedIoNameThatArrivedEscaped) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in\tk","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAliasedIoMagicWithNoAliasArrayFound) { + // TR02 says the metadata carries aliased_io. Here the key is one byte wrong, + // so the walk finds nothing and the header would come back alias-free -- and + // an alias-free header of a threaded .pte binds each aliased output to its own + // storage and stops updating the caller's cache, with nothing to fail on. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliasedXio":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAliasedIoMagicWithTheKeysInSortedOrder) { + // The alias array is searched for past the io_bindings array, so a writer that + // emitted the keys in sorted order would put it out of reach. TR02 is what + // makes that a refusal rather than a silently alias-free header. + const std::string metadata = R"({"aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}],)" + R"("io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesScalarsPastAnAliasedBindingNamedLikeAKey) { + // The two scalar scans search the metadata text, and the alias array sits + // between where io_bindings ends and where those scans used to start, so an + // aliased binding named device_id was matched as the key: the scan then walked + // to the next colon, met the kind string, and failed the whole blob. + const std::string metadata = R"({"io_bindings":[{"name":"device_id","is_input":true},)" + R"({"name":"hardware_compatible","is_input":false}],)" + R"("aliased_io":[{"output":"hardware_compatible","input":"device_id",)" + R"("kind":"kv_cache_update"}],"hardware_compatible":true,"device_id":3})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_TRUE(header.hardware_compatible); + EXPECT_EQ(header.device_id, 3); + EXPECT_EQ(header.aliased_io.size(), 1u); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesTheDeviceIdKeyOutsideAnAliasEntryCarryingOne) { + // The other half of the same window. Being in key position does not tell an + // unknown key inside an alias entry from the real one, which the alias walk + // skips and the scans would otherwise read as the field; starting them past + // the array is what does. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update",)" + R"("device_id":9}],"device_id":3})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_EQ(header.device_id, 3); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsADeviceIdPastTheIntMaximum) { + // Accumulated into an int this wrapped to 1, which is a GPU that exists on + // most machines: cudaSetDevice then succeeds and the engine deserializes on a + // device nobody asked for. + const auto blob = make_blob(R"({"io_bindings":[{"name":"x","is_input":true}],"device_id":4294967297})"); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesTheLargestDeviceIdAnIntHolds) { + // The bound is the int maximum itself, not something short of it, so the + // refusal above is about overflow and not about long-looking values. + const auto blob = make_blob(R"({"io_bindings":[{"name":"x","is_input":true}],"device_id":2147483647})"); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_EQ(header.device_id, 2147483647); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesAStringValueThatIsExactlyAScalarKeyName) { + // A string *value* that reads like the key: quoted the same way, and so a + // match for the same search. This blob carries no device_id of its own, which + // is what an older writer emits, so the value is the only match there is -- + // and reading it as the key meant walking to the next colon, meeting the + // target_platform string, and failing a blob that is perfectly good. What + // separates the two is that a key is followed by its own colon and a value is + // followed by a comma. + const std::string metadata = R"({"io_bindings":[{"name":"x","is_input":true}],)" + R"("serialized_metadata":"device_id","target_platform":"linux"})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_EQ(header.device_id, 0); +} + TEST(ExecuTorchTensorRTBlobHeader, RejectsUnknownFutureMagic) { constexpr char kFutureMagic[4] = {'T', 'R', '0', '3'}; const std::string metadata = R"({"io_bindings":[{"name":"x","is_input":true}]})"; diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py index 8c42437dfa8..588f2c7a824 100644 --- a/tests/py/dynamo/executorch/test_export.py +++ b/tests/py/dynamo/executorch/test_export.py @@ -69,6 +69,10 @@ class FakeEdgeProgramManager: ``export()`` reads back the methods of the manager it is handed, to put each zero-copy method's mutations in the order ExecuTorch finalizes them in. These programs declare no mutation, so that call finds nothing to reorder. + + For a zero-copy method it also replaces ``to_executorch`` on the manager with + one that reads the finalized program back, so the attribute has to exist here + for that to bind to. Nothing in these tests finalizes. """ def __init__(self): @@ -78,6 +82,9 @@ def __init__(self): def exported_program(self, method_name="forward"): return self._programs[method_name] + def to_executorch(self, config=None): + raise AssertionError("these tests do not finalize") + class FakeTensorRTPartitioner: def __init__(self, compile_specs): diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index 0b1f41d4d51..5c8bf78c30c 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -641,6 +641,47 @@ def test_zero_copy_backend_config_reads_the_planning_mode_when_the_pass_runs(): turned_on.to_out_var_pass(graph_module) +@pytest.mark.unit +def test_zero_copy_backend_config_does_not_refuse_a_planner_the_flag_never_reaches(): + """The flag is only a ground to refuse on where ExecuTorch delivers it. + + ``to_executorch`` does not pass ``enable_non_cpu_memory_planning`` to the + memory planner, it assigns it -- and only onto a planner that already has an + attribute of that name. A caller-supplied planner without one, which is what + the user guide tells people to bring for a cache shared between prefill and + decode, never sees the field, so where the caches land is that planner's own + business and ``False`` does not mean the host arena. Refusing on it there + turns the field into a trap that blocks a configuration that would have + worked; the ``.pte`` is still not left unchecked, since ``check_zero_copy_kv`` + reads the arena that planner actually chose. + """ + from executorch.exir import ExecutorchBackendConfig + + def a_planner_of_ones_own(graph_module): + raise AssertionError("memory planning does not run in this test") + + assert not hasattr(a_planner_of_ones_own, "enable_non_cpu_memory_planning") + config = Z.zero_copy_backend_config( + ExecutorchBackendConfig( + enable_non_cpu_memory_planning=False, + memory_planning_pass=a_planner_of_ones_own, + ) + ) + graph_module, _, _ = _direct_delegate_graph() + + config.to_out_var_pass(graph_module) + + # The same field, with a planner ExecuTorch does hand it to, is refused -- + # so what is carried above is the planner and not the flag. + stock = Z.zero_copy_backend_config( + ExecutorchBackendConfig(enable_non_cpu_memory_planning=False) + ) + assert hasattr(stock.memory_planning_pass, "enable_non_cpu_memory_planning") + graph_module, _, _ = _direct_delegate_graph() + with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): + stock.to_out_var_pass(graph_module) + + @pytest.mark.unit def test_zero_copy_backend_config_rebuilt_over_a_derived_config_reads_it(): """Deriving a config with ``dataclasses.replace`` is the case that needs it. @@ -649,11 +690,13 @@ def test_zero_copy_backend_config_rebuilt_over_a_derived_config_reads_it(): config derived from the one this returns carries a pass still reading the original, and the first half below is that known gap, pinned: turning planning off on the derived config alone is not refused by the inherited - pass. It is mostly not silent either -- ``check_zero_copy_kv`` reads the - arena memory planning chose and refuses the program it produces, for any - method holding a host tensor to give the shared arena away -- and the remedy - the docstring gives is the second half: call the function again on the - derived config and the pass it builds is bound to that one. + pass. It is not silent either -- ``check_zero_copy_kv`` reads the arena + memory planning chose and refuses the program it produces, for any method + holding a host tensor to give the shared arena away, and the manager + ``export`` returns runs that check on whatever its ``to_executorch`` + finalizes -- and the remedy the docstring gives is the second half: call the + function again on the derived config and the pass it builds is bound to + that one. """ import dataclasses @@ -780,7 +823,7 @@ def test_unstage_refuses_a_buffer_another_backend_stages_on_the_same_gpu(): ``_h2d_copy_out`` requires a host source and fails ``InvalidArgument`` on a device one, so leaving the copy in place produces a program that does not run. Same GPU is what makes this shape distinct: the device and index both - match, so neither of ``_device_move_is_safe``'s spec comparisons rejects it. + match, so neither of ``_device_placement_is_safe``'s spec comparisons rejects it. """ graph = torch.fx.Graph() k_buffer = graph.placeholder("b_k_0") @@ -819,6 +862,92 @@ def test_unstage_refuses_a_buffer_another_backend_stages_on_the_same_gpu(): assert delegate_other.args[1] is staged_other +@pytest.mark.unit +def test_unstage_refuses_a_left_behind_staging_of_a_buffer_already_on_the_device(): + """The same refusal when the buffer's spec already names the engine's GPU. + + Nothing moves in this shape, so a check asked only about the move skips it + entirely -- and then the pass rewires the delegate anyway and hands the + engine a buffer whose other consumer, another backend's ``_h2d_copy``, reads + device memory as a host source and fails ``InvalidArgument`` on every call. + What the pass has to establish is where the buffer ends up, which is the same + place either way, so the question is asked either way. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_trt = graph.call_function(h2d, (k_buffer,)) + staged_other = graph.call_function(h2d, (k_buffer,)) + lowered_trt = graph.get_attr("lowered_module_0") + lowered_other = graph.get_attr("lowered_module_1") + delegate_trt = graph.call_function( + executorch_call_delegate, (lowered_trt, staged_trt) + ) + delegate_other = graph.call_function( + executorch_call_delegate, (lowered_other, staged_other) + ) + graph.output((k_buffer, delegate_trt, delegate_other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=None + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="CudaBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + # The buffer is already where the staging copy would have put it, which is + # the one thing that separates this from the test above. + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_trt.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + staged_other.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="consumer this pass leaves in place"): + Z._unstage_aliased_buffers(graph_module) + + assert delegate_trt.args[1] is staged_trt + + +@pytest.mark.unit +def test_unstage_moves_a_buffer_two_tensorrt_delegates_stage_from(): + """Two TensorRT delegates staging one cache is not a surviving consumer. + + Both staging copies go, so what is left reading the buffer is two delegates + taking it directly, which is the shape this pass exists to produce. The walk + reaches the second one after the first has already been rewired, so a + surviving-consumer test that did not allow a rewired delegate would refuse + the program the pass had just built. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + h2d = torch.ops.et_copy._h2d_copy.default + staged_first = graph.call_function(h2d, (k_buffer,)) + staged_second = graph.call_function(h2d, (k_buffer,)) + lowered_first = graph.get_attr("lowered_module_0") + lowered_second = graph.get_attr("lowered_module_1") + first = graph.call_function(executorch_call_delegate, (lowered_first, staged_first)) + second = graph.call_function( + executorch_call_delegate, (lowered_second, staged_second) + ) + graph.output((k_buffer, first, second)) + root = torch.nn.Module() + for name in ("lowered_module_0", "lowered_module_1"): + setattr( + root, + name, + SimpleNamespace(backend_id="TensorRTBackend", compile_specs=None), + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) + for staged in (staged_first, staged_second): + staged.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + assert Z._unstage_aliased_buffers(graph_module) == 2 + assert first.args[1] is k_buffer + assert second.args[1] is k_buffer + + @pytest.mark.unit def test_unstage_raises_when_the_staging_copy_is_not_on_cuda(): """Following the staging to the CPU would put the buffer out of the engine's @@ -915,7 +1044,7 @@ def test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus(): un-staged for either: a spec carries one device index, so whichever engine lost the race would be handed an address on the other's GPU. ``spec.device`` is only CUDA/CPU, so it is the device-index comparison in - ``_device_move_is_safe`` that refuses the first delegate here -- both + ``_device_placement_is_safe`` that refuses the first delegate here -- both stagings feed a TensorRT delegate, which is what separates this from the two-backends shapes and leaves the index the only comparison that can catch it. @@ -958,7 +1087,7 @@ def test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus(): @pytest.mark.unit def test_unstage_refuses_to_move_a_buffer_a_second_consumer_stages_to_the_host(): - """The device-*type* half of ``_device_move_is_safe``'s spec comparison. + """The device-*type* half of ``_device_placement_is_safe``'s spec comparison. Its sibling, the device index, is pinned by ``test_unstage_refuses_to_move_a_buffer_staged_to_two_gpus``. Here the second @@ -1008,7 +1137,7 @@ def test_unstage_refuses_a_buffer_another_backend_stages_to_a_different_gpu(): Un-staging skips the other backend's delegate, so its staging copy keeps reading the buffer while staging it to cuda:1, and re-homing the buffer onto the TensorRT engine's cuda:0 would move the source of that read to the wrong - GPU. Two of ``_device_move_is_safe``'s comparisons refuse this shape -- the + GPU. Two of ``_device_placement_is_safe``'s comparisons refuse this shape -- the device index, which runs first, and the surviving copy's non-TensorRT user -- so this test does not discriminate between them. The index comparison on its own is pinned by @@ -1143,6 +1272,7 @@ def _planned( on_device=True, host_tensor_arena=None, device_type=DeviceType.CUDA, + device_index=None, ): """Add what memory planning leaves behind, which the checker reads. @@ -1159,15 +1289,21 @@ def _planned( the output node's spec, which is where a real finalized program carries one. Leaving it unset gives the two shapes above: a separate arena when the program records devices, and the buffer's own when it does not. + ``device_index`` is the GPU the arena is recorded for; unset, it is the one + the graph's own CUDA specs ask for, which is what a planner that honoured + them would record. Passing a different one is the multi-GPU mistake. """ from executorch.exir.schema import NonConstBufferDevice if host_tensor_arena is None: host_tensor_arena = HOST_ARENA if on_device else arena + asked_for = [] for node in graph_module.graph.nodes: spec = node.meta.get("spec") if node.op == "placeholder" and spec is not None: spec.mem_id = arena if spec.device == DeviceType.CUDA else host_tensor_arena + if spec.device == DeviceType.CUDA: + asked_for.append(spec.device_index) if node.op == "output": node.meta["spec"] = [ SimpleNamespace( @@ -1177,7 +1313,11 @@ def _planned( if on_device: graph_module.meta["non_const_buffer_device"] = [ NonConstBufferDevice( - buffer_idx=arena, device_type=device_type, device_index=0 + buffer_idx=arena, + device_type=device_type, + device_index=( + next(iter(asked_for), 0) if device_index is None else device_index + ), ) ] return graph_module @@ -1300,6 +1440,54 @@ def test_check_zero_copy_kv_rejects_an_arena_the_program_records_as_non_cuda(): ) +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_an_arena_recorded_for_another_gpu(): + """A CUDA arena is not enough; it has to be the GPU the cache asks for. + + The runtime allocates the cache out of the arena the program records, so an + arena recorded for another device hands the engine an address on a GPU it is + not running on -- which fails exactly as a host pointer does, and which the + device *type* on its own cannot tell from a correct program. + """ + graph_module, k_buffer, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + assert k_buffer.meta["spec"].device_index == 0 + + with pytest.raises(RuntimeError, match="asks for cuda:0 and was planned"): + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module, device_index=1))) + + +@pytest.mark.unit +def test_check_zero_copy_kv_counts_one_buffer_in_two_slots_once(): + """Two argument slots holding one buffer are one cache, not two. + + The delegate's spec names two elided aliased outputs, so it owes two caches + written in place, and it has one. Counting slots satisfies that count with + the same buffer twice -- which is the arrangement the count exists to catch, + a delegate whose second mark did not survive. The un-staging pass counts the + same way, so the two are pinned together: they must not start disagreeing + about one graph. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + lowered = graph.get_attr("lowered_module_0") + delegate = graph.call_function( + executorch_call_delegate, (lowered, k_buffer, k_buffer) + ) + graph.output((k_buffer, delegate)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs("out_k", "out_v") + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=0) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="takes only 1 of the 2 buffers"): + Z._unstage_aliased_buffers(graph_module) + with pytest.raises(RuntimeError, match="takes 1 marked buffer"): + Z.check_zero_copy_kv(_finalized_program(_planned(graph_module))) + + @pytest.mark.unit def test_check_zero_copy_kv_rejects_a_buffer_staged_at_its_own_delegate(): """An unstamped TensorRT delegate must not stand in for the one that elided. @@ -1539,6 +1727,36 @@ def test_check_zero_copy_kv_rejects_a_stamped_delegate_whose_method_lost_its_mar ) +@pytest.mark.unit +def test_export_manager_checks_the_program_its_to_executorch_returns(): + """The two-call API cannot hand back a silently non-updating program. + + Export removed the copy-back before it returned, so every way of finalizing + that does not also un-stage the caches writes a ``.pte`` whose caches never + update -- and ``to_executorch()`` with ExecuTorch's defaults is one of those + and is the call the documentation spells out. The manager reads its own + finalized program back through the same check ``save`` runs, so the mistake + is an error rather than a file. A program that is right is handed back + untouched, which is the second half here: the check must not be a toll on + the working path. + """ + graph_module, k_buffer, _, _ = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + staged = _finalized_program(_planned(graph_module)) + edge = SimpleNamespace(to_executorch=lambda config=None: staged) + Z._check_zero_copy_kv_when_finalized(edge) + + with pytest.raises(RuntimeError, match="do not reach the TensorRT delegate"): + edge.to_executorch() + + unstaged = _finalized_program(_unstaged_graph()) + good = SimpleNamespace(to_executorch=lambda config=None: unstaged) + Z._check_zero_copy_kv_when_finalized(good) + assert good.to_executorch(config=object()) is unstaged + + @pytest.mark.unit def test_zero_copy_backend_config_defaults_to_executorch_defaults(): """Called with no config it starts from ExecuTorch's defaults, and the one @@ -2073,6 +2291,42 @@ def test_aliased_buffer_mark_survives_real_lowering(generate_etrecord): assert marked == {"k_cache", "v_cache"} +def test_finalizing_a_real_export_with_executorch_defaults_raises(): + """The documented two-call path, with the second call left as the default. + + ``to_executorch()`` with ExecuTorch's own defaults runs no un-staging, and + export has already removed the copy-back, so what it produces is a whole + ``.pte`` whose caches never update -- wrong output for a KV cache, and + nothing about the call says so. This is the real-engine end of + ``test_export_manager_checks_the_program_its_to_executorch_returns``: the + manager export returns refuses it, on a program that really was lowered and + really was finalized. + """ + _require_real_engine() + with torch.no_grad(): + torch.manual_seed(0) + model = _KVDecodeStep().eval().cuda() + tokens = torch.zeros(1, 1, dtype=torch.long).cuda() + input_pos = torch.tensor([0], dtype=torch.long).cuda() + + exported_program = torch.export.export(model, (tokens, input_pos)) + trt_gm = torch_tensorrt.dynamo.compile( + exported_program, + arg_inputs=(tokens, input_pos), + min_block_size=1, + truncate_double=True, + ) + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=(tokens, input_pos), + retrace=False, + zero_copy_kv=True, + ) + + with pytest.raises(RuntimeError, match="do not reach the TensorRT delegate"): + edge.to_executorch() + + class _MixedDecodeStep(torch.nn.Module): """A decode step with an engine-aliased KV cache and a copy-back buffer. @@ -2814,6 +3068,39 @@ def test_save_zero_copy_kv_true_installs_the_real_unstaging_pass(monkeypatch, tm ) +@pytest.mark.unit +def test_save_refuses_skip_h2d_before_it_compiles_anything(monkeypatch, tmp_path): + """The one refusal that reads nothing but the config fires before the compile. + + Reached only through ``zero_copy_backend_config``, it lands after export has + partitioned the graph and built every engine, so a caller who set a field + that was never going to be allowed pays the whole compile to find out. + ``save`` already validates the weight-streaming budget up front for that + reason; this is the same rule applied to the same kind of field. + """ + from executorch.exir import ExecutorchBackendConfig + from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig + + calls = _install_save_stubs(monkeypatch) + + with pytest.raises(ValueError, match="skip_h2d_for_method_inputs"): + torch_tensorrt.save( + _trivial_exported_program(), + str(tmp_path / "model.pte"), + output_format="executorch", + zero_copy_kv=True, + backend_config=ExecutorchBackendConfig( + propagate_device_config=PropagateDeviceConfig( + skip_h2d_for_method_inputs=True + ) + ), + ) + + # Refused before the lowering ran, which is the whole of this: the same + # ValueError comes out either way. + assert calls.export_kwargs is None + + @pytest.mark.unit def test_save_defaults_leave_kv_staged(monkeypatch, tmp_path): """Default save() (zero_copy_kv omitted) never wraps the config, so the KV From 7b723f0693d98c792dd10ebe4bb522a75d6264f2 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Tue, 8 Sep 2026 16:48:45 -0700 Subject: [PATCH 21/22] test(executorch): cover the unkilled zero-copy clauses and drop a dead one Nothing here changes what a caller gets. Six clauses had no test that bites them, one clause is dead and is gone, one is dead and is kept with the reason written down, and four claims in comments and docs were wrong in a way a reader would act on. Twelve mutants were run: eleven are killed by a named test, and the twelfth is reported below rather than counted as coverage. The six untested clauses. Each new test was checked by deleting the code it covers, watching it fail, and restoring. * `partitioner.py`'s extraction-failure handler and its single-engine guard. Three tests reuse the file's existing engine-node fixtures so the real per-engine derivation runs unmocked, and induce the failure by patching `_get_engine_info_for_node` to raise -- inside the `try` and after the `if not aliased: return set()` exit, so both halves of the handler are reachable. Deleting the guarded `raise`, making it unconditional, and deleting `if len(engine_nodes) != 1` each fail exactly one of the three, and each of the three kills exactly one. The two-engine test asserts the same partitioner answers `{"out_k"}` for one of those engines alone, so the empty set is attributable to the second engine and not to a graph with nothing in it. * `test_public_api_symbols_present` read `__all__`, which is written out verbatim in both branches of the `_has_executorch_exir()` guard, so it said nothing about whether the names are bound. It now compares `__all__` against a module-level tuple so the two cannot drift silently, and a second test resolves each name off the package and asserts the stub lane's `ImportError` names the symbol, so it bites in both branches instead of being vacuous in one. Deleting the `_zero_copy` import block from `__init__.py` leaves the old test green and fails the new one. * `export`'s two post-lowering wirings. One test substitutes a `to_executorch` on the stub manager, records the reorder and the finalize-time check, and asserts the reorder ran once per method against the manager's own programs by identity, that the check has not run at that point, and that it runs on the program the manager's `to_executorch` returns. The real `_check_zero_copy_kv_when_finalized` is used, so the wiring and the hook are covered together. Deleting either wiring fails it; both used to leave the stub suite at exactly baseline. * Both GraphModule branches of `save`. A test parametrized over `retrace` hands `save` a `torch.fx.symbolic_trace`d module with `use_legacy_exporter=False`. Dropping `zero_copy_kv=` from the `retrace=False` branch fails only `[retrace-false]`, from the `retrace=True` branch only `[retrace-true]`, and the control -- dropping it from the `ExportedProgram` branch -- fails only the three pre-existing tests. All three call sites are now pinned separately. * Both early returns of the reorder's `gets_a_copy`. A `BUFFER_MUTATION` whose target matches no input spec, so it is absent from the lifted map, and a mutation whose value is the literal `7` spec'd as a `ConstantArgument`. The first is shaped so the lineage test alone would answer the other way, which is what makes the resulting permutation show which predicate was used; under the non-`Node` mutant the second dies in upstream's write-back pass, which is the failure the comment beside that return predicts. * The cannot-decode fallbacks in `_delegate_elided_output_names`. The eight-way unreadable-value parametrization is now a module-level decorator applied to both decoders, so the twin cannot drift over different value sets, and the zero-copy side asserts the empty set where the backend side asserts a raise. A third test covers the no-spec case with a positive control in the same test. That last one is where the surviving mutant is. Deleting `if spec is None: return set()` kills nothing: `getattr(None, "value", None)` is `None`, which the type check two lines down refuses, so the function returns the same empty set. The line does execute now -- confirmed with a `sys.settrace` line probe on a spec-less delegate -- but no test can distinguish its loss. It stays: unlike the clause dropped below it cannot weaken a bound, and it is the documented "no spec at all" outcome, which is the one shape the backend twin treats differently. `TensorRTBlobHeader.cpp`'s `!saw_name ||` clause is dead and is gone. `saw_name` is set at the `name` key and tested there, so by the entry-level check it can only be false when the key never appeared, and then `name` is still empty and `usable_as_binding_name` refuses it on `!name.empty()`. Measured rather than argued: head against the file as of the parent commit over a 7,776-case metadata corpus -- 18 name-key shapes covering absent, empty, repeated, non-string, escaped, NUL, unterminated and nested-decoy names, times the `is_input`, extra-key, tail, second-entry and magic axes -- gives 0 verdict differences, against 864 for the control that drops the `usable_as_binding_name` half instead. Two lines beside the surviving predicate say why the absent-name case needs no flag of its own, since `is_input` immediately below does have one. The first extent check is dead in the same sense and is kept, with a comment saying why so the next mutation run does not rediscover it as a finding. Over 1,946,720 header-field combinations -- 23 values each for `metadata_offset`, `metadata_size` and `engine_offset` including `0x7fffffff`, `0xffffffff` and the exact body length, 10 for `engine_size` including the 2^63 and 2^64-16 wrap cases, 8 file sizes, both magics, 496 of them accepts -- deleting it gives 0 verdict differences, against 3,876 for deleting the metadata-vs-`engine_offset` check and 904 for the engine-vs-`size` one. It is kept because the extent this function then dereferences is `bytes + metadata_offset` for `metadata_size` bytes, two checks below: dropping the check routes that memory-safety argument through a check about the engine, and because the mutant is verdict-equivalent nothing would fail when a later edit to either of the other two breaks it. Four claims a reader would act on: * `serialization.py`'s field-ordering comment described a scan start that has moved. Driving the parser directly, a scalar written between `io_bindings` and `aliased_io`, or before both, parses clean and keeps its C++-side default (`parse=OK hw=0 dev=0`) where the same blob with the scalars last reads `hw=1 dev=6`. The comment now states the rule the code follows -- the scan starts past whichever of the two arrays was walked last -- and that silent defaulting is what breaking it costs. * The save-options paragraph in the user guide undercounted which options reach `to_edge_transform_and_lower`. Counting them is itself the trap, so the sentence is rescoped to the six options the section bullets, which cannot be undercounted by an option documented elsewhere, and every clause is re-checked against `_export.py` and `_compile.py` -- including `generate_etrecord`, which is both forwarded there and read by `save` itself to write the record beside the `.pte`. * The example re-ran `check_zero_copy_kv` after `_save_zero_copy` had already run it through the manager's hook, and its docstring said the updates "are dropped, with no error". Omitting the config makes the manager raise. The duplicate call and its import are gone and both docstrings say what happens; the module docstring now attributes the check to the manager rather than to this file. * `_zero_copy.py` cited `propagate_device_pass.py:216` into a third-party file -- the condition is at 218 in the ExecuTorch checkout on this box -- and spelled its module logger `_LOGGER` where the rest of the package uses `logger`. The citation is now the enclosing method, the form every other upstream reference in the module already uses, and the package and the test directory were swept for both. Two of those measurements falsify a sentence in an earlier commit of this stack. The PR description is built from these messages, so both are corrected where they were written rather than contradicted here: the "only surviving mutant" claim in the commit that added the NUL refusal, and the same `propagate_device_pass.py:216` citation in the commit that closed three remaining guard holes. Python suite 1 failed, 325 passed (`-m unit`), against 1 failed, 307 passed before; the failure is the pre-existing wheel-pin test from #4635, unchanged. `TensorRTBlobHeader.cpp` compiles standalone clean under `-Wall -Wextra`. `black`, `ruff`, `mypy 1.15.0` and `clang-format` clean. --- .../executorch/TensorRTBlobHeader.cpp | 14 +- .../runtime_performance/saving_models.rst | 10 +- .../export_kv_cache_decode.py | 26 +- py/torch_tensorrt/executorch/_zero_copy.py | 12 +- py/torch_tensorrt/executorch/serialization.py | 8 +- tests/py/dynamo/executorch/test_api.py | 32 ++- tests/py/dynamo/executorch/test_backend.py | 73 +++++- tests/py/dynamo/executorch/test_export.py | 61 ++++- .../py/dynamo/executorch/test_zero_copy_kv.py | 222 ++++++++++++++++++ 9 files changed, 418 insertions(+), 40 deletions(-) diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp index a113e0cfa0d..b56e7f49b40 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp @@ -247,7 +247,6 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso std::string name; bool is_input = false; - bool saw_name = false; bool saw_is_input = false; bool name_escaped = false; @@ -278,8 +277,7 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso if (key == "name") { pos = parse_string(json, pos, name, &name_escaped); - saw_name = pos != std::string::npos; - if (!saw_name) { + if (pos == std::string::npos) { return false; } } else if (key == "is_input") { @@ -307,7 +305,9 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso // argument list keeps its full length, and the two are only inferred from // the engine when both are empty -- so one real name beside a blank leaves // a short list that no longer lines up with the engine's bindings. - if (!saw_name || !usable_as_binding_name(name, name_escaped)) { + // An absent "name" key needs no flag of its own, unlike is_input below: + // the string is still empty here, which this predicate already refuses. + if (!usable_as_binding_name(name, name_escaped)) { return false; } // is_input has no safe default, so an entry without it is refused beside @@ -519,6 +519,12 @@ bool TensorRTBlobHeader::parse(const void* data, std::size_t size, TensorRTBlobH // The two metadata extents cannot wrap on a 64-bit size_t -- both operands are // 32-bit fields, so their sum is at most 2^33 -- but they are written the same // way so that the form, not the width of each field, is what makes them safe. + // The metadata-against-size check below cannot be the sole reason a blob is + // refused: the other two imply it, because the metadata extent is held inside + // engine_offset and engine_offset inside size. It is kept anyway, because it + // is the metadata extent this function goes on to dereference, and bounding + // that against the file size where it is read does not depend on a chain + // through a check about the engine. if (out.metadata_offset > size || out.metadata_size > size - out.metadata_offset) { return false; } diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 1d9b86b8bed..f06d9244a02 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -561,10 +561,12 @@ the synchronization is then the only thing making a host read see the new values ``torch_tensorrt.save`` takes these extra keyword arguments. They are only consulted for the ``executorch`` format; passing them with any other -``output_format`` logs a warning and is otherwise ignored. ``constant_methods``, -``transform_passes`` and ``compile_config`` are forwarded to ExecuTorch's -``to_edge_transform_and_lower(...)``; the rest are consumed at other points in -``save``. +``output_format`` logs a warning and is otherwise ignored. Of the six below, +``constant_methods``, ``transform_passes`` and ``compile_config`` are forwarded +to ExecuTorch's ``to_edge_transform_and_lower(...)``, and so is +``generate_etrecord``, which ``save`` also reads itself to write the record +beside the ``.pte``. ``backend_config`` goes to ``to_executorch(...)`` instead, +and ``zero_copy_kv`` is read by ``save`` on both sides of that boundary. * ``constant_methods`` — a ``dict`` of extra constant methods to embed in the ``.pte`` (e.g. ``{"get_max_seq_len": 2048}`` for an LLM runner). diff --git a/examples/torchtrt_executorch_example/export_kv_cache_decode.py b/examples/torchtrt_executorch_example/export_kv_cache_decode.py index 5c1c0d8b3ce..365f65bd645 100644 --- a/examples/torchtrt_executorch_example/export_kv_cache_decode.py +++ b/examples/torchtrt_executorch_example/export_kv_cache_decode.py @@ -20,8 +20,9 @@ here: zero-copy removes the copy that was making the update stick, so if the engine's in-place write is not reaching the caller's buffer the run fails. What it cannot see is a ``--zero_copy`` export that degenerated into an ordinary -staged ``.pte`` -- the two are indistinguishable to it -- so -``check_zero_copy_kv`` refuses to write one. +staged ``.pte`` -- the two are indistinguishable to it -- so the manager +``export()`` returns runs ``check_zero_copy_kv`` on the program its +``to_executorch()`` produces and raises rather than let one be written. Prerequisites ------------- @@ -96,18 +97,16 @@ def _save_zero_copy(trt_gm: torch.fx.GraphModule, inputs: tuple, path: str) -> N Zero-copy needs both ends of the Edge boundary: ``zero_copy_kv`` before lowering, and ``zero_copy_backend_config`` on the config the program is - finalized with. Without the second the cache is still staged and its updates - are dropped, with no error. ``torch_tensorrt.save(output_format="executorch", - zero_copy_kv=True)`` owns both steps and is the shorter way to the same .pte; - this spells them out because both halves are shown, and because a program - that reaches ``to_executorch()`` by any other route has to install the config - itself. + finalized with. Omitting the second leaves the cache staged and its updates + dropped, and the manager ``export()`` returns catches that: it runs + ``check_zero_copy_kv`` on whatever its own ``to_executorch()`` produced, so + the wrong config raises rather than writing a silently staged .pte. + ``torch_tensorrt.save(output_format="executorch", zero_copy_kv=True)`` owns + both steps and is the shorter way to the same .pte; this spells them out + because both halves are shown, and because a program that reaches + ``to_executorch()`` by any other route has to install the config itself. """ - from torch_tensorrt.executorch import ( - check_zero_copy_kv, - export, - zero_copy_backend_config, - ) + from torch_tensorrt.executorch import export, zero_copy_backend_config # retrace=True here, retrace=False for the plain save() below, so the two # exporters are both covered. Which way round matters: the legacy exporter @@ -118,7 +117,6 @@ def _save_zero_copy(trt_gm: torch.fx.GraphModule, inputs: tuple, path: str) -> N # cache is told from an ordinary copy-back buffer. edge = export(trt_gm, arg_inputs=inputs, retrace=True, zero_copy_kv=True) program = edge.to_executorch(zero_copy_backend_config()) - check_zero_copy_kv(program) with open(path, "wb") as output: program.write_to_file(output) if program._tensor_data: diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index f6e902c820e..9c1f7a7357a 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -47,7 +47,7 @@ if TYPE_CHECKING: from executorch.exir import ExecutorchBackendConfig -_LOGGER = logging.getLogger(__name__) +logger = logging.getLogger(__name__) def _aliased_inputs_by_output_index( @@ -239,7 +239,7 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> List[str]: signature = exported_program.graph_signature mutations = _aliased_buffer_mutations(exported_program) if not mutations: - _LOGGER.debug("no aliased buffer mutations to rewire") + logger.debug("no aliased buffer mutations to rewire") return [] engines_with_elided_outputs: Set[Node] = set() @@ -300,7 +300,7 @@ def rewire_aliased_mutations_to_buffers(exported_program: Any) -> List[str]: exported_program._graph_signature = ExportGraphSignature( input_specs=list(signature.input_specs), output_specs=output_specs ) - _LOGGER.debug( + logger.debug( "rewired %d aliased mutation(s) to their buffers, eliding outputs %s", len(mutations), elided_output_names, @@ -428,7 +428,7 @@ def gets_a_copy(index: int) -> bool: input_specs=list(signature.input_specs), output_specs=new_specs ) moved = sum(1 for slot, index in zip(slots, source) if slot != index) - _LOGGER.debug("moved %d mutation(s) so the copy-back ones come first", moved) + logger.debug("moved %d mutation(s) so the copy-back ones come first", moved) return moved @@ -927,7 +927,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> Any: unstaged = _unstage_aliased_buffers( graph_module, device_memory_planning=planning ) - _LOGGER.debug("un-staged %d aliased delegate buffer(s)", unstaged) + logger.debug("un-staged %d aliased delegate buffer(s)", unstaged) return inner(graph_module) return _UnstageThenToOutVar() @@ -1357,7 +1357,7 @@ def _refuse_skip_h2d(config: "ExecutorchBackendConfig") -> None: What is refused is every value that pass reads as on, which is every truthy one rather than only ``True``. ``PropagateDevicePass`` is handed the field - whole and only tests it for truth (``propagate_device_pass.py:216``), never + whole and only tests it for truth (in ``_insert_h2d_copies``), never resolving it per method, so it reads any non-empty dict as on for every method -- one whose entries are all ``False`` included. Refusing only the true entries would carry such a dict through and hand back a config that diff --git a/py/torch_tensorrt/executorch/serialization.py b/py/torch_tensorrt/executorch/serialization.py index ba7ccac10c2..70599c4cfb2 100644 --- a/py/torch_tensorrt/executorch/serialization.py +++ b/py/torch_tensorrt/executorch/serialization.py @@ -50,8 +50,12 @@ class TensorRTBlobMetadata: target_platform: str = "" def to_json(self) -> bytes: - # Keep field order stable because the C++ parser is intentionally small - # and searches forward after io_bindings for the scalar fields. + # Keep field order stable because the C++ parser is intentionally small. + # It walks io_bindings, then aliased_io, then searches forward from the + # end of whichever of those two it last walked for the scalar fields -- + # so a scalar written before either array is not found and keeps its + # C++-side default while the parse still succeeds. Any field added here + # that the C++ side reads by key must go after both arrays. data = { "io_bindings": [ { diff --git a/tests/py/dynamo/executorch/test_api.py b/tests/py/dynamo/executorch/test_api.py index 2e0b849aeb8..0768b1c8eba 100644 --- a/tests/py/dynamo/executorch/test_api.py +++ b/tests/py/dynamo/executorch/test_api.py @@ -92,20 +92,40 @@ def test_load_executorch_dispatches_to_delegate(monkeypatch): ) +_PUBLIC_API_SYMBOLS = ( + "get_edge_compile_config", + "TensorRTPartitioner", + "TensorRTBackend", + "export", + "zero_copy_backend_config", + "check_zero_copy_kv", +) + + @pytest.mark.unit def test_public_api_symbols_present(): module = importlib.import_module("torch_tensorrt.executorch") - assert "get_edge_compile_config" in module.__all__ - assert "TensorRTPartitioner" in module.__all__ - assert "TensorRTBackend" in module.__all__ - assert "export" in module.__all__ - assert "zero_copy_backend_config" in module.__all__ - assert "check_zero_copy_kv" in module.__all__ + assert set(module.__all__) == set(_PUBLIC_API_SYMBOLS) assert "Program" not in module.__all__ assert "load" not in module.__all__ assert "to_executorch" not in module.__all__ +@pytest.mark.unit +def test_public_api_symbols_are_bound_not_just_advertised(): + # __all__ is a literal written out in both branches of the + # _has_executorch_exir() guard, so reading it cannot tell whether the + # package binds what it advertises. Resolve each name instead. + module = importlib.import_module("torch_tensorrt.executorch") + if module._has_executorch_exir(): + for name in _PUBLIC_API_SYMBOLS: + assert getattr(module, name) is not None + else: + for name in _PUBLIC_API_SYMBOLS: + with pytest.raises(ImportError, match=name): + getattr(module, name) + + _REPO_ROOT = Path(__file__).resolve().parents[4] _SETUP_PY = _REPO_ROOT / "setup.py" _RUNTIME_SETUP_PY = _REPO_ROOT / "py/torch-tensorrt-executorch-runtime/setup.py" diff --git a/tests/py/dynamo/executorch/test_backend.py b/tests/py/dynamo/executorch/test_backend.py index 791221ea6c1..6a3a04cc78e 100644 --- a/tests/py/dynamo/executorch/test_backend.py +++ b/tests/py/dynamo/executorch/test_backend.py @@ -676,8 +676,10 @@ def test_preprocess_rejects_a_delegate_with_no_outputs_at_all(): TensorRTBackend.preprocess(edge_program, [spec]) -@pytest.mark.unit -@pytest.mark.parametrize( +# Values a hand-built zero-copy spec may carry that neither decoder can read. +# Applied to both decoders below: what each does with an unreadable value is the +# whole of the documented difference between them. +_unreadable_spec_values = pytest.mark.parametrize( "value", [1, True, None, b"", b"kv0", b'"kv0"', b"[1, 2", b"\xff\xfe"], ids=[ @@ -691,6 +693,25 @@ def test_preprocess_rejects_a_delegate_with_no_outputs_at_all(): "invalid-utf8", ], ) + + +def _delegate_graph_with_specs(compile_specs): + """A one-node lowered graph whose delegate carries ``compile_specs``.""" + from executorch.exir.delegate import executorch_call_delegate + + graph = torch.fx.Graph() + lowered = graph.get_attr("lowered_module_0") + delegate = graph.call_function(executorch_call_delegate, (lowered,)) + graph.output((delegate,)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=compile_specs + ) + return torch.fx.GraphModule(root, graph), delegate + + +@pytest.mark.unit +@_unreadable_spec_values def test_elided_output_names_refuses_a_value_it_cannot_read(value): """A hand-built spec may carry anything, and every shape has to name the key. @@ -711,6 +732,54 @@ def test_elided_output_names_refuses_a_value_it_cannot_read(value): _elided_output_names(specs) +@pytest.mark.unit +@_unreadable_spec_values +def test_zero_copy_twin_reads_an_unreadable_value_as_cannot_tell(value): + """``_zero_copy``'s decoder on the same values the backend twin refuses. + + Where the backend raises, this returns the empty set, which its callers read + as "cannot tell" and answer by demanding at least one marked buffer. Drifting + to a raise here would abort an export these cross-checks only weaken. + """ + from torch_tensorrt.executorch._zero_copy import _delegate_elided_output_names + from torch_tensorrt.executorch.backend import ZERO_COPY_KV_COMPILE_SPEC_KEY + + graph_module, delegate = _delegate_graph_with_specs( + [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, value)] + ) + + assert _delegate_elided_output_names(graph_module, delegate) == set() + + +@pytest.mark.unit +def test_zero_copy_twin_reads_a_missing_spec_as_cannot_tell(): + """A delegate carrying no zero-copy spec at all is the other "cannot tell". + + The one shape the two decoders do not share: for the backend a missing spec + is ``None`` and keeps a missing output an error, while here it joins the + values above. The second half is the control, so the empty sets are the + fallbacks and not a decoder that never answers. + """ + from torch_tensorrt.executorch._zero_copy import _delegate_elided_output_names + from torch_tensorrt.executorch.backend import ( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names, + ) + + graph_module, delegate = _delegate_graph_with_specs([]) + assert _delegate_elided_output_names(graph_module, delegate) == set() + + graph_module, delegate = _delegate_graph_with_specs( + [ + CompileSpec( + ZERO_COPY_KV_COMPILE_SPEC_KEY, + _serialize_elided_output_names(["out_k"]), + ) + ] + ) + assert _delegate_elided_output_names(graph_module, delegate) == {"out_k"} + + @pytest.mark.unit def test_elided_output_names_reads_the_list_the_partitioner_writes(): """The control for the refusals above, so they cannot pass vacuously.""" diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py index 588f2c7a824..31d3ce36593 100644 --- a/tests/py/dynamo/executorch/test_export.py +++ b/tests/py/dynamo/executorch/test_export.py @@ -72,7 +72,9 @@ class FakeEdgeProgramManager: For a zero-copy method it also replaces ``to_executorch`` on the manager with one that reads the finalized program back, so the attribute has to exist here - for that to bind to. Nothing in these tests finalizes. + for that to bind to. A test that wants to watch that happen substitutes its + own on the instance; this default is here to make an accidental finalization + loud rather than to be called. """ def __init__(self): @@ -83,7 +85,10 @@ def exported_program(self, method_name="forward"): return self._programs[method_name] def to_executorch(self, config=None): - raise AssertionError("these tests do not finalize") + raise AssertionError( + "this stub manager does not finalize; a test that needs to must " + "substitute its own to_executorch on the instance" + ) class FakeTensorRTPartitioner: @@ -616,6 +621,58 @@ def test_export_zero_copy_kv_rewires_every_method(monkeypatch): ) +@pytest.mark.unit +def test_export_wires_the_reorder_and_the_finalize_check(monkeypatch): + """The two things export does after lowering, on the lane without a GPU. + + Both are wiring rather than computation, so the tests of the pieces + themselves say nothing about whether export calls them. The reorder has to + run after lowering because ``to_edge`` re-derives the graph signature, and + the check has to be installed on the manager export hands back because that + is the only object that knows zero-copy was asked for. + """ + import torch_tensorrt.executorch._zero_copy as zero_copy + + export_module, lower = _patch_lowering(monkeypatch) + _patch_declare(monkeypatch) + _patch_rewire(monkeypatch) + + manager = FakeEdgeProgramManager() + manager._programs = { + "prefill": FakeExportedProgram(), + "decode": FakeExportedProgram(), + } + manager.methods = set(manager._programs) + finalized = object() + manager.to_executorch = lambda config=None: finalized + lower.return_value = manager + + reordered = [] + monkeypatch.setattr( + zero_copy, + "order_copyback_mutations_first", + lambda program: (reordered.append(program), 0)[1], + ) + checked = [] + monkeypatch.setattr(zero_copy, "check_zero_copy_kv", checked.append) + + result = export_module.export( + {"prefill": FakeExportedProgram(), "decode": FakeExportedProgram()}, + partitioners={"prefill": [object()], "decode": [object()]}, + zero_copy_kv=True, + ) + + assert result is manager + assert {id(program) for program in reordered} == { + id(program) for program in manager._programs.values() + } + # Nothing is finalized yet, so the check must not have run. + assert checked == [] + + assert result.to_executorch() is finalized + assert checked == [finalized] + + @pytest.mark.unit def test_export_does_not_exempt_a_method_that_kept_all_its_outputs(monkeypatch): """The exemption is per method and only where an output was actually elided. diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index 5c8bf78c30c..79feb922b97 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -15,6 +15,7 @@ """ import json +import logging import operator import re from types import SimpleNamespace @@ -2196,6 +2197,95 @@ def test_a_mixed_alias_engine_derives_the_narrower_set_and_is_then_refused(): ) +def _one_aliased_engine_partition(marked): + """One engine with one aliased output, its buffer input marked or not.""" + from torch_tensorrt.executorch.partitioner import TensorRTPartitioner + + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + engine = _no_op_engine_node( + graph, + [k_buffer], + aliased_io={"out_k": ("k_in", "kv_cache_update")}, + input_names=["k_in"], + output_names=["out_k"], + ) + graph.output((engine,)) + if marked: + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + program = SimpleNamespace( + graph_module=torch.fx.GraphModule(torch.nn.Module(), graph), + graph_signature=SimpleNamespace(buffers_to_mutate={}, inputs_to_buffers={}), + constants={}, + ) + return TensorRTPartitioner(), program, SimpleNamespace(id=0, nodes=[engine]) + + +def _make_engine_info_unreadable(monkeypatch): + def _boom(*args, **kwargs): + raise RuntimeError("engine record unreadable") + + monkeypatch.setattr( + "torch_tensorrt.executorch.partitioner._get_engine_info_for_node", _boom + ) + + +@pytest.mark.unit +def test_unreadable_engine_propagates_when_a_buffer_was_rewired(monkeypatch): + """A method with a rewired buffer must not fall back to eliding nothing. + + The aliased outputs of a rewired buffer left the graph before partitioning, + so an empty set stamps no delegate and the export dies further down blaming a + lost aliased-buffer mark -- a failure that names neither this partition nor + the record that would not read. + """ + partitioner, program, partition = _one_aliased_engine_partition(marked=True) + _make_engine_info_unreadable(monkeypatch) + + with pytest.raises(RuntimeError, match="engine record unreadable"): + partitioner._partition_elided_output_names(program, partition) + + +@pytest.mark.unit +def test_unreadable_engine_elides_nothing_when_no_buffer_was_rewired( + monkeypatch, caplog +): + """With nothing rewired the delegate really does carry every binding. + + The pair with the test above is the whole point of the handler: the same + failure is survivable here and not there, and only the graph says which. + """ + partitioner, program, partition = _one_aliased_engine_partition(marked=False) + _make_engine_info_unreadable(monkeypatch) + + with caplog.at_level( + logging.WARNING, logger="torch_tensorrt.executorch.partitioner" + ): + assert partitioner._partition_elided_output_names(program, partition) == set() + assert "could not resolve elided outputs" in caplog.text + + +@pytest.mark.unit +def test_a_partition_holding_two_engines_elides_nothing(): + """Elision is derived from one engine's aliased_io, so two is not answerable. + + ``TensorRTBackend.preprocess`` refuses a multi-engine partition outright, but + this runs first, and guessing here would stamp one engine's binding names + onto a delegate that also carries another's. + """ + from torch_tensorrt.executorch.partitioner import TensorRTPartitioner + + program, engine_a, engine_b = _two_engine_program() + both = SimpleNamespace(id=0, nodes=[engine_a, engine_b]) + + partitioner = TensorRTPartitioner() + assert partitioner._partition_elided_output_names(program, both) == set() + # The same partitioner does answer for engine_a alone, so the empty set above + # is the two-engine shape and not a graph that had nothing to elide. + single = SimpleNamespace(id=0, nodes=[engine_a]) + assert partitioner._partition_elided_output_names(program, single) == {"out_k"} + + # -------------------------------------------------------------------------- # GPU integration: the mark set during rewiring must survive real lowering, or # the un-staging pass has nothing to act on and every KV update is lost. Only a @@ -2545,6 +2635,102 @@ def test_reorder_leaves_an_already_correct_order_alone(): assert [spec.arg.name for spec in program.graph_signature.output_specs] == before +@pytest.mark.unit +def test_reorder_treats_an_unlifted_mutation_target_as_needing_no_copy(): + """A mutation upstream cannot resolve to a lifted input gets no copy. + + ``insert_write_back_for_buffers_pass`` only copies mutations whose target is + in the map it builds from the input specs, so one that is not belongs with + the copy-free mutations however its value was produced -- and its value here + is an ordinary functional result, which is what the lineage test reads as + needing a copy. Asking only the lineage test would leave this slot ahead of + the real copy-back and cross the finalized pairing. + """ + from torch.export.exported_program import ExportGraphSignature + from torch.export.graph_signature import InputKind, InputSpec + + graph = torch.fx.Graph() + b_cb = graph.placeholder("b_cb") + x = graph.placeholder("x") + unlifted_value = graph.call_function(torch.ops.aten.add.Tensor, (x, x)) + cb_value = graph.call_function(torch.ops.aten.add.Tensor, (b_cb, x)) + graph.output((unlifted_value, cb_value)) + program = _StubProgram( + torch.fx.GraphModule(torch.nn.Module(), graph), + ExportGraphSignature( + input_specs=[ + InputSpec(InputKind.BUFFER, TensorArgument("b_cb"), "cb", False), + InputSpec(InputKind.USER_INPUT, TensorArgument("x"), None), + ], + output_specs=[ + # No input spec targets "ghost", so it is not in the lifted map. + OutputSpec( + OutputKind.BUFFER_MUTATION, + TensorArgument(unlifted_value.name), + "ghost", + ), + OutputSpec( + OutputKind.BUFFER_MUTATION, TensorArgument(cb_value.name), "cb" + ), + ], + ), + ) + + assert Z.order_copyback_mutations_first(program) == 2 + assert [spec.target for spec in program.graph_signature.output_specs] == [ + "cb", + "ghost", + ] + + +@pytest.mark.unit +def test_reorder_groups_a_non_node_mutation_value_with_the_copies(): + """A mutation slot holding a literal is ordered as though it were copied. + + Upstream reads a non-Node value as needing a copy and then raises walking it, + so putting it anywhere else would make this reorder the thing that raises and + hide the program upstream is actually complaining about. + """ + from torch.export.exported_program import ExportGraphSignature + from torch.export.graph_signature import ConstantArgument, InputKind, InputSpec + + graph = torch.fx.Graph() + b_first = graph.placeholder("b_first") + b_cb = graph.placeholder("b_cb") + x = graph.placeholder("x") + inplace_value = graph.call_function(torch.ops.aten.add_.Tensor, (b_first, x)) + graph.output((inplace_value, 7)) + program = _StubProgram( + torch.fx.GraphModule(torch.nn.Module(), graph), + ExportGraphSignature( + input_specs=[ + InputSpec(InputKind.BUFFER, TensorArgument("b_first"), "first", False), + InputSpec(InputKind.BUFFER, TensorArgument("b_cb"), "cb", False), + InputSpec(InputKind.USER_INPUT, TensorArgument("x"), None), + ], + output_specs=[ + OutputSpec( + OutputKind.BUFFER_MUTATION, + TensorArgument(inplace_value.name), + "first", + ), + OutputSpec( + OutputKind.BUFFER_MUTATION, + ConstantArgument(name="literal", value=7), + "cb", + ), + ], + ), + ) + + assert Z.order_copyback_mutations_first(program) == 2 + assert [spec.target for spec in program.graph_signature.output_specs] == [ + "cb", + "first", + ] + assert program.graph_module.graph.output_node().args[0][0] == 7 + + def _assert_each_mutation_names_its_own_value(program): """The finalized signature pairs every mutated buffer with its own new value. @@ -3120,3 +3306,39 @@ def test_save_defaults_leave_kv_staged(monkeypatch, tmp_path): assert calls.wrap_args == [] assert calls.to_executorch_config is user_cfg assert calls.checked == [] + + +@pytest.mark.unit +@pytest.mark.parametrize( + "retrace", [False, True], ids=["retrace-false", "retrace-true"] +) +def test_save_forwards_zero_copy_kv_from_a_graph_module(monkeypatch, tmp_path, retrace): + """A compiled module is a GraphModule, and save reaches ExecuTorch by a + different branch for each value of ``retrace``. + + The tests above hand save an ``ExportedProgram``, which is the third branch. + Dropping the option from either of these two leaves a caller who asked for + zero-copy with an ordinary staged ``.pte`` and no error, since export is + never told and so nothing is rewired for the checker to miss. + """ + calls = _install_save_stubs(monkeypatch) + + class _Add(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + 1 + + torch_tensorrt.save( + torch.fx.symbolic_trace(_Add()), + str(tmp_path / "model.pte"), + output_format="executorch", + arg_inputs=[torch.randn(3)], + retrace=retrace, + # Neither branch's default exporter reads a plain traced GraphModule: the + # legacy one wants the engine-node shape a real compile produces. + use_legacy_exporter=False, + zero_copy_kv=True, + ) + + assert calls.export_kwargs["zero_copy_kv"] is True + assert calls.to_executorch_config is calls.wrapped_sentinel + assert calls.checked == [calls.program] From 1b7c21e6b1ddf16d6f359274b3c8b1b71e66c48b Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Tue, 8 Sep 2026 18:58:17 -0700 Subject: [PATCH 22/22] fix(executorch): make each zero-copy guard refuse what its own prose claims Eleven behaviour findings in the zero-copy guards. Who each one changes is enumerated at the end; two of them should be read first, because one widens the blast radius of this feature's machinery to every export and the other reverses an acceptance an earlier commit in this stack chose on purpose. `order_copyback_mutations_first` now runs for every method of every export, gated on nothing. The crossing it repairs is not zero-copy's, and that is measured rather than argued. A stock-ExecuTorch probe -- no torch_tensorrt, no TensorRT, no `zero_copy_kv` in it at all -- with two buffers, one written from another buffer and one from a user input, comes out of `to_edge().to_executorch()` with each mutation naming the other's value: `dst` against the `copy__default` that writes `log`, and `log` against `b_src`. Run the reorder first and the pair is correct; it reports `moved=2` on that program. Upstream's own `_inplace_lineage`, which is what the reorder asks, answers True for one of those slots and False for the other, so this is exactly the copy-free-before-copy order the repair exists for. Gating it on `zero_copy_methods` therefore left a silently wrong pairing in place for exports that never asked for anything, and made whether a method is repaired depend on whether some sibling method opted in. A program whose pairing is already right is untouched -- the reorder moves nothing and returns 0 -- and it cannot introduce a crossing, because it moves toward the order upstream's own classifier demands at that point, and where a later `reinplace_pass` invalidates that classification the pair was already crossed with or without the move. Two costs, taken knowingly: the module's private import of `_inplace_lineage` now breaks every `torch_tensorrt.executorch.export()` rather than only zero-copy ones if upstream renames it, and the loop reads `edge_manager.methods` on every export, so the two stand-ins for `to_edge_transform_and_lower` -- a bare `object()` and a class without `.methods`, nine test cases between them -- were given `methods = ()` rather than re-gating; the real `EdgeProgramManager` always has both. The finalize-time check keeps a gate of its own, now `if zero_copy_kv:`. `check_zero_copy_kv` now refuses a program whose planner recorded no device for the cache's arena. This reverses a deliberate acceptance, documented as such where it was made: the runtime fact was already in that docstring, and absence was read as "cannot tell" so as not to block a bring-your-own planner, which the user guide tells people to use for a cache shared between prefill and decode. It is reversed because `MethodMeta::memory_planned_buffer_device` answers `Device{CPU, 0}` both when `non_const_buffer_device()` is null and when the arena is absent from the sparse list (`method_meta.cpp:381-383` and `399-400`, both read at source), so a `.pte` with no record cannot tell any runner where the cache lives and the reference runner in this repo backs it with host memory. Who that refuses who is correct today, stated plainly: a caller pairing a custom planner that does place the cache on the device with a custom runner that ignores `memory_planned_buffer_device` and backs every arena on the device. That program runs correctly now and `save` and the manager's `to_executorch` both refuse it, with no opt-out. It is not a hypothetical caller, since a prefill/decode shared arena needs a custom runtime by design. The refusal names the remedy -- `apply_algo` with `enable_non_cpu_memory_planning=True` -- and the docstring and the user guide now say the record is required rather than optional. The smallest escape hatch, if one is wanted, is a keyword on `check_zero_copy_kv`. Beside it, an absent `mem_id` is no longer reported as host-planned. `_is_host_planned` answered True for `mem_id=None` against `{3: 0}` while `_planned_on_another_gpu` answered None on the same value, so an unplanned cache was refused in words that blamed the planner for something it had not done. The three comprehensions are one `if/elif` chain now, ordered so the positive host-arena ground is still read first where it applies -- otherwise the common `enable_non_cpu_memory_planning=False` case, which both puts the cache among the host tensors and writes no record, would have lost its more specific message. The test for the new refusal asserts the message does not say "host tensors", and carries the recorded-CUDA-arena program as a positive control. Four in the blob parser, all of them the two ends reading one blob differently: * An escaped name is refused only where the escape is one this parser reads differently from a JSON reader. `parse_string` drops the backslash and keeps what follows, which is exactly correct decoding for `\"`, `\\` and `\/` and wrong for every other escape, so the predicate is `escape_reads_as_written` and the out-parameter is `saw_misread_escape`. The blanket refusal added earlier in this stack regressed names the merge-base parser loads: `a"b`, `a\b` and `a/b` go from refused back to accepted and come back byte-identical, while non-ASCII, `\n`, `\t`, `\r`, `\b`, `\f` and a bogus `\q` stay refused. `\/` is included because it is provably the same decoding both ends do -- `json.dumps` writes `/` plain, but a hand-assembled blob may escape it, and refusing that would refuse a name both ends agree on. * The same flag now reaches the callers that compare the text: the `io_bindings` entry key, the `aliased_io` entry key, and `kind`. The three binding names already had it. A backslash before any character of `is_input` produces the collision, not just one position -- all eight read as the real key at the parent commit and are refused now, with the unescaped control still parsing -- and `\user`, `u\ser`, `us\er` and `use\r` all came back as the plain string `user`. The comment claiming escape handling "does not matter" for the key callers is corrected, since that comment is the reason a reader would not look there. * Both array keys are found the way the scalars are, through `value_pos_after_key`, so the array has to be that key's value rather than the next `[` anywhere in the text. This fixes the consistency of what the four keys count as a key and nothing else: the silent defaulting of a scalar written before `aliased_io` is deliberately left alone, since closing it needs the parser to search the whole metadata or to refuse a key found behind the cursor, which is a larger change than this. The ordering rule is instead pinned from both ends against the same stated rule, by a gtest over the exact key order `TensorRTBlobMetadata.to_json` emits and a Python test that the writer keeps every by-key scalar after both arrays; each names the other, because the C++ parser is not reachable from Python and the writer is not reachable from a gtest. Three good blobs the parent commit refuses, because a metadata value happens to spell an array key and the old `find('[')` then walked the wrong array, are accepted now. * The two disproven sub-triggers of that one are not re-tried as defects: the nested `build_info.device_id` decoy and the sorted-keys blob sit in the corpus as controls and show 0 differences at both revisions. `_unstage_aliased_buffers` and `check_zero_copy_kv` no longer accept different graphs. `satisfied_placeholders.add(...)` sat outside the `declares_zero_copy` branch, so a marked buffer whose only TensorRT delegate carries no zero-copy spec was accepted by the pass and refused by the checker two steps later -- reproduced on the project's own direct-delegate fixture with `compile_specs=None`, the pass un-staging 0 and the checker raising "marked for in-place update but do not reach the TensorRT delegate". The identical split sits one branch down on the staged route and measures the same way, so both moved; fixing one half of an invariant is worse than fixing neither. End to end this changes which error fires and when, since both documented entry points already run the checker, but a caller who finalizes through `zero_copy_backend_config` by hand and never calls it now gets an error where they previously got a silently non-updating `.pte`. Eleven existing tests built delegates with `compile_specs=None`, which a real `export(zero_copy_kv=True)` never produces -- the partitioner stamps the delegate whose engine elided -- so nine were given real specs and two had their message matcher updated. The refusal is pinned on both routes, each with a stamped positive control, so it is attributable to the missing stamp. The surviving-consumer check is asked on the direct route too. One graph built twice, with another backend's `_h2d_copy` reading the marked buffer on `cuda:1` in both: staged refused, direct accepted with nothing un-staged. The direct branch now asks `_device_placement_is_safe` with the buffer's own spec device and index before counting it satisfied, which makes true the two docstrings that already promised the question is asked whether or not the spec names the target device. What it refuses is a program whose every `execute()` fails with `InvalidArgument`, so no caller getting a correct result is affected. The finalize-time check follows what was asked for, not what was rewired. `_apply_zero_copy_kv` returns `{}` when nothing was elided, so `export(ep, zero_copy_kv=True)` on a model with no aliased buffer mutation installed no hook and finalized a plain staged `.pte`, while `save(..., zero_copy_kv=True)` calls the checker unconditionally and refused that same model -- the two entry points disagreeing about one program. It now raises at `to_executorch()`, and the export-time warning says finalization will refuse rather than leaving the caller to meet the raise later. The `_compile.py` comment that conceded the asymmetry goes with it. The hook refuses before the manager is spent. `EdgeProgramManager.to_executorch` runs the passes over `self.exported_program.graph_module` and copies the finalized graph back, so a manager that has finalized once cannot do it again, and a hook that only reads the result leaves the caller nothing to retry. `_UnstageThenToOutVar` is hoisted to module level, taking `inner` and `device_memory_planning` as constructor arguments -- it was defined inside `unstage_aliased_buffers_pass`, so every call built a fresh class object and nothing could ask a config what its pass was -- and the wrapper now reads `config.to_out_var_pass` from `kwargs["config"]` or the first positional and raises before delegating, saying in the message why it fires early. The post-check stays, for the two things the config cannot answer: the arenas memory planning picks afterwards, and a lost mark. That program already raised; what changes is that the failure is now recoverable on the same manager. `test_zero_copy_backend_config_defaults_to_executorch_defaults` compared the pass by class name as a workaround for the class not being reachable, and uses `isinstance` now, which is what the hook does. One in the runtime, at one of the two sites that share its shape. In `TensorRTBackend.cpp`'s `execute()`, `inflight_pending` was cleared on the line before the `cudaEventSynchronize` return code was tested, so a failed synchronize disarmed the only wait `~EngineHandle` has while the enqueue might still be running; the assignment now sits below the test. At the second site there is nothing to fix: the flag is provably already false there -- the only `= true` in the file is the last statement of the other branch, the wait at the top of `execute()` clears whatever a previous call left, and the must-sync branch this sits in records no event -- so that assignment cannot disarm anything. The reorder is applied there for shape, with a comment stating the invariant so the asymmetry with the first site is not a puzzle, and it is reported as a no-op rather than a fix. The residual hazard at that site is real and out of scope: if `cudaStreamSynchronize` fails, the enqueue may still be running with nothing armed, which needs a mechanism and not a reorder. Both sites ship unmeasured beyond `clang-format`: no gtest target in this repo depends on `//cpp:tensorrt_executorch_backend`, so nothing here compiles that file, and reaching the path needs the backend library, a runner, a real `.pte` and an induced CUDA event-synchronize failure. Who this changes, against "no behaviour change for a caller getting a correct result today": * a non-zero-copy export whose pairing is crossed -- a silently wrong result fixed; a correct one is untouched; * blobs named `a"b`, `a\b` or `a/b` -- loading restored, which this stack had regressed; * blobs with a misread escape in an entry key or in `kind` -- refused instead of silently read as something else, with the unescaped controls unaffected; * blobs whose metadata value spells an array key -- three shapes accepted that were refused; * a marked buffer on an unstamped delegate, and a shared buffer reaching its delegate directly -- both already fail end to end, refused earlier and by name; * `zero_copy_kv=True` that rewired nothing, and `to_executorch()` with the wrong config -- both already raised, now sooner and recoverably; * a custom planner paired with a custom runner that ignores the arena record -- the one caller who is correct today and is refused now, described above; * a failed `cudaEventSynchronize` -- strictly safer. Python suite 1 failed, 335 passed (`-m unit`), the failure the same pre-existing wheel-pin test from #4635. Blob-header gtests 39 -> 48. The C++ differential over a 186-case metadata corpus, head against the parser as of the parent commit, changes 51 verdicts: 12 newly accepted, 39 newly refused. Seventeen mutants were run and all seventeen are killed by a named test; two tests were rewritten mid-flight because the first draft did not discriminate -- one asserted a refusal both revisions already made, and the escaped alias-entry key had no test at all until its mutant killed nothing. `black`, `ruff`, `mypy 1.15.0` and `clang-format` clean. Five claims in earlier messages in this stack are corrected where they were written, since the PR description is built from them: the escaped-name refusal, the hook-installation rule, the scope of the satisfaction-count agreement and the documented arena gap in the commit that made the documented path refuse its own bad program, and the acceptance of an unrecorded arena in the commit that corrected four guards and bounded the blob extents. --- .../executorch/TensorRTBackend.cpp | 14 +- .../executorch/TensorRTBlobHeader.cpp | 132 ++++--- .../runtime_performance/saving_models.rst | 26 +- .../export_kv_cache_decode.py | 8 +- py/torch_tensorrt/_compile.py | 6 +- py/torch_tensorrt/executorch/_export.py | 45 ++- py/torch_tensorrt/executorch/_zero_copy.py | 325 +++++++++++++----- .../test_executorch_blob_header.cpp | 163 ++++++++- tests/py/dynamo/executorch/test_api.py | 4 + tests/py/dynamo/executorch/test_export.py | 74 +++- .../dynamo/executorch/test_serialization.py | 44 +++ .../test_weight_streaming_budget.py | 4 +- .../py/dynamo/executorch/test_zero_copy_kv.py | 243 ++++++++++--- 13 files changed, 871 insertions(+), 217 deletions(-) diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index 73646867918..337b7e5bd7a 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -676,13 +676,18 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* // on the shared exec_ctx. Wait for it before reconfiguring the context below: // TensorRT forbids mutating a context while one of its enqueues is in flight, and // setInputShape/setTensorAddress run on the host, so this must be a host-side wait. + // The flag is cleared only once the wait has actually succeeded: it is the + // destructor's only reason to wait, and a failed synchronize is precisely the + // case where the enqueue may still be running. Clearing it first would hand + // the destructor a handle it believes is idle and let it free the staging + // buffers and reset exec_ctx underneath a live enqueue. if (engine->inflight_pending) { cuda_err = cudaEventSynchronize(engine->inflight_event); - engine->inflight_pending = false; if (cuda_err != cudaSuccess) { ET_LOG(Error, "TensorRTBackend::execute: cudaEventSynchronize failed: %s", cudaGetErrorString(cuda_err)); return Error::InvalidProgram; } + engine->inflight_pending = false; } const auto caller_stream = ::executorch::extension::cuda::getCallerStream(); const bool caller_stream_set = caller_stream.has_value(); @@ -1054,11 +1059,16 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* } } cuda_err = cudaStreamSynchronize(stream); - engine->inflight_pending = false; if (cuda_err != cudaSuccess) { ET_LOG(Error, "TensorRTBackend::execute: cudaStreamSynchronize failed: %s", cudaGetErrorString(cuda_err)); return Error::InvalidProgram; } + // Same shape as the wait at the top of execute(), and for the same reason, + // though nothing here can be armed yet: this branch never records the + // event, and the wait at the top has already cleared any flag a previous + // call left. Arming happens as the last statement of the other branch, so + // the only writer that can make this clear anything is one added later. + engine->inflight_pending = false; if (copy_err != Error::Ok) { return copy_err; } diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp index b56e7f49b40..a7959c9e805 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBlobHeader.cpp @@ -30,18 +30,28 @@ std::size_t skip_ws(const std::string& s, std::size_t pos) { return pos; } -// A backslash is dropped and the character after it kept, which is not what a -// JSON escape means: "\n" comes out as the letter n, and a \u escape as the -// letter u and its four digits. The writer is json.dumps, which escapes every -// control character and everything non-ASCII, so those are exactly the strings -// the two ends would read differently. Rather than decode them, the two callers -// that record a name refuse one that arrived escaped (see -// usable_as_binding_name); saw_escape is how they are told. Every other caller -// -- the keys, and the values skip_value walks past -- only has to find the end -// of the string, which this does correctly either way. -std::size_t parse_string(const std::string& s, std::size_t pos, std::string& out, bool* saw_escape = nullptr) { - if (saw_escape != nullptr) { - *saw_escape = false; +// The three escapes whose JSON meaning is exactly what this parser does with +// every escape: drop the backslash, keep the character after it. json.dumps +// emits the first two whenever a name holds a quote or a backslash, and the +// third is one a hand-written blob may carry, so these are strings this parser +// and a JSON reader agree on and there is nothing to refuse. +bool escape_reads_as_written(char after_backslash) { + return after_backslash == '"' || after_backslash == '\\' || after_backslash == '/'; +} + +// A backslash is dropped and the character after it kept. For the three escapes +// above that is the correct decoding; for every other one the two ends read +// these bytes differently. "\n" comes out as the letter n where a JSON reader +// sees a newline, a \u escape as the letter u and its four digits where a JSON +// reader sees one codepoint, and "\q" comes out as the letter q where a JSON +// reader refuses the string outright. Rather than decode them, every caller +// that *compares* the text it gets -- against the engine's binding names, +// against a key, or against an alias kind -- refuses a string carrying one, and +// saw_misread_escape is how they are told. The values skip_value walks past +// need only their end, which this finds either way. +std::size_t parse_string(const std::string& s, std::size_t pos, std::string& out, bool* saw_misread_escape = nullptr) { + if (saw_misread_escape != nullptr) { + *saw_misread_escape = false; } if (pos >= s.size() || s[pos] != '"') { return std::string::npos; @@ -50,8 +60,8 @@ std::size_t parse_string(const std::string& s, std::size_t pos, std::string& out out.clear(); while (pos < s.size() && s[pos] != '"') { if (s[pos] == '\\' && pos + 1 < s.size()) { - if (saw_escape != nullptr) { - *saw_escape = true; + if (saw_misread_escape != nullptr && !escape_reads_as_written(s[pos + 1])) { + *saw_misread_escape = true; } ++pos; } @@ -185,20 +195,23 @@ bool parse_int_after_key(const std::string& json, std::size_t search_from, const // name can hold a NUL: two names differing only after one are distinct here and // are the same tensor to TensorRT, which is exactly what those refusals exist to // stop, and a name that is only a NUL is non-empty here and empty to TensorRT. -// A name that arrived escaped is the same problem one step earlier: parse_string -// does not decode escapes, so the name recorded here is not the one the writer -// wrote and so not one the engine has either, and two names differing only in an -// escape collapse into one and are refused as a repeat. Refusing both outright -// makes what the parser compares be what the engine will compare. No writer -// emits a NUL -- json.dumps escapes it -- and json.dumps escapes only what a -// binding name has no business holding, so this refuses only a blob assembled -// some other way. +// A name carrying an escape this parser reads differently from a JSON reader is +// the same problem one step earlier: the name recorded here is not the one the +// writer wrote and so not one the engine has either, and two names differing +// only in such an escape collapse into one and are refused as a repeat. +// Refusing both outright makes what the parser compares be what the engine will +// compare. An escaped quote, backslash or forward slash is *not* refused: those +// three come out of parse_string exactly as the writer wrote them, and the first +// two are what json.dumps emits for a name holding a quote or a backslash. No +// writer emits a NUL -- json.dumps escapes it -- and the escapes this does +// refuse are the ones json.dumps reserves for what a binding name has no +// business holding, so it refuses only a blob assembled some other way. // Like every refusal in this parser it is silent: parse() returns false and // the caller reports its own generic parse failure, so a NUL is not // distinguishable at load time from a bad magic or a wrapped extent. Refusing // uniformly is the deliberate choice here, not a missing diagnostic. -bool usable_as_binding_name(const std::string& name, bool escaped) { - return !escaped && !name.empty() && name.find('\0') == std::string::npos; +bool usable_as_binding_name(const std::string& name, bool misread_escape) { + return !misread_escape && !name.empty() && name.find('\0') == std::string::npos; } bool parse_metadata_json(const std::string& json, bool expects_aliased_io, TensorRTBlobHeader& out) { @@ -208,12 +221,14 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso out.hardware_compatible = false; out.device_id = 0; - const std::size_t bindings_pos = json.find("\"io_bindings\""); - if (bindings_pos == std::string::npos) { - return false; - } - const std::size_t arr_start = json.find('[', bindings_pos); - if (arr_start == std::string::npos) { + // Found the same way as the two scalars below, so all four keys agree about + // what a key is: an occurrence followed by its own colon, and the array right + // after that colon rather than the next '[' anywhere in the text. A string + // value that spells io_bindings is then passed over instead of being taken + // for the key, and a blob whose io_bindings is not an array is refused rather + // than walked from some unrelated bracket further on. + const std::size_t arr_start = value_pos_after_key(json, 0, "\"io_bindings\""); + if (arr_start == std::string::npos || arr_start >= json.size() || json[arr_start] != '[') { return false; } @@ -248,7 +263,7 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso std::string name; bool is_input = false; bool saw_is_input = false; - bool name_escaped = false; + bool name_misread = false; while (true) { pos = skip_ws(json, pos); @@ -264,9 +279,19 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso continue; } + // A key is compared against name and is_input below, so it is held to the + // same rule those comparisons need: a key carrying an escape this parser + // reads differently is refused rather than matched. Without that, + // "is_i\nput" arrives here as is_input and decides which list the binding + // goes in, while a JSON reader sees a key with a newline in it, ignores + // it, and leaves is_input at its own default -- the two readers of one + // blob putting one binding in opposite lists, which is what the + // saw_is_input refusal below exists to stop for the plainly misspelled + // key. std::string key; - pos = parse_string(json, pos, key); - if (pos == std::string::npos) { + bool key_misread = false; + pos = parse_string(json, pos, key, &key_misread); + if (pos == std::string::npos || key_misread) { return false; } pos = skip_ws(json, pos); @@ -276,7 +301,7 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso pos = skip_ws(json, pos + 1); if (key == "name") { - pos = parse_string(json, pos, name, &name_escaped); + pos = parse_string(json, pos, name, &name_misread); if (pos == std::string::npos) { return false; } @@ -307,7 +332,7 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso // a short list that no longer lines up with the engine's bindings. // An absent "name" key needs no flag of its own, unlike is_input below: // the string is still empty here, which this predicate already refuses. - if (!usable_as_binding_name(name, name_escaped)) { + if (!usable_as_binding_name(name, name_misread)) { return false; } // is_input has no safe default, so an entry without it is refused beside @@ -339,10 +364,9 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso // Search from pos (past the io_bindings array) so a model input literally // named "aliased_io" isn't matched as the array key. std::size_t scalars_from = pos; - const std::size_t alias_key = json.find("\"aliased_io\"", pos); - if (alias_key != std::string::npos) { - std::size_t apos = json.find('[', alias_key); - if (apos == std::string::npos) { + std::size_t apos = value_pos_after_key(json, pos, "\"aliased_io\""); + if (apos != std::string::npos) { + if (apos >= json.size() || json[apos] != '[') { return false; } ++apos; @@ -367,8 +391,9 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso ++apos; AliasedBinding ab; - bool output_escaped = false; - bool input_escaped = false; + bool output_misread = false; + bool input_misread = false; + bool kind_misread = false; while (true) { apos = skip_ws(json, apos); if (apos >= json.size()) { @@ -382,9 +407,13 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso ++apos; continue; } + // Held to the same rule as the io_bindings keys above, for the same + // reason: these are compared, so an escape the two ends read + // differently would match a key here that a JSON reader does not see. std::string key; - apos = parse_string(json, apos, key); - if (apos == std::string::npos) { + bool key_misread = false; + apos = parse_string(json, apos, key, &key_misread); + if (apos == std::string::npos || key_misread) { return false; } apos = skip_ws(json, apos); @@ -393,11 +422,11 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso } apos = skip_ws(json, apos + 1); if (key == "output") { - apos = parse_string(json, apos, ab.output, &output_escaped); + apos = parse_string(json, apos, ab.output, &output_misread); } else if (key == "input") { - apos = parse_string(json, apos, ab.input, &input_escaped); + apos = parse_string(json, apos, ab.input, &input_misread); } else if (key == "kind") { - apos = parse_string(json, apos, ab.kind); + apos = parse_string(json, apos, ab.kind, &kind_misread); } else { apos = skip_value(json, apos); } @@ -411,7 +440,16 @@ bool parse_metadata_json(const std::string& json, bool expects_aliased_io, Tenso // argument list, so a dropped entry surfaces as an argument-count error at // every execute, which never mentions aliasing, instead of at parse, which // the blob-header tests reach without a GPU. - if (!usable_as_binding_name(ab.output, output_escaped) || !usable_as_binding_name(ab.input, input_escaped)) { + if (!usable_as_binding_name(ab.output, output_misread) || !usable_as_binding_name(ab.input, input_misread)) { + return false; + } + // kind is not a binding name, but it is compared -- init() reads "user" + // as the kind validated on shape alone rather than confirmed against the + // engine's own aliasing -- so an escape the two ends read differently is + // refused here too. "\user" and "use\r" both arrive as the plain string + // user, while a JSON reader refuses the first and reads the second as + // "use" and a carriage return. + if (kind_misread) { return false; } // An output binding may be claimed by at most one entry. A second entry diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index f06d9244a02..026dd7edaea 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -401,8 +401,9 @@ that name never receives the field at all -- ``to_executorch`` assigns it onto the planner rather than passing it -- so where the caches land is that planner's own business. Neither is left to the runtime to discover: ``check_zero_copy_kv`` reads the arena planning actually chose, and the manager ``export`` returns runs -it for you -- so long as that planner records its arenas, which is its own -responsibility and is covered below. And +it for you. That planner does have to record the device of each arena it places, +or the check refuses the program for saying nothing about where the cache +lives -- see below. And ``propagate_device_config.skip_h2d_for_method_inputs`` is refused outright: ``PropagateDevicePass`` refuses to un-stage a method input whose placeholder does not have exactly one user, and a zero-copy cache always @@ -424,11 +425,14 @@ one silently would break a runner built before this feature. the copy-back; finalizing without ``zero_copy_backend_config`` leaves the engine writing a per-call staging copy that is discarded, so the cache never updates. For a KV cache that is wrong output, not a crash. The manager - ``export`` returns is what stops that: its ``to_executorch`` reads its own - finalized program through + ``export`` returns is what stops that. Its ``to_executorch`` refuses a + config that does not carry the un-staging pass *before* it finalizes + anything -- finalization rewrites the manager's edge programs in place, so a + refusal after the fact could only be acted on by exporting again -- and it + reads the finalized program back through ``torch_tensorrt.executorch.check_zero_copy_kv``, which refuses one whose - caches are still staged, or planned somewhere the engine cannot write them, - before there is a ``.pte`` to write. + caches ended up planned somewhere the engine cannot write them, before there + is a ``.pte`` to write. ``transform()`` and ``to_backend()`` return a *new* manager, which that check does not travel to, so a program finalized off one of those needs it @@ -455,7 +459,8 @@ one as ``backend_config`` as well: that installs the pass twice, which is redundant rather than an error -- the second run finds the buffers already un-staged. The two entry points are alternatives, not a pair. -Three further responsibilities are the caller's, and none raises: +Three further responsibilities are the caller's. The first two nothing checks; +the third is refused where it is visible in the ``.pte``: * **One CUDA stream for every delegate**, if the ``.pte`` is coalesced. Getting this wrong is a race, not a deterministic error: it is intermittent and can @@ -486,9 +491,10 @@ Three further responsibilities are the caller's, and none raises: ``False`` and with it off ``apply_algo`` plans every spec into one CPU bucket and records nothing either. Without that record the ``.pte`` reports every planned buffer as CPU, and a runner that honours it backs the cache with host - memory the engine cannot write. A missing record is not something - ``check_zero_copy_kv`` refuses on: it is also what a device-aware planner of - your own leaves, so it is read as "cannot tell" rather than as the host. + memory the engine cannot write. ``check_zero_copy_kv`` refuses a program in + that state, under a message of its own rather than the host-arena one: what + it can read is that nothing says where the cache lives, not that your planner + put it among the host tensors. **Coalesced TensorRT + CUDA .pte** diff --git a/examples/torchtrt_executorch_example/export_kv_cache_decode.py b/examples/torchtrt_executorch_example/export_kv_cache_decode.py index 365f65bd645..9110efcaa84 100644 --- a/examples/torchtrt_executorch_example/export_kv_cache_decode.py +++ b/examples/torchtrt_executorch_example/export_kv_cache_decode.py @@ -98,9 +98,11 @@ def _save_zero_copy(trt_gm: torch.fx.GraphModule, inputs: tuple, path: str) -> N Zero-copy needs both ends of the Edge boundary: ``zero_copy_kv`` before lowering, and ``zero_copy_backend_config`` on the config the program is finalized with. Omitting the second leaves the cache staged and its updates - dropped, and the manager ``export()`` returns catches that: it runs - ``check_zero_copy_kv`` on whatever its own ``to_executorch()`` produced, so - the wrong config raises rather than writing a silently staged .pte. + dropped, and the manager ``export()`` returns catches that: it refuses a + config that does not carry the un-staging pass before it finalizes anything, + and runs ``check_zero_copy_kv`` on whatever its own ``to_executorch()`` did + produce, so the wrong config raises rather than writing a silently staged + .pte. ``torch_tensorrt.save(output_format="executorch", zero_copy_kv=True)`` owns both steps and is the shorter way to the same .pte; this spells them out because both halves are shown, and because a program that reaches diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index fcf3408472d..d11504d4650 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -1508,9 +1508,9 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None # arena it can write. Both halves of zero-copy no-op quietly when they # find nothing to do, so without this a save() that asked for zero-copy # could still write an ordinary staged .pte. export() installs the same - # check on the manager it returns, but only for a method it rewired - # something in: the case this call adds is the one where it rewired - # nothing, which is a zero_copy_kv=True that bought the caller nothing. + # check on the manager it returns whenever zero-copy was asked for, so + # the two entry points refuse the same models; save() runs it here rather + # than through that manager because it finalizes the program itself. check_zero_copy_kv(executorch_program) with open(file_path, "wb") as f: executorch_program.write_to_file(f) diff --git a/py/torch_tensorrt/executorch/_export.py b/py/torch_tensorrt/executorch/_export.py index 093ec0eb37b..a0dace3191f 100644 --- a/py/torch_tensorrt/executorch/_export.py +++ b/py/torch_tensorrt/executorch/_export.py @@ -325,7 +325,9 @@ def _apply_zero_copy_kv( if not any(elided.values()): logger.warning( "zero_copy_kv=True, but no aliased buffer mutation was found in %s, " - "so zero-copy KV was not applied.", + "so zero-copy KV was not applied. The returned manager still checks " + "its finalized program, so to_executorch() will refuse it; export " + "without zero_copy_kv to get an ordinary staged program.", ", ".join(sorted(program_map)), ) return {} @@ -455,11 +457,14 @@ def export( change. Finalize such a program with ``to_executorch(torch_tensorrt.executorch.zero_copy_backend_config(config))`` -- without it the buffer is still staged and its updates would be discarded. - The returned manager reads its own finalized program back through - :func:`torch_tensorrt.executorch.check_zero_copy_kv` and raises on that - rather than handing back a program whose caches never update. That reaches - the manager returned here and no other: ``transform()`` and ``to_backend()`` - build a new one, so a program finalized off either owes that call by hand. + The returned manager refuses that rather than handing back a program whose + caches never update: it reads the config before it finalizes anything, so a + config without the pass is refused while the manager can still be finalized + again with one, and it reads the finalized program back through + :func:`torch_tensorrt.executorch.check_zero_copy_kv` for what the config + alone cannot say. That reaches the manager returned here and no other: + ``transform()`` and ``to_backend()`` build a new one, so a program finalized + off either owes that call by hand. Only a buffer the engine declares aliased is affected. A method may hold both kinds at once: a mutable buffer with no aliasing available -- a convolution @@ -749,15 +754,31 @@ def export( ), generate_etrecord=generate_etrecord, ) - if zero_copy_methods: - # After lowering, not beside the rewiring: to_edge re-derives the graph - # signature, so an order set earlier does not reach the finalizer. See - # order_copyback_mutations_first. - for name in edge_manager.methods: - order_copyback_mutations_first(edge_manager.exported_program(name)) + # After lowering, not beside the rewiring: to_edge re-derives the graph + # signature, so an order set earlier does not reach the finalizer. See + # order_copyback_mutations_first. + # + # Every method, not only the ones zero-copy rewired, because the crossing + # this repairs is not zero-copy's. A plain nn.Module writing one buffer from + # another buffer and a second buffer from a user input -- no TensorRT, no + # zero_copy_kv -- comes out of stock to_edge().to_executorch() with the first + # buffer's mutation spec naming the copy that writes the second. Gating this + # on zero_copy_methods would make a general repair depend on an unrelated + # opt-in, and would leave a sibling method repaired or not according to + # whether some other method asked for zero-copy. + for name in edge_manager.methods: + order_copyback_mutations_first(edge_manager.exported_program(name)) + if zero_copy_kv: # The copy-back is gone from here on, so finalizing without the matching # un-staging writes a .pte whose caches never update. This makes the # manager's own to_executorch read the finalized program back and refuse # that, which the two-call API otherwise leaves to the caller. + # + # Gated on what the caller asked for, not on what was rewired. A method + # with no aliased buffer mutation gets the warning above and a plain + # staged program, which is precisely one of the shapes the check refuses + # -- and save(..., zero_copy_kv=True) refuses that same model. Installing + # the hook only where something was rewired would leave the one entry + # point that cannot catch it the one that never looks. _check_zero_copy_kv_when_finalized(edge_manager) return edge_manager diff --git a/py/torch_tensorrt/executorch/_zero_copy.py b/py/torch_tensorrt/executorch/_zero_copy.py index 9c1f7a7357a..a47ff4273a2 100644 --- a/py/torch_tensorrt/executorch/_zero_copy.py +++ b/py/torch_tensorrt/executorch/_zero_copy.py @@ -22,7 +22,9 @@ partitioning, drops the copy-back by declaring that the buffer *is* the mutation's result. The aliased output then has no user and disappears from the partition. :func:`order_copyback_mutations_first` then repairs, on the Edge - program, the mutation-spec pairing that declaration disturbs downstream. + program, the mutation-spec pairing that declaration disturbs downstream -- + a crossing this is one cause of and not the only one, which is why + ``export()`` runs that repair over every method rather than the rewired ones. * :func:`unstage_aliased_buffers_pass`, as a ``to_out_var_pass``, drops the staging so the engine writes the caller's buffer rather than a copy. @@ -42,6 +44,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Set import torch +from executorch.exir.pass_base import PassBase from torch.fx import Node if TYPE_CHECKING: @@ -357,6 +360,13 @@ def order_copyback_mutations_first(exported_program: Any) -> int: graph's leading results into the state dict in that order and so updates each buffer from another one's value. + Zero-copy is not the only way in, which is why ``export()`` runs this over + every method rather than only the rewired ones. A plain ``nn.Module`` that + writes one buffer from another buffer -- whose mutation value is then that + other buffer's placeholder -- and a second buffer from a user input comes out + of stock ``to_edge().to_executorch()``, with no TensorRT anywhere, with the + first buffer's mutation spec naming the copy that writes the second. + Putting the mutations that still get a copy first restores the correspondence. Which ones those are is decided by asking upstream's own predicate (``_inplace_lineage``, imported rather than reimplemented) rather @@ -606,17 +616,29 @@ def _unstage_aliased_buffers( makes handing the buffer straight to the engine valid at all: a host-arena pointer is not something the engine can write. It is refused when the buffer has a consumer this pass leaves behind that does not survive being handed - the buffer from there (see :func:`_device_placement_is_safe`) -- asked - whether or not the spec already names that device, since it is the placement - and not the change of device that such a consumer does not survive. + the buffer from there (see :func:`_device_placement_is_safe`) -- asked on + both routes below and whether or not the spec already names that device, + since it is the placement and not the change of device that such a consumer + does not survive. A buffer already reaching its delegate directly is in the + same position as one this pass moves there, and is the shape this pass's own + first run leaves for its second. What the pass has to establish is the *post-condition*: every marked buffer - is a direct argument of a TensorRT delegate *and* ends up planned in device + is a direct argument of a TensorRT delegate *declaring zero-copy KV* -- one + whose own engine elided an aliased output -- *and* ends up planned in device memory. Removing a staging copy is only the usual way of getting there, not the goal, and a marked buffer that already satisfies both is left alone and counts as satisfied -- which is what running this pass a second time over a program it has already un-staged finds. + The zero-copy declaration is what narrows the delegates that count, on both + routes, because the mark says an engine writes the buffer in place and only a + stamped delegate's engine did. An unstamped TensorRT delegate taking the + buffer proves nothing about the stamped one, whose write would still go to a + staging copy that is discarded. :func:`check_zero_copy_kv` narrows the same + way over the finalized program, so the two halves of this post-condition + accept and refuse the same graphs. + Being a direct argument is not on its own enough, so it is not on its own accepted. Two things decide where the buffer is planned, and the second is not visible in the graph: the spec's own device, and whether memory planning @@ -725,8 +747,26 @@ def _unstage_aliased_buffers( "already been removed, so the update would be lost. " f"{remedy}" ) - satisfied_placeholders.add(arg) + if not _device_placement_is_safe( + graph_module, + arg, + h2d_copy, + direct_spec.device, + direct_spec.device_index, + orphaned_stagings, + ): + raise RuntimeError( + "TensorRT zero-copy KV: buffer " + f"'{arg.name}' reaches a TensorRT delegate directly " + "and is read by a consumer that does not survive it " + "being planned in this engine's device memory -- a " + "staging copy left behind reads it as a host source " + "and fails InvalidArgument. Export this method " + "without zero_copy_kv, or stop sharing the aliased " + "buffer." + ) if declares_zero_copy: + satisfied_placeholders.add(arg) satisfied_per_delegate[node].add(arg) continue if arg.target is not h2d_copy: @@ -785,9 +825,9 @@ def _unstage_aliased_buffers( source_spec.device_index = staged_spec.device_index new_args[i] = source unstaged += 1 - satisfied_placeholders.add(source) orphaned_stagings[arg] = None if declares_zero_copy: + satisfied_placeholders.add(source) satisfied_per_delegate[node].add(source) node.args = tuple(new_args) @@ -799,11 +839,15 @@ def _unstage_aliased_buffers( raise RuntimeError( "TensorRT zero-copy KV: buffer(s) " f"{names} were marked for in-place update but no TensorRT delegate " - "takes them, either directly or through a staging copy this pass " - "could remove. Export removed their copy-back on the promise that a " - "TensorRT engine writes them in place, so as this program stands " - "nothing updates them. Export this method without zero_copy_kv, or " - "keep the aliased buffer on a TensorRT delegate." + f"declaring zero-copy KV (compile spec " + f"'{ZERO_COPY_KV_COMPILE_SPEC_KEY}') takes them, either directly or " + "through a staging copy this pass could remove. Export removed their " + "copy-back on the promise that a TensorRT engine writes them in " + "place, so as this program stands nothing updates them. An unstamped " + "TensorRT delegate taking the buffer does not count: the stamp is " + "what records that an engine elided its aliased output for it. " + "Export this method without zero_copy_kv, or keep the aliased buffer " + "on the delegate whose engine elided it." ) for delegate in zero_copy_delegates: # One marked buffer per aliased output the spec says this delegate @@ -874,6 +918,40 @@ def _config_plans_on_devices(config: "ExecutorchBackendConfig") -> bool: return bool(config.enable_non_cpu_memory_planning) +class _UnstageThenToOutVar(PassBase): # type: ignore[misc] + """The ``to_out_var_pass`` :func:`unstage_aliased_buffers_pass` builds. + + A named type rather than a closure so that a config can be *recognised*: + :func:`_check_zero_copy_kv_when_finalized` has to answer, before it lets + finalization start, whether the config it was handed carries this pass, and + a class made afresh inside the builder gives a different type object on + every call for nothing to match against. + + ``inner`` is the ``to_out_var_pass`` that would otherwise have run. + ``device_memory_planning`` is what an unbound pass reads; once + ``finalization_config`` is set, the flag is resolved off that config on + every call instead, as :func:`_config_plans_on_devices` describes. + """ + + def __init__(self, inner: Any, device_memory_planning: bool) -> None: + self.inner = inner + self.device_memory_planning = device_memory_planning + self.finalization_config: Optional["ExecutorchBackendConfig"] = None + + def call(self, graph_module: torch.fx.GraphModule) -> Any: + config = self.finalization_config + planning = ( + self.device_memory_planning + if config is None + else _config_plans_on_devices(config) + ) + unstaged = _unstage_aliased_buffers( + graph_module, device_memory_planning=planning + ) + logger.debug("un-staged %d aliased delegate buffer(s)", unstaged) + return self.inner(graph_module) + + def unstage_aliased_buffers_pass( inner_pass: Optional[Any] = None, *, device_memory_planning: bool = True ) -> Any: @@ -906,31 +984,13 @@ def unstage_aliased_buffers_pass( binding would override it on every call. """ from executorch.exir import ExecutorchBackendConfig - from executorch.exir.pass_base import PassBase inner = ( inner_pass if inner_pass is not None else ExecutorchBackendConfig().to_out_var_pass ) - - class _UnstageThenToOutVar(PassBase): # type: ignore[misc] - finalization_config: Optional["ExecutorchBackendConfig"] = None - - def call(self, graph_module: torch.fx.GraphModule) -> Any: - config = self.finalization_config - planning = ( - device_memory_planning - if config is None - else _config_plans_on_devices(config) - ) - unstaged = _unstage_aliased_buffers( - graph_module, device_memory_planning=planning - ) - logger.debug("un-staged %d aliased delegate buffer(s)", unstaged) - return inner(graph_module) - - return _UnstageThenToOutVar() + return _UnstageThenToOutVar(inner, device_memory_planning) def _device_planned_arenas( @@ -1004,11 +1064,15 @@ def _is_host_planned( Sharing an arena with a host tensor settles it whatever the program records; otherwise, when the program does record its arena devices, an arena missing from that record is not a CUDA one. + + Asked only for a buffer that has a ``mem_id``. One that does not was never + planned, which is not a placement this can read and is refused as its own + thing by the caller. """ - mem_id = getattr(node.meta.get("spec"), "mem_id", None) - if mem_id in host_arenas: - return True - return device_arenas is not None and mem_id not in device_arenas + mem_id = node.meta["spec"].mem_id + return mem_id in host_arenas or ( + device_arenas is not None and mem_id not in device_arenas + ) def _planned_on_another_gpu( @@ -1059,14 +1123,15 @@ def check_zero_copy_kv(program: Any) -> None: error nor the speedup. The wrong-output case is the one below: a rewiring that did happen and then lost its mark. - Five shapes are refused: a marked buffer that is not a direct argument of a + Six shapes are refused: a marked buffer that is not a direct argument of a TensorRT delegate carrying the zero-copy compile spec, a stamped delegate that takes fewer marked buffers than its own spec says it elided aliased outputs -- including one in a method with no marked buffer at all, which is the lost-mark case -- a marked buffer that reaches such a delegate directly - but is not planned in device memory, one planned in a device arena the - program records for another GPU, and a program carrying neither a marked - buffer nor a stamped delegate in any of its methods. The first three are + but is planned in a host arena, one whose placement the program records + nowhere, one planned in a device arena the program records for another GPU, + and a program carrying neither a marked buffer nor a stamped delegate in any + of its methods. The first three are what finalizing without :func:`zero_copy_backend_config` leaves behind -- all but the lost-mark case folded into the second, which no finalization choice produces. That config's pass gets to all three earlier, off the @@ -1106,25 +1171,30 @@ def check_zero_copy_kv(program: Any) -> None: host tensors is one whatever the program records (see :func:`_host_planned_arenas`); failing that, an arena missing from the CUDA ones the program *does* record in ``non_const_buffer_device`` is one too. - Absence of that record is not itself read as the host: only ``apply_algo`` - writes it, so a caller-supplied ``memory_planning_pass`` that does not go - through it records nothing whatever it planned, and refusing on that would - block the one thing the guide says to bring your own planner for. That - acceptance is about what this function can *tell*, not about what the - runtime does with such a program: ``MethodMeta::memory_planned_buffer_device`` - answers ``CPU`` for a buffer the ``.pte`` records nothing for, so a runner - that honours it -- ``examples/executorch_reference_runner`` does -- backs - that arena with host memory and the engine then fails the alias-target guard - on every call. A planner for a zero-copy cache has to leave the record - behind, which means calling ``apply_algo`` with - ``enable_non_cpu_memory_planning=True``: that parameter defaults to + + A placement the program does not record at all is a refusal of its own + rather than either of those, because it is not an accusation about the + planner's choice -- a buffer with no ``mem_id`` was not planned, and a + program with no ``non_const_buffer_device`` entries says nothing about any + of its arenas. It is refused because of what the runtime makes of it: + ``MethodMeta::memory_planned_buffer_device`` answers ``CPU`` for an arena + the ``.pte`` records nothing for, so a runner that honours it -- + ``examples/executorch_reference_runner`` does -- backs that arena with host + memory and the engine then fails the alias-target guard on every call. Only + ``apply_algo`` writes that record, so a caller-supplied + ``memory_planning_pass`` that does not go through it leaves a ``.pte`` in + exactly that state whatever it planned; a planner for a zero-copy cache has + to leave the record behind, which means going through ``apply_algo`` with + ``enable_non_cpu_memory_planning=True``. That parameter defaults to ``False``, and with it off ``apply_algo`` plans every spec into one CPU - bucket and writes no record either. Where the - record does name a GPU it is read as one: an arena recorded for ``cuda:1`` - holding a cache whose own spec asks for ``cuda:0`` is an address on the - wrong device, which fails the same way a host one does, so the index is - compared and not only the type. Either index left unrecorded says nothing - and is accepted, as an unrecorded arena is. + bucket and writes no record either. + + Where the record does name a GPU it is read as one: an arena recorded for + ``cuda:1`` holding a cache whose own spec asks for ``cuda:0`` is an address + on the wrong device, which fails the same way a host one does, so the index + is compared and not only the type. An unrecorded index on either side says + nothing and is accepted, since the arena's device type is already settled by + then. Every method is read, not only ``forward``. ``export()`` rewires each method on its own, so a check that stopped at ``forward`` would pass a program whose @@ -1147,9 +1217,12 @@ def check_zero_copy_kv(program: Any) -> None: Both entry points run it for you: ``save(..., zero_copy_kv=True)`` before it writes the file, and the ``EdgeProgramManager`` that ``export(..., zero_copy_kv=True)`` returns on whatever its ``to_executorch`` - produces. Call it by hand for a program that reached finalization some other - way -- one derived from that manager by ``transform()`` or ``to_backend()``, - which are new managers without the hook. + produces -- that manager also refuses a config without the un-staging pass + before finalizing at all, which is the one refusal here that can be answered + without exporting again. Call this by hand for a program that reached + finalization some other way -- one derived from that manager by + ``transform()`` or ``to_backend()``, which are new managers without the + hook. Arguments: program (executorch.exir.ExecutorchProgramManager): The finalized program @@ -1165,6 +1238,7 @@ def check_zero_copy_kv(program: Any) -> None: method_names = sorted(program.methods) staged_by_method: Dict[str, List[str]] = {} short_by_method: Dict[str, List[str]] = {} + unrecorded_by_method: Dict[str, List[str]] = {} host_planned_by_method: Dict[str, List[str]] = {} wrong_gpu_by_method: Dict[str, List[str]] = {} marked_anywhere = False @@ -1227,22 +1301,37 @@ def check_zero_copy_kv(program: Any) -> None: device_arenas = _device_planned_arenas(graph_module) host_arenas = _host_planned_arenas(graph_module) reaching = [node for node in marked if node in zero_copy_delegate_args] - host_planned = [ - node.name - for node in reaching - if _is_host_planned(node, device_arenas, host_arenas) - ] + unrecorded: List[str] = [] + host_planned: List[str] = [] + wrong_gpu: List[str] = [] + # One classification per buffer. The two unrecorded shapes are not + # evidence about the host arena and get a refusal of their own, but the + # positive host-arena ground is read first where it applies, since + # host-only planning both puts the cache among the host tensors and + # writes no arena record, and naming the arena it is actually in tells + # the caller more than saying nothing was recorded. + for node in reaching: + mem_id = getattr(node.meta.get("spec"), "mem_id", None) + if mem_id is None: + unrecorded.append( + f"'{node.name}' carries no mem_id, so memory planning left " + "it unplanned and nothing in the program says where it lives" + ) + elif _is_host_planned(node, device_arenas, host_arenas): + host_planned.append(node.name) + elif device_arenas is None: + unrecorded.append( + f"'{node.name}' is planned in arena {mem_id}, and the " + "program records no CUDA arena at all" + ) + else: + detail = _planned_on_another_gpu(node, device_arenas) + if detail is not None: + wrong_gpu.append(detail) + if unrecorded: + unrecorded_by_method[method_name] = unrecorded if host_planned: host_planned_by_method[method_name] = host_planned - wrong_gpu = [ - detail - for detail in ( - _planned_on_another_gpu(node, device_arenas) - for node in reaching - if not _is_host_planned(node, device_arenas, host_arenas) - ) - if detail is not None - ] if wrong_gpu: wrong_gpu_by_method[method_name] = wrong_gpu if staged_by_method: @@ -1295,6 +1384,27 @@ def check_zero_copy_kv(program: Any) -> None: "buffer among the host tensors, and it has to give the delegate's " "device an arena of its own." ) + if unrecorded_by_method: + raise RuntimeError( + "TensorRT zero-copy KV: buffer(s) " + + "; ".join( + f"{detail} in method '{method}'" + for method, details in unrecorded_by_method.items() + for detail in details + ) + + ". These buffers reach their TensorRT delegate directly, so the " + "engine writes them through the pointer the runtime allocates for " + "them, and the .pte has to say that pointer is device memory. " + "MethodMeta::memory_planned_buffer_device answers CPU for an arena " + "the program records nothing for, so a runner that honours it backs " + "the arena with host memory and every execute() fails on the " + "alias-target guard. The memory_planning_pass in use has to plan " + "these buffers and leave the record behind, which means going " + "through ExecuTorch's apply_algo with " + "enable_non_cpu_memory_planning=True -- that parameter defaults to " + "False, and with it off apply_algo plans every spec into one CPU " + "bucket and writes no record either." + ) if wrong_gpu_by_method: raise RuntimeError( "TensorRT zero-copy KV: buffer(s) " @@ -1312,23 +1422,35 @@ def check_zero_copy_kv(program: Any) -> None: def _check_zero_copy_kv_when_finalized(edge_manager: Any) -> None: - """Make one manager's own ``to_executorch`` run :func:`check_zero_copy_kv`. + """Guard one manager's own ``to_executorch``, before and after it runs. Export removes the copy-back of the rewired caches before it returns, so from that point on every way of finalizing the program that does not also un-stage them produces a ``.pte`` whose caches never update -- silently, and for a KV cache that is wrong output rather than a crash. ``to_executorch()`` with ExecuTorch's defaults is one such way, and it is the call the two-step - API documents. Reading the finalized program back is what tells the two - apart, and this is the only place holding it that knows the export asked for - zero-copy, so the check runs here rather than being left for the caller to - remember. It is the same check ``save(..., zero_copy_kv=True)`` runs at the - same point for the same reason. - - Nothing about the program changes: a correct one is returned untouched, and - a broken one raises instead of being handed back. The refusals name what to - do -- finalize through :func:`zero_copy_backend_config` -- so this needs no - message of its own. + API documents. This is the only place holding the manager that knows the + export asked for zero-copy, so the refusal lives here rather than being left + for the caller to remember. + + It asks two questions at two moments, because only the first can be answered + while there is still something to be done about it. ``to_executorch`` does + not copy the edge programs: it runs the write-back and the device passes over + the manager's own graph modules and copies the finalized graph back into + them, so once finalization has run the manager is spent -- calling it a + second time dies inside ``insert_write_back_for_buffers_pass``. A refusal + raised after the fact could therefore only be acted on by exporting again and + rebuilding every engine. So the config is read *first*: one that does not + carry the un-staging pass cannot produce a zero-copy program whatever the + graph looks like, and saying so before delegating leaves the manager + untouched and the remedy -- the same call with + :func:`zero_copy_backend_config` -- available. + + The finalized program is still read back afterwards, through + :func:`check_zero_copy_kv`, because a config can be the right one and the + program still come out wrong: memory planning chooses the arenas after this + pass has run, and a lost mark is a property of the graph. Those refusals + name what to do themselves. Bound to the instance rather than to a type, because ``EdgeProgramManager`` is ExecuTorch's. That reaches the manager ``export()`` hands back and no @@ -1340,6 +1462,24 @@ def _check_zero_copy_kv_when_finalized(edge_manager: Any) -> None: @functools.wraps(finalize) def to_executorch(*args: Any, **kwargs: Any) -> Any: + config = kwargs.get("config", args[0] if args else None) + if not isinstance( + getattr(config, "to_out_var_pass", None), _UnstageThenToOutVar + ): + raise RuntimeError( + "TensorRT zero-copy KV: this program was exported with " + "zero_copy_kv=True, which removed the copy-back of its aliased " + "buffers, and this configuration does not carry the pass that " + "un-stages them. Finalizing it would write a .pte whose engine " + "writes a per-call staging copy that is discarded, so the caches " + "never update -- for a KV cache, wrong output rather than a " + "crash. Finalize with " + "to_executorch(torch_tensorrt.executorch.zero_copy_backend_config" + "(config)) instead. This is raised before finalizing rather than " + "after, because to_executorch rewrites the manager's own edge " + "programs in place and a manager that has finalized once cannot " + "do it again." + ) program = finalize(*args, **kwargs) check_zero_copy_kv(program) return program @@ -1431,10 +1571,11 @@ def zero_copy_backend_config( planner's own business. Neither is left to the runtime to discover: :func:`check_zero_copy_kv` reads the arena memory planning actually chose and refuses the program either mistake produces, and the manager - ``export(..., zero_copy_kv=True)`` returns runs that check itself. Two - placements it cannot settle are a method holding no host tensor to give - the shared arena away, and a planner that leaves no arena-device record - for it to read; :func:`check_zero_copy_kv` describes both. + ``export(..., zero_copy_kv=True)`` returns runs that check itself. The one + placement it cannot settle is a method holding no host tensor to give the + shared arena away and a program that records its arena devices; a planner + that records nothing is refused rather than passed over, as + :func:`check_zero_copy_kv` describes. * ``propagate_device_config.skip_h2d_for_method_inputs`` is *refused*, wherever it is written -- in the single ``PropagateDeviceConfig`` or in a per-method dict of them -- and on every value ``PropagateDevicePass`` @@ -1452,8 +1593,10 @@ def zero_copy_backend_config( Finalizing a ``zero_copy_kv=True`` program *without* this config leaves the engine writing a per-call staging copy that is then discarded, so the buffer never updates -- for a KV cache, wrong output rather than a - crash. The manager ``export()`` hands back refuses that itself, reading - its own finalized program through :func:`check_zero_copy_kv`, and + crash. The manager ``export()`` hands back refuses that itself: it reads + the config before finalizing and refuses one that does not carry this + pass, and it reads its own finalized program through + :func:`check_zero_copy_kv` for what the config cannot say. ``torch_tensorrt.save(..., zero_copy_kv=True)`` runs the same check before it writes the file. A program finalized off some *other* manager -- one ``transform()`` or ``to_backend()`` derived from that diff --git a/tests/cpp/executorch/test_executorch_blob_header.cpp b/tests/cpp/executorch/test_executorch_blob_header.cpp index 1eaa5abaf42..95556f6673e 100644 --- a/tests/cpp/executorch/test_executorch_blob_header.cpp +++ b/tests/cpp/executorch/test_executorch_blob_header.cpp @@ -457,10 +457,9 @@ TEST(ExecuTorchTensorRTBlobHeader, RejectsBindingEntryWithAMisspelledIsInputKey) TEST(ExecuTorchTensorRTBlobHeader, RejectsABindingNameThatArrivedEscaped) { // parse_string drops the backslash and keeps what follows, so this name is - // recorded as the three characters anb and TensorRT has no such tensor. The - // writer is json.dumps, which escapes every control character and everything - // non-ASCII, so an escape is exactly where the two ends read a name - // differently. + // recorded as the three characters anb and TensorRT has no such tensor. Every + // escape json.dumps reserves for a control character or a non-ASCII codepoint + // is a place the two ends read a name differently. const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" R"({"name":"a\nb","is_input":false}]})"; const auto blob = make_blob(metadata); @@ -469,6 +468,89 @@ TEST(ExecuTorchTensorRTBlobHeader, RejectsABindingNameThatArrivedEscaped) { EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); } +TEST(ExecuTorchTensorRTBlobHeader, ParsesABindingNameEscapedTheWayThisParserReadsIt) { + // The three escapes whose JSON meaning is "the character after the + // backslash", which is what parse_string produces for every escape. A quote + // and a backslash are what json.dumps emits for a name holding one, so + // refusing these would refuse names the writer really produces and the + // merge-base parser recorded exactly as written. The forward slash is the + // third of the three; json.dumps writes it plain, but a hand-assembled blob + // may escape it and both ends read it the same way. + const std::string metadata = R"({"io_bindings":[{"name":"a\"b","is_input":true},)" + R"({"name":"c\\d","is_input":false},)" + R"({"name":"e\/f","is_input":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + ASSERT_EQ(header.input_binding_names.size(), 1u); + EXPECT_EQ(header.input_binding_names[0], "a\"b"); + ASSERT_EQ(header.output_binding_names.size(), 2u); + EXPECT_EQ(header.output_binding_names[0], "c\\d"); + EXPECT_EQ(header.output_binding_names[1], "e/f"); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAnEntryKeySpelledThroughAnEscape) { + // The key comparisons see the text parse_string produced, so a backslash + // before any letter of is_input still reads as is_input here while a JSON + // reader sees a key with a newline in it, ignores it, and leaves is_input at + // its own default -- which serialization.py's TensorRTIOBinding makes an + // input and the initializer here makes an output. That is the two readers of + // one blob putting one binding in opposite lists, which is exactly what + // RejectsBindingEntryWithAMisspelledIsInputKey stops for the plainly + // misspelled key. + const std::string metadata = R"({"io_bindings":[{"name":"x","is_i\nput":false}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesAnEntryKeyCarryingAnEscapeBothEndsAgreeOn) { + // The control for the refusal above: an escaped backslash in a key is decoded + // the same way here and by a JSON reader, so it names the key dty\pe for + // both, which neither matches. Nothing depends on it and the blob is good. + const std::string metadata = R"({"io_bindings":[{"name":"x","dty\\pe":"f32","is_input":true}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + ASSERT_EQ(header.input_binding_names.size(), 1u); + EXPECT_EQ(header.input_binding_names[0], "x"); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAnAliasEntryKeySpelledThroughAnEscape) { + // The alias entry's keys are compared the same way the binding entry's are. + // A backslash before any letter of output leaves this parser reading the + // output name out of it, while a JSON reader sees a key it does not know and + // an entry with no output at all -- which it would then have to refuse. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"\output":"out_k","input":"in_k","kind":"kv_cache_update"}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsAnAliasKindSpelledThroughAnEscape) { + // kind is not a binding name, but init() compares it: "user" is the kind + // validated on shape alone rather than confirmed against the engine's own + // aliasing, so reading a different string as user picks the weaker check. + // Here both entries come out as the plain string user, while a JSON reader + // refuses "\user" outright and reads "use\r" as use plus a carriage return. + for (const char* kind : {R"(\user)", R"(use\r)"}) { + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":")" + + std::string(kind) + R"("}]})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)) << kind; + } +} + TEST(ExecuTorchTensorRTBlobHeader, RejectsAnAliasedIoNameThatArrivedEscaped) { const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" R"({"name":"out_k","is_input":false}],)" @@ -577,6 +659,79 @@ TEST(ExecuTorchTensorRTBlobHeader, ParsesAStringValueThatIsExactlyAScalarKeyName EXPECT_EQ(header.device_id, 0); } +TEST(ExecuTorchTensorRTBlobHeader, ParsesEveryScalarFromTheWriterKeyOrder) { + // The key order TensorRTBlobMetadata.to_json emits, with every field present + // and both scalars set away from their defaults. The scalar scans start past + // whichever array they last walked, so a field moved ahead of one of them is + // not found and keeps its C++-side default while the parse still succeeds -- + // this is what fails if that order changes. The Python half of the same rule + // is test_serialization.py::test_to_json_writes_every_scalar_after_both_arrays. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","dtype":"float32","shape":[1,2],"is_input":true},)" + R"({"name":"out_k","dtype":"float32","shape":[1,2],"is_input":false}],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}],)" + R"("hardware_compatible":true,"device_id":6,)" + R"("serialized_metadata":"","target_platform":"linux_x86_64"})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_EQ(header.input_binding_names, std::vector{"in_k"}); + EXPECT_EQ(header.output_binding_names, std::vector{"out_k"}); + ASSERT_EQ(header.aliased_io.size(), 1u); + EXPECT_EQ(header.aliased_io[0].output, "out_k"); + EXPECT_EQ(header.aliased_io[0].input, "in_k"); + EXPECT_EQ(header.aliased_io[0].kind, "kv_cache_update"); + EXPECT_TRUE(header.hardware_compatible); + EXPECT_EQ(header.device_id, 6); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesAnArrayKeySpelledByAnEarlierStringValue) { + // The two array keys are found the way the two scalars are: an occurrence + // followed by its own colon. Taking the first occurrence anywhere and then + // the next '[' reads the value below as the key and walks the shape array + // that follows it, which refuses a blob that is perfectly good. + const std::string metadata = R"({"serialized_metadata":"io_bindings","shape":[9],)" + R"("io_bindings":[{"name":"x","is_input":true}],"device_id":6})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + EXPECT_EQ(header.input_binding_names, std::vector{"x"}); + EXPECT_EQ(header.device_id, 6); +} + +TEST(ExecuTorchTensorRTBlobHeader, ParsesTheAliasArrayKeySpelledByAnEarlierStringValue) { + // The alias key gets the same treatment, and its window is narrower: the + // search already starts past io_bindings, so only a value between the two + // arrays can stand in for it -- which is where serialized_metadata sits. + const std::string metadata = R"({"io_bindings":[{"name":"in_k","is_input":true},)" + R"({"name":"out_k","is_input":false}],)" + R"("serialized_metadata":"aliased_io","shape":[9],)" + R"("aliased_io":[{"output":"out_k","input":"in_k","kind":"kv_cache_update"}],)" + R"("device_id":6})"; + const auto blob = make_blob(metadata, 4, TENSORRT_MAGIC_ALIASED_IO); + + TensorRTBlobHeader header; + ASSERT_TRUE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); + ASSERT_EQ(header.aliased_io.size(), 1u); + EXPECT_EQ(header.aliased_io[0].output, "out_k"); + EXPECT_EQ(header.device_id, 6); +} + +TEST(ExecuTorchTensorRTBlobHeader, RejectsIoBindingsWhoseValueIsNotAnArray) { + // The array has to be the value of the key, not the next '[' in the text: a + // blob whose io_bindings is an object is otherwise walked from an unrelated + // bracket further on, and the entries found there are recorded as this + // engine's bindings. The array below is shaped like the real one so that + // walking it succeeds, which is what makes the wrong answer a silent one. + const std::string metadata = R"({"io_bindings":{"name":"x"},)" + R"("elsewhere":[{"name":"y","is_input":true}]})"; + const auto blob = make_blob(metadata); + + TensorRTBlobHeader header; + EXPECT_FALSE(TensorRTBlobHeader::parse(blob.data(), blob.size(), header)); +} + TEST(ExecuTorchTensorRTBlobHeader, RejectsUnknownFutureMagic) { constexpr char kFutureMagic[4] = {'T', 'R', '0', '3'}; const std::string metadata = R"({"io_bindings":[{"name":"x","is_input":true}]})"; diff --git a/tests/py/dynamo/executorch/test_api.py b/tests/py/dynamo/executorch/test_api.py index 0768b1c8eba..191a93ea85f 100644 --- a/tests/py/dynamo/executorch/test_api.py +++ b/tests/py/dynamo/executorch/test_api.py @@ -563,6 +563,10 @@ def get_etrecord(self): return _FakeETRecord() class _FakeEdge: + # export() reorders each method's mutations after lowering, over every + # method the manager holds. No method here holds a program to reorder. + methods = () + def to_executorch(self, config=None): captured["backend_config"] = config return _FakeExec() diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py index 31d3ce36593..0bf58f40dd4 100644 --- a/tests/py/dynamo/executorch/test_export.py +++ b/tests/py/dynamo/executorch/test_export.py @@ -669,10 +669,82 @@ def test_export_wires_the_reorder_and_the_finalize_check(monkeypatch): # Nothing is finalized yet, so the check must not have run. assert checked == [] - assert result.to_executorch() is finalized + assert result.to_executorch(zero_copy.zero_copy_backend_config()) is finalized assert checked == [finalized] +@pytest.mark.unit +def test_export_installs_the_finalize_check_even_when_nothing_was_rewired(monkeypatch): + """A zero_copy_kv=True that rewired nothing is the shape the check refuses. + + Export warns and carries on, so the manager it hands back would otherwise + finalize a plain staged program without a word, while + ``save(..., zero_copy_kv=True)`` refuses that same model. Gating the hook on + what was rewired would leave the case that most needs reading back as the + one case nobody reads back. + """ + import torch_tensorrt.executorch._zero_copy as zero_copy + + export_module, lower = _patch_lowering(monkeypatch) + _patch_declare(monkeypatch) + _patch_rewire(monkeypatch, elided_names=()) + + manager = FakeEdgeProgramManager() + finalized = object() + manager.to_executorch = lambda config=None: finalized + lower.return_value = manager + monkeypatch.setattr(zero_copy, "order_copyback_mutations_first", lambda program: 0) + checked = [] + monkeypatch.setattr(zero_copy, "check_zero_copy_kv", checked.append) + + result = export_module.export(FakeExportedProgram(), zero_copy_kv=True) + + assert checked == [] + assert result.to_executorch(zero_copy.zero_copy_backend_config()) is finalized + assert checked == [finalized] + + +@pytest.mark.unit +def test_export_reorders_mutations_for_a_method_that_never_asked_for_zero_copy( + monkeypatch, +): + """The reorder is not gated on zero-copy, because the crossing is not either. + + Stock ``to_edge().to_executorch()`` crosses the mutation pairing on a plain + module that writes one buffer from another buffer and a second from a user + input, with no TensorRT and no ``zero_copy_kv``. Gating the repair on + ``zero_copy_kv`` would leave that caller with the crossed program and would + repair a sibling method only because another method opted in. + """ + import torch_tensorrt.executorch._zero_copy as zero_copy + + export_module, lower = _patch_lowering(monkeypatch) + _patch_declare(monkeypatch) + + manager = FakeEdgeProgramManager() + lower.return_value = manager + + reordered = [] + monkeypatch.setattr( + zero_copy, + "order_copyback_mutations_first", + lambda program: (reordered.append(program), 0)[1], + ) + checked = [] + monkeypatch.setattr(zero_copy, "check_zero_copy_kv", checked.append) + + result = export_module.export(FakeExportedProgram(), zero_copy_kv=False) + + assert result is manager + assert [id(program) for program in reordered] == [ + id(manager.exported_program("forward")) + ] + # No zero-copy was asked for, so no finalize-time check is installed: the + # hook binds to the instance, so an unhooked manager has nothing there. + assert "to_executorch" not in vars(result) + assert checked == [] + + @pytest.mark.unit def test_export_does_not_exempt_a_method_that_kept_all_its_outputs(monkeypatch): """The exemption is per method and only where an output was actually elided. diff --git a/tests/py/dynamo/executorch/test_serialization.py b/tests/py/dynamo/executorch/test_serialization.py index c025b11c024..7fc45f9b24e 100644 --- a/tests/py/dynamo/executorch/test_serialization.py +++ b/tests/py/dynamo/executorch/test_serialization.py @@ -120,3 +120,47 @@ def test_deserialize_accepts_both_magics(aliased_io): engine, parsed = deserialize_engine(serialize_engine(b"engine-bytes", metadata)) assert engine == b"engine-bytes" assert parsed.aliased_io == aliased_io + + +@pytest.mark.unit +def test_to_json_writes_every_scalar_after_both_arrays(): + """The ordering rule the C++ parser's scalar scans depend on. + + ``TensorRTBlobHeader.cpp`` walks ``io_bindings``, then ``aliased_io``, then + searches forward from the end of whichever of those it last walked for the + scalar fields, so a scalar written ahead of either array is not found and + keeps its C++-side default while the parse still succeeds -- no error, and a + ``device_id`` of 0 means the engine deserializes on a GPU nobody named. The + C++ half of this rule is + ``ParsesEveryScalarFromTheWriterKeyOrder`` in + ``tests/cpp/executorch/test_executorch_blob_header.cpp``; this is the half + that fails when a field is added to ``to_json`` in the wrong place. + + Every key that side reads by key is listed, not only the two scalars: a new + scalar it learns to read is only safe in the same position. + """ + metadata = TensorRTBlobMetadata( + io_bindings=[ + TensorRTIOBinding(name="in_k", dtype="float32", shape=[1, 2]), + TensorRTIOBinding(name="out_k", dtype="float32", is_input=False), + ], + aliased_io={"out_k": ("in_k", "kv_cache_update")}, + hardware_compatible=True, + device_id=6, + target_platform="linux_x86_64", + ) + + text = metadata.to_json().decode("utf-8") + arrays_end = max( + text.index("]", text.index(key)) for key in ('"io_bindings"', '"aliased_io"') + ) + for key in ('"hardware_compatible"', '"device_id"'): + assert text.index(key) > arrays_end, f"{key} is written before an array" + + # Read back through the writer's own reader as well, so the ordering + # assertion is made about a payload that is otherwise correct. + restored = TensorRTBlobMetadata.from_json(metadata.to_json()) + assert restored.hardware_compatible is True + assert restored.device_id == 6 + assert restored.aliased_io == {"out_k": ("in_k", "kv_cache_update")} + assert [b.name for b in restored.io_bindings] == ["in_k", "out_k"] diff --git a/tests/py/dynamo/executorch/test_weight_streaming_budget.py b/tests/py/dynamo/executorch/test_weight_streaming_budget.py index 42b10ea97d8..eff03a3e0c0 100644 --- a/tests/py/dynamo/executorch/test_weight_streaming_budget.py +++ b/tests/py/dynamo/executorch/test_weight_streaming_budget.py @@ -70,7 +70,9 @@ def _patch_lowering(monkeypatch, engine_counts=None): ) export_module = importlib.import_module("torch_tensorrt.executorch._export") engine_counts = engine_counts or {} - lower = MagicMock(return_value=object()) + # export() reorders each method's mutations over every method the manager + # holds after lowering, so the stand-in has to answer that much. + lower = MagicMock(return_value=SimpleNamespace(methods=())) monkeypatch.setattr(executorch.exir, "to_edge_transform_and_lower", lower) monkeypatch.setattr(executorch_api, "TensorRTPartitioner", FakeTensorRTPartitioner) monkeypatch.setattr(executorch_api, "get_edge_compile_config", lambda: "default") diff --git a/tests/py/dynamo/executorch/test_zero_copy_kv.py b/tests/py/dynamo/executorch/test_zero_copy_kv.py index 79feb922b97..2d88379cc75 100644 --- a/tests/py/dynamo/executorch/test_zero_copy_kv.py +++ b/tests/py/dynamo/executorch/test_zero_copy_kv.py @@ -465,7 +465,9 @@ def test_unstage_feeds_the_buffer_straight_to_the_delegate(): Moving the spec is not cosmetic: memory planning reads it, and a buffer left in a host arena is somewhere the engine cannot write. """ - graph_module, k_buffer, staged_k, delegate = _staged_delegate_graph() + graph_module, k_buffer, staged_k, delegate = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True assert Z._unstage_aliased_buffers(graph_module) == 1 @@ -517,23 +519,17 @@ def _direct_delegate_graph(*, compile_specs=None, device=DeviceType.CUDA): @pytest.mark.unit -@pytest.mark.parametrize( - "compile_specs", - [None, [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")]], - ids=["plain", "declares-zero-copy"], -) -def test_unstage_accepts_a_buffer_that_never_had_a_staging_copy(compile_specs): +def test_unstage_accepts_a_buffer_that_never_had_a_staging_copy(): """A marked buffer already handed straight to its delegate needs no work. What the pass has to leave behind is a marked buffer that is a delegate argument planned in device memory; removing a staging copy is only the usual route there. Keying success on having removed one instead rejects this program, which is already in the shape zero-copy wants -- and it is the shape - the pass's own second run sees. The delegate's own zero-copy declaration is - cross-checked against the same count, so it is parametrized here too. + the pass's own second run sees. """ graph_module, k_buffer, delegate = _direct_delegate_graph( - compile_specs=compile_specs + compile_specs=[CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")] ) assert Z._unstage_aliased_buffers(graph_module) == 0 @@ -542,6 +538,39 @@ def test_unstage_accepts_a_buffer_that_never_had_a_staging_copy(compile_specs): assert k_buffer.meta["spec"].device == DeviceType.CUDA +@pytest.mark.unit +@pytest.mark.parametrize("route", ["direct", "staged"]) +def test_unstage_refuses_a_marked_buffer_whose_only_delegate_is_unstamped(route): + """The two halves of one post-condition have to answer one graph the same way. + + A TensorRT delegate carrying no zero-copy compile spec is one whose engine + elided no aliased output, so a marked buffer reaching it says nothing about + the engine that does write it in place -- that engine's write is still going + to a staging copy nothing reads back. :func:`check_zero_copy_kv` counts only + stamped delegates over the finalized program, on either route the buffer + took, so both routes are pinned here. + """ + stamped = [CompileSpec(ZERO_COPY_KV_COMPILE_SPEC_KEY, b"[]")] + if route == "direct": + graph_module, _, _ = _direct_delegate_graph(compile_specs=None) + else: + graph_module, k_buffer, _, _ = _staged_delegate_graph(compile_specs=None) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="declaring zero-copy KV"): + Z._unstage_aliased_buffers(graph_module) + + # The same graph with the delegate stamped is accepted, so what the refusal + # reads is the missing stamp and not something else about these graphs. + if route == "direct": + graph_module, _, _ = _direct_delegate_graph(compile_specs=stamped) + assert Z._unstage_aliased_buffers(graph_module) == 0 + else: + graph_module, k_buffer, _, _ = _staged_delegate_graph(compile_specs=stamped) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + assert Z._unstage_aliased_buffers(graph_module) == 1 + + @pytest.mark.unit @pytest.mark.parametrize("spec", ["absent", "host"]) def test_unstage_refuses_a_direct_buffer_that_is_not_on_the_device(spec): @@ -630,7 +659,7 @@ def test_zero_copy_backend_config_reads_the_planning_mode_when_the_pass_runs(): turned_off = Z.zero_copy_backend_config() turned_off.enable_non_cpu_memory_planning = False - graph_module, _, _ = _direct_delegate_graph() + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): turned_off.to_out_var_pass(graph_module) @@ -638,7 +667,7 @@ def test_zero_copy_backend_config_reads_the_planning_mode_when_the_pass_runs(): ExecutorchBackendConfig(enable_non_cpu_memory_planning=False) ) turned_on.enable_non_cpu_memory_planning = True - graph_module, _, _ = _direct_delegate_graph() + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) turned_on.to_out_var_pass(graph_module) @@ -668,7 +697,7 @@ def a_planner_of_ones_own(graph_module): memory_planning_pass=a_planner_of_ones_own, ) ) - graph_module, _, _ = _direct_delegate_graph() + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) config.to_out_var_pass(graph_module) @@ -678,7 +707,7 @@ def a_planner_of_ones_own(graph_module): ExecutorchBackendConfig(enable_non_cpu_memory_planning=False) ) assert hasattr(stock.memory_planning_pass, "enable_non_cpu_memory_planning") - graph_module, _, _ = _direct_delegate_graph() + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): stock.to_out_var_pass(graph_module) @@ -704,11 +733,11 @@ def test_zero_copy_backend_config_rebuilt_over_a_derived_config_reads_it(): derived = dataclasses.replace( Z.zero_copy_backend_config(), enable_non_cpu_memory_planning=False ) - graph_module, _, _ = _direct_delegate_graph() + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) derived.to_out_var_pass(graph_module) rebuilt = Z.zero_copy_backend_config(derived) - graph_module, _, _ = _direct_delegate_graph() + graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) with pytest.raises(RuntimeError, match="enable_non_cpu_memory_planning=False"): rebuilt.to_out_var_pass(graph_module) @@ -722,7 +751,9 @@ def test_unstage_runs_a_second_time_without_raising(): second run finds the buffer already wired to the delegate and returns without un-staging anything. """ - graph_module, k_buffer, _, delegate = _staged_delegate_graph() + graph_module, k_buffer, _, delegate = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True assert Z._unstage_aliased_buffers(graph_module) == 1 @@ -741,7 +772,9 @@ def test_unstage_raises_for_a_marked_buffer_on_another_backends_delegate(): ) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - with pytest.raises(RuntimeError, match="no TensorRT delegate takes them"): + with pytest.raises( + RuntimeError, match="no TensorRT delegate declaring zero-copy KV" + ): Z._unstage_aliased_buffers(graph_module) @@ -757,7 +790,9 @@ def test_unstage_raises_when_a_marked_buffer_is_never_unstaged(): k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True - with pytest.raises(RuntimeError, match="no TensorRT delegate takes them"): + with pytest.raises( + RuntimeError, match="no TensorRT delegate declaring zero-copy KV" + ): Z._unstage_aliased_buffers(graph_module) @@ -936,7 +971,9 @@ def test_unstage_moves_a_buffer_two_tensorrt_delegates_stage_from(): setattr( root, name, - SimpleNamespace(backend_id="TensorRTBackend", compile_specs=None), + SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs() + ), ) graph_module = torch.fx.GraphModule(root, graph) k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=0) @@ -1000,7 +1037,7 @@ def test_unstage_allows_a_buffer_that_is_also_its_mutation_output(): graph.output((k_buffer, delegate)) root = torch.nn.Module() root.lowered_module_0 = SimpleNamespace( - backend_id="TensorRTBackend", compile_specs=None + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs() ) graph_module = torch.fx.GraphModule(root, graph) k_buffer.meta["spec"] = SimpleNamespace(device=DeviceType.CPU, device_index=3) @@ -1012,6 +1049,56 @@ def test_unstage_allows_a_buffer_that_is_also_its_mutation_output(): assert k_buffer.meta["spec"].device == DeviceType.CUDA +@pytest.mark.unit +@pytest.mark.parametrize("route", ["staged", "direct"]) +def test_unstage_refuses_a_buffer_a_surviving_consumer_reads(route): + """The surviving-consumer refusal belongs to the placement, not to the move. + + A buffer already reaching its delegate directly is in the same position as + one this pass un-stages: it is planned in the engine's device memory, and a + consumer that reads it as a host source fails ``InvalidArgument`` on the + first call. Asking only on the branch that removes a staging copy would let + the identical graph through by the other door -- including the graph this + pass's own first run produces, which its second run reads directly. + """ + graph = torch.fx.Graph() + k_buffer = graph.placeholder("b_k_0") + lowered = graph.get_attr("lowered_module_0") + h2d = torch.ops.et_copy._h2d_copy.default + delegate_arg = ( + graph.call_function(h2d, (k_buffer,)) if route == "staged" else k_buffer + ) + # Another backend's host copy of the same buffer, which nothing removes. + foreign = graph.call_function(h2d, (k_buffer,)) + lowered_other = graph.get_attr("lowered_module_1") + other = graph.call_function(executorch_call_delegate, (lowered_other, foreign)) + delegate = graph.call_function(executorch_call_delegate, (lowered, delegate_arg)) + graph.output((k_buffer, delegate, other)) + root = torch.nn.Module() + root.lowered_module_0 = SimpleNamespace( + backend_id="TensorRTBackend", compile_specs=_zero_copy_specs() + ) + root.lowered_module_1 = SimpleNamespace( + backend_id="CudaBackend", compile_specs=None + ) + graph_module = torch.fx.GraphModule(root, graph) + k_buffer.meta["spec"] = SimpleNamespace( + device=DeviceType.CPU if route == "staged" else DeviceType.CUDA, + device_index=0, + ) + if route == "staged": + delegate_arg.meta["spec"] = SimpleNamespace( + device=DeviceType.CUDA, device_index=0 + ) + # On the GPU the engine is not on, so it is the placement and not merely the + # presence of a second reader that this refuses. + foreign.meta["spec"] = SimpleNamespace(device=DeviceType.CUDA, device_index=1) + k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True + + with pytest.raises(RuntimeError, match="does not survive|leaves in place"): + Z._unstage_aliased_buffers(graph_module) + + @pytest.mark.unit def test_unstage_refuses_to_move_a_shared_buffer(): """A buffer read by a consumer other than its TensorRT delegate staging @@ -1230,7 +1317,9 @@ def test_zero_copy_backend_config_keeps_the_callers_config(): assert config.memory_planning_pass is base.memory_planning_pass assert config.to_out_var_pass is not inner # The caller's to_out_var_pass is not dropped, it is run after the un-staging. - graph_module, k_buffer, _, delegate = _staged_delegate_graph() + graph_module, k_buffer, _, delegate = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True seen = [] base = ExecutorchBackendConfig(to_out_var_pass=lambda gm: seen.append(gm)) @@ -1393,30 +1482,58 @@ def test_check_zero_copy_kv_rejects_a_direct_buffer_planned_on_the_host(): @pytest.mark.unit -def test_check_zero_copy_kv_accepts_a_device_arena_a_custom_planner_did_not_record(): - """An absent arena-device record is "cannot tell", not "on the host". +def test_check_zero_copy_kv_rejects_an_arena_no_planner_recorded(): + """An unrecorded arena is what the runtime reads as the host, whoever planned it. ``apply_algo`` is the only thing in ExecuTorch that writes ``non_const_buffer_device``, and ``to_executorch`` takes any callable as - ``memory_planning_pass`` -- which the user guide tells people to supply for a - cache shared between prefill and decode. Reading the absent record as a host - placement refuses that program, and ``save`` runs this check with no way to - opt out, so nothing is written at all. This is the previous test's graph with - the host tensors moved out of the buffer's arena, which is the whole - difference between the two. + ``memory_planning_pass``, so a caller-supplied planner can put the cache in + device memory and leave the ``.pte`` saying nothing. It is refused all the + same: ``MethodMeta::memory_planned_buffer_device`` answers ``CPU`` for an + arena with no entry, so the runner backs it with host memory and the engine + fails the alias-target guard on the first call. + + This is the host-arena test's graph with the host tensors moved out of the + buffer's arena, so the two refusals are told apart by which of them fires: + the arena here is not one the program's host tensors are in, and the message + says the record is absent rather than accusing the planner of putting the + cache among the host tensors. """ graph_module, _, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) - Z.check_zero_copy_kv( - _finalized_program( - _planned( - graph_module, - arena=CUDA_ARENA, - on_device=False, - host_tensor_arena=HOST_ARENA, + with pytest.raises(RuntimeError, match="records no CUDA arena at all"): + Z.check_zero_copy_kv( + _finalized_program( + _planned( + graph_module, + arena=CUDA_ARENA, + on_device=False, + host_tensor_arena=HOST_ARENA, + ) ) ) - ) + + +@pytest.mark.unit +def test_check_zero_copy_kv_rejects_an_unplanned_buffer_without_blaming_the_planner(): + """A buffer with no ``mem_id`` was not planned, which is not the host arena. + + A planner that excludes mutable buffers, or one built with graph-input + allocation off, leaves the cache with no ``mem_id`` at all. Nothing then + says where it lives, so it is refused -- but reporting it among the host + tensors would tell the caller their planner made a placement it never made, + and the two are separated here by the message. The control is the same + program with the cache in the recorded CUDA arena, which is accepted, so the + refusal is the missing ``mem_id`` and nothing else about this graph. + """ + graph_module, k_buffer, _ = _direct_delegate_graph(compile_specs=_zero_copy_specs()) + program = _finalized_program(_planned(graph_module, host_tensor_arena=HOST_ARENA)) + Z.check_zero_copy_kv(program) + + del k_buffer.meta["spec"].mem_id + with pytest.raises(RuntimeError, match="carries no mem_id") as raised: + Z.check_zero_copy_kv(program) + assert "host tensors" not in str(raised.value) @pytest.mark.unit @@ -1741,6 +1858,7 @@ def test_export_manager_checks_the_program_its_to_executorch_returns(): untouched, which is the second half here: the check must not be a toll on the working path. """ + zero_copy_config = torch_tensorrt.executorch.zero_copy_backend_config() graph_module, k_buffer, _, _ = _staged_delegate_graph( compile_specs=_zero_copy_specs() ) @@ -1750,12 +1868,47 @@ def test_export_manager_checks_the_program_its_to_executorch_returns(): Z._check_zero_copy_kv_when_finalized(edge) with pytest.raises(RuntimeError, match="do not reach the TensorRT delegate"): - edge.to_executorch() + edge.to_executorch(config=zero_copy_config) unstaged = _finalized_program(_unstaged_graph()) good = SimpleNamespace(to_executorch=lambda config=None: unstaged) Z._check_zero_copy_kv_when_finalized(good) - assert good.to_executorch(config=object()) is unstaged + assert good.to_executorch(config=zero_copy_config) is unstaged + + +@pytest.mark.unit +@pytest.mark.parametrize("how", ["omitted", "positional", "keyword"]) +def test_export_manager_refuses_a_config_without_the_pass_before_finalizing(how): + """The refusal has to come before the manager is spent, not after. + + ``to_executorch`` rewrites the manager's own edge programs in place, so a + manager that has finalized once cannot finalize again -- the second call + dies inside ``insert_write_back_for_buffers_pass``. A refusal raised on the + finalized program could therefore only be acted on by exporting again and + rebuilding every engine, while the config alone already settles it. The + config reaches ``to_executorch`` either way round, so both are read. + """ + from executorch.exir import ExecutorchBackendConfig + + program = _finalized_program(_unstaged_graph()) + finalized = [] + edge = SimpleNamespace( + to_executorch=lambda config=None: (finalized.append(config), program)[1] + ) + Z._check_zero_copy_kv_when_finalized(edge) + + with pytest.raises(RuntimeError, match="does not carry the pass"): + if how == "omitted": + edge.to_executorch() + elif how == "positional": + edge.to_executorch(ExecutorchBackendConfig()) + else: + edge.to_executorch(config=ExecutorchBackendConfig()) + + # Nothing was finalized, so the caller can follow the remedy on this manager. + assert finalized == [] + zero_copy_config = torch_tensorrt.executorch.zero_copy_backend_config() + assert edge.to_executorch(zero_copy_config) is program @pytest.mark.unit @@ -1768,8 +1921,10 @@ def test_zero_copy_backend_config_defaults_to_executorch_defaults(): defaults = ExecutorchBackendConfig() # The un-staging pass specifically, not merely "some object that is not the - # default" -- which is all any wrapper would have to be. - assert type(config.to_out_var_pass).__name__ == "_UnstageThenToOutVar" + # default" -- which is all any wrapper would have to be. Compared by type + # rather than by name, which is also what the finalization hook does when it + # asks a config whether it is one of these. + assert isinstance(config.to_out_var_pass, Z._UnstageThenToOutVar) assert type(config.memory_planning_pass) is type(defaults.memory_planning_pass) assert type(config.sym_shape_eval_pass) is type(defaults.sym_shape_eval_pass) assert config.emit_stacktrace == defaults.emit_stacktrace @@ -1863,7 +2018,9 @@ def test_zero_copy_backend_config_carries_skip_h2d_left_falsy(skip): @pytest.mark.unit def test_unstage_pass_runs_the_inner_pass_after_unstaging(): """A caller's own to_out_var_pass has to survive being composed with.""" - graph_module, k_buffer, _, delegate = _staged_delegate_graph() + graph_module, k_buffer, _, delegate = _staged_delegate_graph( + compile_specs=_zero_copy_specs() + ) k_buffer.meta["_torch_tensorrt_aliased_buffer"] = True seen = []