feat(executorch): copy-back for non-KV mutable buffers in the TensorRT delegate - #4459
feat(executorch): copy-back for non-KV mutable buffers in the TensorRT delegate#4459Conarnar wants to merge 13 commits into
Conversation
github-actions
Bot
added
component: tests
|
#4459 widens the #4445 ordering blocker into a three-way reversal. lift_mutated_buffers appends copy-back values to the graph output: 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. |
| # index_put KV write correctly falls through to copy-back rather than being | ||
| # dropped in the false expectation of aliasing. (A dedicated index_put -> | ||
| # IKVCacheUpdateLayer converter would let it use zero-copy aliasing instead.) | ||
| _KV_WRITE_TARGETS = { |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
| self.assertEqual(len(lifted), 2) | ||
| self.assertEqual(new_gm.meta["_copyback_mutation_buffers"], ["state"]) | ||
|
|
||
|
|
There was a problem hiding this comment.
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.
Same with #4445, I was not able to reproduce this. |
6a88214 to
a8b0d61
Compare
|
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 |
|
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 |
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. |
|
A KV cache owned by a submodule is skipped, so the copy-back path never runs. I tried this PR on a decoder whose per-layer caches live inside the attention modules, which is how a transformer normally stores them, and the export fails. Two separate problems, both about fully qualified buffer names. The second one is arguably yours to fix since this PR adds the consumer; the first is older but this PR is where it starts to matter. Problem 1: nested buffers are never lifted
buffer_name = get_attr_node.target
if not hasattr(gm, buffer_name):
logger.warning(
"lift_mutated_buffers: get_attr target %s not found on gm; skipping",
buffer_name,
)
continue
buffer_tensor = getattr(gm, buffer_name)A hasattr(gm, "layer.state") # False
gm.get_buffer("layer.state") # the tensorEvery mutation is skipped with a warning, nothing is lifted, and the cache is later frozen as constant data. On my model that was 16 warnings and a silently dead cache.
try:
buffer_tensor = gm.get_buffer(buffer_name)
except AttributeError:
...Problem 2: the copy-back list keeps names that no longer existOnce problem 1 is fixed, the mutations are found and classified, and then the export fails in the verifier:
if "." in buf_name:
attr_name = "lifted_buf_" + buf_name.replace(".", "_")
...
buf_to_attr[buf_name] = attr_nameBut A minimal reproducer, no TensorRT needed: class Attn(torch.nn.Module):
def __init__(self):
super().__init__()
self.register_buffer("state", torch.zeros(4, 8))
def forward(self, idx, val):
updated = self.state.index_put([idx], val)
self.state.copy_(updated)
return updated.sum(-1)
class Model(torch.nn.Module): # cache one level down
def __init__(self):
super().__init__()
self.layer = Attn()
def forward(self, idx, val):
return self.layer(idx, val)
ep = torch.export.export(Model().eval(), (torch.tensor([0]), torch.zeros(1, 8)), strict=True)
gm, lifted = lift_mutated_buffers(ep.module())
gm = inline_lifted_buffers_into_gm(gm, lifted)
print(gm.meta["_copyback_mutation_buffers"]) # ['layer.state']
print(list(dict(gm.named_buffers()))) # ['lifted_buf_layer_state']Note the ordering: on this PR as it stands, that snippet stops at problem 1 and prints With problem 1 fixed it gets further and prints the mismatch above. So problem 1 hides problem 2, and fixing only the lookup turns a silently frozen cache into a hard verifier error. Suggested fixRemap the copy-back list through the same mapping that renamed the buffers, right after gm.meta["_lifted_buffer_attr_names"] = dict(buf_to_attr)
copyback = gm.meta.get("_copyback_mutation_buffers")
if copyback:
gm.meta["_copyback_mutation_buffers"] = [
buf_to_attr.get(name, name) for name in copyback
]Recording the mapping in The KV aliasing path already avoids this, because it reads the name off the Why CI does not catch itEvery test in self.register_buffer("cache", torch.zeros(1, 4, 16, 8))A flat name has no dots, so VerifiedWith both fixes applied I get a working multi-method program: Worth noting the failure mode differs per problem, and only one is loud. Problem 2 is caught by the verifier with a clear message. Problem 1 only logs a warning and then produces a program whose cache silently does not persist, which reads as a model quality issue rather than an export bug. If a discovered mutation target cannot be resolved, raising instead of warning would turn that into an immediate, obvious failure. Environment: torch |
|
Two problems when a buffer write is excluded from TensorRT, which together block a two-delegate program. I hit these while exporting a model that splits across the TensorRT and CUDA delegates. Both are in Problem 1: the aliasing prediction ignores
|
a8b0d61 to
da17e5b
Compare
|
@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). |
da17e5b to
718d0c8
Compare
…ng for hybrid graphs
torch_tensorrt.save(retrace=False) uses the legacy dynamo exporter, which inlines the
partitioned _run_on_gpu (non-TensorRT) submodules back into the graph before building an
ExportedProgram. For a hybrid graph interleaving TensorRT engines with a CUDA/pytorch
delegated op, inline_torch_modules wired each submodule's inputs by MATCHING placeholder
names to graph nodes (get_duplicate_nodes). Name matching binds an input to a same-named
but unrelated node on a collision (e.g. a submodule input placeholder name-matching a
different engine's getitem), which:
- rewires a consumer to the wrong producer and orphans the real one; the orphan is then
pruned by dead-code elimination, leaving a delegate short an output at runtime (an
aliased engine reports "expected N args, got N-1"); and
- for a submodule mixing graph-input and computed-intermediate inputs, leaks the
computed intermediates as spurious graph placeholders (misclassified USER_INPUTs).
Wire submodule inputs POSITIONALLY from the call_module args (gm_node.args, which is
authoritative) instead of by name: let graph_copy create a fresh placeholder for each
submodule input, then rewire each to submodule_inputs[i] by position and erase it. Drop
get_duplicate_nodes (now unused).
Also fix two torch-version-compat gaps this path hits on recent torch:
- lift(): pass an explicit persistent= flag on BUFFER InputSpecs (required since 2.3).
- create_trt_exp_program(): an inlined GraphModule may carry a plain fx.CodeGen (no
pytree_info); fall back to specs rebuilt from the example inputs + graph outputs.
With these, retrace=False export of a hybrid TensorRT+CUDA program is bit-identical to
retrace=True (validated on a 2-layer int4 MoE decode: per-step argmax + logits match).
Tests: tests/py/dynamo/models/test_exporter_inlining.py -- positional input wiring under a
name collision, and multi-output preservation (GPU-free fx unit tests).
Adds end-to-end caller-owned KV-cache support to the ExecuTorch TensorRT
delegate: the KV buffers are owned by the caller above the delegate and threaded
in as mutable-buffer delegate args, instead of being self-allocated inside a
(stateless) TensorRT engine.
Runtime + serialization (delegate):
- serialize each engine's aliased (KV-cache / in-place) I/O into the delegate blob
(serialization.py, backend.py, TensorRTBlobHeader.{h,cpp});
- at runtime bind each aliased TRT output binding to its aliased input's
caller-provided pointer (in-place) and reflect the result into the delegate
output EValue -- a no-op when the memory planner already aliased the two
(TensorRTBackend.{h,cpp}).
Export/lowering (torch_tensorrt):
- expose each engine's aliased outputs as graph-level BUFFER_MUTATIONs so
ExecuTorch keeps the KV buffers as caller-owned mutable buffers: at transform
time for the legacy exporter (retrace=False), and via a post-export pass
(_declare_aliased_kv_mutations_on_ep) for torch.export (retrace=True), which
otherwise truncates the aliased outputs at the fx boundary;
- keep delegate-mutated buffers above the delegate in TensorRTPartitioner
(tag_constant_data would otherwise freeze them as constants).
The retrace=True pass runs for exported_program as well as executorch. The
truncation happens at the fx boundary for every output format, so declaring only
on the executorch path left an exported_program saved with the mutation absent
from its signature while the engine still updated the cache in place. It is
declared before _normalize_engine_constants_to_python, which rewrites the engine
constants the pass reads aliased_io from. retrace=False was already correct for
every format via create_trt_exp_program. aot_inductor stays undeclared and
warns: whether an aliased in-place mutation survives functionalization under
inductor is unverified.
Tests cover serialization round-trip, the exposure-flag dispatch across both
retrace modes, the buffer-mutation declaration, and the partitioner un-tagging.
…mutations
_declare_aliased_kv_mutations_on_ep rebuilds the ExportedProgram to attach the
new output specs, but reconstructed only root/graph/signature/state_dict/
range_constraints/module_call_graph/constants. example_inputs and verifiers are
not recoverable from the graph and reset to their defaults when omitted, so the
pass was not a drop-in replacement for the program it rewrites.
That was invisible while the pass ran only on the executorch path, since to_edge
does not read either. Declaring mutations for exported_program as well makes it
reachable: torch.export.save then persists a program whose example inputs are
silently gone. AOTI refuses such a program outright ("exported_program.
example_inputs is required to be set in order for AOTInductor compilation"), so
this also has to be fixed before that format can ever declare mutations.
Carry both through. The stub programs in test_kv_cache_export.py now model them,
and the capturing test asserts they reach the constructor.
Reported by cehongwang in review.
Whether an aliased KV output gets declared a BUFFER_MUTATION is decided in two places: the exporter (the legacy one declares at transform time, via create_trt_exp_program) and save()'s per-output-format branch. Nothing reconciles them, so a program that arrives already declared is declared a second time. The duplicate spec then fails the ExportedProgram verifier's output ordering check. Seed already_exposed from the incoming signature's BUFFER_MUTATION targets rather than an empty set, so the pass skips buffers that are already declared and returns the program untouched when nothing new remains. That covers every combination that reaches it, including exported_program with use_legacy_exporter=True, which the preceding commit made reachable. The added test drives the pass on an already-declared program with ExportedProgram monkeypatched to raise, so a regression fails on "rebuilt the program" rather than on some later verifier complaint. The no-op fixture grows an output_specs field, which a real ExportGraphSignature always has. Reported by shoumikhin in review.
Four comments describe a zero-copy path where the memory planner places the delegate's output slot on the aliased input, so the reflect is skipped. That never happens: the delegate input and its aliased output are live at the same time, so the planner cannot co-locate them. Instrumenting the branch over a 30-layer export confirms it -- dst == bind_ptr was false for all 120 aliased outputs across a prefill and a decode step. Two of them go further and call the skipped case the common fast path, which inverts what the code does: every aliased output reflects, and a model with aliased outputs therefore always syncs before returning. Say what actually happens in all four (TensorRTBackend.h, and the reflect list, the reflect loop and the must_sync rationale in TensorRTBackend.cpp). The dst != bind_ptr check itself stays -- it is unreachable today, but it is what stops a self-copy if planning ever changes -- and its comment now says so instead of advertising a fast path. Also two comment fixes found in the same pass: "These two" in _exporter.py referred forward to checks the reader had not reached yet, and a test comment said "previously-dropped" where it meant the output torch.export truncates. No behaviour change. Reported by shoumikhin in review.
The aliased-output branch discarded the resize_tensor result and skipped the resize altogether when the rank was out of range, while the sibling non-aliased branch treats both as fatal. et_alias_out.nbytes() is read on the next line and sizes both the reflect D2D copy and, through the delegate output EValue, ExecuTorch's write-back copy_. A dynamic aliased output that outgrows its planned size would therefore move the stale planned byte count in both, silently truncating the cache update rather than failing. Mirror the sibling branch: reject an out-of-range rank and propagate a resize_tensor error. Reported by shoumikhin in review.
…th too Whether an aliased KV output is declared a BUFFER_MUTATION depends on two independent switches. save() picks the exporter (retrace, plus an optional use_legacy_exporter override) and only the legacy exporter exposes the mutations, at transform time; save()'s per-format branch declares them for everything else. The retrace=False branch did neither, so retrace=False with use_legacy_exporter=False produced a program that silently omits an update the engine performs -- no declaration and no diagnostic. Run the declaration pass on the exported_program and executorch branches there as well. The preceding commit made the pass skip buffers that already carry a spec, so this is correct for either exporter: the legacy one keeps declaring at transform time and the pass returns its program untouched. aot_inductor stays undeclared on both paths, as before. The added test drives save() over both formats and both exporters and asserts the pass runs exactly once; against the previous commit all four parametrizations fail. Reported by shoumikhin in review.
The runtime binds delegate output i to output_binding_names[i]. That holds because the partition's outputs are getitem(engine_node, i) in index order, and nothing verifies it: inputs are checked in _reorder_input_names_for_executorch, outputs were not. arrange_graph_outputs moves buffer mutations ahead of user outputs and is a no-op here only while the mutated buffers stay above the delegate, so a regression there would swap the serialized names silently. Validate the correspondence in preprocess. A single-output engine returned unwrapped is accepted -- one binding has no order to get wrong -- and anything else must be one getitem per binding, in index order. _build_edge_program only ever emitted `output((engine_node,))`, including for its three-output-binding case, which is not a shape that can occur: a three-tuple cannot be consumed as one value. It now emits one getitem per output binding, so the fixtures model what the backend actually receives. Reported by shoumikhin in review.
…d_io aliased_io changes what a blob means. A parser that predates it binds each aliased output to its own allocation instead of the input it aliases, so it does not fail -- it returns wrong results. The magic is the only field that parser validates, so it is the only thing that can make the skew fail closed. Emit TR02 when metadata.aliased_io is non-empty and keep TR01 otherwise, rather than bumping unconditionally: a blob with no alias map means exactly what it meant before, so it stays loadable by an older runtime. Both magics are accepted on read, so new runtimes still load existing artifacts. Verified against a real older build rather than a simulated one: a TR02 blob loads and produces the expected KV-persistence result on a runtime built from this branch, and a runtime built before the change rejects the same blob with "failed to parse TensorRT blob". Reported by shoumikhin in review.
…r verify The caller-owned KV path had no CI coverage: kv_cache_decode_check ships in the release tarball and is defined in the packaged CMake project, but nothing built or ran it, so a regression in the aliased binding would only surface downstream. Export a decode .pte alongside the static-shape one and pass it to the verify script as an optional second argument. When present the script builds kv_cache_decode_check from the unpacked tarball, runs it, and requires the persistence assertion to pass; the same no-libtorch link check the example runner gets is applied to it. Without the argument the script behaves as before. Also assert the tarball ships kv_cache_decode_check.cpp, next to the existing entries, so the packaging contract is checked rather than assumed. Reported by shoumikhin in review.
…ering claim Four `continue`s in _declare_aliased_kv_mutations_on_ep left an aliased output undeclared without saying so, and the mismatch only surfaced later as a delegate arity error at execute. Log at each, at the level the case warrants: warn when the persisted alias map disagrees with the engine's bindings (unknown input name, or an index past the delegate args) and when the aliased input is not a registered buffer, since all three leave the engine with an output binding the delegate cannot satisfy; debug when the buffer already carries a spec, which is the expected idempotent skip. The non-aliased output path stays silent -- it is the common case, not a fault. _reorder_input_names_for_executorch's docstring also justified skipping the output reordering by claiming a TensorRT partition has no mutation outputs. That has not been true since aliased I/O landed. The order does survive lowering, but for a different reason: _keep_mutated_buffers_above_delegate keeps mutated buffers out of the delegate, so ExecuTorch records the mutation as a USER_OUTPUT and arrange_graph_outputs computes the identity permutation. Say that, and note the guarantee is conditional -- _validate_output_binding_order is what enforces it. Reported by cehongwang in review.
…T delegate A mutable buffer with no engine aliasing -- a convolution-state ring buffer, say -- is mutated outside the engine, so the KV aliasing path does not cover it. lift_mutated_buffers appends its new value as a trailing user output; reclassifying that output as a BUFFER_MUTATION of its buffer makes ExecuTorch copy it back after the delegate runs. Thread the buffer FQNs from gm.meta['_copyback_mutation_buffers'] into _declare_aliased_kv_mutations_on_ep on the retrace=True executorch path, and have create_trt_exp_program do the equivalent at transform time on the legacy path. The reclassification is positional: it takes the trailing len(copyback_buffers) outputs. That cannot skip an already-declared buffer the way the KV path's already_exposed set does, because after a first declaration the mutations sit at the front and the trailing outputs are the user's -- a second pass would retarget those, dropping a user output and leaving the buffer with two mutation specs. Reject that combination instead of silently corrupting the signature.
…program Copy-back was threaded only on the executorch branch, so saving as exported_program produced a signature that omits an update the program performs. Hoist _copyback_bufs to the retrace=True block and pass it on both format branches. retrace=False is left alone deliberately. Only the legacy exporter declares copy-back there, at transform time, and the pass cannot pick it up afterwards: its reclassification is positional, so running it on an already-declared program would retarget the wrong outputs. That leaves retrace=False with use_legacy_exporter=False unable to declare copy-back at all, which now warns instead of silently dropping it, mirroring the aot_inductor warning above.
718d0c8 to
1d2ed27
Compare
3 participants
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'sIKVCacheUpdateLayeraliasing to write the new value back to the caller-owned storage (zero-copy).That assumption only holds for KV-cache writes. The
slice_scatterandindex_copyconverters have a fast path that emits anIKVCacheUpdateLayerwhose output is aliased in-place to the cache input. Any other in-place mutable buffer has no such aliasing — for example theconv_state/recurrent_statering-buffers of a Gated DeltaNet (GDN) layer. For those, erasing thecopy_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:slice_scatter/index_copy) keep the existing zero-copy aliasing path — thecopy_is erased and the write-back is handled by the engine'sIKVCacheUpdateLayer.BUFFER_MUTATIONof 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_copythe converter cannot turn into anIKVCacheUpdateLayerfalls to copy-back instead of being dropped.assert_predicted_kv_aliasedthen checks at compile time that every write predicted as KV really does appear in an engine'saliased_io, turning a mis-prediction into a loud error rather than a silently lost write-back.How it works
The classification happens once, in
lift_mutated_buffers, and is threaded to both save paths:dynamo/lowering/_buffer_lifting.py— for each lifted buffer, a KV write is left to aliasing; a non-KV write has its new value appended as a trailing graph output (so it survives DCE now that thecopy_is gone) and its buffer name recorded ingm.meta["_copyback_mutation_buffers"].dynamo/_compiler.py— forwards that list onto the compiled module's meta so it reaches the exporters.dynamo/_exporter.pycreate_trt_exp_program(retrace=False): tags the trailing outputs with their buffer target and moves all mutation outputs ahead of the user outputs (verifier requirement)._declare_aliased_kv_mutations_on_ep(retrace=True): reclassifies the lastlen(copyback_buffers)user outputs fromUSER_OUTPUTtoBUFFER_MUTATIONand rebuilds the top-levelout_specsoto_edge's unflatten sees the right leaf count.Where copy-back is declared
retraceexported_programexecutorchaot_inductoruse_legacy_exporter=FalseUnder
retrace=Truethe declaration pass now runs on theexported_programbranch as well asexecutorch, so both formats carry a correctBUFFER_MUTATION.aot_inductorstays undeclared, matching #4445's KV scoping.retrace=Falsedeclares copy-back only through the legacy exporter, at transform time. The pass cannot pick it up afterwards: its reclassification is positional, so running it on a program the legacy exporter already declared would retarget outputs that are no longer the trailing ones.retrace=Falsewithuse_legacy_exporter=Falsetherefore cannot declare copy-back at all, and now emits a warning rather than silently saving a signature that omits the update.The change is a no-op for any model without a non-KV in-place mutable buffer — including all standard models and all KV-cache models. A KV-only export is byte-identical to before.
Behaviour change worth flagging
Declaring copy-back twice now raises. The positional reclassification cannot detect an already-declared buffer the way the KV path's
already_exposedset does: after a first declaration the mutations sit at the front and the trailing outputs are the user's, so a second pass retargets those — dropping a user output and leaving the buffer with two mutation specs. That combination is reachable today viaretrace=True+use_legacy_exporter=True+ dynamic shapes, where the legacy exporter declares at transform time and the save path would declare again. It previously produced a corrupted signature silently; it now fails with a clear error.Testing
Unit tests added in this PR (CPU-only, gated on
executorch.exir):tests/py/dynamo/lowering/test_buffer_lifting.py— classification: KVslice_scatter/index_copywrites 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_putfalls to copy-back; a mixed KV + non-KV graph records only the non-KV buffer; an ineligibleindex_copy(2-D cache) and an ineligibleslice_scatter(wrong dim) both fall to copy-back rather than being dropped. Plus theassert_predicted_kv_aliasedbackstop: passes when the prediction is aliased, raises when it is not, aggregatesaliased_ioacross multiple engines, and no-ops with no prediction.tests/py/dynamo/executorch/test_kv_cache_export.py— both exporter passes reclassify a trailing copy-back output toBUFFER_MUTATIONahead of the user outputs; a write-only copy-back buffer still reachesliftas aget_attr(drives the reallift/ExportedProgramconstructor so a regression reproduces the verifier error rather than passing vacuously); theget_attrre-add does not duplicate an existing one or invent a dangling one; nested (dotted) buffer names are remapped to their flattened form; declaring copy-back twice raises; andsave()warns exactly on theretrace=False+ non-legacy combination.The copy-back path was also exercised end-to-end during development — a KV-only model as a no-regression sentinel, and a hybrid TensorRT + CUDA decode model with a convolution-state copy-back buffer (both retrace modes) — but those runs are not part of this PR's automated tests.
Stacking
Stacked on #4445 (caller-owned KV-cache), which is stacked on #4446 (retrace=False legacy-exporter fix). Should land after both.
Follow-ups / interaction with in-flight PRs
_compile.pyintotorch_tensorrt.executorch.export(). The retrace=True re-declaration (including thecopyback_buffersthreading added here) currently lives in_compile.py's save path and will need re-homing into the newexport()— this applies to fix(executorch): support KV-cache aliased I/O in the TensorRT delegate #4445's KV re-declaration as well. The retrace=False path is unaffected. feat(executorch): expose composable Edge export API #4440 also introduces per-method partitioners, so the re-declaration will need to run per method.