fix(executorch): lift nested and partitioner-excluded mutable buffers - #4470
Open
shoumikhin wants to merge 4 commits into
Open
fix(executorch): lift nested and partitioner-excluded mutable buffers#4470shoumikhin wants to merge 4 commits into
shoumikhin wants to merge 4 commits into
Conversation
…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).
Tests cover serialization round-trip, the exposure-flag dispatch across both
retrace modes, the buffer-mutation declaration, and the partitioner un-tagging.
…T delegate lift_mutated_buffers erased every copy_ that mutates a lifted buffer, assuming the write-back happens through engine-level aliasing. That is true only for KV-cache writes (slice_scatter / index_copy), which the converter lowers to an IKVCacheUpdateLayer with aliased I/O. Any other mutable buffer -- e.g. a Gated DeltaNet conv_state ring-shift -- has no such aliasing, so erasing its copy_ silently dropped the write-back and the runtime delegate received too few args. Distinguish the two kinds: KV writes keep the existing zero-copy aliasing path; a non-KV mutation has its new value re-attached as an ordinary BUFFER_MUTATION graph output so ExecuTorch copies it back to the caller-owned buffer after the delegate runs. Threaded through both save paths -- create_trt_exp_program (retrace=False) and _declare_aliased_kv_mutations_on_ep (retrace=True) -- via a _copyback_mutation_buffers list carried on gm.meta. Tests (CPU-only, gated on executorch.exir): - test_buffer_lifting.py: KV writes stay aliased (no copy-back); a non-KV mutation is recorded and re-attached as the trailing output; index_put falls into copy-back; mixed KV + non-KV records only the non-KV buffer. - test_kv_cache_export.py: both exporter passes reclassify a trailing copy-back output to BUFFER_MUTATION ahead of the user outputs.
Four problems that together stop a model from exporting when its caches live in
submodules or when part of the graph runs on another backend. The later ones only
become visible once the earlier ones are fixed.
1. A get_attr target is fully qualified, so a buffer owned by a submodule arrives
as "layers.0.self_attn.kv_cache.k_cache". hasattr and getattr do not walk a
dotted path, so every nested buffer is reported missing and skipped, and the
copy-back path never runs. Use get_buffer, which resolves through submodules.
2. register_buffer rejects dots, so a nested buffer is renamed to lifted_buf_*,
but the recorded mutation targets keep the original dotted names. The verifier
then rejects the program:
SpecViolationError: Buffer output getitem_1 does not point to a buffer
that exists
Remap the recorded targets through the same mapping, and keep the mapping in
meta so a later consumer does not have to reconstruct the renaming rule.
3. The aliasing prediction never sees torch_executed_ops, so a write the caller
excluded from TensorRT is still predicted to alias and its copy_ is dropped.
compile() then fails its own cross-check, reporting that the write-back would
be silently dropped. An op that never reaches a converter cannot emit an
IKVCacheUpdateLayer, so honor the exclusion list, matched the way the
partitioner matches it so the two cannot disagree.
4. Even with the prediction corrected, lifting an excluded write turns the
mutation into a delegate output that feeds the other delegate. ExecuTorch
cannot express a buffer mutation consumed across a delegate boundary:
RuntimeError: Tried to erase Node getitem_145 but it still had 1 users
in the graph: {executorch_call_delegate_1: None}
If TensorRT is not converting the write, leave the buffer alone entirely so
the copy_ stands for whichever backend claims it.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this fixes
Four problems that stop a model from exporting when its caches live inside submodules, or
when part of the graph runs on a different backend. They are in one commit because they
block the same thing and are in the same function, and because the later ones only become
visible once the earlier ones are fixed.
1. A buffer owned by a submodule is skipped
A
get_attrtarget is fully qualified, so a cache owned by a submodule arrives aslayers.0.self_attn.kv_cache.k_cache. Neitherhasattrnorgetattrwalks a dotted path:Every nested buffer is reported missing and skipped, so the copy-back path never runs. Any
real transformer owns its caches per layer, so this is the common case rather than an edge
case.
2. Renamed buffers leave the mutation targets dangling
register_bufferrejects dots, so a nested buffer is renamed tolifted_buf_*. Therecorded mutation targets keep the original dotted names, which no longer resolve:
Remapped through the same mapping, which is also recorded in
metaso a later consumerdoes not have to reconstruct the renaming rule.
3. The aliasing prediction ignores
torch_executed_ops_kv_write_will_aliasdecides whether a write becomes engine-level aliasing by looking atthe op, and never sees the ops the caller excluded from TensorRT. An excluded write is still
predicted to alias, its
copy_is dropped, and then the cross-check incompile()fails:An op that never reaches a converter cannot emit an
IKVCacheUpdateLayer, so the predictionnow honors the exclusion list, matched with
ConverterRegistry.qualified_name_or_strthesame way the partitioner matches it, so the two cannot disagree about the same op.
That guard is worth calling out: it turned a cache that would silently stop updating into a
clear error naming every affected buffer.
4. An excluded write should not be lifted at all
With the prediction corrected the export gets further, then fails inside ExecuTorch:
Lifting the buffer turns the mutation into a delegate output. When the write runs on the
other backend, that output feeds the second delegate, and ExecuTorch cannot express a buffer
mutation consumed across a delegate boundary. If TensorRT is not converting the write, it
now leaves the buffer alone so the
copy_stands for whichever backend claims it.Testing
The existing suite passes unchanged:
tests/py/dynamo/lowering/test_buffer_lifting.py,19 passed.
Those tests all use a cache owned by the top-level module (a name with no dots), which is
why none of them catch problem 1. Verified the fixes on a two-layer model that owns its
caches at
layers.N.attn.k_cache:End to end, with all four applied, a multi-method program carries both delegates and runs
correctly from C++:
Without them, excluding an op from TensorRT either drops a cache write or fails to lower,
so a model needing part of its graph on another backend cannot be exported.
Happy to split this into four commits if that reviews better, or to fold it into #4459
directly if the author prefers.