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}).
• USER aliases are shape-validated at init; kv_cache_update aliases are
  shape-enforced by TensorRT's IKVCacheUpdateLayer.

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.
• Keep delegate-mutated buffers above the delegate in TensorRTPartitioner
  (tag_constant_data would otherwise freeze them as constants).

Dependency

This PR is stacked on #4446 and must land after it — #4446 fixes the legacy
(retrace=False) submodule inlining that the composable ExecuTorch export path
depends on.

Follow-up tests (gated on other PRs)

A cross-delegate prefill/decode acceptance test — decode consuming the KV
cache that prefill wrote through a separate per-method delegate — will be added
once #4440 (per-method TensorRTPartitioner → separate delegate instances)
and #4454 (shared caller CUDA stream, for ordering the dependent GPU work
between the two) land. 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); exposure-flag dispatch across both retrace
  modes; the BUFFER_MUTATION declaration; partitioner keeps only
  mutation-target buffers above the delegate.
• End-to-end: exported and ran a KV model through both retrace=False and
  retrace=True — the cache persists in place across steps (step 0 output
  reproduces at step 1, then diverges as it accumulates). A non-KV model exports
  and runs unchanged (output byte-identical to eager).

[shoumikhin]

Thanks for adding aliased-I/O metadata to the ExecuTorch blob. This fixes a real missing capability, and the serialization changes look reasonable.

I found a blocking issue that applies even when only one method is exported. The patch removes every aliased input from the delegate arguments and replaces it with a private, zero-initialized buffer. If the caller supplies the cache tensor, its existing contents are ignored and the caller cannot observe the update.

For example, the expected behavior is:

caller cache    -> TensorRT input
TensorRT output -> same caller cache

The new behavior is:

caller cache        -> ignored
private empty cache -> TensorRT input and output

Checking only kind == "kv_cache_update" is not enough, because a compiler-generated KV update can still use caller-owned storage. Alias kind describes who enforces the alias, not who owns the buffer, so I think the metadata needs to describe storage ownership separately from alias kind.

Runtime-owned storage also needs a defined lifetime. The current buffer is initialized once, has no reset operation, and is shared by every call using that loaded method. Two conversations, or concurrent requests, would therefore use the same cache.

Could caller-owned aliases stay explicit delegate inputs, with the output bound to the same pointer? If runtime-owned storage is genuinely needed, please add explicit ownership and a stable state identity, plus reset or session-selection behavior.

On testing, it would help to pin the runtime behavior directly rather than through serialization round-trips: caller-visible mutation, repeated calls on one loaded method, a fresh load starting from a known state, and sequence isolation.

For the multi-method case: multi-method Torch-TensorRT ExecuTorch export is still in flight (#4440), so a natural next step is to build a two-method prefill/decode test on top of it and assert that decode observes cache state written by prefill. I want to flag the expected outcome up front: with the current design I believe that test fails, because each method loads as its own delegate with its own private cache, so there is no shared object for the two methods to write through. That is really why I think the cache has to be owned above the delegate and bound into both methods; the shared-cache test is the acceptance criterion for that ownership change rather than something stacking alone will make pass.

For mixed TensorRT and CUDA execution there are two independent requirements worth testing separately: (1) both methods bind the same KV-cache storage, and (2) dependent GPU work from the two delegates is ordered. Ordering needs a shared caller stream (there is separate in-flight work for that, #4421), so a mixed test should run on top of it, but note the shared stream only provides (2). Property (1), shared storage across the TensorRT and CUDA delegates, cannot come from a delegate-private buffer, so it also depends on moving cache ownership above the delegate.

[shoumikhin]

Following up with something concrete I should have led with, plus a correction to my own comment.

The existing runtime already defines the expected behavior

The non-ExecuTorch C++ runtime binds an aliased output to the caller's pointer
(core/runtime/execute_engine.cpp:426):

ctx->setTensorAddress(name.c_str(), in_it->second.data_ptr());

and examples/dynamo/aliased_io_user_inputs.py documents that contract for users:
the caller owns the cache, passes it in on every call, and after the call
cache.data_ptr() is unchanged and the mutation is visible.

So this is not only a question of which design is nicer. It is that the ExecuTorch
delegate would behave differently from the runtime that already ships, for the same
compiled engine. That is the part I would most like to resolve before this lands.

Correction: the alias kind is never consulted

I said checking kind == "kv_cache_update" was "not enough". Looking again, the kind
is not checked at all. output_alias_kind is populated and then never read, and every
aliased input is self-owned unconditionally:

for (int in_idx : handle->output_aliased_input_idx) {
  if (in_idx >= 0) {
    handle->input_is_self_owned[in_idx] = true;   // no kind check
  }
}

That matters for AliasKind::USER, which _ConversionContext.py defines as: "the
runtime must validate shape compatibility and bind both input and output to the same
device pointer." A user-declared alias is caller-owned by construction, so self-owning
it is the opposite of the documented behavior. Either way the conclusion is the same:
ownership needs to be described explicitly rather than inferred from the presence of an
alias.

On the constraint you hit

Your comment in the header explains the real obstacle, and I do not think I gave it
enough credit: with the graph-level copy_ gone, ExecuTorch sees a non-mutated buffer
and freezes it as a constant, so it never arrives as a delegate arg. That is a genuine
blocker, not an oversight.

My concern is where it gets solved. Working around it inside the delegate means the
delegate silently substitutes its own storage for the caller's, which is invisible from
the outside and, as above, disagrees with the existing runtime. Keeping the buffer
mutable through export so ExecuTorch threads it as an argument fixes the arg-count
mismatch and preserves caller-visible mutation at the same time. I realize that is more
work than this patch, and I am happy to help look at the export side if useful.

Two smaller things

• cudaMalloc in init is not freed on the failure paths that follow it (for example
  if initialize_input_profiles returns an error), so a failed init leaks the buffers.
• Self-owned buffers reject dynamic dims with Error::InvalidProgram at init. A
  dynamic-shape model with a KV alias would then fail to load, where today it fails at
  execute(). Worth stating as a known limit if it stays.

The serialization work (carrying aliased_io through the blob, list form for the C++
parser, missing-key backward compatibility) looks right to me and is reusable whichever
ownership model wins. If it helps unblock things, that part could land on its own ahead
of the runtime binding change.

Composition note

Re-checking my earlier point about multi-method, in #4440 each method gets its own
TensorRTPartitioner, so each becomes a separate delegate instance with separate state.
A prefill/decode test where decode reads what prefill wrote would therefore fail under
a delegate-private cache, which is why I think that test is the right acceptance
criterion for the ownership question rather than something that stacking alone
resolves. For a mixed TensorRT and CUDA program the two requirements stay independent:
shared cache storage across delegates (ownership, this PR) and ordering of dependent GPU
work (the shared caller stream, #4454, not #4421 as I mistyped earlier).

[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 logits EValue 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]

These tests verify that mutations get prepended in the ExportedProgram, which is one half of the contract. The other half — that the serialized output_binding_names line up with the delegate arg order the C++ backend indexes positionally — isn't covered here or anywhere else, because to_edge, arrange_graph_outputs, and preprocess are all outside the mocked boundary.

Could you add a test at that seam? Roughly:

1. Build a KV model with both a user output and an aliased cache output.
2. Run it through to_edge_transform_and_lower with the TRT partitioner.
3. Read the delegate's output-arg order from the lowered module's output specs.
4. Assert deserialize_engine(...).io_bindings output names match that order.

This needs no GPU and would fail today.

[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).

RFC 0003 §6.4 and the Python runtime both treat alias sources as required-device-resident. Could you add an explicit check before staging: if input i is the target of any entry in output_aliased_input_idx, require unified_memory || is_cuda_accessible_ptr(...) and nbytes() > 0, and return Error::InvalidArgument with a message naming the binding otherwise? Failing loudly is much better than a silently stale cache.

[Conarnar]

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

[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.