feat(executorch): copy-back for non-KV mutable buffers in the TensorRT delegate

[Conarnar]

Why this PR is needed

Caller-owned KV cache support (#4445) lets a mutable buffer live above the delegate and be updated in place by the engine. It handles this by lifting each mutated buffer to a delegate input and erasing the trailing copy_, relying on TensorRT's IKVCacheUpdateLayer aliasing to write the new value back to the caller-owned storage (zero-copy).

That assumption only holds for KV-cache writes. The slice_scatter and index_copy converters have a fast path that emits an IKVCacheUpdateLayer whose output is aliased in-place to the cache input. Any other in-place mutable buffer has no such aliasing — for example the conv_state / recurrent_state ring-buffers of a Gated DeltaNet (GDN) layer. For those, erasing the copy_ dropped the write-back entirely: the update became dead code and was eliminated, so the engine received fewer args than expected at runtime and the buffer never updated (silently wrong output).

The fix

Distinguish the two kinds of mutation in lift_mutated_buffers:

• KV writes (slice_scatter / index_copy) keep the existing zero-copy aliasing path — the copy_ is erased and the write-back is handled by the engine's IKVCacheUpdateLayer.
• Any other ("copy-back") mutation has its new value re-attached as an ordinary graph output and recorded as a BUFFER_MUTATION of its caller-owned buffer, so ExecuTorch copies it back after the delegate runs.

This uses the standard mutable-buffer representation rather than the engine-enforced aliased-I/O path (which is reserved for zero-copy KV aliasing).

Classification reuses the converters' own eligibility predicates rather than the op target alone, so a slice_scatter / index_copy the converter cannot turn into an IKVCacheUpdateLayer falls to copy-back instead of being dropped. assert_predicted_kv_aliased then cross-checks the classification against the engines that were actually built, in both directions: every write predicted as KV must appear in an engine's aliased_io, and no write filed copy-back may appear there. The first direction catches a silently lost write-back. The second catches a module that raises on every call, because the runtime treats an aliased output as an engine side effect and does not return it, while the surrounding graph still reads the trailing output copy-back added.

One derivation for slice_scatter eligibility

A predicate is only as good as its arguments: deriving _kv_eligible's inputs independently on the two sides mis-predicts as effectively as a different predicate would. resolve_slice_scatter_write in dynamo/conversion/impl/slice_scatter.py is the single derivation both sides call — the converter derives its arguments from the TRT tensors, the predictor from the fx node, and both go through this helper. It fills in the op's defaults, counts negative indices from the end, and reports which of the converter's early exits a write hits through a KVWriteStatus: FULL_OVERWRITE (the converter returns the source and emits no KV layer), DYNAMIC_BOUNDS (a non-int bound, which the converter raises NotImplementedError on), BAD_DIM (a dim that is not a Python int or does not index the cache, which raises IndexError), or OK. Only OK reaches _kv_eligible, with update_len taken as end - start. Every other status is a case in which the converter returns or raises before it can emit an aliasing layer, so none of them can be predicted to alias: a full overwrite keeps its copy-back, and the two raising statuses abort the compile, so nothing is copied back at all.

_kv_eligible itself does not take step. The shared derivation reads step for the full-overwrite shortcut, which requires step == 1, and otherwise passes it through untouched. On the converter side only the scatter fallback reads it: try_emit_kv_cache_update takes no step at all and writes update_len slots consecutively from start. So a strided write the KV fast path accepts is lowered as if it were contiguous — cache[:, :, 0:8:2, :] lands in slots 0, 1, 2, 3 rather than 0, 2, 4, 6, silently. That is confined to the KV fast path: a strided write that falls through to the scatter fallback is lowered correctly, because the fallback builds its scatter indices with step in hand, and test_fallback_step_two drives exactly that shape and compares against eager. The KV-path miscompile behaves the same way on main, and this change neither introduces nor fixes it; the OK return carries a comment recording it as known and unfixed and scoping it to the KV path, so the next reader is not sent at a fallback that is correct. test_step_is_returned_unchanged pins only what the shared derivation owes its two callers — that step comes back as it went in — so freezing the derivation does not freeze the routing built on top of it.

Two converter behaviours differ from main, because the dim check moved into the shared helper so that predictor and converter apply one rule to dim: an out-of-range dim raises IndexError from the converter instead of surfacing as a TensorRT message out of input.shape[dim], and dim must be a Python int, so a numpy.int64 dim is read as a bad dim where on main it works.

The cache has to reach the converter as a network input

emit_kv_cache_update_layer aliases the cache only when the argument it is handed is a direct network input — anything else has no input binding name, and the converter falls back to a plain scatter. After lifting, a mutated buffer is a placeholder, so the ordinary write-only shape satisfies that directly. A clone in between does not, and that is the shape a forward takes when it also reads the post-write cache: the write's cache argument is a call_function, but remove_input_alias_fixing_clones — an ATEN lowering pass inside post_lowering — then erases the clone, and the converter does see a network input after all.

_effective_cache_input therefore predicts what the converter will be handed rather than reading what is in the graph at classification time, reproducing that pass's own condition exactly: the clone is peeled only when its input is a placeholder and that placeholder's sole user is the clone. A placeholder some other node also reads keeps its clone, the converter sees a call_function, and nothing aliases. For the same reason classification runs after the trailing copy_ is erased — the copy_ is the last reader of the buffer placeholder in the write-only shape, so whether it is still there decides whether the clone survives. The peel has to reach both write ops: _index_copy_kv_eligible applies a placeholder check of its own, so it takes an input_node keyword the classifier passes and the partitioner does not; running as a capability validator, after lowering has settled what the input is, it reads node.args[0] itself.

Predicting this wrong in the copy-back direction is worse than a lost write-back: the write is filed copy-back, its new value re-attached as a trailing graph output, the engine aliases the cache anyway, and every call raises. That is what the second direction of the cross-check exists for. remove_input_alias_fixing_clones carries a note back to the classifier, so that loosening the pass's condition, or deleting the pass as its own TODO invites, is done knowing a prediction depends on it.

A KV write nothing else reads

Erasing the trailing copy_ can leave a KV-eligible write with no consumer at all — a forward that updates a cache and returns something unrelated. Dead-code elimination then takes the write before any engine can see it, so a prediction of engine aliasing could never be fulfilled. Such a write is routed to copy-back instead: re-attaching its new value as a graph output makes it live again, so the model compiles at the default min_block_size and the ExecuTorch runtime applies the buffer, rather than the compile failing on a prediction nothing could fulfil.

Copy-back and engine aliasing are mutually exclusive for one buffer, so a re-routed write also has to be kept out of aliased_io. It carries node.meta["_trt_no_kv_alias"], and both eligibility checks honour it: impl.slice_scatter reads it off ctx.current_node, and _index_copy_kv_eligible off the node it is handed. The key is a bare string literal at every site, which is how every other custom meta key in dynamo/ is written — _fp8_softmax_scale, set by a lowering pass and read off ctx.current_node in a converter, is the same shape — and it keeps _buffer_lifting.py free of a module-level import from conversion/.

The marker is deliberately narrow. It goes only on a write the classifier knows the converter would otherwise have aliased and whose result nothing consumes, never on copy-back writes at large: marking every copy-back write would leave the copy-back direction of the cross-check nothing to catch, and catching a genuine mis-prediction there is the whole point of it. The marker rides on node.meta through every pass in post_lowering, where a pass that rebuilds a node without its meta would drop it silently, so assert_no_kv_alias_markers_survived reconciles the names recorded before lowering against the lowered graph before conversion. That reconciliation runs in both directions and reports which of three faults it found: a recorded name that kept its node but lost the marker, an unrecorded node that gained one, or a recorded name gone with an unrecorded node carrying the marker instead.

The marker sits on the write node while the classification is per buffer, so one value node written back into two buffers would put the two in conflict. torch.export rejects that shape with SpecViolationError before the classifier sees it, which makes it a constraint on what can arrive rather than a case to guard against.

Two further cases

• Submodule-owned buffers. A get_attr target is fully qualified (layers.0.self_attn.kv_cache.k_cache), and hasattr/getattr do not walk a dotted path, so nested caches were silently skipped and frozen as constants. lift_mutated_buffers resolves through submodules with get_buffer; register_buffer rejects dots, so the buffer is renamed to a flat attribute (lifted_buf_*) and the recorded copy-back targets are remapped through the same renaming (otherwise the verifier rejects the program).
• Writes excluded from TensorRT. An op in settings.torch_executed_ops never reaches a converter, so it cannot emit an IKVCacheUpdateLayer and the engine will not alias that buffer. The classifier honours the exclusion list (matched the way the partitioner matches it, so the two cannot disagree) and leaves such a write on the copy-back path — which is what lets a program split across the TensorRT and CUDA delegates export.

How it works

The classification happens once, in lift_mutated_buffers, and is threaded to both save paths:

1. dynamo/lowering/_buffer_lifting.py — for each lifted buffer, a KV write is left to aliasing and its input-binding name recorded in gm.meta["_predicted_kv_bindings"]; a non-KV write has its new value appended as a trailing graph output (so it survives DCE now that the copy_ is gone), its buffer name recorded in output order in gm.meta["_copyback_mutation_buffers"], and its input-binding name recorded in gm.meta["_copyback_bindings"] for the other direction of the cross-check. A KV-eligible write re-routed for want of a consumer is marked, and the marked node names recorded in gm.meta["_no_kv_alias_writes"]. Nested buffers are resolved through get_buffer and renamed; excluded writes are left on the copy-back path.
2. dynamo/_compiler.py — reads the markers back against the lowered graph (assert_no_kv_alias_markers_survived) after post_lowering and before conversion, forwards the copy-back list onto the compiled module's meta so it reaches the exporters, runs assert_predicted_kv_aliased against the aliased_io of the compiled submodules with both binding sets, and hides the copy-back values from what the compiled module returns (hide_copyback_outputs).
3. dynamo/_exporter.py
   • create_trt_exp_program (retrace=False, legacy exporter): tags the trailing outputs with their buffer target and moves all mutation outputs ahead of the user outputs (verifier requirement). A write-only buffer has no surviving get_attr after DCE and lift derives BUFFER input specs from get_attr nodes alone, so an unused one is re-added for it.
   • _declare_aliased_kv_mutations_on_ep (run on the resulting ExportedProgram by every save() branch that serializes a signature, and by torch_tensorrt.executorch.export() for every source shape it accepts): detaches the full trailing run of copy-back values, pairs each positionally with its buffer (in copyback_buffers order), and declares as BUFFER_MUTATION only the buffers torch.export did not already declare itself — skipping the rest — then rebuilds the top-level out_spec so to_edge's unflatten sees the right leaf count. Pairing across the full run before filtering is deliberate: dropping already-declared buffers first would shift the run and pair a value with the wrong buffer, which the runtime would then copy into a buffer of a different shape.
4. executorch/_export.py — torch_tensorrt.executorch.export() runs the declaration pass over program_map once _prepare_programs has normalized all three accepted source shapes into it, threading each method's _copyback_mutation_buffers from the source that method came from. It runs on the staged copy rather than the caller's program, so a caller who hands in their own ExportedProgram keeps it usable afterwards.

The KV half of _declare_aliased_kv_mutations_on_ep needs the optional [executorch] extra to read an engine's binding names; on a plain install that import fails, the engine scan is skipped, and the copy-back half still runs.

What compile() returns

A copy-back value is only ever consumed by whatever serializes the module. Nothing between compile() and the ExecuTorch runtime writes it anywhere, so a compiled module called directly from PyTorch leaves its buffer at the value it was compiled with, exactly as on main. Handing that value back as an extra return value would report a mutation the caller cannot act on, so hide_copyback_outputs installs a CodeGen whose generated forward stops short of the trailing copy-back values. The compiled module keeps the arity of the model it was compiled from: a non-KV mutable buffer does not turn one return value into two.

Only the generated forward stops early. The values stay on the graph's output node, so everything that reads the graph rather than calling the module still finds them — the exporters, dead-code elimination, and the torch.fx.Interpreter that a non-strict torch.export runs a GraphModule through, which is why a retrace keeps them too. Non-strict is what both retrace paths get: _compile.save passes strict=False and dynamo._exporter.export takes the default. A caller who exports the compiled module themselves with strict=True gets the module called rather than interpreted, and the hidden values do not reach the program; the _HiddenCopybackOutputs docstring records that at the point of decision. _TorchTensorRTModule.forward does the same with the aliased outputs the interpreter appends: the engine produces them, the caller never sees them.

Where copy-back is declared

source	exporter	exported_program	executorch	aot_inductor
ExportedProgram	n/a	KV only	KV + copy-back	not declared
GraphModule, retrace=True	either	KV + copy-back	KV + copy-back	not declared
GraphModule, retrace=False	legacy (default)	KV + copy-back, at transform time	KV + copy-back, at transform time	KV + copy-back, at transform time
GraphModule, retrace=False	use_legacy_exporter=False	KV only, warns	KV + copy-back	not declared, warns

torch_tensorrt.executorch.export() declares for every source shape it accepts, and save(output_format="executorch") reaches it on every branch, so that column is uniform. save(output_format="exported_program") declares only the KV half of an already-exported program, since no copy-back buffer list is threaded on that branch; save() is therefore not uniform across its two serializing formats.

aot_inductor is left undeclared on the torch.export paths — whether an aliased in-place mutation survives functionalization under inductor is unverified — matching #4445's KV scoping. It warns when the module carries engines with aliased I/O, and, under retrace=False with use_legacy_exporter=False, when a copy-back buffer would go undeclared.

Among the GraphModule paths, retrace=False with use_legacy_exporter=False and output_format="exported_program" is the one combination that leaves a copy-back mutation undeclared: on the retrace=False branch the legacy exporter is what declares copy-back, and torch_tensorrt.executorch.export() covers the executorch format, so this combination falls between the two. It warns rather than silently saving a signature that omits the update.

Idempotency

The declaration pass is safe to call more than once, by two separate mechanisms. The KV half skips any buffer already carrying a BUFFER_MUTATION spec, which is also how the legacy exporter's transform-time declaration survives a later run. The copy-back half cannot skip one buffer at a time, because it detaches its outputs by position rather than by name, and per-buffer state cannot tell the two "already declared" cases apart: torch.export declares a mutation it recognises and leaves the trailing value, while a previous run of this pass declares it and consumes the value. So the consuming step records gm.meta["_copyback_mutations_declared"] on the GraphModule and the slice is skipped when it is set. create_trt_exp_program sets the same marker, since it declares copy-back itself at transform time. Without the marker a second run slices whatever trails the graph — by then a genuine user output — and the saved program fails to load with treespec.unflatten ... leaves has length 0.

Uninitialized copy-back buffers

Declaring a write as an ExecuTorch BUFFER_MUTATION puts the buffer under ExecuTorch's rule that a mutated buffer has a meaningless initial state: only its shape and dtype are serialized. A buffer read before it is written — a GDN conv_state or recurrent_state, an EMA, a running statistic — therefore starts from whatever the allocation held. ExecuTorch warns about this itself, but generically, naming no buffer. torch_tensorrt.executorch.export() names the copy-back buffers and states the rule they fall under, with the condition attached rather than stated flat, since a caller who has already supplied the pass below does have the value in the .pte: they "are declared mutated, so ExecuTorch serializes only their shape and dtype, not their value, unless a pass marks them meta["et_init_buffer"]; without that, a buffer read before it is written starts from whatever the allocation held."

The remedy is named, but only together with the one thing upstream does not say about it. ExecuTorch's own way to put an initial value back is InitializedMutableBufferPass, which sets et_init_buffer on the placeholders its patterns match; the emitter then marks the spec const and serializes the buffer by reading its storage host-side through ctypes.cast(spec.storage.data_ptr(), ...). Whether that read is safe is decided by where the buffer lives: it works while the buffer is on CPU, and it takes the export down with a segfault, rather than raising, once the buffer is CUDA-resident — which is what exporting a model on CUDA leaves behind, and exporting on CUDA is the ordinary way to reach a TensorRT engine. ExecuTorch's emitter names the pass to the same user later in the same export and says nothing about where the buffer has to live, so withholding the name here would withhold the caveat and not the advice. The warning carries both, and the test that covers it requires the two to travel together.

Engine-aliased KV buffers lose their serialized initial value the same way and are deliberately excluded from the warning. An aliased write is a cache-position write on the sequence axis — the only form IKVCacheUpdateLayer expresses — and such a model is assumed to read the cache under that same position, through a bounded slice or a position mask, so slots no step has written should not reach the output. Nothing verifies that; it is a property of the model. On that assumption initializing a cache buys nothing and would cost what is typically the model's largest tensor, written whole into the .pte.

What copy-back costs

Only an engine-aliased (KV) write is free. A BUFFER_MUTATION output is defined as the buffer's new contents rather than the slice that changed, so a copy-back write re-attaches the whole post-write buffer as a graph output. Each call therefore pays, per copy-back buffer, a full-buffer materialization of that output plus a full-buffer copy from ExecuTorch's own write-back pass after the delegate returns. For a decode step writing one position into a multi-megabyte KV cache, that easily dominates the step. It is the right correctness tradeoff — the alternative is silently losing the write — but it is why the classifier routes a cache write to engine aliasing instead, and why a model that unexpectedly lands in copy-back comes out slow rather than wrong. This is recorded in the _buffer_lifting.py module docstring, with gm.meta['_copyback_mutation_buffers'] named as the way to find out which buffers are paying it.

The engine-converter entry point

convert_exported_program_to_serialized_trt_engine(..., lift_mutable_buffers=True) hands the caller raw engine bindings and performs no write-back of its own. A copy-back write there would append an extra engine output that the entry point reports nothing about, so the buffer would come back stale; it is rejected right after lifting, with the offending buffers named and the working route pointed at: torch_tensorrt.dynamo.compile and saving the result, since it is the save that declares the buffer mutation for the runtime to apply — compile() on its own performs no write-back either. The cross-check runs there too, against interpreter_result.aliased_io.

Diagnosing a failed cross-check

assert_predicted_kv_aliased takes the aliased-input binding set rather than a module, because the two callers hold ground truth in different shapes — compile() reads it off the compiled submodules, the engine converter off the interpreter result — and aliased_input_bindings assembles it for both. Both directions are keyed on the buf_* input-binding name, which is stable across the buffer rename that inlining does later and is exactly what aliased_io records on the input side.

Classification runs before lowering and partitioning, so an unfulfilled KV prediction means something in between kept the write out of the engine it was predicted into. The error names the buffers and then names a likely cause, citing min_block_size only where that setting can be the reason: above 1 it reports the value in force and that min_block_size=1 rules the setting out, and at 1 it points at a converter or a capability validator rejecting the op instead, since naming the setting there would blame the value it would have to recommend.

Under compile(dryrun=True) no engine is built, so aliased_io is empty for reasons that say nothing about the predictions and every one of them would look unfulfilled. The skip is keyed on an explicit engines_built keyword rather than on settings.dryrun, because the two entry points do not mean the same thing by that flag: convert_exported_program_to_serialized_trt_engine accepts dryrun, puts it in CompilationSettings and then never consults it, so an engine comes out regardless. Reading dryrun inside the shared helper would hand that entry point an opt-out kwarg for a check it is meant to have none for. Only compile() passes engines_built=False, and only where it returns before conversion.

Reading engine metadata without re-serializing the engine

__getstate__ on a TensorRT engine is the pickling hook, not a metadata accessor: it re-serializes the whole ICudaEngine and base64-encodes it. get_engine_info_from_state(..., metadata_only=True) takes the record from serialize_metadata_only instead, and _get_engine_info_for_node forwards the flag. Two callers read metadata only and opt in. The aliased-mutation declaration pass reads ALIASED_IO, INPUT_BINDING_NAMES and OUTPUT_BINDING_NAMES and runs its engine loop before it can know whether any engine has aliased I/O at all, so a program with no aliased KV paid a full serialization per engine node and then returned unchanged — measured at 0.4 s and a 67 MB transient string for one ~50 MB engine. TensorRTPartitioner._resolve_target_device_for_partition reads a single field, DEVICE_IDX, once per partition, on every export that does not pin target_device, which is the default: 967.6 ms against 3.1 ms on a ~100 MB engine, plus a 134 MB base64 string built and dropped. The accessor's caveat that ENGINE_IDX is unreliable under metadata_only applies to neither, since neither reads that slot.

Scope

The classification is a no-op for any model with no non-KV in-place mutable buffer and no KV write whose result nothing consumes — including all standard models, and any KV-cache model whose writes the converter turns into IKVCacheUpdateLayers. The dim handling above is the one converter behaviour that differs for programs with no mutable buffer at all.

Testing

Unit tests, CPU-only unless noted, and gated on executorch.exir where they touch it:

• tests/py/dynamo/lowering/test_buffer_lifting.py — classification: KV slice_scatter / index_copy writes stay aliased with no copy-back; a non-KV mutation is recorded and its new value re-attached as the trailing output (verified numerically equal to the updated buffer); index_put falls to copy-back; a mixed KV + non-KV graph records only the non-KV buffer; an ineligible index_copy (2-D cache) and an ineligible slice_scatter (wrong dim) fall to copy-back rather than being dropped; a write excluded through torch_executed_ops is left on the copy-back path; nested buffers are lifted, and inlining remaps the recorded copy-back targets to their flattened attribute names. TestSliceScatterDerivationIsShared pins the corners where an independent derivation would part company with the converter — full overwrite, open-ended end == INT64_MAX, negative start, non-int start — plus each KVWriteStatus, the bounds its contract promises, and step coming back as it went in. TestCacheMustReachTheConverterAsANetworkInput covers what the converter will be handed: a direct placeholder and a sole-use clone of one classify as KV, for slice_scatter and for index_copy, while a shared clone, a clone of a non-placeholder, a cache read through another op and a cache still read from a get_attr do not — and the index_copy validator still reads node.args[0] when nobody overrides it, so the partitioner's view of the same node is unchanged. TestDeadKvWriteRouting covers the write nothing else reads: it is filed copy-back rather than predicted KV, the re-routed write carries the marker, a live KV write and a non-KV write do not, and the marker turns _index_copy_kv_eligible from true to false on the same node. TestNoKvAliasMarkerSurvival drives the reconciliation: a surviving marker passes, while a stripped marker, a marker copied onto an unmarked node, a marker appearing where none was recorded, and a marked node replaced by another carrying the marker each raise. TestPredictedKvAssertion covers the backstop in both directions: it passes when a predicted-KV write is aliased and raises when it is not, raises when a copy-back write is aliased and passes when it is not, aggregates aliased_io across multiple engines, no-ops with no prediction, skips under engines_built=False for either direction, does not skip on dryrun alone, and names min_block_size only where that setting can be the cause. TestHiddenCopybackOutputs pins that hiding a value leaves the graph's output node untouched, that hiding nothing is a no-op, and that a graph consumer still sees the hidden value. TestCompileSeam (CUDA) drives compile() end to end and pins the wire between the two halves: the copy-back list reaching trt_gm.meta and still resolving to a live buffer, the cross-check being called with the predictions lifting made, each direction of it failing a compile, compile() telling the check that a dryrun built no engines, a cloned slice_scatter and a cloned index_copy cache write each compiling and writing back, a dead slice_scatter write compiling with an empty aliased_io at min_block_size 1 and at the default 5 and its mutation arriving in a saved program, the same for a dead index_copy write, for a dead cloned write where both mechanisms meet, and for a dead write whose buffer still has another reader, the marker read-back running on the lowered graph rather than before lowering, and the copy-back value staying out of the compiled module's return — beside an engine-aliased KV buffer in the same module, and while a retrace and a saved program still find it.
• tests/py/dynamo/executorch/test_kv_cache_export.py — both exporter passes reclassify a trailing copy-back output to BUFFER_MUTATION ahead of the user outputs; a buffer torch.export already declared is skipped rather than declared twice, and with a mix of declared and undeclared buffers each remaining copy-back value is paired with the buffer it was appended for; a second run of the pass is a no-op; the copy-back half still runs when the [executorch] extra is unimportable; the declaration pass reads engine metadata only; a write-only copy-back buffer still reaches lift as a get_attr (driving the real lift and ExportedProgram constructor so a regression reproduces the verifier error rather than passing vacuously), and the re-added get_attr neither duplicates an existing one nor invents a dangling one; a saved copy-back program reloads, returns exactly the outputs the source module returned, and updates its buffer; torch_tensorrt.executorch.export() declares copy-back for each source shape it accepts (both GraphModule retrace modes, an ExportedProgram, and a method mapping) and leaves the caller's own program intact; every save() branch that serializes a signature reaches the declaration pass, swept over module type by serializing format; save() warns on exactly the one combination that cannot declare; and the engine-converter entry point rejects a copy-back buffer by name, matched in its rendered form so a substring of an unrelated word cannot satisfy the assertion.
• tests/py/dynamo/executorch/test_export.py — the uninitialized-buffer warning names the buffer, states what ExecuTorch keeps of it and under what condition, and names InitializedMutableBufferPass only in a message that also carries the CUDA caveat; an export with nothing to copy back stays silent.
• tests/py/dynamo/conversion/test_slice_scatter_aten.py — TestSliceScatterFallback drives the scatter fallback against eager across ranks and dims, including a step=2 write and a dynamic-shape case; TestSliceScatterEarlyExits drives the converter's raising exits: dynamic bounds, and both forms of bad dim.
• tests/py/dynamo/executorch/test_api.py — the partitioner asks for the engine record with metadata_only=True.

The copy-back path was also exercised outside the test suite, by running .ptes on a real ExecuTorch runner: a synthetic copy-back model; an L=2 Gemma4 MoE with an added non-KV buffer, which puts copy-back and engine-aliased KV in one hybrid TensorRT + CUDA model; a Qwen3.5 GDN decode whose four per-layer conv_* / rec_* state buffers are all copy-back and whose recurrent state has no KV aliasing anywhere; and a clone / cache-position-write / write-back decode step. In each, a step observes the previous step's write. Those runs are not part of this PR's automated tests.

Stacking

Not stacked. #4446 (retrace=False legacy-exporter fix) and #4445 (caller-owned KV cache) have both landed, and this branch is based on main above them.

Follow-ups

• feat(executorch): expose composable Edge export API #4440 (composable Edge export API) has landed, and the declaration for output_format="executorch" now lives in torch_tensorrt.executorch.export(), where it runs per method. The output_format="exported_program" branches still declare from save() in _compile.py, and save(exported_program, output_format="exported_program") still declares only the KV half; making save() uniform across its two serializing formats is a follow-up.
• refactor(executorch)!: share one caller stream, and ship the CUDA delegate in the runtime wheel #4454 (shared caller stream) has landed. A copy-back is an ExecuTorch-level write performed after the delegate; on the shared caller stream it is naturally ordered before the next delegate reads the buffer, so cross-delegate copy-back is correct without additional code. A cross-delegate prefill/decode acceptance test is now unblocked and will be added as a follow-up.
• _write_op_is_torch_executed and the dead-write re-route are the two reasons a write can miss an engine that the classifier can see coming, and both file the write as copy-back up front. The others — min_block_size, a missing converter, a failed capability validator — are not knowable before partitioning, so the cross-check catches them after it instead. Making the prediction non-destructive, by always keeping the copy-back output and dropping it after conversion once aliased_io has settled which buffers the engine aliased, would remove the need for the cross-check entirely; that is a larger change and belongs in its own PR.

[cehongwang]

#4459 widens the #4445 ordering blocker into a three-way reversal. lift_mutated_buffers appends copy-back values to the graph output:

out_args = list(output_node.args[0])
out_args.extend(nv for nv, _ in copyback)
output_node.args = (tuple(out_args),)

and the interpreter appends aliased KV outputs after that. So the engine binding order becomes [user…, copyback…, kv_aliased…]. But _declare_aliased_kv_mutations_on_ep sets the graph output to kv_getitems + copyback_getitems + out_args, giving delegate args [kv_aliased…, copyback…, user…] — the three groups in reverse. Since preprocess still passes output names through in engine order, a model with all three kinds gets a full permutation mismatch rather than the two-way swap in #4445.

[cehongwang]

This classification is op-level, but aliasing isn't an op-level property — so this can still drop a write-back in exactly the way the PR sets out to fix.

The comment above says the set must "stay in sync with the ops the converters actually turn into an IKVCacheUpdateLayer." The problem is that no set of ops can stay in sync with that, because whether an IKVCacheUpdateLayer gets emitted depends on shapes and network position, not just the target. In slice_scatter.py, _kv_eligible requires:

• a static s_max
• a 4-D [b, d, s_max, h] cache
• dim == 2
• a non-dynamic batch dim

and on top of that emit_kv_cache_update_layer bails when the cache isn't a direct network input (input_binding_name returns None) or when add_kv_cache_update returns None. index_copy.py has the same eligible/index_copy_fallback split.

When any of those fail, the converter emits a plain scatter with no aliasing — but _is_kv_cache_write already returned True, so the copy_ is erased with no copy-back appended and the write-back is silently lost. A 3-D cache or dim != 2 is enough to hit it.

Before this PR the erase was unconditional, so this failure mode existed for every mutated buffer; this PR narrows it to these two ops but keeps the same "trust the op" assumption for them.

Could the classification be derived from what was actually emitted rather than predicted? The engine's aliased_io map is the ground truth, and _declare_aliased_kv_mutations_on_ep already reads it. Deriving copy-back as "mutated buffer not present in aliased_io" post-conversion would make the two sides agree by construction.

If you'd rather keep the pre-conversion classification to avoid restructuring, the minimum would be a post-conversion assertion: every buffer classified as KV must appear in the engine's aliased_io, otherwise error (or fall back to copy-back) rather than silently dropping the write.

[Conarnar]

By the time aliased_io is generated, the mutation write-backs have already been removed by DCE. Deriving copy-back from it would require re-attaching the write-back for ALL buffers and removing the ones present in aliased_io afterward.

Applied your suggestion to the classification and added the post-conversion assertion.

[cehongwang]

The KV cases here use slice_scatter / index_copy shapes that are KV-eligible, so they only cover the happy path of the classification.

Could you add a case where the op is slice_scatter but the converter's fast path would not fire — e.g. a 3-D cache, or dim != 2, or a dynamic s_max? Today that lands in the KV bucket and loses its write-back. Whatever behavior you settle on (copy-back, or a hard error), a test pinning it would keep the two sides from drifting as _kv_eligible evolves.

[Conarnar]

#4459 widens the #4445 ordering blocker into a three-way reversal. lift_mutated_buffers appends copy-back values to the graph output:

out_args = list(output_node.args[0])
out_args.extend(nv for nv, _ in copyback)
output_node.args = (tuple(out_args),)

and the interpreter appends aliased KV outputs after that. So the engine binding order becomes [user…, copyback…, kv_aliased…]. But _declare_aliased_kv_mutations_on_ep sets the graph output to kv_getitems + copyback_getitems + out_args, giving delegate args [kv_aliased…, copyback…, user…] — the three groups in reverse. Since preprocess still passes output names through in engine order, a model with all three kinds gets a full permutation mismatch rather than the two-way swap in #4445.

Same with #4445, I was not able to reproduce this.

[cehongwang]

I am not 100% sure, but it seems like, in the decode step the KV cache is written in place by the engine, then copied again by ExecuTorch here: https://github.com/pytorch/executorch/blob/7811c69f41db6cc84abfd0b4a32f2e49393958a5/exir/passes/insert_write_back_for_buffers_pass.py#L163 because it is also a BUFFER_MUTATION output node. Please correct me if I am wrong

[cehongwang]

This fixes the bug for executorch but leaves Torch-TensorRT runtime wrong. @narendasan Should we use the same mechanism to support in-place operation for Torch-TensorRT runtime by appending a copy node at the end of graph and run it in python as graph break? If we want to do that, we should move that part from export.py to compiler or interpreter

[Conarnar]

I am not 100% sure, but it seems like, in the decode step the KV cache is written in place by the engine, then copied again by ExecuTorch here: https://github.com/pytorch/executorch/blob/7811c69f41db6cc84abfd0b4a32f2e49393958a5/exir/passes/insert_write_back_for_buffers_pass.py#L163 because it is also a BUFFER_MUTATION output node. Please correct me if I am wrong

You are correct, there is a redundant copy added by ET. Unfortunately, removing it would cause ExecuTorch to stop tracking and persisting the mutation. Fixing this properly would require delegate-boundary aliasing support on the ExecuTorch side.

[Conarnar]

@shoumikhin Thanks for taking a look and the thorough writeup. I have taken your suggestion regarding KV caching in submodules and applied the fix to problem 2 here. Problem 1 is a pre-existing bug on main and can be landed separately (#4472).

Regarding excluded buffer writes from TRT, I believe those are out of scope for this PR and can remain in yours (#4470).

[shoumikhin]

I had left too many long comments here, several of them wrong or already fixed, so I deleted them and replaced them with this one list. Anything below that I did not personally reproduce is marked as such.

Inline threads with your replies are untouched, since your answers there are worth keeping.

Needs fixing before merge

1. Saving with retrace=True and the legacy exporter destroys a user output.

_declare_aliased_kv_mutations_on_ep drops the last N outputs, but on that path the legacy exporter has already declared the copy-back and moved it to the front, so the outputs it drops are ordinary user outputs. On a model whose user output differs from the written value:

BEFORE  specs [('USER_OUTPUT', 'mul', None)]        run -> tensor([4., 6.])
AFTER   specs [('BUFFER_MUTATION', 'mul', 'state_0')] run -> ()

The user output is not just dropped, it is relabeled as the new buffer contents. Saving succeeds silently and the saved program then fails on load with treespec.unflatten ... leaves has length 0 but the spec refers to a pytree that holds 1 items. Your own comment at _compile.py:1142-1146 describes this hazard, and the retrace=False branch already avoids it by passing no copyback_buffers. Simplest fix is to do the same on the retrace=True branch, or skip any buffer that already has a BUFFER_MUTATION spec.

2. Copy-back is skipped entirely when the optional executorch extra is missing, with no warning.

_exporter.py:624-626 catches ImportError and returns early, before any copy-back work, and copy-back needs nothing from executorch. Measured with the package genuinely hidden from the import system:

extra present : [USER_OUTPUT mul] -> [BUFFER_MUTATION mul state]
extra absent  : [USER_OUTPUT mul] -> [USER_OUTPUT mul]

buffer before [1.0, 2.0, 3.0]  after [1.0, 2.0, 3.0]   updated? False
log records   : NONE

Fix is to move that import below the copy-back handling, or narrow the try to the KV block.

3. convert_exported_program_to_serialized_trt_engine adds an engine output but not the mapping.

It calls lift_mutated_buffers at _compiler.py:2015 but never reads _copyback_mutation_buffers and never calls assert_predicted_kv_aliased, both of which compile() does. Copy-back appends a trailing graph output, and _TRTInterpreter.py:803-815 indexes caller-supplied names by position, so a caller who sized that list from the original program is one short. Raising a clear error is fine if this API is not meant to support copy-back; silently returning an engine with an extra output and stale state is worse.

Worth a look, but I could not confirm these myself

4. A buffer with a trained initial value may be wrong from the first call. Declaring a buffer BUFFER_MUTATION puts it under ExecuTorch's rule that mutated buffers have a meaningless initial state, so only shape and dtype are serialized. ExecuTorch warns about this itself, and nothing in this PR applies InitializedMutableBufferPass or warns. I confirmed the warning and that no mitigation exists, but a local toolchain problem stopped me from reproducing the wrong numbers end to end. Any model with a trained conv_state, an EMA, or a running statistic would be affected.

5. torch_tensorrt.executorch.export(..., retrace=True) may never declare copy-back. Every call site of the declaration pass is in _compile.py, and that entry point passes use_legacy_exporter=False, so nothing tags the value. I checked the call sites by reading, not by running it.

Two things I got wrong, so you do not chase them

• I said a rebase would clear the red CI. It will not. The same jobs are red on main, because constant_fold_exclusions landed as a new package directory but was never added to the dynamo_packages list in setup.py, so the built wheel omits it and the runtime import fails. Not your change, and not fixable by rebasing.
• I reported that returning the same value you write back deletes the user output, and that two buffers sharing one value loses a name. Both were artifacts of my own test setup skipping run_decompositions, which the real pipeline runs before lifting. Through the actual path these behave the same on your branch and on the base, so please ignore them.

Notes

The strided slice_scatter case (step > 1) is real and silent, but it predates this PR, so it is not a blocker here. If you want to close it anyway, returning False from the predicate when step != 1 routes the write to copy-back, and it breaks nothing: all 22 tests in test_buffer_lifting.py still pass, and no test in either new file uses a step argument.

The design is right. The classifier's split between engine-aliased KV writes and copy-back writes is the correct call, and assert_predicted_kv_aliased catches a class of silent write loss that exists on the base today. The issues above are all about the output layout assumptions, not the approach.

[shoumikhin]

One more, and this one is on merged main rather than something new here. I am raising it on this PR because this is where the save paths are still being worked on, and because it is not fixed by the current head.

When save() is handed an already-exported program, the executorch branch never declares the mutation, while its two siblings do.

elif module_type == _ModuleType.ep:
    ...
    if output_format == "exported_program":
        module = _declare_aliased_kv_mutations_on_ep(module)   # declares
    ...
    elif output_format == "executorch":
        _save_as_executorch(                                   # does not
            module,

That is _compile.py:1112-1113 on main. The equivalent branches in the other two module-type blocks, at :1210 and :1321, both call the declaration first.

Measured by instrumenting the call, same model and same input both times:

save(ep, output_format="exported_program") -> declaration called? True
save(ep, output_format="executorch")       -> declaration called? False

I ran that against merged main and then against this PR's head, and got the same result both times, so this is not something the current head fixes. Diffing _compile.py between main and this head, none of the 55 changed lines touch that branch.

Effect: for save(exported_program, output_format="executorch") the buffer is never marked mutated, so the aliased I/O the merged change exists to provide is quietly not declared for that one entry point. No error, no warning.

I think it is a two-line fix, adding the same call the sibling branches make right before _save_as_executorch. Worth a regression test too, since all three branches doing the same thing is exactly the kind of symmetry that silently drifts.

Separately, and not yours: executorch-runtime-test is gated on needs.executorch-runtime-build.result == 'success', and that build is currently failing on main for the setup.py packaging reason I mentioned earlier. So the ExecuTorch tests are being skipped rather than run, on this PR and on main. Until that is sorted, none of this is being checked by CI, which is probably how the asymmetry above survived review in the first place.

[Conarnar]

Changes since the last push, mapped against your numbered items.

1. Cross-check turns working compiles into hard failures at the default min_block_size

Both of the minimum items are in.

Skip when no engine was built. The skip is an explicit engines_built keyword on assert_predicted_kv_aliased, not a dryrun check inside the helper. compile() passes engines_built=not settings.dryrun, and it is the only caller that passes anything: convert_exported_program_to_serialized_trt_engine accepts dryrun, puts it in CompilationSettings and then never consults it — interpret_module_to_result runs and an engine comes out — so reading dryrun inside the shared helper would hand that entry point an opt-out kwarg for a check it is meant to have none for. It keeps the default engines_built=True. test_dryrun_alone_does_not_skip pins that: assert_predicted_kv_aliased with CompilationSettings(dryrun=True) and an unfulfilled prediction still raises.

Name min_block_size. The message now appends ", most often because min_block_size (5) rejected the subgraph it landed in; min_block_size=1 rules that out", with the value read off the settings in force. test_message_names_min_block_size asserts both min_block_size (5) and min_block_size=1 appear.

Tests: test_no_engines_built_does_not_raise, test_dryrun_alone_does_not_skip, test_message_names_min_block_size, and TestCompileSeam.test_compile_does_not_cross_check_a_dryrun for the compile() call site. Each fails with its half of the change disabled.

Separately, one of the ways to reach your failure is now gone: a full-extent slice_scatter is predicted FULL_OVERWRITE and routed to copy-back rather than predicted to alias — see item 3.

The larger rework you prefer — always keep the copy-back output, drop the prediction, and let the engine's own aliased_io settle it after conversion — may well be possible. A prototype did not hit the obstacle I expected, but it reaches past the export path into how the runtime decides which engine outputs belong to the caller, and I would want that properly tested before proposing it. I would rather work it up as its own PR than fold it in here.

2. The seam between the two halves has no test

TestCompileSeam in tests/py/dynamo/lowering/test_buffer_lifting.py, four tests, all through compile(), CUDA-gated. Your verification recipe, by deletion:

deleted from _compiler.py	fails
trt_gm.meta["_copyback_mutation_buffers"] = _copyback_mutation_buffers	test_compile_threads_copyback_buffers_into_trt_gm_meta
the assert_predicted_kv_aliased(...) call	test_compile_raises_when_a_predicted_kv_write_is_not_aliased, and test_compile_runs_the_predicted_kv_cross_check, which patches the symbol and asserts it was called with ["buf_cache"]
the engines_built=not settings.dryrun argument at that call site	test_compile_does_not_cross_check_a_dryrun

The model in test_compile_raises_when_a_predicted_kv_write_is_not_aliased is your repro — zeros(2, 4, 16, 8) cache, cache[:, :, 3:4, :] = x, min_block_size=5 — so the raise is pinned as behaviour rather than left to the assertion being present. The meta test also checks that the recorded name still resolves to a buffer after inlining and that the trailing output carries the post-write value, not a stale read.

3. The classifier re-derives _kv_eligible's arguments

Factored into resolve_slice_scatter_write in dynamo/conversion/impl/slice_scatter.py, returning (start, end, step, KVWriteStatus). The converter and _kv_write_will_alias both call it; the converter's argument derivation is now literally the same code the predictor runs. Your four:

• update_len from the source shape. Both sides now take end - start. test_open_ended_slice_is_not_kv drives cache[:, :, 3:, :] with end == INT64_MAX, which fails start + update_len <= s_max in the predictor exactly as it does in the converter.
• A non-int start treated as 0. Now KVWriteStatus.DYNAMIC_BOUNDS: the predictor returns False, the converter raises NotImplementedError. test_non_int_start_is_not_kv.
• No negative-index normalisation. The helper counts a negative start or end from dim_size before anything reads it. test_negative_start_is_normalised covers both sides of the bound — -4 in a 16-slot cache normalises to 12, so a 4-slot write is eligible and an 8-slot write is not.
• No full-overwrite shortcut. Now KVWriteStatus.FULL_OVERWRITE, which the predictor treats as "nothing aliases the cache" and routes to copy-back. test_full_overwrite_is_not_kv, both with explicit 0, 16 bounds and with the bounds omitted.

step is deliberately unchanged. _kv_eligible does not consult it, and the helper consults it only where the converter does — the full-overwrite shortcut, which requires step == 1 — so a strided write with concrete bounds still takes the KV path. That matches your note that the case predates this PR. test_strided_write_still_takes_the_kv_path pins the current answer so the shared derivation cannot change it by accident, without claiming it is the right one.

Two behaviour changes ride along, both from the dim check moving into the shared helper so predictor and converter apply one rule to dim: an out-of-range dim raises IndexError from the converter rather than surfacing out of input.shape[dim] with a TensorRT message, and dim must be a Python int, so a numpy.int64 dim is read as a bad dim where it used to work. The IndexError prints the offending value and its type name, because those two causes share one status and an in-range numpy.int64 would otherwise read as an out-of-range complaint about an in-range number. TestSliceScatterEarlyExits in tests/py/dynamo/conversion/test_slice_scatter_aten.py drives both raising exits through the converter and matches the type name rather than the rendered value, which numpy 2.0 changed; the dynamic-bounds exit had no coverage anywhere, so restoring the helper's bounds pass-through left that guard unreachable with every test still green.

4. The copy-back cost is not written down

Now in the _buffer_lifting.py module docstring, under "What copy-back costs": only an engine-aliased write is free; a copy-back write re-attaches the whole post-write buffer as a graph output, because a BUFFER_MUTATION output is defined as the buffer's new contents rather than the slice that changed. Each call therefore pays, per copy-back buffer, the graph materializing the full buffer as that output plus ExecuTorch's write-back pass copying it into the caller-owned buffer after the delegate returns. For a decode step writing one position into a multi-megabyte cache that is a full-cache-sized materialization plus a full-cache-sized copy to record a single-slot update, and easily dominates the step. It states that this is the right correctness tradeoff, that it is why the classifier routes a cache write to engine aliasing instead, and that a model landing unexpectedly in copy-back comes out slow rather than wrong — with gm.meta['_copyback_mutation_buffers'] named as the way to find out which buffers are paying. lift_mutated_buffers points at the docstring rather than restating it.

Your 15:12 comment: save(ep, output_format="executorch") never declares the mutation

Fixed, though not with the two lines in _compile.py. The declaration now runs inside torch_tensorrt.executorch.export(), over program_map once _prepare_programs has normalized all three source shapes into it, threading each method's _copyback_mutation_buffers from the source that method came from. save(ep, output_format="executorch") reaches it through _save_as_executorch. Doing it there also covers the two source shapes save() cannot reach at all — an ExportedProgram or a method mapping handed straight to export(), neither of which was declared either. The pass is idempotent, so a program that already carries the declaration (the legacy exporter's, or a caller's) comes through unchanged, and it runs on the staged copy so a caller's own program is left intact.

Regression test, since you asked: test_save_declares_aliased_mutations_on_every_serializing_branch sweeps (module type ∈ {ep, fx_retrace, fx_no_retrace}) × (format ∈ {exported_program, executorch}), spying on the pass while delegating to the real one so the pipeline underneath stays real. aot_inductor is excluded — it warns instead of declaring. Removing the declaration from torch_tensorrt.executorch.export() fails the (ep, executorch) case and leaves the other five passing, which is the case that was missing. test_executorch_export_declares_copyback_for_every_source_shape checks the resulting Edge program's output specs once per source shape export() accepts: both GraphModule retrace modes, an ExportedProgram, and a method mapping.

Knock-on: save()'s retrace=False / use_legacy_exporter=False warning became false for output_format="executorch", since that branch now declares, so the warning is gated on the format. test_save_warns_when_copyback_cannot_be_declared takes an output_format parameter and pins all four combinations together.

Still asymmetric in one place: save(exported_program, output_format="exported_program") declares only the KV half, because no copy-back buffer list is threaded on that branch. This does not make save() uniform across its two serializing formats.

Changes you did not ask for

Idempotency marker. _declare_aliased_kv_mutations_on_ep claimed idempotency on the strength of its already_exposed check, which holds for the KV half only. The copy-back half detaches the trailing N outputs positionally before consulting that set, so a second run sliced whatever trailed the graph — by then a genuine user output — and the saved program failed to load with treespec.unflatten ... leaves has length 0. Per-buffer state cannot separate the two cases: torch.export declares a mutation it recognises and leaves the trailing value, while a previous run of this pass declares it and consumes the value. So the consuming step records gm.meta["_copyback_mutations_declared"] and the slice is skipped when it is set; create_trt_exp_program records it too. Latent through today's entry points, but the docstring asserted the property, which is how a future caller would get it wrong.

Metadata-only engine read. _get_engine_info_for_node ended in a hand-rolled engine_obj.__getstate__() — the pickling hook, which re-serializes the whole ICudaEngine and base64-encodes it. Two callers read metadata and nothing else, and both now pass metadata_only=True: the declaration pass, which reads ALIASED_IO and the binding names and runs its engine loop before it can know whether any engine has aliased I/O, so a program with no aliased KV paid a full serialization per engine node and then returned unchanged (0.4 s and a 67 MB transient string for one ~50 MB engine); and TensorRTPartitioner._resolve_target_device_for_partition, which reads DEVICE_IDX once per partition on every export that does not pin target_device, the default, behind no early-out at all (967.6 ms against 3.1 ms on a ~100 MB engine, plus a 134 MB base64 string built and dropped).

Uninitialized copy-back buffers. Declaring a copy-back write a BUFFER_MUTATION puts the buffer under ExecuTorch's rule that a mutated buffer has a meaningless initial state — only shape and dtype are serialized — so a buffer read before it is written starts from whatever the allocation held. ExecuTorch warns generically and names no buffer; export() now names them and states the rule they fall under. It offers no remedy, and must not: ExecuTorch's own InitializedMutableBufferPass sets et_init_buffer, after which the emitter marks the spec const and serializes the buffer by reading its storage host-side through ctypes.cast(spec.storage.data_ptr(), ...) — which segfaults the process, rather than raising, whenever the buffer is CUDA-resident, and exporting a model on CUDA is the ordinary way to reach a TensorRT engine. test_export_warns_that_copyback_buffers_load_uninitialized asserts the pass stays unnamed, and fails against both earlier drafts of this string. Engine-aliased KV buffers are excluded on purpose: an aliased write is a cache-position write on the sequence axis, such a model is assumed to read the cache under that same position, and initializing it would write the model's largest tensor whole into the .pte for nothing.

Three things worth knowing

Two tests fail on main today, and both are test defects rather than production defects. tests/py/dynamo/executorch/test_api.py::test_public_save_forwards_lowering_kwargs_graphmodule_no_retrace stubs _exporter.export to return a bare object() sentinel; #4445 added a _declare_aliased_kv_mutations_on_ep(exp_program) call on that same branch, and the sentinel flows straight into it. tests/py/dynamo/executorch/test_edge_cases.py::test_save_as_executorch_uses_public_lowering_and_persists_data asserts export.assert_called_once_with(source, partitioners=..., compile_specs=...), while _save_as_executorch forwards eight — #4433 added transform_passes, compile_config, constant_methods and generate_etrecord, #4336 added weight_streaming_budget_per_engine. Both reproduce at this head unchanged. A fix is prepared and will be sent as its own PR rather than folded in here.

origin/main does not pass its own isort pre-commit hook. With the pinned isort 6.0.0 and the repo's profile = "black", py/torch_tensorrt/executorch/_export.py and tests/py/dynamo/executorch/test_weight_streaming_budget.py both fail --check-only, for the same import: WEIGHT_STREAMING_BUDGET_COMPILE_SPEC_KEY is ordered after normalize_weight_streaming_budget_per_engine in a from torch_tensorrt.executorch.partitioner import (...) block, and isort wants the constant first. Traced to #4336. Not fixed here, but this PR touches both files, so a PR-shaped hook run reports them.

CI is not checking any of this. Noted from your comment: executorch-runtime-test is gated on executorch-runtime-build, which is failing on main, so the ExecuTorch tests are skipped rather than run, here and on main. Nothing above changes that.

[shoumikhin]

Re-reviewed at 697cb65756. The three earlier items look closed to me: the seam tests fail when either half is deleted, the shared resolve_slice_scatter_write makes the docstring true, and the cost note is in. Thanks for the writeup mapping them.

Five new things, ordered by what I would fix first. Everything below I ran on this head and on the merge base 08276e78ae.

1. A buffer can be classified copy-back and still be engine-aliased, and the model then raises on every call

py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py:115

When the write reads the buffer through a clone, classification sees a call_function at args[0], files the write as copy-back and appends the new value as a trailing graph output. remove_input_alias_fixing_clones runs later inside post_lowering and deletes that clone, so the converter sees a direct network input and emits the KV layer. The runtime treats the trailing aliased output as a side effect and truncates it, while the outer graph still reads index 1.

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("k", torch.zeros(1, 2, 16, 4).cuda())
    def forward(self, x):
        k = self.k.clone()
        k[:, :, 3:4, :] = x
        self.k.copy_(k)
        return k.sum(dim=-1)

HEAD  copyback ['k'], aliased_io {'output1': ('buf_k', 'kv_cache_update')}
      IndexError: index 1 is out of bounds for dimension 0 with size 1
BASE  1 output, matches eager, buffer matches eager

Same failure on an 8 head decoder step with a 128 slot cache at the default min_block_size=5, so it is not confined to toy graphs or to min_block_size=1. Clone, scatter, write back is the ordinary shape when the post-write cache is also read in the same forward.

assert_predicted_kv_aliased cannot catch it. predicted_kv_bindings is empty, so it returns at the early exit. It asserts that a predicted KV write must alias; the direction that fails here is that a copy-back buffer must not alias.

Two suggestions, and I would take both:

• Make the cross-check bidirectional. If a buffer in _copyback_mutation_buffers turns up in aliased_io, raise there rather than returning a module that cannot run. That is a few lines and it turns a broken module into a clear error for whatever the next divergence turns out to be.
• Peel a clone of a placeholder in _kv_write_will_alias so the case is classified KV and works. I checked that this single change makes head match base on the repro above.

Either way the guard on line 115 needs a test. I removed it and the suite still passed.

2. test_strided_write_still_takes_the_kv_path pins a miscompile

tests/py/dynamo/lowering/test_buffer_lifting.py:613

I agree the strided case predates this PR. I measured it on the merge base and it is identical: cache[:, :, 0:8:2, :] = x writes slots 0, 1, 2, 3 rather than 0, 2, 4, 6, and the user output sum is unchanged so nothing surfaces it.

My concern is the test rather than the bug. As written it asserts the wrong answer is the expected answer, so the eventual fix arrives looking like a regression and has to delete a passing test to land. If the intent is to freeze the shared derivation rather than the classification, a test on resolve_slice_scatter_write returning step unchanged does that without blessing the KV routing. Otherwise I would drop it and leave a comment on the step line saying the case is known wrong and unfixed.

3. The serialized-engine error recommends an API that does not do what the error says

py/torch_tensorrt/dynamo/_compiler.py:2121 ends with "Use torch_tensorrt.dynamo.compile, which performs the write-back itself." It does not write back, see item 5. The raise itself is right, only that sentence is wrong. Something like "Use torch_tensorrt.dynamo.compile and save the result, which declares the buffer mutation for the runtime to apply" would be accurate.

4. The cross-check message points at min_block_size when that is not the cause

py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py:214

A KV shaped write whose result nothing else consumes never reaches an engine, so it raises at head at both min_block_size=1 and the default 5, while the merge base compiles and runs:

def forward(self, x):
    self.cache[:, :, 3:4, :] = self.lin(x)
    return x.sum() * 2.0

The message says min_block_size most likely rejected the subgraph and that min_block_size=1 rules that out. Here the write is dead rather than partitioned out, so the advice leads nowhere. Raising is still better than the silent drop on base. The suggestion is to check whether the write survived into any partition before naming min_block_size, and to say so when it did not, or to route the dead case to copy-back so it does not raise at all.

5. compile() changes the public output arity, and the extra value is not usable in process

py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py:389

This is the arity half of what @cehongwang raised on 7 Aug about the Torch-TensorRT runtime being left wrong. Measured:

HEAD  copyback ['cache'], 2 return values, module buffer unchanged
BASE  1 return value, module buffer unchanged

Any torch_tensorrt.dynamo.compile or torch_tensorrt.compile(ir="dynamo") user with one non-KV mutable buffer goes from one return value to two. The value is numerically right, but nothing writes it back, so in process the buffer is stale exactly as it was before. torch.compile(backend="tensorrt") is unaffected.

This is a design call rather than a patch and I would not rush it. My preference is to keep the value in gm.meta and out of the return, since the exporters are the only consumers and a public extra output reads like the mutation was handled when it was not. Appending a real copy at the end of the graph, as @cehongwang suggested, fixes it properly for the runtime too. Either is fine. What I would ask is that the decision land with this PR rather than after it, because once the extra output ships, taking it back is itself a breaking change. If it stays, please say in the description that the arity changes and that the buffer is not updated.

Smaller

• _buffer_lifting.py:135 says returning False keeps the copy_. It is erased at line 377 on both branches.
• Docstrings say an ineligible write is lowered to a non-aliasing scatter. Two of the four statuses raise instead.
• The description quotes a warning string I cannot find in the code.

Not from this change

The red checks are red on main with the same failure names. Open-ended slices, returning the written value directly, and one value written into two buffers all fail the same way on the merge base.

One coverage note: tests/py/dynamo/executorch/** runs only under trt_tier_executorch, reached only from executorch-runtime-test, which is gated on executorch-runtime-build. That build is red on main, so the tier executes nowhere and the copy-back path has no CI coverage even though the tests exist.

[Conarnar]

Changes since 697cb65756, against the five numbered items.

1. Copy-back classification that the engine aliases anyway

Both suggestions are in.

The cross-check is bidirectional. A buffer in _copyback_mutation_buffers that turns up in aliased_io now raises, alongside the existing direction. lift_mutated_buffers records the copy-back writes' input bindings in gm.meta["_copyback_bindings"], and compile() hands them to assert_predicted_kv_aliased as copyback_bindings=. The message states the consequence — the runtime treats an aliased output as an engine side effect and does not return it, so the surrounding graph reads past the end of the results and the module raises on every call — and names remove_input_alias_fixing_clones erasing a clone of the cache as the known way to get there. convert_exported_program_to_serialized_trt_engine has nothing to pass in that direction: it rejects copy-back buffers outright a few lines after lifting.

The clone is peeled. _effective_cache_input returns the node the converter will be handed. It peels an aten.clone.default only when the clone's own input is a placeholder and that placeholder's sole user is the clone — remove_input_alias_fixing_clones's own condition, reproduced exactly. A placeholder some other node also reads keeps its clone, the converter sees a call_function, and the write stays copy-back. Classification also had to move to after the trailing copy_ is erased: the copy_ is the last reader of the buffer placeholder in the ordinary write-only shape, so with it still in place the clone always looks shared and the peel never fires.

On the item-1 repro at min_block_size=1, and on an 8-head/128-slot decoder step at the default 5: copyback empty, both caches in aliased_io, one output, output and buffer both matching eager — the merge base's answer.

The guard has tests. TestCacheMustReachTheConverterAsANetworkInput, nine of them: three positive — a direct placeholder, and a sole-use clone of one, classify as KV, for slice_scatter and for index_copy; five negative — a shared clone (once through each op), a clone of a non-placeholder, a cache read through another op, and a cache still read from a get_attr do not; and one pinning that the index_copy validator still reads node.args[0] when nobody overrides it. Two more, in TestCompileSeam, build real engines from a cloned slice_scatter write and a cloned index_copy write and check the buffer afterwards. Those two cannot pass without the peel by construction: drop it and the write is filed copy-back, the pass erases the clone anyway, the engine aliases the buffer, and the new copy-back direction of the cross-check fails the compile.

The item-1 repro used slice_scatter, and that turned out to be half of it. _index_copy_kv_eligible applies its own args[0].op == "placeholder" check, so peeling inside _kv_write_will_alias fixed slice_scatter and left index_copy — the per-step decode write, the shape this path exists for — still filed copy-back while the converter aliased it. The same failure, reached through the second op. Fixed by giving _index_copy_kv_eligible a keyword-only input_node override that the classifier passes and the partitioner does not: by the time it runs as a capability validator, lowering has settled what the input is, so its behaviour there is unchanged, and test_index_copy_validator_still_reads_args0_by_default pins that. Found while building the end-to-end run rather than from the review.

2. test_strided_write_still_takes_the_kv_path

Dropped. test_step_is_returned_unchanged takes its place: resolve_slice_scatter_write returns step unchanged for step in 1, 2, 3 — the shared derivation's contract to its two callers, frozen without blessing the routing built on it.

One correction to the framing, in the direction that matters: the miscompile is confined to the KV fast path, not to strided slice_scatter in general. slice_scatter() reads step when it builds the fallback's scatter indices (np.arange(start, end, step)), so a strided write that misses the KV path is lowered correctly. test_fallback_step_two drives step=2 on a 4-D cache at dim=2, forced off the KV path, and compares against eager; it passes. What ignores step is try_emit_kv_cache_update, which takes no step parameter at all and writes its slots consecutively from start. The comment on the OK return says that now — known wrong, unfixed, and only on the KV path, with test_fallback_step_two named — so the next reader is not sent at code that is correct, and does not "fix" the OK return in a way that breaks the fallback. The test docstring and the PR description were corrected the same way.

Behaviour is unchanged either way: the decision that step != 1 stays KV-eligible is a pre-existing one, and nothing on the conversion path moved.

3. The serialized-engine error's remedy

Corrected to the wording item 3 proposes, after checking both halves of it against the code — compile() plus save() does record a BUFFER_MUTATION for the buffer, and something that loads the saved program does apply it. It says "the runtime" rather than naming ExecuTorch because ep.module() re-lifts the declaration into a copy_ and applies it in PyTorch too. The comment above the raise had the same defect in a second form: it said the buffer is updated afterwards without saying by what, and said this entry point does not perform that update, implying some other entry point does. It now names the serialization step.

test_serialized_engine_rejects_copyback_buffers asserts the buffer name in its rendered form — ['state'], because a bare state also matches stateless and state_dict — and deliberately does not assert the remedy's wording, since pinning the words would break on a correct reword and pass on a differently wrong one. What holds the advice to account is following it end to end: test_executorch_export_declares_copyback_for_every_source_shape, and TestCompileSeam::test_saved_program_declares_the_hidden_copyback_mutation, which compiles, saves, and asserts the BUFFER_MUTATION arrives.

4. The cross-check message and min_block_size

Both remedies item 4 offers are in. The second is the substantive one.

The dead case is routed to copy-back, so it no longer raises. The item-4 repro compiles at the default min_block_size and at 1, with an empty aliased_io, and the ExecuTorch runtime applies the buffer — run end to end on a real runner, and pinned by test_compile_routes_a_dead_slice_scatter_write_to_copyback (both block sizes) and test_saved_program_declares_the_dead_write_mutation. lift_mutated_buffers re-routes a write it classified KV when the erased copy_ was that write's only consumer: copy-back re-attaches the value as a graph output, which makes the write live again instead of leaving it for dead-code elimination.

That is only safe if the converter then leaves the buffer alone, because a buffer that is both copy-back and engine-aliased is exactly the every-call failure item 1 describes. So the re-routed write carries node.meta["_trt_no_kv_alias"], and both eligibility checks honour it — impl.slice_scatter off ctx.current_node, _index_copy_kv_eligible off the node it is handed, which is a second check and a second call site rather than the same one twice.

The marker is deliberately narrow: it goes only on a write the classifier knows the converter would otherwise have aliased and whose result nothing consumes. Marking every copy-back write would be simpler and would disarm the very cross-check direction item 1 asked for, whose job is to catch a write filed copy-back that the engine aliased anyway. test_a_live_kv_write_is_not_marked and test_a_non_kv_write_is_not_marked hold that boundary from the other side. The marker also rides on node.meta through every pass in post_lowering, where a pass that rebuilt the node without its meta would drop it silently and hand back the every-call failure, so compile() reconciles the names recorded before lowering against the lowered graph before conversion (assert_no_kv_alias_markers_survived), in both directions, and names which of three faults it found: a marker lost, a marker gained, or a marker that travelled to a replacement node.

The message no longer recommends a setting the caller already has. Above 1 it reports the value in force and that min_block_size=1 rules the setting out; at 1 it names a converter or a capability validator rejecting the op instead, since blaming the value it would have to recommend is no advice. test_message_names_min_block_size and test_message_does_not_offer_min_block_size_1_when_it_is_already_1 pin the two branches.

That is the whole of item 4's first suggestion that survives. The other half — checking whether the write reached a partition and saying so when it did not — was built and then dropped as unreachable: once a dead write is filed copy-back, the case that branch existed to report cannot raise here at all.

5. compile()'s public output arity

Option A, the preference item 5 states. hide_copyback_outputs installs a torch.fx.graph.CodeGen subclass whose generated forward stops short of the trailing copy-back values. The values stay on the graph's output node, so the exporters still find them and still declare the BUFFER_MUTATION; only what the module returns is shortened. Arity matches the merge base again — the item-5 repro goes from two return values back to one.

The strongest evidence that nothing else moved with it: the .pte exported from the same model is byte-identical (matching md5) to the one exported at 697cb65756, while the public return went from two values to one. Unchanged alongside it: the output specs from save(retrace=False), save(retrace=True) and torch_tensorrt.executorch.export(), and torch.compile(backend="tensorrt")'s arity.

Two defensive calls were drafted and then removed, because the property they were guarding does not exist on this torch: re-exposing the outputs inside _exporter.transform() and around _compile.py's retrace, on the theory that a retrace reads the outputs off the module's return. It does not — a non-strict torch.export runs a GraphModule through torch.fx.Interpreter, which reads the output node and ignores the codegen, and non-strict is what both retrace paths get. Deleting each call in turn left every test green, including one that pins the surviving USER_OUTPUT. Rather than ship two undemonstrable calls, the property is pinned by tests instead — test_graph_consumers_still_see_the_hidden_value, test_retracing_exporter_sees_the_hidden_copyback_value, and test_saved_program_declares_the_hidden_copyback_mutation under retrace=True — so a future torch that honours the codegen on export fails loudly rather than losing the declaration quietly.

One case the truncation arithmetic depends on now has its own test. hide_copyback_outputs drops the last N values, which is right only while the copy-back values really are the trailing ones, and two mechanisms append outputs: lift appends copy-back last, and the interpreter appends an aliased output per KV write inside each submodule. test_compile_hides_copyback_beside_an_engine_aliased_kv_buffer compiles one module holding both kinds of buffer and checks that the return is the model's arity, the KV buffer is aliased and written in place, and the one extra value on the graph is the copy-back buffer's new contents rather than a KV output.

What has not changed is the in-process staleness item 5 measured: a module called directly from PyTorch leaves its buffer where the merge base left it. That is now stated in the _buffer_lifting.py module docstring rather than hinted at by an extra return value.

Smaller

• _kv_write_will_alias's comment that returning False keeps the copy_: corrected. The copy_ is erased on both branches; what False buys is the new value being re-attached as a graph output instead.
• The docstrings saying an ineligible write is lowered to a non-aliasing scatter: corrected in _kv_write_will_alias, in lift_mutated_buffers and in the TestCopyBackClassification class docstring. For slice_scatter it is one of four outcomes — a non-aliasing scatter under OK but failing _kv_eligible, the source returned under FULL_OVERWRITE, NotImplementedError under DYNAMIC_BOUNDS, IndexError under BAD_DIM — and index_copy adds a fifth: its fallback is a scatter but itself raises for a dynamic index or an unsupported dim. The two per-test docstrings that say "scatter" are left alone, since each describes one specific model and both of those models do take the scatter fallback.
• The description quote item 3's list could not find: it was stale, and so was the paragraph built on it, which argued the warning names no remedy on purpose. The shipped warning does name InitializedMutableBufferPass — attributed to ExecuTorch's own warning, and never separated from the caveat that the pass reads the buffer host-side to serialize it, so it works while the buffer is on CPU and takes the export down without raising once the buffer is CUDA-resident. The description now quotes what the code says, and the test asserts the name and the caveat travel together.

Changed without being asked

• Classification now runs after the trailing copy_ is erased rather than before. It is a prerequisite for item 1's peel rather than a change of behaviour on its own, but it changes which graph the classifier reads, so it is worth naming separately: with the copy_ still in place the buffer placeholder has two readers and every clone of it looks shared.
• remove_input_alias_fixing_clones carries a note back to the classifier that reproduces its condition. Loosening that condition, or deleting the pass as its own TODO invites, silently invalidates the prediction, and the failure mode is the module that raises on every call rather than a lost optimization.

Worth knowing

The CI observation is right, and by two routes rather than one. tests/py/dynamo/executorch/ is reached by trt_tier_executorch in tests/py/utils/ci_helpers.sh, invoked only from executorch-runtime-test, which is gated on needs.executorch-runtime-build.result == 'success' and that build is red on main; and by the executorch suite in tests/ci/suites.py, which is tier l2 with lanes=("nightly",), so it is out of the fast lane a PR push resolves to and out of the full lane an approval resolves to. Either way the exporter and save-path tests execute nowhere on this PR. The other half is better off: tests/py/dynamo/lowering/ is the dynamo-lowering suite, tier l0 with lanes=("fast", "full", "nightly"), so the classifier, the cross-check, the marker reconciliation and the compile-seam tests do run on every push.

Not a substitute for that, and not claimed as one, but here is what was actually run. A wide sweep over tests/py/dynamo/lowering/, tests/py/dynamo/conversion/ and tests/py/dynamo/runtime/ gives 2 failed, 2601 passed, 69 skipped; both failures were reproduced at 697cb65756 and are unrelated to this branch (test_aten_lowering_passes.py::TestComplexSubgraph::test_complex_subgraph and test_hf_static_cache_xfail.py::TestHFStaticCacheCurrentLimitations::test_compile_fails_with_known_error). The sweep needs -n 8 --dist=loadfile; the repo's default -n auto puts one CUDA context per CPU on the single visible GPU and errors out before it collects anything.

Beyond the suite, these changes were exercised end to end by running .ptes on a real ExecuTorch runner: a synthetic copy-back model; an L=2 Gemma4 MoE with an added non-KV buffer, which puts copy-back and engine-aliased KV in one model across the TensorRT and CUDA backends; a Qwen3.5 GDN decode, unmodified, whose four per-layer conv_* / rec_* state buffers are all real copy-back with no KV aliasing anywhere in the model — run with input_pos held constant, so the only possible source of step-to-step variation in the logits is the copy-back buffers, and the logits do vary and converge as a decaying recurrence does; and the item-1 model as an index_copy decode step, which produces no .pte at all at 697cb65756 and runs to a correct cache at the current head. The item-4 dead-write model runs the documented ramp at the default min_block_size.