fix(executorch): lift nested and partitioner-excluded mutable buffers

[shoumikhin]

Depends on #4459. These fixes are on top of that branch, so the file diff here
includes its commits and CI will stay red until it lands. Only the last four commits
are mine (af29266c7ed4..HEAD), at +280/-20 across five files, three of which are
tests.

What this fixes

Two problems that stop a model from exporting when part of its graph runs on a different
backend, plus the nested-buffer follow-ups that only become visible once the graph gets
that far.

1. An op excluded from TensorRT was predicted to alias

_kv_write_will_alias decides whether a buffer write becomes engine-level aliasing by
looking at the op, and never saw the ops the caller excluded via torch_executed_ops. An
excluded write was still predicted to alias, its copy_ was dropped, and the cross-check
in compile() then failed:

RuntimeError: lift_mutated_buffers classified these buffer writes as KV-cache
(engine-aliased) and dropped their copy_, but the compiled engine did not alias them
(absent from aliased_io): [...]. Their write-back would be silently dropped.

An op that never reaches a converter cannot emit an IKVCacheUpdateLayer, so the
prediction now honors the exclusion list, matched with
ConverterRegistry.qualified_name_or_str the same way the partitioners match it. The
write is still lifted and still copied back, like any other write the engine does not
alias.

That guard is worth calling out: it turned a cache that would silently stop updating into
a clear error naming every affected buffer.

2. A copy-back output could be paired with the wrong buffer

lift_mutated_buffers appends each copy-back buffer's new value as a trailing output, in
copyback_buffers order, and the exporter reclassified them by taking the last N.

That breaks once torch.export recognises a mutation itself, because it then records its
own BUFFER_MUTATION and moves it ahead of the user outputs. Filtering the buffer list
first and slicing the same number of trailing outputs shifts the two lists against each
other:

buffers [a, b, c], b already declared
pending [a, c] -> takes the last two values, which are b's and c's
result a is declared as the mutation of b's value

With different shapes the runtime fails copying that value into the buffer; with matching
shapes it silently writes the wrong state. On a transformer whose cache write stays inside
the engine, an earlier form of this bug declared the model's logits as a cache update and
export failed with expand: attempting to expand a dimension of length 16 -> 4096.

The fix detaches the full trailing run first and pairs it positionally with the full
buffer list, then keeps only the buffers not already declared. When every buffer was
already declared, nothing used to be detached and the appended values stayed as user
outputs, changing the saved model's output arity.

3. Renamed nested buffers left the mutation targets dangling

register_buffer rejects dots, so a nested buffer is renamed to lifted_buf_*. The
recorded mutation targets kept the original dotted names, which no longer resolve:

SpecViolationError: Buffer output getitem_1 does not point to a buffer that exists.
mutated target : 'layers.0.self_attn.kv_cache.k_cache'
buffer available : 'lifted_buf_layers_0_self_attn_kv_cache_k_cache'

They are remapped through the same mapping.

What changed since the first version

The first version of this PR also skipped an excluded buffer entirely, leaving the
get_attr and copy_ in place, on the theory that another backend would claim the write.
That was wrong outside ExecuTorch. lift_mutated_buffers runs for every compile()
caller, and post_lowering's remove_num_users_is_0_nodes erases the orphaned copy_
because nothing reads it, so the buffer froze at its compile-time contents and the cache
silently stopped updating. Worse, that input used to fail loudly at the cross-check above,
so the skip turned a clear error into wrong numbers.

The delegate-boundary error that motivated the skip has a different cause. ExecuTorch only
rejects a mutation consumed after a delegate when the delegate also absorbed the
buffer; returning a mutation value out of a delegate is a supported, tested path
(exir/backend/test/test_partitioner.py::test_not_delegate_mutable_buffers). Absorption is
already prevented, by tag_constant_data skipping mutated buffers and by
_keep_mutated_buffers_above_delegate covering the case it misses. So the copy-back route
handles this correctly and no skip is needed.

Also dropped gm.meta["_lifted_buffer_attr_names"], which had no reader.

Testing

tests/py/dynamo/lowering/test_buffer_lifting.py and
tests/py/dynamo/executorch/test_kv_cache_export.py, plus three new regression tests,
each of which fails on the behavior it replaces:

• test_torch_executed_write_is_copyback_not_kv: an excluded write is lifted, recorded as
  copy-back rather than predicted-aliased, and leaves no dead nodes for post_lowering to
  erase.
• test_declare_aliased_kv_mutations_pairs_copyback_by_position: with [a, b, c] and only
  b already declared, each buffer is paired with its own value.
• test_declare_aliased_kv_mutations_skips_already_declared_copyback: an already-declared
  buffer is not declared twice, and its trailing value still leaves the user outputs.

Note the nested-buffer lookup itself (gm.get_buffer) landed separately in #4472 and is
inherited here through this branch's base, so it is no longer part of this PR's changes.

[shoumikhin]

Pushed an amendment to this PR.

What changed

A copy-back value is now paired with its buffer by name rather than by position.

The old code assumed those values were always the last outputs:

copyback_getitems=list(out_args[-num_copyback:])

That assumption does not always hold. torch.export usually leaves them at the end, but
when it recognises the mutation itself it declares them BUFFER_MUTATION and moves them
to the front, because the verifier requires mutations to come before user outputs.
Taking the tail unconditionally then grabs a real user output and relabels it as a buffer
mutation.

How it showed up

On a transformer whose cache write stays inside the engine, it declared the model's
logits as a cache update. Five mutation specs for four buffers, one of them pointing at a
tensor of a completely different shape:

layers_1 v_cache <- (1, 16, 202048) the logits
layers_0 k_cache <- declared twice

Export then failed at runtime lowering:

RuntimeError: expand: attempting to expand a dimension of length 16 -> 4096

16 is the sequence length and 4096 is the cache length, which is the shape mismatch you
would expect from copying a user output into a cache buffer.

The fix

Reclassify only the buffers the program does not already declare as mutated:

already_mutated= {
spec.targetforspecinorig_specsifspec.kind==OutputKind.BUFFER_MUTATIONandisinstance(spec.target, str)
}
pending_copyback= [bufforbufincopyback_buffersifbufnotinalready_mutated]

This makes the pass idempotent with respect to what torch.export already recorded. The
companion if num_copyback and module_call_graph guard became
if copyback_specs and module_call_graph, so the output spec is only rebuilt when this
pass actually moved outputs. Rebuilding it after reclassifying nothing would drop a leaf
count that already matched.

After the change the same model exports and produces four mutation specs for four
buffers, each cache appearing exactly once, with the logits left as a user output.

[shoumikhin]

Rebased onto your updated #4459 (5e7b37d). Two add/add conflicts, both in the copy-back remap comment, resolved in favor of your wording since #4459 now carries that fix itself.

One collision worth your call. Your new commit raises RuntimeError when a copy-back buffer already carries a BUFFER_MUTATION spec, because positional reclassification would retarget the trailing outputs. My commit on top fixes that positional pairing: it detaches the full num_copyback run and pairs it against the full copyback_buffers list before filtering, so an already-declared buffer is skipped rather than shifting the run.

So we now have two contracts for the same situation. Your test_declare_aliased_kv_mutations_rejects_redeclared_copyback expects the raise; my test_declare_aliased_kv_mutations_skips_already_declared_copyback expects a graceful skip. They do not always both fire, since your guard reads _kv_exposed_mutation_targets while mine reads graph_signature.output_specs, but the intent differs.

My read is that the raise is still correct for the KV path your guard watches, and the graceful skip is correct for the torch.export-declared case, so both can stand if we scope the guard to the KV targets explicitly. Happy to do that, or to drop my overlapping test and let the raise be the single contract, whichever you prefer for #4459.

[shoumikhin]

Closing this in favour of #4459, which is the same work with one owner.

This branch was #4459 plus four commits on top, so there is nothing here that #4459 cannot carry directly. I rebased it onto today's main first, so the conflict resolutions are done and #4459 can take them rather than redo them. The details and the four commits to adopt are written up in #4459 (comment).

The branches stay up so the commits remain fetchable:

• shoumikhin:for-4459-nested-buffer-lifting, the full rebased series
• shoumikhin:rebased-4459-base, feat(executorch): copy-back for non-KV mutable buffers in the TensorRT delegate #4459's current content rebased on its own

One note for whoever rebases next: the commit fix(dynamo): correct legacy exporter (retrace=False) submodule inlining for hybrid graphs already landed on main as #4446, and the landed version is a superset of the branch version, so it should be dropped rather than replayed.