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