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 3 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.
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.
| 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.
2 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).
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.Scope and behavior on non-executorch output formats
Caller-owned mutable buffers (KV via #4445, and copy-back here) are supported for
output_format="executorch"in both retrace modes. 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 (KV writes append nothing and take the aliasing path unchanged; a KV-only export is byte-identical to before).For the niche case of a model that does have a non-KV in-place mutable buffer and is exported to a non-executorch format:
exported_program/aot_inductor):create_trt_exp_programruns for all formats, so the mutation is emitted as a correctBUFFER_MUTATION.exported_program/aot_inductor): the re-declaration pass runs only on the executorch path, so the copy-back value is surfaced as an extraUSER_OUTPUTrather than aBUFFER_MUTATION.Importantly, no previously-correct path is changed: before this PR that same combination silently dropped the mutation, so this either fixes it (executorch, and retrace=False other formats) or replaces one unsupported behavior with another (retrace=True other formats). This matches #4445's own executorch scoping.
Testing
Unit tests added in this PR (CPU-only, gated on
executorch.exir):tests/py/dynamo/lowering/test_buffer_lifting.py— the classification: KV writes (slice_scatter/index_copy) stay aliased with no copy-back; a non-KV mutation (add_) is recorded and its new value re-attached as the trailing output (verified numerically equal to the updated buffer); anindex_putwrite falls into copy-back; a mixed KV + non-KV graph records only the non-KV buffer.tests/py/dynamo/executorch/test_kv_cache_export.py— both exporter passes (create_trt_exp_programfor retrace=False and_declare_aliased_kv_mutations_on_epfor retrace=True) reclassify a trailing copy-back output toBUFFER_MUTATION, ordered ahead of the user outputs.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
Two open PRs will require follow-up once they land (neither blocks the mechanism here; both affect where/how it is wired):
_compile.pyintotorch_tensorrt.executorch.export(). The retrace=True re-declaration (_declare_aliased_kv_mutations_on_ep, including thecopyback_buffersthreading added here) currently lives in_compile.py's save path and will need to be re-homed 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 (it lives increate_trt_exp_program, which feat(executorch): expose composable Edge export API #4440 does not touch). feat(executorch): expose composable Edge export API #4440 also introduces per-method partitioners (separate prefill/decode delegates); the re-declaration will need to run per method.conv_state, a separate delegate reads it) once feat(executorch): expose composable Edge export API #4440 (separate delegates) and refactor(executorch): use shared caller stream #4454 (shared stream) are both available.