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

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

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]

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

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

[shoumikhin]

Missing stream drain if a reflect copy fails, so the next call can touch a live context.

enqueueV3() has already succeeded by the time we get here, so the engine is running on the stream. If one of these cudaMemcpyAsync calls fails we return straight away, and engine->inflight_pending is only set further down (line 811), so it stays false.

Both later guards are gated on that flag:

• the next execute() (line 470) waits only if (engine->inflight_pending), and it does that before calling setInputShape / setTensorAddress
• the destructor (line 83) waits only if (inflight_pending) before destroying the context and freeing the staging buffers

So the flag says "nothing in flight" while TensorRT is still working. The next call is then free to reconfigure the context, which the comment at line 466 correctly says TensorRT forbids.

You already have exactly the right pattern a few lines below, at 802-808:

// Could not arm the completion marker; drain now so a later execute() or the
// destructor never reconfigures or frees exec_ctx while this enqueue runs.
(void)cudaStreamSynchronize(stream);
engine->inflight_pending = false;
return Error::InvalidProgram;

Same two lines would fix this branch.

Worth considering the more robust shape too: set the pending marker as soon as the first async work is submitted, rather than at the end of the happy path. Then any early return added later between enqueue and the end of the function inherits the protection instead of quietly reintroducing this.

Note the D2H branch at 786-793 looks like it has the same gap, but that one pre-dates this PR, so it is probably a separate fix.

[shoumikhin]

Aliased models saved in the other two formats get no equivalent handling, and no warning.

This normalization runs only in the output_format == "executorch" branch. The exported_program branch (1116) and aot_inductor branch (1126) receive the same retrace=True program and save it as-is.

if output_format == "exported_program":     # no alias handling
elif output_format == "aot_inductor":       # no alias handling
elif output_format == "executorch":
    exp_program = _declare_aliased_kv_mutations_on_ep(exp_program)   # only here

I understand why it lives here: ExecuTorch has tag_constant_data, which freezes an undeclared buffer, and that pass is the reason this function exists. The other formats have no equivalent, so the harm is not the same.

But the outcome still differs per format in a way a user cannot see:

• exported_program: the engine still mutates the buffer in place at runtime, so the cache probably does update, but the saved signature does not declare the mutation. Anything downstream that trusts the signature gets a wrong answer.
• aot_inductor: this goes through inductor, and whether an in-place mutation on a custom op survives functionalization is not obvious. I could not find test coverage for that combination.

If aliased I/O is only intended to be supported for executorch, could the other two reject or warn when the module has aliased engines? A user who saves as aot_inductor today gets no signal either way, which is the part that seems worth closing regardless of which behavior you pick.