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. getAliasedInputTensor is a TensorRT 10.15+ API, so that cross-check is version-gated (NV_TENSORRT_MAJOR/MINOR) and skipped on older TensorRT — e.g. Jetson's 10.13, which would otherwise fail to compile — trusting the persisted map instead (no functional change on ≥ 10.15; mirrors fix: stop calling TensorRT 10.15 aliasing APIs unconditionally #4468 in core/runtime). 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

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.

[cehongwang]

The alias I/O should be supported in both executorch and standard Torch-TensorRT runtime and AOT inductor

[cehongwang]

You might want to add _declare_aliased_kv_mutations_on_ep to exp program. AoT inductor is more complicated and we will fix that later

[Conarnar]

Will rebase and enable exported_program on #4459 as well.

[cehongwang]

Brief warning here for better debugging

[Conarnar]

This one in particular is an ordinary case. It fires for every engine output
that isn't aliased, which is most of them, so a warning would be hundreds of
lines on a larger model. The legacy path leaves the equivalent skip silent for
the same reason. If you do want logging here, logger.debug rather than a
warning.

[cehongwang]

here

[cehongwang]

here

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

memory_planned_buffer_device does not exist in the ExecuTorch commit this repo pins (MODULE.bazel:50, 6118688a09, dated 2026-05-18); it landed on release/1.3 after that pin, so //examples/executorch_reference_runner:kv_cache_decode_check does not compile. Bump the pin to a commit that has the API.

[Conarnar]

Looks like it's already fixed on main: e2f18eb23c (executorch==1.3.1) landed
there in #4398, so this stack picks it up in the rebase. Verified it has the API
but not whether it builds. Will do that.

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

This discards the resize_tensor result and skips the resize entirely on a bad rank, while the non-aliased branch at lines 693 to 706 treats both as fatal. On a dynamic aliased output that outgrows its planned size, nbytes() on the next line is then the stale planned size and drives both the reflect and ExecuTorch's write-back copy_, so handle the return value and the rank check the way the sibling branch does.

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

aliased_io changes what a blob means, but the magic stays TR01 and the header is unchanged, so a pre-PR runtime still loads a new blob, ignores the alias map, and binds the KV output to its own allocation instead of the aliased input. I compiled the merge-base parser against head-format metadata and it accepted it, so bump the magic or add a required-feature field the old parser rejects, and let version skew fail closed instead of returning wrong results.

[shoumikhin]

The order these names are serialized in is correct only because _keep_mutated_buffers_above_delegate (partitioner.py:165) keeps the mutated buffer above the delegate, and nothing checks that: inputs get an arity check, outputs get none. An assertion that the lowered delegate's output specs are in engine binding order would turn a future regression from silently swapped outputs into a loud failure.

[shoumikhin]

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.

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