fix(executorch): support KV-cache aliased I/O in the TensorRT delegate

[Conarnar]

Description

Adds end-to-end caller-owned KV-cache support to the ExecuTorch TensorRT delegate. The KV buffers are owned by the caller above the delegate and threaded through as mutable-buffer delegate args (both the input and the engine's aliased output), so a TensorRT engine updates them in place and the cache persists across decode steps — matching the contract the non-ExecuTorch TensorRT runtime already exposes.

Runtime + serialization (delegate)

• Serialize each engine's aliased (KV-cache / user) I/O into the delegate blob (serialization.py, backend.py, TensorRTBlobHeader.{h,cpp}).
• At runtime, bind each aliased TRT output binding to its aliased input's caller-provided pointer (in-place) and reflect the result into the delegate output EValue — a no-op when the memory planner already aliased the two (zero-copy) (TensorRTBackend.{h,cpp}).
• Validate the persisted alias map at init: kind must be one we understand (kv_cache_update / user); kv_cache_update entries are cross-checked against the engine's own getAliasedInputTensor (TensorRT is the source of truth), and user entries are shape-checked before two tensors are bound to the same storage. A caller-owned aliased input must be device-resident — otherwise its in-place update would be staged through host scratch and silently lost, so it's rejected loudly.
• A non-zero-copy aliased reflect is ordered to complete before execute() returns, since ExecuTorch's buffer-mutation copy_ reads the delegate output EValue afterward; the zero-copy caller-owned KV fast path records no reflect and is untouched.

Export / lowering (torch_tensorrt)

• Surface each engine's aliased outputs as graph-level BUFFER_MUTATIONs so ExecuTorch keeps the KV buffers as caller-owned mutable buffers instead of freezing them: at transform time for the legacy exporter (retrace=False), and via a post-export pass (_declare_aliased_kv_mutations_on_ep) for torch.export (retrace=True), which otherwise drops the aliased outputs at the fx boundary. The retrace=True pass runs for both executorch and exported_program; aot_inductor is left undeclared (and warns), since whether an aliased in-place mutation survives functionalization under inductor is unverified.
• Keep delegate-mutated buffers above the delegate in TensorRTPartitioner (tag_constant_data would otherwise freeze them as constants).

Dependency

Depends on #4446 (the legacy retrace=False submodule-inlining fix the composable ExecuTorch export path relies on), which has landed — this PR is now based on main.

Follow-up tests

Now that #4440 (per-method TensorRTPartitioner → separate delegate instances) and #4454 (shared caller CUDA stream, for ordering the dependent GPU work between two delegates) have both landed, a cross-delegate prefill/decode acceptance test — decode consuming the KV cache that prefill wrote through a separate per-method delegate — is unblocked and will be added as a follow-up. That configuration is what exercises cross-delegate cache sharing, which single-delegate tests cannot cover.

Testing

• Unit:aliased_io serialization round-trip; blob-header parse (present / empty / missing-key, and a model input literally named aliased_io not mistaken for the array key); exposure-flag dispatch across both retrace modes; the BUFFER_MUTATION declaration; partitioner keeps only mutation-target buffers above the delegate.
• Runnable example:examples/executorch_reference_runner/kv_cache_decode_check (exported by export_kv_cache_decode.py) drives multi-step decode against the delegate and asserts the KV cache persists in place across steps, starting from an explicitly-zeroed cache (ExecuTorch drops a mutated buffer's initial value) with a NaN guard so the check can't pass vacuously.

[Conarnar]

What changed vs the earlier (delegate-owned) version

The earlier revision made the delegate own the KV cache: because ExecuTorch
saw the KV buffers as non-mutated and froze them as constants, the delegate
cudaMalloc'd a persistent device buffer per aliased KV input, bound both the
input and its aliased output there, and accumulated the cache internally across
execute() calls. The aliased outputs were not threaded as delegate args.

This revision makes the cache caller-owned, above the delegate.

Export / lowering (new):

• Each engine's aliased outputs are exposed as graph-level BUFFER_MUTATIONs
  (transform-time for retrace=False; a post-export pass for retrace=True), and
  TensorRTPartitioner strips the delegation tag from mutation-target buffers, so
  ExecuTorch keeps the KV buffers as caller-owned mutable buffers and threads
  them as delegate args instead of freezing them.

Runtime (changed):

• Removed all delegate-internal allocation — the per-input cudaMalloc/persistent
  buffer and the input_is_self_owned / num_self_owned_inputs bookkeeping.
• Every input and aliased output is now a delegate arg (1:1). Each aliased output
  binding is bound to its aliased input's caller-provided pointer (in-place),
  and the result is reflected into the delegate output EValue — a no-op when the
  memory planner already aliased the two (zero-copy).

[narendasan]

@cehongwang Please review usage of the aliased i/o feature

[cehongwang]

line 223:

Only inputs need this. Outputs are also bound positionally by the runtime,
but they are getitem(engine_node, idx) nodes whose index order equals the
engine output-binding order. ExecuTorch lowering can reorder delegate outputs
(arrange_graph_outputs moves buffer-mutation outputs ahead of user
outputs), but a TensorRT delegate partition is a functional inference engine
with no mutation outputs, so that pass is a no-op here and the output order is
preserved. If a TRT partition ever produced mutation outputs, outputs would
need the same node-identity reordering as inputs.

We need to account for the output order, or otherwise there is a mismatch

[cehongwang]

The engine appends the buffer mutation at the end of output, while the delegate prepends it

[cehongwang]

Runtime walkthrough

Assume:

TRT inputs: [tokens, k_cache_in]
TRT outputs: [logits, k_cache_out]

ExecuTorch passes:

args[0] = tokens
args[1] = k_cache_in
args[2] = k_cache_mutation
args[3] = logits

After consuming the inputs, arg_idx == 2.

First output iteration

o = 0
name = output_binding_names[0] = logits
arg = args[2] = k_cache_mutation

The backend binds TensorRT's logits output to the cache-mutation output storage.

Second output iteration

o = 1
name = output_binding_names[1] = k_cache_out
out_arg = args[3] = logits

TensorRT correctly binds k_cache_out to k_cache_in for the in-place update, but the backend treats the logitsEValue as its mutation output slot. Its reflect copy therefore writes the cache result into the logits output.

The result is effectively:

TensorRT output	Lands in
logits	cache mutation slot
cache update	logits slot

If shapes or capacities differ, execution may fail during resize/binding/enqueue. If they are compatible, it can run successfully while returning incorrect logits and corrupting the observable cache state.

The required fix is to reorder serialized output_binding_names into actual delegate-output order, analogous to _reorder_input_names_for_executorch.

[Conarnar]

I was not able to reproduce this. Seems like arrange_graph_outputs reorders the submodule outputs, output_specs, and the parent getitems, but it doesn't touch the call node's meta["val"]. node.args is permuted by fusion (so _reorder_input_names_for_executorch is still needed), but outputs go through meta["val"], which arrange_graph_outputs leaves alone.

If meta["val"] is being rearranged somewhere that would be a problem, but the fix would be different from how _reorder_input_names_for_executorch handles it.

[cehongwang]

Ok I spend a ton of time on this and it is quite surprising that the order was actually correct. But here is some findings:

_keep_mutated_buffers_above_delegate(exported_program)
https://github.com/Conarnar/TensorRT/blob/b9dbb306a4597eb4fb0c4ef0fdb546235d3924b1/py/torch_tensorrt/executorch/partitioner.py#L165
This function lifes mutated buffer to the executorch program level, therefore arrange_graph_outputs did not have any buffer_mutation, and therefore the order is perserved.

[cehongwang]

One testing gap worth closing before merge: there's currently no test anywhere that exercises a full to_edge_transform_and_lower → to_executorch → execute path. Grepping tests/ for to_edge_transform_and_lower, to_executorch, and load_for_executorch returns nothing, and executorch-static-linux.yml only builds the backend plus the bazel C++ tests — it never lowers a KV model or runs a delegate.

That means the export side and the blob side are each tested in isolation and each is individually correct, while the bug lives in their composition. Could we get:

1. A CPU test asserting blob output-binding order matches delegate output-arg order for a model with both a mutation and a user output (details in the inline comment).
2. A two-step decode test on GPU or via the reference runner: same cache storage across two execute() calls, asserting step 1 observes step 0's update.

(2) is what RFC 0003 §7.4 asked for, and it would also cover the device-residency and reflect-path questions raised elsewhere in this review. Happy to gate it on GPU availability, but it should exist as a runnable target.

[cehongwang]

In tests/py/dynamo/executorch/test_backend.py, test_preprocess_preserves_output_binding_order
This test's premise no longer holds after this PR, and as written it locks in the bug.

The comment says output order is "stable by construction (getitem index order == engine output-binding order)." That was true when a TRT partition was purely functional. This PR introduces BUFFER_MUTATION outputs, and ExecuTorch's arrange_graph_outputs reorders delegate outputs to [mutations..., user_outputs...] while preprocess still serializes them in engine order ([user_outputs..., aliased_outputs...]).

The fixture here has no mutation outputs, so it passes — but it asserts "preprocess must pass output names through unchanged," which is the behavior that needs to change.

Could you update this test to cover the mutation case? Something like a fixture with one aliased output and one user output, asserting the serialized output_binding_names match the delegate's output-arg order rather than engine order.

[cehongwang]

Aliased inputs can still take the host-staging path, which breaks caller-owned semantics.

If an aliased input's data pointer isn't CUDA-accessible, this falls through to the staging branch and binds engine->cached_input_ptrs[i] — a delegate-owned scratch buffer. The aliased output then binds to input_bind_ptrs[alias_in], i.e. that same scratch buffer, so the in-place KV update lands in delegate scratch rather than the caller's storage. On the next execute() the staging copy re-reads the caller's unchanged host buffer, so the update is silently lost. Recovery depends entirely on the reflect copy, which has its own problem (see the reflect comment).

[Conarnar]

Can you clarify what RFC 0003 §6.4 (and §7.4) is? I was not able to find any references for that.

[cehongwang]

Oh, here it is:
Require the caller's tensor itself to be device-resident. Aliasing only
works when the pointer bound to the input is the caller's real storage. Two
existing branches break that: the H2D fallback copies a host tensor into a
device staging buffer and binds that, and the zero-byte branch binds a
cached scratch allocation. Either would make the engine write the update
somewhere the caller never reads. For an alias-source input, gate on the
caller tensor being device/unified memory (is_cuda_accessible_ptr) with
non-zero size, and otherwise return Error::InvalidArgument ("aliased input
'%s' must be on GPU or unified memory") — i.e. reject before both the
staging and zero-byte paths, not only the H2D one.

[cehongwang]

Do we need to account for aliase I/O? Is there a race possible?

[Conarnar]

Yes, but only on the non-zero-copy reflect path: with a caller stream active and no end sync, a pending reflect into the delegate output could still be in flight when ExecuTorch's buffer-mutation copy_ reads it. Will handle it.

[cehongwang]

This doesn't cross-check the persisted alias map against the engine, and it accepts unknown kind values.

The Python runtime's _TRTEngine._reconcile_aliased_io treats getAliasedInputTensor as the source of truth for kv_cache_update aliases and preserves user ones as metadata-trusted. Here, a kind that is neither "kv_cache_update" nor "user" — a typo in the wire format, or a future kind written by a newer exporter — skips the shape check and gets registered as if it were a KV alias, which then binds two tensors to the same storage.

Could you mirror the Python behavior:

1. Reject unknown kinds with Error::InvalidProgram.
2. For kv_cache_update, compare the persisted ab.input against engine->getAliasedInputTensor(ab.output.c_str()) and error on disagreement.
3. Keep user as metadata-trusted after the shape check (TRT can't see those aliases).

Related: the parser leaves ab.kind empty when the "kind" key is absent, while the Python side defaults to "kv_cache_update". Once unknown kinds are rejected, that mismatch turns an old blob into a hard failure — worth defaulting to "kv_cache_update" in the parser to match.

[cehongwang]

This'd better be warnings or errors

[Conarnar]

Guessing you mean the two continues below.

[cehongwang]

here

[cehongwang]

The previous comment no longer holds, and debugging can be painful if anything comes wrong in the future. Change the comment to something like

def_reorder_input_names_for_executorch(
edge_program: ExportedProgram, engine_node: Any, input_names: List[str]
) ->List[str]:
"""Reorder TRT binding names into executorch_call_delegate argument order. The runtime binds positionally (``execute()`` arg ``i`` -> input_binding_names ``[i]``), but ExecuTorch fusion may permute the delegate placeholders relative to the TRT-submodule order that produced ``input_binding_names``. The names can't be matched (TRT names are semantic, lowered placeholders are generic ``arg_N``), so recover the permutation by node identity: the engine node's first arg lists its input nodes in binding order, so sort the names by each node's slot among the graph placeholders (its runtime delegate-arg position). Only inputs need this. Outputs are also bound positionally, but they are ``getitem(engine_node, idx)`` nodes whose index order equals the engine output-binding order, and that order survives lowering -- though not because the partition is mutation-free. With aliased-I/O (KV-cache) support a TensorRT partition *does* produce mutation outputs, and ``arrange_graph_outputs`` does move buffer-mutation outputs ahead of user outputs. It stays a no-op here because ``_keep_mutated_buffers_above_delegate`` (``partitioner.py``) strips the ``delegation_tag`` from mutated buffer placeholders, so they stay out of the delegate's state dict and constants; ExecuTorch's ``_get_new_signature`` then records the mutation as a plain ``USER_OUTPUT`` rather than a ``BUFFER_MUTATION`` (it uses the latter only when the delegate itself consumes the buffer). The lowered submodule therefore has no mutation specs, so ``arrange_graph_outputs`` computes the identity permutation and the getitem indices still line up with the engine's output bindings. That guarantee is conditional, not structural. If a mutated buffer is ever tagged into a delegate, its spec becomes ``BUFFER_MUTATION``, the delegate's outputs are permuted, and they would need the same node-identity reordering as the inputs below. """

[cehongwang]

The reconstruction drops example_inputs and verifiers, so the pass isn't a drop-in replacement for its input.

ExportedProgram.init also takes example_inputs and verifiers; neither is carried over, so the returned program has them unset. No impact on the current to_edge path, which is presumably why it hasn't surfaced, but it bites in two places:

If the pass moves to the exported_program branch, the saved artifact silently loses its example inputs.
AOTI rejects such a program outright: RuntimeError: exported_program.example_inputs is required to be set in order for AOTInductor compilation. I hit this while testing whether the pass fixes AOTI and had to restore _example_inputs by hand before I could even reach the real failure.
Passing both through costs two keyword arguments.

[Conarnar]

Fixed in the next push. Both example_inputs and verifiers are threaded
through now.

I should mention that the retrace=False path has a related gap:
create_trt_exp_program constructs its ExportedProgram without
example_inputs either, so AOTI would refuse a program saved that way for the
same reason. That one predates this PR and isn't a dropped value. There's no
source program to carry from, though it does receive arg_inputs it could use.
Let me know if you want that fixed as well.

[shoumikhin]

On the output ordering question, it does not reproduce at this head and I think the author's reply is right. Running the real ExecuTorch to_edge_transform_and_lower over the head partitioner and preprocess gives delegate args [cache, tokens, logits, cache_out] against blob names [cache_in, tokens, logits, cache_out], with the parent mapping delegate output 1 to the buffer mutation. Deleting _keep_mutated_buffers_above_delegate reorders the subprogram to [cache_out, logits] while the blob stays [logits, cache_out], which is exactly the swap that was described, so the mechanism is real but the helper prevents it.

The comments below are the blockers and majors I found. On the other side, the device-residency rejection is correctly placed before both the host-staging and zero-byte paths, the inflight_pending drain gap I raised earlier is fixed, and the TensorRT version gate around getAliasedInputTensor is correct on both sides of 10.15 (checked by preprocessing it against several SDK header sets).

[shoumikhin]

Whether the mutation gets declared is decided twice, once by use_legacy_exporter inside _exporter.export and once by the outer retrace branch here, so two of the four combinations are wrong: retrace=True, use_legacy_exporter=True declares it twice and raises SpecViolationError: User output getitem_3 is not in the correct order, and retrace=False, use_legacy_exporter=False never declares it and prints no warning. Decide exposure once from the exporter that actually ran, and make _declare_aliased_kv_mutations_on_ep skip a buffer that already has a BUFFER_MUTATION spec.

[Conarnar]

Turns out declaring on the exported_program branch created a third bad combination on top of your two that you found.

Fixed in the next push, using your suggestion and that covers all three combinations rather than special-casing any of them.

[shoumikhin]

Two cases drop an aliased output here while the runtime still requires one delegate arg per engine output (TensorRTBackend.cpp:426): a second engine aliasing the same buffer is deduped, and a caller-supplied cache that is not a registered buffer gives buf_fqn is None and is skipped, both ending in Error::InvalidArgument at execute. The legacy path skips the same two at lines 765 to 793, so either emit one output arg per engine output, or record in the blob which outputs are not threaded so the runtime can adjust its expected arg count.

[Conarnar]

I was able to confirm the second case, but I wasn't able to test the first.

I couldn't get a two-engine model that aliases one buffer through the retrace=True path. The export fails earlier, inside torch.export, on the hybrid TRT+eager graph. So I can't say whether that branch is reachable in practice. If you have a model in mind that hits it, that would help.

Regarding the fix, your second option looks workable for the non-buffer case specifically, since there's no ET buffer to write back to. I'd be wary of applying it generally though. An un-threaded aliased output also gets no write-back, and for a real buffer that turns an execute-time error into a silently stale cache. So scoping the blob-side elision to the non-buffer case and keeping one-arg-per-output elsewhere seems safer.

[shoumikhin]

dst == bind_ptr cannot happen: the delegate input and output are live at the same time, so the memory planner never places them at the same address, and a mutable buffer is not in the planned arena at all. That makes the zero-copy path described in the comment and in the PR description dead, and every decode step pays a full cache-size device copy here plus ExecuTorch's write-back copy_ on a buffer the engine already updated in place, so it is worth stating whether that cost is inherent to expressing this as BUFFER_MUTATION.

[Conarnar]

Confirmed dead. I instrumented the branch on a 30-layer export: 120 aliased outputs across prefill and decode, dst == bind_ptr false every time. The comment should go.

One correction: in our .pte the mutable buffers are planned values (mem_id 1, no mutable data segment). The reason dst != bind_ptr is the one you give, plus a second one: the delegate never receives the buffer at all. PropagateDevicePass wraps every delegate input in _h2d_copy, so the engine binds a per-call staging copy, and the reflect plus ExecuTorch's write-back exist to carry the update back to the buffer.

On whether the cost is inherent to BUFFER_MUTATION: it isn't. Pointing the mutation at the buffer placeholder does make ExecuTorch emit no copy_, but that alone breaks persistence, because the staging copy is re-made from the buffer at the top of every call, so the engine's in-place write is overwritten before it is read. Removing the staging too fixes it: the buffer lands in the device arena, the delegate binds it directly, and every full-cache copy goes away. On a single-layer KV model that is 16 -> 8 values and 10 -> 4 instructions, with the persistence check matching eager exactly (0.235384) where the elision-only build reported 0.

That is prototype-grade: a toy model, and the un-staging rule still needs a guard that every consumer of the buffer is device-capable.

The natural home looks like PropagateDevicePass itself, so the copy is never inserted for a delegate-mutated buffer in the first place, though that is more your call than mine. We could also carry it as a Torch-TensorRT pass that undoes the staging for our delegate only, if that sequences better, but it would have to hook in between PropagateDevicePass and memory planning, which is more fragile than doing it properly upstream.

Worth pursuing separately. For this PR I just corrected the comments. I kept the dst != bind_ptr check itself as a self-copy guard rather than deleting it, since it costs a pointer compare and is what stops a self-copy if planning ever changes.

[shoumikhin]

Fixed, verified on the current head (c87d3bbd8). The dst == bind_ptr comparison is gone; bind_ptr now appears only on the input-binding side, so the dead zero-copy branch and the comment claiming it no longer exist.

Please resolve this thread; I do not have write access to do it myself.

[shoumikhin]

Correcting my previous reply, which was wrong. I grepped for bind_ptr with a truncated output and only saw the input-binding uses, so I said the comparison was gone. It is not:

void* dst = et_alias_out.nbytes() > 0 ? et_alias_out.mutable_data_ptr() : nullptr;
// dst != bind_ptr guards against issuing a self-copy. The memory planner does// not currently place the delegate's output slot on the aliased input -- the// two are live at the same time -- so this holds for every aliased output.if (dst != nullptr && dst != bind_ptr) {

What has changed is the comment, which now states the same thing I raised: the planner does not place them together, so the branch is always taken. So the misleading claim of a zero-copy fast path is gone, which was half my point.

The other half stands and is not fixed: every decode step still pays a full cache-size device copy here, plus ExecuTorch's write-back copy_. Please leave this thread open for that.

[Conarnar]

None of this is exercised by CI. kv_cache_decode_check is not built or run anywhere: .github/scripts/verify-executorch-reference-runner.sh builds --target example_executorch_runner only, and the CMake target has no add_test. Every Python test mocks the seam too, so nothing fails if the binding order is wrong or if _keep_mutated_buffers_above_delegate is deleted, and retrace=True has no end-to-end coverage at all.

Separately, no build or test job has run at this head. _decide.yml maps a non-approval pull_request_review event to lane=skip, and the one pull_request run that started was cancelled by the later review submissions, so the three red gate checks are that cancellation rather than a real failure. Worth getting a real run before merge. There is also a small conflict with main in examples/executorch_reference_runner/BUILD, where both sides add one line to the same filegroup.

Got it. The setup for the reference runner will be mirrored for kv_cache_decode_check.
Also, the conflict will be fixed after rebasing.

[shoumikhin]

Cleaning up after myself: I had left a lot of long comments here, several already fixed and a couple simply wrong, so I deleted them and replaced them with this note. I left the inline threads where you replied, since your answers there are worth keeping.

Nothing on this PR blocks merge from my side.

Everything I raised is either fixed on the current head, which I re-checked, or parked with a reason I agree with. The two open ones are the arity mismatch, where half the case is still unproven and neither of us could build a two-engine repro, and the double-declare, which is really a #4459 concern.

Two corrections so nobody chases them:

• The stale 1:1 bindings comment I complained about is not hiding a runtime failure. preprocess validates first, at backend.py:363 and :371, and unknown alias names are rejected during backend initialization, so for any accepted program that comment is accurate. If anything the misleading comments are the Python ones at _exporter.py:613-616 and :638-640, which point at the execute-time arity check. Minor either way.
• I said a rebase would clear the red checks. It will not. The same jobs are red on main: constant_fold_exclusions landed as a new package directory but was never added to the dynamo_packages list in setup.py, so the built wheel omits it and the runtime import fails. Not caused by this PR, and not fixed by rebasing.

The remaining review comments are on #4459, which is where the copy-back work lives.