Skip to content

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

Open
shoumikhin wants to merge 4 commits into
pytorch:mainfrom
shoumikhin:fix-nested-buffer-lifting
Open

fix(executorch): lift nested and partitioner-excluded mutable buffers#4470
shoumikhin wants to merge 4 commits into
pytorch:mainfrom
shoumikhin:fix-nested-buffer-lifting

Conversation

@shoumikhin

@shoumikhin shoumikhin commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Depends on #4459. These fixes are on top of that branch, so the file diff here
includes its commits and CI will stay red until it lands. Only the final commit
(fix(executorch): lift nested and partitioner-excluded mutable buffers) is mine,
at +70/-7 across two files.

Overlaps #4472. That pull request fixes the same dotted-name buffer lookup against
main, standalone and with a test. This one carries the fix too, because #4459's branch
does not have it and the remaining fixes here would be unreachable without it. If #4472
lands first, drop that hunk from this one; if this lands first, #4472 becomes a no-op.
Either order works, they just should not both land unreconciled.

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_attr target is fully qualified, so a cache owned by a submodule arrives as
layers.0.self_attn.kv_cache.k_cache. Neither hasattr nor getattr walks a dotted path:

hasattr(gm, "layer.state")     # False, even though the buffer exists
gm.get_buffer("layer.state")   # the tensor

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_buffer rejects dots, so a nested buffer is renamed to lifted_buf_*. The
recorded mutation targets keep the original dotted names, which no longer resolve:

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

Remapped through the same mapping, which is also recorded in meta so a later consumer
does not have to reconstruct the renaming rule.

3. The aliasing prediction ignores torch_executed_ops

_kv_write_will_alias decides whether a write becomes engine-level aliasing by looking at
the 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 in compile() fails:

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

An op that never reaches a converter cannot emit an IKVCacheUpdateLayer, so the prediction
now honors the exclusion list, matched with ConverterRegistry.qualified_name_or_str the
same way the 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:

RuntimeError: Tried to erase Node getitem_145 but it still had 1 users
in the graph: {executorch_call_delegate_1: None}

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:

lifted buffers: 2                        (0 before this change, the lookup failed)
name mapping recorded: 2
  layers.0.attn.k_cache -> lifted_buf_layers_0_attn_k_cache
  layers.1.attn.k_cache -> lifted_buf_layers_1_attn_k_cache
copy-back targets dangling: 0

with the write excluded from TensorRT: lifted 0   (correctly left alone)
without exclusion:                    lifted 2

End to end, with all four applied, a multi-method program carries both delegates and runs
correctly from C++:

prefill   delegates=['CudaBackend', 'TensorRTBackend']
decode    delegates=['CudaBackend', 'TensorRTBackend']
generated tokens vary per step, and logits match eager to 9.6e-05

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.

Conarnar and others added 4 commits July 31, 2026 19:31
…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.
@meta-cla meta-cla Bot added the cla signed label Aug 8, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests component: lowering Issues re: The lowering / preprocessing passes component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: api [C++] Issues re: C++ API component: runtime component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 8, 2026
@github-actions
github-actions Bot requested a review from narendasan August 8, 2026 23:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [C++] Issues re: C++ API component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes component: runtime component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants