feat(executorch): let a TensorRT engine write the caller's KV cache in place - #4667
Open
Conarnar wants to merge 19 commits into
Open
feat(executorch): let a TensorRT engine write the caller's KV cache in place#4667Conarnar wants to merge 19 commits into
Conarnar wants to merge 19 commits into
Conversation
…in place An engine with aliased I/O writes its aliased output *through* the aliased input's pointer, so by the time the engine returns the caller's buffer is already updated. Nothing in the ExecuTorch pipeline knows that, and it pays for the update twice on every execution: PropagateDevicePass hands the delegate an `_h2d_copy` staging copy instead of the buffer, and the aliased output is threaded back out so ExecuTorch can copy it into the buffer afterwards. For a KV cache that is two cache-sized copies per token. Add the two passes that remove them. They are separate functions because they run at different points and there is no single point that works: `rewire_aliased_mutations_to_buffers` runs on the exported program, before partitioning, because the partition boundary is what fixes the delegate's outputs. Pointing each aliased BUFFER_MUTATION at its own buffer placeholder says "the buffer is the result", so there is nothing to copy back and the aliased output -- now userless -- leaves the graph and the delegate with it. `unstage_aliased_buffers_pass` runs as a `to_out_var_pass`, the last hook in the window after PropagateDevicePass and before memory planning -- the window where the staging copies exist and the buffers' placement is still open. (`sym_shape_eval_pass` is a caller-supplied hook in that window too, but it runs first.) Which mutations are aliased is read from the engine's own `aliased_io`, not inferred from the graph. A copy-back mutation and a mutation of a buffer that is also an engine input are both a getitem off the engine node, indistinguishable by shape, and rewiring either would delete a real update with no error. Both silent-corruption shapes raise instead. Eliding every output of an engine would leave a delegate with no outputs, which nothing downstream reports: the runtime infers elision from a single argument count, and a delegate nothing reads is a pure node a later dead-code elimination can erase. And a marked buffer whose staging cannot be followed to CUDA would be written by the engine and then discarded, since its copy-back is already gone. Neither function is public: they are one feature and applying half of it is worse than applying none. The next commits reach them through an opt-in on export().
The output-binding validator assumed every engine output binding is a delegate output. With zero-copy KV that stops being true: the aliased outputs are gone from the delegate, because the engine's in-place write through the aliased input already is the buffer update. Accept that shape, but only when the caller asked for it. A delegate missing its aliased outputs because nothing declared them as buffer mutations looks exactly like a zero-copy one -- here and in the runtime, which reads elision off a single argument count -- and today that case is a loud export-time error. Relaxing the check unconditionally would turn it into a .pte that runs and quietly never updates its cache. So the permission travels down explicitly, on a CompileSpec that only export() sets, over the partitioner's DelegationSpec: the one channel from the export call to `preprocess`. Even with permission, a partial drop stays an error. It would be a buffer update lost without a word, and the argument count cannot express it in any case.
With zero-copy KV the aliased outputs are not delegate arguments at all: the engine's write through the aliased input's pointer is the buffer update, so there is no mutation slot to fill and no copy to reflect. Detect that from the argument count rather than a serialized flag. Export elides either all of an engine's aliased outputs or none, so the two shapes differ in arity by exactly the aliased-output count the handle already knows -- and a .pte written before zero-copy existed keeps taking the threaded branch with no version check and no new blob field. (Elision is all-or-nothing only because every aliased output today is a `kv_cache_update` on a buffer; a future `kind="user"` alias beside an elided one would break the identity, and the arity check would reject the .pte at execute.) `setTensorAddress` has already pointed the binding at the caller's buffer by the point the branch is taken, so eliding only skips consuming an argument and recording a reflect. The end-to-end check for this path comes with the next commit, which adds the export-side option that can produce such a .pte.
…nfig
Until now the only way to get zero-copy KV was to monkey-patch a private
function, because the rewiring has to land between two things export() does
internally -- after the aliased mutations are declared, before the program is
staged -- and there was no hook there. Give it a real one.
export(..., zero_copy_kv=True)
edge.to_executorch(zero_copy_backend_config(config))
Opt-in rather than automatic: a .pte whose aliased outputs are elided needs a
runtime that understands that shape, so producing one unasked would break a
runner built before this feature, and the option changes nothing for existing
caller-owned KV users.
Two calls rather than one because to_executorch() is ExecuTorch's, not ours --
export() returns at the Edge boundary and never sees the config the program is
finalized with. zero_copy_backend_config composes onto the caller's config
rather than replacing it, so their memory planning and their own to_out_var_pass
survive. It is the module's only public name; the two passes stay private,
since applying one without the other is worse than applying neither.
The backend's permission to accept an elided delegate is granted per method and
only where a mutation was actually rewired, so a method that lost an output for
some other reason is still rejected.
That leaves one thing the caller must not forget, and it is documented as such
along with the other two quiet contracts of this feature: finalizing without the
config, running the delegates on separate CUDA streams, and expecting one cache
to be shared across methods without a memory-planning pass that says so. All
three produce wrong values rather than an error.
The persistence check gets a zero-copy variant in CI: the existing runner
asserts a decode step sees the KV the previous step wrote, which is exactly the
property zero-copy has to preserve after removing the copy that used to provide
it.
The two-call contract is unavoidable on the direct export()+to_executorch()
path, but not through torch_tensorrt.save(output_format="executorch"): that path
owns both steps -- it calls export() and then to_executorch() itself -- so it
can make zero-copy foolproof. Give save() a single `zero_copy_kv=` flag that
threads the opt-in into export() and installs zero_copy_backend_config before
to_executorch(), so forgetting the second call is not possible there. It is
single-method only, like the rest of save(); multi-method stays on the direct
export path. Wrapping preserves a caller-supplied backend_config, so passing
both is fine. The one exclusion is handing save() a backend_config that already
carries the pass: save() wraps it again and finalization raises. The two entry
points are mutually exclusive.
A method may hold both kinds of mutable buffer: a KV cache the engine writes through an aliased binding, and a non-KV buffer -- a convolution state -- whose new value comes back as a trailing delegate output for ExecuTorch to copy back. The copy-back path predates `zero_copy_kv`, but the combination of the two was never decided or covered, only mechanically tolerated. Treat it as supported. `_aliased_buffer_mutations` already discriminates the two by reading the engine's own `aliased_io` rather than the graph, in which they are identical -- both a `getitem` off the engine node whose buffer is also an engine input. The aliased caches go zero-copy; the copy-back buffer keeps its staging copy and the output that writes it. Refusing the combination instead would give up the feature for every model carrying one non-KV mutable buffer beside its cache, and rewiring the copy-back would delete a real update with no error. Neither is acceptable, and the discriminator that avoids both is already load-bearing, so pin it: a stub-level test on one method holding both mutations, and a real-engine export of a decode step with `k_cache`/`v_cache` beside a ring-shifted `conv_state`, asserting the caches end up bound to their own placeholders while `conv_state` stays bound to a delegate output. That real-engine export runs on both exporters, because they reach the discriminator by different routes. The legacy exporter (`retrace=False`) declares all three mutations as it builds the program, leaving `_declare_aliased_kv_mutations_on_ep` nothing to do. Under `retrace=True`, which is `save()`'s default, the retraced program arrives with no mutation declared at all -- `torch.export` drops the aliased outputs at the fx boundary and leaves the copy-back value as a plain return -- so that post-export pass is what separates the two kinds, and it is exercised only on that parameter.
`zero_copy_kv` was pinned on one real engine only: a single method, a single TensorRT delegate, the aliased caches and a `conv_state` copy-back side by side. Two shapes that shape does not reach are covered here, both on real engines. A method whose copy-back rides on a *different* TensorRT delegate than the aliased caches. `TensorRTPartitioner` derives the elided binding names per engine and stamps `zero_copy_kv` only on the delegate that lost an output, so the plain compute delegate beside it must carry no spec; if it did, `_unstage_aliased_buffers`'s cross-check would demand an aliased buffer it never had and the export would die. The test asserts the stamping directly and then finalizes, which is where that cross-check runs. A method that also holds an ExecuTorch CUDA (AOTI) delegate. `erfinv` has no TensorRT converter, so with a `CudaPartitioner` catch-all the method lowers to TensorRT, CudaBackend and TensorRT in sequence. Un-staging must reach the TensorRT delegate only -- `_is_tensorrt_delegate` gates it, because no other backend promises the in-place write through an aliased binding. Both models assert their own shape before asserting the behaviour: a partitioner change that collapsed either back to one delegate would otherwise leave the test passing while covering nothing. The trailing `Linear` in the CUDA model exists for that reason -- ending on `erfinv` leaves exactly one TensorRT delegate, and "only the KV delegate is stamped" is then true no matter what the partitioner does. Both of these run on both exporters. The split-delegate one is the only shape where `_declare_aliased_kv_mutations_on_ep` has to pick the aliased engine out of several -- it scans every `execute_engine` node and skips the ones whose `aliased_io` is empty, and the copy-back value it detaches comes off a different engine than the caches it declares. The legacy exporter declares all of that as it builds the program, so that scan runs only under `retrace=True`, which is `save()`'s default and had no real-engine `zero_copy_kv` exercise before.
github-actions
Bot
added
documentation
shoumikhin
reviewed
Sep 3, 2026
shoumikhin
left a comment
shoumikhin
left a comment
Contributor
There was a problem hiding this comment.
Some notes inline.
These are about lines the change does not touch, so they could not be attached to one.
In examples/executorch_reference_runner/kv_cache_decode_check.cpp, around line 166:
The reference runner check covers the argument count branch, but not the new part that skips the stream sync. The KV decode check never installs a caller stream guard, so the caller stream is absent and the code always takes the synchronizing branch no matter what. The zero copy no sync path never runs in CI.
That is the riskier half. A wrong argument count fails loudly and every time. A missing sync fails sometimes, with numbers that look fine.
Could the KV decode check wrap its decode loop in a caller stream guard on its own stream, the way the main runner already does, and run each model both ways? That would make the branch real coverage instead of a comment.
One related point on wording. Before this change a model with aliased outputs always synchronized at the end of execute, because the reflect forced it. The commit message calls the shared stream contract a source of wrong values. The docs in this repo already say that situation is a race that can surface as wrong results or an illegal memory access. Worth saying it at that strength where the contract is written down.
In cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp, around line 518:
In execute, the aliased output count now decides the delegate arity, but init increments it once per aliased_io entry rather than once per aliased output binding. A blob that lists the same entry twice passes every check in init, because both copies resolve to the same output and the same engine alias, and the count still ends up at two.
Two things then go wrong. If the count is above the number of outputs actually flagged, the elision test can be true while the loops consume more arguments than the span holds, and the span index is not bounds checked. If the count is above the number of outputs, the subtraction underflows, the sum in the length check wraps, and the check passes for any arity.
Your own exporter cannot produce a duplicate, so this needs a crafted or corrupt file. It is worth two lines anyway, because before this change the counter only fed a log. Increment only when the slot was still unset, and reject a repeated entry the way a duplicate weight streaming spec is already rejected. Refusing a count larger than the number of output bindings makes the underflow unreachable too.
| # un-staging pass runs after lowering, where the engine's aliased_io is no | ||
| # longer reachable from the graph: it has become an opaque delegate blob. | ||
| mutation.placeholder.meta["_torch_tensorrt_aliased_buffer"] = True | ||
| output_args[spec_index] = mutation.placeholder |
Contributor
There was a problem hiding this comment.
Repointing the mutation at the buffer placeholder makes ExecuTorch's write-back pass skip the copy for that buffer, which is the point. But that pass then patches the signature with a counter over mutation specs while indexing a list that only holds the copy nodes it created. A mutation with no copy shifts every later one by one.
For a method holding both kinds this crosses them. The KV cache mutation ends up naming the copy that writes the conv state, and the conv state mutation ends up naming the KV placeholder. A KV only method is fine because there is no copy at all, and the same model without zero copy is fine because both get one, so nothing you have measured so far would show it.
The file itself looks like it still runs, because copy carries its own destination and the emitter records outputs by position. What is wrong is the finalized program's mutation map, which the exir execution path and the ETRecord both read.
This was checked against the ExecuTorch version the workflow installs. Moving the rewired mutations behind the un-rewired ones in the output specs and args gives the correct pairing.
Contributor
Author
There was a problem hiding this comment.
Fixed by ordering the rewired mutations last, though the mechanism is not quite the one described.
The list upstream indexes is not only the copies — _insert_copy ends with
buffer_output_nodes.extend(user_output_nodes), so it returns the copy_ nodes it created
followed by every output it copied nothing for, in graph order. That matters for the fix rather
than being a detail: under the "only the copies" reading a mutation spec past the number of copies
would index off the end, so reordering would not be what saves it. What happens instead is that
ordering the rewired mutations last gives them spec positions C..M-1 and puts them at the head
of that tail in the same order, so the copy-producing specs index the copy prefix one-for-one and
the rewired specs index the tail entries that are those same outputs.
On stock ExecuTorch with run_reinplace_pass=True and two mutated buffers where only the first is
reinplaced, the first buffer's mutation names the copy_ that writes the second and the second's
names the index_put_ that wrote the first — a swap, not an index error, which only the tail
explains.
That configuration is also the fix's limit: reinplace_pass runs after the ordering and can swap
them again. There is no hook between the two, so it is scoped and stated rather than fixed.
| and every cross-boundary dependency is satisfied, while execution stays | ||
| asynchronous. | ||
| (each backend exposes a caller-stream hook: scope both | ||
| ``torch_tensorrt::executorch_backend::CudaStreamGuard`` and |
Contributor
There was a problem hiding this comment.
The single stream section now says to scope the Torch-TensorRT CUDA stream guard as well as the ExecuTorch caller stream guard. That first class is gone. The backend README says it was removed on purpose, with no deprecated alias, so this code would not compile. The reason given is also backwards: one caller stream guard already covers every CUDA delegate, because the shared CUDA extension means all of them read the same caller-stream storage, which that README states directly. The old wording, "each backend exposes a caller-stream hook", was correct. Please name only the ExecuTorch guard here.
| Finalizing a ``zero_copy_kv=True`` program *without* this config does | ||
| not raise. The engine writes a per-call staging copy that is then | ||
| discarded and the buffer never updates, which for a KV cache is wrong | ||
| output rather than a crash. Nothing downstream can detect the omission, |
Contributor
There was a problem hiding this comment.
The warning says nothing downstream can detect the omission, so pairing the two calls is the caller's responsibility. The example in this same change detects it: it walks the finalized program for marked buffers that are still staged and refuses to write the file. So the check is possible, and it is already written twice, in the example and in the tests. The single-call path in save holds that finalized program right before it writes, and does not run it, which is the one path that promised the mistake was impossible. Please move that check into the library and call it there, and drop the "nothing downstream can detect" sentence.
Contributor
Author
There was a problem hiding this comment.
Fixed. check_zero_copy_kv() is public, save() runs it before writing, and the example calls it
rather than carrying its own copy.
It walks every method, not just forward — the first version took exported_program() with no
name, which raised KeyError on a multi-method program and skipped the other methods where a
forward did exist. "Nothing marked" raises only when no method has marks, matching the warning
_apply_zero_copy_kv already emits; a staged buffer raises per-method and names it.
| # single argument count, which a zero-output delegate satisfies, and a | ||
| # delegate nothing reads is a pure node that a later graph-wide dead-code | ||
| # elimination can erase, taking the computation with it. Stop here instead. | ||
| # The eliminate_dead_code() below does not erase this engine node: unlike |
Contributor
There was a problem hiding this comment.
In the rewiring pass, the comment above the every output elided check says an execute_engine node is impure to FX and so survives dead code elimination with no users. I do not think that is true. The operator is declared with no write annotation, and the Python fallback registers it with an empty mutates_args, so FX sees a pure op and will erase it once nothing reads it. Nothing in this repository adds it to FX's side effect list either.
That also leaves a small hole in the check above. It only fires when every user of the engine is an elided getitem. If one non elided getitem is still there but is itself dead, dead code elimination erases that getitem first, then the engine, and the computation goes with it.
Either count live users only, or drop the sentence and say plainly that the guard is what keeps the engine alive.
Contributor
Author
There was a problem hiding this comment.
The premise does not hold: EffectHolder._set_default_effect gives it EffectType.ORDERED, so
is_impure() is true and DCE keeps it with no users. The comment was right.
The hole you inferred from it is real, though, and is fixed — the check now counts live users, so a
surviving-but-dead non-elided getitem no longer lets the engine and its computation be erased.
Contributor
Author
Fixed, and instrumented to confirm it: the guarded mode takes the skip path on all three of its Two limits worth stating. The staged
Fixed at parse rather than in Reproducing it turned up a second path the It needs no crafted graph — the same engine, argument list and arity come from a correct |
The single-stream section told runners to scope `torch_tensorrt::executorch_backend::CudaStreamGuard` alongside `executorch::extension::cuda::CallerStreamGuard`. That class no longer exists -- `cpp/src/torch_tensorrt/executorch/README.md` records its removal, deliberately with no deprecated alias -- so the snippet named a type that does not compile. The reason given for scoping both was backwards as well. One caller-stream guard already reaches every CUDA-capable delegate, because they resolve a single shared `libextension_cuda` and so read the same caller-stream storage; that is why the old class could be dropped rather than aliased. Say that instead.
The warning on `zero_copy_backend_config` said nothing downstream could detect a program exported with `zero_copy_kv=True` and then finalized without it. The example in this same change detects exactly that: it walks the finalized program for marked buffers that no longer reach a delegate directly and refuses to write the `.pte`. So the claim was false, and the check was written twice -- in the example and in the tests -- while the one path that owned both ends, `save()`, held the finalized program and did not run it. Move it into the library as `torch_tensorrt.executorch.check_zero_copy_kv`, call it from `save(..., zero_copy_kv=True)` before the `.pte` is written, and have the example call the library version. It refuses two shapes: a marked buffer that is not a direct argument of a TensorRT delegate (finalized without the config, so the engine writes a staging copy that is discarded), and nothing marked for in-place update at all (`zero_copy_kv=True` only warns when it finds no aliased buffer mutation). Where the un-staging pass does run it has already raised on the first shape; this catches the case where it never ran. Only a TensorRT delegate counts as that direct argument, which is the same filter the un-staging pass applies. A buffer is marked because a TensorRT engine writes it in place, so another backend's delegate holding it says nothing about whether the engine got it un-staged: a program can hand the buffer straight to a `CudaBackend` delegate while the TensorRT engine beside it still reads a staging copy that is thrown away. Every method is read, not only `forward`. `ExecutorchProgramManager.exported_program()` defaults to `forward`, and `export()` rewires each method separately, so a check that took the default would raise a bare `KeyError: 'forward'` on the prefill/decode program the user guide's own zero-copy example builds, and on a program that does have a `forward` beside other methods it would pass one whose decode had degenerated to staged. The failure names the method the buffer is in. The "nothing marked" refusal stays about the program rather than about each method, matching the warning `export()` emits for the same condition: a method with no aliased buffer mutation of its own is not an error, so a model that rewires only its decode step is accepted, and only a program where no method rewired anything is refused. The docs and the docstring now point at it instead of asserting the mistake is undetectable. The three real-engine tests hand it their finalized program before running their own stricter graph assertion. Every other test of it builds the program itself -- a `SimpleNamespace` for the check's own cases, a monkeypatch for `save()`'s -- so nothing else puts `methods` or `exported_program(name)` in front of a real `ExecutorchProgramManager`, and an upstream rename would leave all of them green while `save(zero_copy_kv=True)` raised `AttributeError` for every caller.
…liminated The guard against eliding every output of an engine fired only when *every* user of the engine was one of the elided getitems. A non-elided getitem that is itself dead defeated it: the check saw an output, the `eliminate_dead_code()` right below erased that getitem anyway, and the partition got the zero-output delegate the guard exists to refuse. Counting only the users that something reads does not settle it either, because it looks one step past the engine and a dead chain can be longer than that. The first link of a two-node dead chain does have a user, so the guard stays silent and the elimination then erases the whole chain. Measured on hand-built graphs with chains of 1, 2 and 3 nodes: only the one-node chain raised, and at 2 and 3 the rewiring returned normally leaving the engine node with no users at all. So run the elimination first and then ask whether the engine still has a user: what survives it is what the delegate will have, whatever the chain length. The test is parametrized over a one- and a two-node chain, the second being the length a rule reading only the engine's immediate users misses. The comment above it claimed an `execute_engine` node is impure to FX and survives DCE with no users. That is true, and it now says why, because the reason is not local to this repository: PyTorch defaults any operator taking a ScriptObject argument to an ORDERED effect (`torch._library.effects.EffectHolder._set_default_effect`), and `execute_engine` takes the engine as one. Confirmed on the op this code builds against -- `_get_effect(torch.ops.tensorrt.execute_engine.default)` is `EffectType.ORDERED`, `Node.is_impure()` is `True`, and a userless engine node survives `Graph.eliminate_dead_code()`; the delegate that replaces it after lowering has no effect registered and does not. The comment also no longer leaves the reader to infer that the guard, and not the DCE, is what stops the bad shape.
…parsed `init` incremented the aliased-output count once per `aliased_io` entry rather than once per output binding it claimed. A blob listing the same entry twice passed every check -- both copies resolve to the same output and the same engine alias -- and left the count at two for one aliased output. `execute` reads that count to decide the delegate arity, so an inflated one is not just a log line. With the count above the number of flagged outputs, the elision test can be true while the output loop consumes more arguments than the span holds, and nothing bounds its index into `args`. With the count above the number of output bindings, `num_outputs - num_aliased_outputs` underflows, the sum in the length check wraps, and the check passes for any arity. A repeat is malformed for every reader of the blob, not only for `init`, and it is visible from the bytes alone -- so refuse it in `TensorRTBlobHeader::parse`, where the blob-header unit tests reach it without a GPU or a real engine. It is refused the way the parser refuses any other malformed metadata, which means the diagnostic is `init`'s generic parse failure rather than a message naming the output. `execute` keeps the `num_aliased_outputs <= num_outputs` bound. The parser cannot establish it: a blob may carry an empty `io_bindings` array, and then the output bindings are inferred from the deserialized engine, which the parser has not seen. Both subtractions there are unsigned and the length check is the only thing bounding how far the loops index into `args`, so the bound is still worth checking for a header that reached the backend some other way. Our own exporter cannot emit a duplicate; this needs a crafted or corrupt file.
…tream The KV persistence check never installed a caller stream guard, so `getCallerStream()` was always empty and the backend always took the synchronizing branch at the end of `execute()`. The skip-the-sync path zero-copy KV depends on -- no host staging, no aliased reflect, a caller stream set -- had therefore never run under this check, and neither had the machinery that exists to make it safe: the `inflight_event` record, the wait on it at the top of the next `execute()`, and the drain in `~EngineHandle`. Run both scenarios twice, once with no caller stream and once with a `CallerStreamGuard` scoped over the decode loop on a stream the check owns, the way the main runner does. The guard is constructed only in the second mode: an explicitly null selection is still a selection (see `tests/cpp/executorch/test_caller_stream.cpp`), so an unconditional guard over a null stream would cover one branch twice. The guarded run synchronizes and destroys its stream before reading the logits. What that buys is branch coverage, and only on the zero-copy `.pte`. Instrumenting the branch decision shows the zero-copy model in the guarded mode taking the skip path on all three of its `execute()` calls, waiting on the previous enqueue at the top of the second `execute()` of the two-step scenario, and draining a pending enqueue at each of the two teardowns; the unguarded mode reproduces the old always-synchronize behaviour in the same binary. The staged `.pte` still synchronizes in both modes, because its aliased outputs are delegate output args and so an aliased reflect is always pending -- running it under the guard moves the engine onto a caller-supplied stream and nothing else. What it does not buy is a check that can fail when a synchronization is missing. The stream comes from `cudaStreamCreate`, which is a blocking stream, and the logits are read with a synchronous `cudaMemcpy` on the legacy default stream, which implicitly waits for every blocking stream in the context. Two controls say so on this model: deleting the check's own `cudaStreamSynchronize` still passes, and a `cudaStreamQuery` placed where that sync was reports the stream already complete every time. So repeated runs agreeing to the last bit are not evidence that the ordering is right. The check exercises the path and would fail on a hard error in it -- a rejected event record, a context reconfigured under a live enqueue, a mis-elided argument -- but it is not a race detector. Making it one would mean a `cudaStreamNonBlocking` stream and dropping the explicit sync that is redundant with the default-stream copy. `verify-executorch-reference-runner.sh` now greps for both modes by name, so dropping one cannot leave the lane green with the branch uncovered. Also: where the shared-stream contract is written down, the zero-copy section described getting it wrong as failing quietly. The coalesced-`.pte` section, a few paragraphs down, already calls it a race that can surface as wrong results or an illegal memory access. Say it at that strength in both places.
…ng like Two defects in `_unstage_aliased_buffers`, both from deciding by the edit it made rather than by the shape it has to leave behind. **It keyed success on having deleted a staging copy.** The property zero-copy needs is that the marked buffer is a direct argument of a TensorRT delegate; removing an `_h2d_copy` is only the usual route there. `ExecutorchBackendConfig(enable_non_cpu_memory_planning=False)` is a supported configuration -- the field defaults to True and `zero_copy_backend_config` preserves whatever the caller set -- and under it `PropagateDevicePass` inserts no staging copies at all. Measured against that pass directly: with the flag True the delegate's argument is an `_h2d_copy` and the placeholder stays on the host; with it False there is no copy node, the placeholder *is* the argument, and its spec is already `cuda:0`. That program is in the shape zero-copy wants and the pass rejected it, naming a caller two causes -- the buffer never reached a delegate, or the pass was installed twice -- neither of which had happened, and telling them to install it once. `check_zero_copy_kv` accepted the same graph, so the two halves disagreed about one program. The pass now counts a marked buffer as satisfied when a TensorRT delegate takes it, whether or not this pass is what put it there, and refuses only when no TensorRT delegate takes it at all -- directly or through a staging copy the pass can remove. That is the condition `check_zero_copy_kv` reads, so the two accept and refuse the same graphs. A consequence is that installing the pass twice is now a no-op rather than an error: the second run finds the first run's work in place. The docstring, the `save()` comment and the user guide each promised that finalization raises there, and now describe what it does. **It allowed a second staging copy to the same GPU to survive the move.** A marked buffer feeding two same-GPU `_h2d_copy` nodes, one to a TensorRT delegate and one to another backend, had the TensorRT staging erased and its own spec flipped host->CUDA, leaving the other copy reading a source that is now device memory. `_h2d_copy_out` requires a host source and fails `InvalidArgument` on a device one (portable kernel; the ATen branch has no such check). The rule is now that any staging copy this pass does not itself remove blocks the move, same GPU or not, and a copy is removed exactly when every user of it is a TensorRT delegate whose argument gets rewired. That reverses a decision earlier in this stack. `test_unstage_leaves_another_backends_same_gpu_staging_in_place` was rewritten to pin the allowance -- that the buffer still moves and the other backend's copy stays -- and now pins the refusal, under a name that says so. Two neighbouring tests describe which comparison catches which shape and are corrected with it: the different-GPU-other-backend shape is no longer refused by the device index alone, and the two-TensorRT-GPUs shape, where every staging does feed a TensorRT delegate, now asserts that nothing moved rather than only that something raised -- without that, dropping the index comparison would still leave it green.
Four pieces of prose that a reader can check against the code and find disagreeing. All four sit in commits already pushed, so they are collected here rather than folded back into them. `_aliased_inputs_by_output_index` said the skipped case is "an output binding absent from the delegate's inputs". What the two `continue`s skip is an `aliased_io` entry whose *input* does not resolve -- a name that is not one of the engine's input bindings, or an index past the delegate's argument list. `_declare_aliased_kv_mutations_on_ep` warns on both, which is the sentence's point and is why neither is reported twice. `rewire_aliased_mutations_to_buffers` said "an engine mixing the two is caught". An engine carrying a rewired aliased output beside an un-rewired user alias is not caught and must not be: it exports cleanly, because the un-rewired output is still a delegate output and that is what the binding check wants. What is caught is a delegate that dropped the un-rewired one as well. The user guide's zero-copy example passed `CudaPartitioner([])` for each of two methods. `export()`'s own docstring says a partitioner whose specs name no method leaves a backend that reads its method name from them -- the CUDA backend -- unable to find it, so the snippet as written raises during lowering. Nothing in the example needs a second backend; drop the argument. `test_validate_output_binding_order_still_accepts_aliased_outputs_threaded` was described as covering "the pre-existing shape", which tells a later reader nothing: there is no before. It covers that naming a binding elidable permits the drop without requiring it. Also two smaller ones: a comment in `test_edge_cases.py` explained the `zero_copy_kv=False` assertion in terms of a KV buffer that model does not have; and the `**ExecuTorch lowering options**` heading had no blank line above it, so it ran into the paragraph before. That last one is on `main` and predates this stack, but the paragraph it runs into is one this stack rewrote.
Conarnar
force-pushed
the
feat/executorch-zero-copy-kv
branch
from
September 6, 2026 04:02
60b008b to
aa7fe51
Compare
…emory `_unstage_aliased_buffers` counted a marked buffer as satisfied the moment a TensorRT delegate took it directly, on the strength of the mark and the delegate edge alone. Being a direct argument is half of what zero-copy needs. The other half is that memory planning puts the buffer in a device arena, and nothing checked it. `ExecutorchBackendConfig(enable_non_cpu_memory_planning=False)` produces exactly that shape and does not satisfy it. `PropagateDevicePass` inserts no staging copy under that flag and writes the delegate's device straight onto the placeholder's own spec, so the buffer looks placed; memory planning with the flag off then ignores every spec device and puts the whole program in one host arena. Measured on a real export, the marked placeholders carry `DeviceType.CUDA` either way, and land in a CUDA arena with the flag on and in the single host arena with it off. The `.pte` that came out exported clean, passed `check_zero_copy_kv`, and failed its first `execute()` on the runtime's alias-target guard -- "aliased input 'buf_k_cache' must be device-resident", `Error::InvalidArgument`. That was run on device rather than inferred. The spec's device is therefore not on its own the answer either. Two things decide where the buffer is planned and only one of them is in the graph, so `zero_copy_backend_config` reads `enable_non_cpu_memory_planning` off the configuration it is building from and hands it to the pass. The pass refuses a direct argument that is either off-CUDA or host-planned, naming the buffer, since that buffer is the one whose update is lost. The docstring named that configuration as one of the two supported ways to reach the direct-argument shape. It is not one, and the sentence now says what being a direct argument is and is not enough for; the pass's own second run over its output is the route that remains. Both tests pinning the shape set the spec to CUDA by hand, which is what the real configuration does too, so neither could have caught this. The new ones cover each conjunct on its own and the wiring that carries the planning mode in.
…ly elided Export decides elision per binding name; the runtime decides it with one subtraction. The two agree only while every aliased output of an engine is a rewired buffer mutation. An engine that also aliases onto a plain input -- one buffer KV cache beside an argument cache whose updated value is returned -- has the narrower set elided, passes the output-binding check, passes `check_zero_copy_kv`, and writes a `.pte` that cannot run: the runtime reads elision by subtracting the engine's whole aliased-output count from the argument count it was handed, so a delegate short of only some of them reads as not elided at all. Loaded through the ExecuTorch runtime with this backend, such a file opens, binds both aliases at init, and then fails every `execute()` on the argument count -- "expected at least 7 args, got 6" for the two-alias engine measured -- with `Error::InvalidArgument`. That was run on device rather than inferred from the arity code. `preprocess` already holds the elidable names and the engine's whole `aliased_io` ten lines apart. It now compares them and refuses, so the failure lands where the export can still be re-run instead of on the device, and it says plainly that partial elision is not expressible by the runtime. Teaching the runtime to read the elided names from the compile spec it is already handed is the other way out of this, and is deliberately not taken: refusing is much smaller, and going from refused to supported later is a compatible progression. Deriving the narrower set is not what is wrong here, so the test that pins that derivation keeps its assertion and gains the refusal beside it.
The un-staging pass froze `enable_non_cpu_memory_planning` into itself at
config-build time and read it on only one of its two branches.
`ExecutorchBackendConfig` is a plain mutable dataclass, so a caller who turns
the flag off on the config `zero_copy_backend_config()` returned -- the config
`to_executorch` then uses -- got a program the pass accepted and the finalizer
planned into the host arena. The pass now reads the flag off that config when
it runs, and refuses a host-only planning mode once for the whole graph rather
than on one branch, so the staged shape is covered too.
`check_zero_copy_kv`, which the docs and `save()` tell people to rely on, unioned
the arguments of every TensorRT delegate and never looked at placement. Measured
at head: it accepts a program finalized with
`ExecutorchBackendConfig(enable_non_cpu_memory_planning=False)` whose caches are
direct delegate arguments planned in the single host arena -- the `.pte` whose
first `execute()` fails on the alias-target guard. It now narrows that union to
the delegates carrying the zero-copy spec, so an unrelated engine reading the
same buffer cannot stand in for the one whose write was elided, and requires the
buffer's `mem_id` to name one of the CUDA arenas memory planning recorded in
`non_const_buffer_device`. In the same function, the completeness check accepted
partial marker loss: one zero-copy delegate taking two staged buffers with only
one still marked returned 1 and raised nothing, leaving the second cache wired
through its `_h2d_copy`. It now compares against the number of aliased outputs
the delegate's compile spec says it elided.
A method mixing zero-copy caches with a copy-back buffer finalized with a crossed
mutation map -- reproduced on the real pipeline as
`{'copy__default': 'k_cache', 'b_k_cache': 'v_cache', 'b_v_cache': 'conv_state'}`
against the unrewired control's correct pairing. Upstream's write-back pass
inserts a copy only for a mutation whose value is not already its buffer
placeholder, moves those copies to the front of the output tuple, and reassigns
each mutation spec's argument by position; rewiring a cache to its own
placeholder takes it out of that leading run. `order_copyback_mutations_first`
puts the mutations that still get a copy first. It runs on the Edge program, not
beside the rewiring: `to_edge_transform_and_lower` re-derives the signature, so
an order set earlier is discarded (measured -- the pre-edge program comes out
`conv_state, k_cache, v_cache` and the Edge program `k_cache, v_cache,
conv_state`). The `.pte` was written correctly either way; what was wrong is the
finalized signature and ExecuTorch's eager call path, which walks
`buffers_to_mutate` writing the graph's leading results into the state dict in
that order.
The blob parser refused a repeated `aliased_io` entry but accepted a repeated
binding name in `io_bindings`, which is the same hazard one step earlier: a
TensorRT engine has one name space, so both slots resolve to one tensor, `init`
records the alias on the first, and `execute` re-binds the same name for the
later one. Confirmed on TensorRT 11.1.0.106 that a second `setTensorAddress` for
one name replaces the first and the engine's write lands entirely in the second
buffer, so the address replaced is the caller's cache. The parser now refuses the
repeat.
Tests: the duplicate-`aliased_io` test repeated the entry exactly, so a parser
keyed on the output/input pair -- the regression that reopens the inflated-count
bug -- passed the whole file; the same-output-different-input case is added
alongside three for duplicate binding names. On the Python side the mixed test
now asserts the finalized signature, which it did not before, and each new
refusal has a test that fails when its clause is mutated away.
One hazard the aliased-output elision looks like it introduces, and does not:
eliding the aliased reflects does drop the term that used to force the
end-of-execute sync, but the engine's other outputs are ordered by the
caller-stream contract instead. Every host consumer ExecuTorch inserts for a
device delegate output -- including the write-back of a copy-back buffer beside
the caches -- is an `et_copy::_d2h_copy`, whose kernel issues its copy on
`getCallerStream()`, the same stream the engine enqueued on, and then
synchronizes it. With no caller stream that kernel is a blocking `cudaMemcpy`
and `must_sync` is true anyway. Written down at the `must_sync` computation.
…tents Three of the fixes below are in the machinery that threads the planning flag into the un-staging pass and reorders the copy-back mutations; the other two are in the partitioner's alias-resolution fallback and the blob parser. `zero_copy_backend_config` preserved `skip_h2d_for_method_inputs`, and handed back a config that cannot finalize at all. Measured on the real mixed-KV model: `to_executorch` raises `skip_h2d_for_method_inputs=True requires placeholder 'b_k_cache' to have exactly one user, but it has 2 users`. That option is ExecuTorch's own un-staging of method inputs and it demands a single user, while a rewired cache always has two -- the delegate, and the graph output it is its own mutation result for -- so every zero-copy graph has the shape it refuses. This is the path the guide tells people to use, so the wrapper now refuses the option itself, naming it: the bool form, and any method a per-method dict asks for. `PropagateDevicePass` only tests that field for truth rather than resolving it per method, so a dict whose entries are all `False` is carried through here and still read as on a layer down; refusing that one too is left open. `order_copyback_mutations_first` did not do what it claimed. It split the mutations on this feature's own mark instead of on the predicate upstream keys on, and `run_reinplace_pass` and `reinplace_extra_ops` are supported `ExecutorchBackendConfig` fields whose pass runs just before the write-back and produces the same copy-free shape from an ordinary mutation. Measured against ExecuTorch's real `insert_write_back_for_buffers_pass`: an in-place-lineage mutation ahead of a copy-back one comes out with each buffer named against the other's value, and the reorder reports zero moved. It now asks upstream's own `_inplace_lineage`, imported rather than reimplemented, and returns the number of slots whose value changed. `check_zero_copy_kv` read an absent `non_const_buffer_device` as proof of a host placement. `apply_algo` is the only thing in ExecuTorch that writes that key and `to_executorch` takes any callable as `memory_planning_pass` -- which the guide tells people to bring for a cache shared between prefill and decode -- so a correct program was refused and `save` wrote nothing. Absence now means "cannot tell", and what still catches the hole the check was added for is that host-only planning puts every tensor in one bucket: the cache's arena also holds the program's host tensors, which no device-aware planner would do. All three shapes are measured on the real model -- the good one, the hole, and a caller-supplied planner that records no devices. A resolver failure inside `_partition_elided_output_names` did not degrade safely. Injecting one, the export died with "the aliased-buffer mark did not survive lowering", which is not what happened: the aliased outputs were removed before partitioning, so returning an empty set stamps nothing and guarantees a downstream failure that names the wrong cause. It re-raises now when any buffer in the method carries the mark, and keeps the fallback for the case where the claim is true. Two in the blob parser, both memory safety. The engine extent check added a 64-bit attacker-controlled `engine_size` to the offset before comparing against the blob length, so the sum wraps: an 8 KiB blob declaring 2^64-4086 parses clean and that length reaches `deserializeCudaEngine`. Every extent is now checked by subtracting the offset from the bound instead; the two metadata extents cannot wrap on a 64-bit `size_t` but are written the same way so the form, not the width of each field, is what makes them safe. Separately, the repeated-entry refusal covered only the output side of `aliased_io`, so two entries naming different outputs and one input both parsed and both resolved to one caller pointer. Confirmed by driving TensorRT 11.1 directly: it accepts two output bindings at one address, runs without error, and only the second write survives. The `kv_cache_update` kind is ruled out by the engine cross-check, but the `user` kind is compared only on shape. An empty or absent binding name is now refused in the same place rather than silently skipped, which used to shorten the recorded list while the delegate's argument list kept its length. Also: the two public helpers are in the API reference; the user guide's snippet no longer passes an unbound `backend_config` and no longer scopes the caller-stream synchronization duty to a coalesced `.pte`; the two readers of the zero-copy compile spec take the same values, so a hand-built one raises naming the key instead of decoding a JSON string into its own characters; a delegate with no outputs at all is refused after lowering, which naming every binding elidable used to satisfy; `save()`'s options message is built from the option table rather than hand-maintained; and the spec-missing refusal names whichever node is bare. Ten tests were added or widened for clauses that had none: the reorder against the real upstream pass, the skip-H2D refusal, both halves of the placement inference, the CUDA filter on the recorded arenas, the no-spec arms on both sides of the device move, the device-type half of the surviving-consumer check, the two unresolvable-alias skips, the empty delegate output list, and the compile-spec value shapes. Four C++ parser tests fail against the old parser. The multi-delegate CUDA-neighbour test is parametrized over the exporter, the mark-survival test asserts both cache names rather than a non-empty list, and the declaration-before-rewiring test records both passes in one tagged log so reversing them fails.
`zero_copy_backend_config` refused `skip_h2d_for_method_inputs` only where its
value was true, and only in a single `PropagateDeviceConfig`. Both are narrower
than the pass they guard. `PropagateDevicePass` is handed the field whole and
only tests it for truth (`propagate_device_pass.py:216`), so it reads any
non-empty dict as on for every method; and `propagate_device_config` is itself
typed as one `PropagateDeviceConfig` or a dict of them keyed by method, with
ExecuTorch handing the pass the entry for the method being finalized. Measured
on the real mixed-KV model, both
`PropagateDeviceConfig(skip_h2d_for_method_inputs={"decode": False})` and
`{"forward": PropagateDeviceConfig(skip_h2d_for_method_inputs=True)}` were
accepted here and then died inside `to_executorch` with `requires placeholder
'b_k_cache' to have exactly one user, but it has 2 users`, which is the failure
this refusal exists to prevent. It now refuses the option wherever it is
written, and on every value that pass reads as on rather than only on `True`.
`False` and `{}` are still carried, because those are exactly what that pass
reads as off.
`check_zero_copy_kv` read the stamped delegates as one set. A method that lowers
to two of them can hand both caches to one and leave the other reading a staging
copy of the cache whose copy-back export had already removed: every marked
buffer does reach a stamped delegate, so nothing refused it, and that engine's
write was discarded. Each stamped delegate is now counted against its own
compile spec -- one marked buffer per aliased output the spec says it elided,
falling back to demanding at least one when the spec names none -- which is the
cross-check `_unstage_aliased_buffers` already made before planning, now made
again after it. This is the function `save()` and the user guide tell people to
rely on, so it is where the last stand-in had to be closed. What neither count
can separate is an exact swap, since the mark is a bare flag naming no engine;
that is stated in the docstring and under Known gaps.
`zero_copy_backend_config` passed `unstage_aliased_buffers_pass` a
`device_memory_planning` that the next line then overrode for every call, by
binding `finalization_config` to the config it returns. Proved dead by poisoning
it to `False` and re-running the suite unchanged. The argument is gone from that
call site; the parameter stays on the builder, where a pass constructed by hand
and left unbound does read it.
What that binding does not cover is documented rather than fixed. The planning
flag is a bool copied by value and the pass an object copied by reference, so a
config derived from the returned one with `dataclasses.replace` carries a pass
still reading the original, and turning planning off on the derived config alone
is not refused there. Making the pass follow the derivation was considered and
rejected: nothing on the config points back at the pass, and nothing on the pass
can see which config `to_executorch` is finalizing, so the only complete fixes
are returning a dynamically built subclass of the caller's config class -- which
changes the returned type and breaks pickling -- or moving the guard behind a
wrapped `memory_planning_pass`, which relocates a tested refusal to a later
point. Both are out of proportion to a defect whose measured outcome is already
a later refusal: calling `zero_copy_backend_config` again on the derived config
is the remedy -- the pass it builds is bound to that one -- and
`check_zero_copy_kv` refuses the program the mistake produces for any method
holding a host tensor to give the shared arena away. Measured, documented, and
pinned by a test.
Five tests: the all-`False` dict, the per-method config dict, a crossed pair of
zero-copy delegates against the matched pair that is accepted, a stamped
delegate whose spec names no aliased output, and the derived-config remedy. Each
new clause was mutated away and kills only the tests belonging to it.
The blob metadata is copied out with an explicit length, so a binding name can hold a NUL byte, while every consumer resolves the name through `c_str()` and TensorRT stops at the first one. The parser has three repeat refusals and one emptiness check, and a NUL walks past every one of them. Past the repeat refusals in four ways: two names differing only after a NUL are two entries here and one tensor to the engine -- the exact collision those refusals exist to stop -- and the set covering `io_bindings` spans both lists, so an input colliding with an output defeats it as well as two outputs do. Past the emptiness check because that check draws its line where `size()` does rather than where `c_str()` does, so a name that is only a NUL is non-empty here and empty to the engine. One predicate beside the emptiness check closes all five routes, since every name the engine is asked for has to be an `io_bindings` name first. Measured against a standalone build of the parser: two output bindings `"x\0y"`/`"x\0z"`, an input and an output colliding across the two lists, both `aliased_io` shapes and the NUL-only name all parsed clean and reported the same `c_str()`; all five are refused now and a plain blob is unaffected. The same predicate on the two `aliased_io` names stops no collapse of its own -- `init()` compares those whole, so a NUL there matches no binding and `init()` refuses the blob -- but it keeps the invariant that every recorded name is one the engine can be asked for, and moves that failure to parse. No writer emits a NUL -- `json.dumps` escapes it -- so this refuses only a blob assembled some other way. An `aliased_io` entry naming only one of its two bindings was silently skipped and now shares that refusal. `init()` counts the entries it accepts and `execute()` subtracts that count from the delegate's argument list, so a dropped entry surfaces as an argument-count error on every call, in a message that never mentions aliasing, instead of at parse where the blob-header tests reach it without a GPU. The writer always emits both keys (`serialization.py`) and an older blob carries no `aliased_io` array at all. `check_zero_copy_kv` passed over any method with no marked buffer before it looked at the delegates, so a method still carrying a stamped zero-copy delegate whose mark was lost was invisible to it while `_unstage_aliased_buffers` refuses that same graph -- measured on a two-method program, the pass raising and the check returning clean. That skip is one side of a two-sided question: the function exists to catch the marks and the stamped delegates disagreeing, and driving the walk from the marks alone leaves anything recorded only on the delegate side outside it. Both records are now enumerated in every method and a method is passed over only when it carries neither; the per-delegate count already in place then refuses the lost-mark case. The program-wide "nothing is marked" refusal moved below the two that fire on such a disagreement, because on a single-method program in that state it answered "probably not exported with zero_copy_kv=True" while the compile spec said it was. `order_copyback_mutations_first`'s docstring claimed that asking upstream's predicate rather than this module's mark is what covers a mutation `run_reinplace_pass` turns in place. It does not: `reinplace_pass` runs inside `to_executorch` (`_program.py:1706`), after the reorder, so such a mutation is ordinary when the predicate is asked and in-place when the write-back reads it. Measured on a two-buffer model with no TensorRT in it, the pair finalizes crossed with and without the reorder and `moved` is 0 either way. The docstring and the test that cited the case now claim only what they cover -- any mutation already in place at the Edge boundary -- and name the reinplace ordering as upstream's. Coverage the mutation runs showed missing. Dropping the clause that compares the metadata extent against `engine_offset` left every blob-header test green, so a blob whose metadata sits past the engine parsed; one case pins it. The partial-elision refusal in `preprocess` was pinned only from the zero-copy file, so replacing it with a constant false left `test_backend.py` green; a case now sits beside its neighbours. The two new public names were in `__all__` and on the API page but not in the test that pins the package's surface. The four tests that build a real engine gated CUDA with a decorator `skipif`, which resolves during collection -- off the GPU host on a remote-GPU runner, so the skip freezes in before a GPU is attached and the only real-engine coverage this feature has disappears with the lane still green. They use the runtime gate the rest of this directory documents three times over. A note on the bazel target for the KV decode check records that nothing depends on it and the CMake build beside it is what CI compiles, so its dependency list is kept in step by hand. Blob-header gtests 22 -> 28, Python suite 325 -> 327 collected, one pre-existing wheel-pin failure unchanged. Every behavioural clause was mutated away and the test that bites it named; the only surviving mutant is the metadata extent rewritten in addition form, which the comment there already says cannot differ.
Conarnar
force-pushed
the
feat/executorch-zero-copy-kv
branch
from
September 8, 2026 05:59
af5ef7e to
4634ea2
Compare
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
2 participants
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.
Description
Let a TensorRT engine write a caller's KV cache in place, instead of routing every update through a staging copy and a copy-back.
An engine with aliased I/O writes its aliased output through the aliased input's pointer, so running the engine over the buffer already is the update. Nothing in the ExecuTorch pipeline knows that, so by default the buffer makes a full round trip on every execution:
PropagateDevicePasswraps each delegate input inet_copy._h2d_copy, so the engine writes a per-call staging copy rather than the caller's buffer; and the aliased output is threaded back out as a delegate output for ExecuTorch to copy in afterwards.For a cache-sized buffer that is two copies per execution of something the engine could have written directly. On a 30-layer Gemma4-MoE at prefill 2048 / max_len 2560 it is 576.7 MB in and 577.2 MB out every decode step, at 6.5–7.0 GB/s pageable.
What this changes
export(..., zero_copy_kv=True)opts in. The feature needs two changes at different points in the pipeline, and the export flag performs only the first:to_out_var_pass, installed byzero_copy_backend_config(), the staging is removed so the engine is handed the caller's buffer.Applying only the first would leave the engine writing a discarded staging copy with the copy-back already gone — the buffer would silently never update. Neither pass is public on its own.
torch_tensorrt.save(output_format="executorch", zero_copy_kv=True)owns both steps and needs only the one flag. It installszero_copy_backend_config()itself, so the two are alternatives rather than a pair: handingsave()a config built with that function applies the un-staging twice, and the second run finds the buffers already wired straight to their delegates and changes nothing.The discriminator, which everything else depends on
An aliased KV mutation and a #4459 copy-back mutation are identical in the graph: both are
BUFFER_MUTATION target=X <- getitem(engine, i)whereXis also an engine input. They are told apart only by reading the engine's own serializedaliased_io. Both ways of getting it wrong are silent — rewiring a copy-back loses a real update, and leaving an aliased cache staged has the engine write scratch that is discarded.A method may hold both kinds at once, and that combination is supported rather than refused. It is measured on a real engine: a model with
index_copy_KV caches plus a ring-shiftedconv_stateexports with the caches on their own placeholders andconv_statestill bound to a delegate output, in one method, with no code change beyond thealiased_ioread. #4459's_trt_no_kv_aliasmakes that signal cleaner still, by excluding non-KV mutable buffers from aliasing at the converter.Measured
30-layer Gemma4-MoE 26B-A4B, real weights, int4 MoE + bf16, full 262144 vocab, prefill 2048 / max_len 2560, hybrid TensorRT+CUDA on one 80GB A100, TensorRT 11.1.0.106:
The 1,730× collapse in staging is the mechanism; the rest follows. Host-side gaps fall with it, 114.39 → 4.81 ms, and compute is untouched — the int4 MoE GEMV runs at 298.7 µs against 299.5 before. Generated token ids are byte-identical to the staged run.
Reading it commit by commit
The first six commits are the layering, in dependency order; the rest answer the review. Every commit stands alone, with no failure at any of them but the pre-existing
#4635one noted below:_zero_copy.py).backend.py).TensorRTBackend.cpp, +35/−13. The two shapes differ in arity by exactly the aliased-output count, so the argument count identifies which one a.pteis; no serialized flag, and a.ptewritten before zero-copy still takes the threaded-output branch.zero_copy_kvopt-in,zero_copy_backend_config(), thesave()flag, docs, the KV example, and the CI reference-runner check. The widest one.check_zero_copy_kv()moved into the library, walking every method, run bysave(), and handed a real finalized program by the engine tests; the elided-output guard run after the dead-code elimination, so a dead chain of any length cannot pass for a surviving output; a repeatedaliased_ioentry refused when the blob is parsed; the KV decode check run with and without a caller stream; the un-staging keyed on the shape the program ends up in rather than on the edit the pass made, which is what makes a second application a no-op and what refuses any staging copy the pass leaves behind; a prose pass over the comments, docstrings and user guide; and the two refusals below.Two shapes that exported cleanly and then failed on device are now refused at export. Both came out of review, which could not run anything, so before either was changed both were reproduced by loading the offending
.ptethrough a real ExecuTorch C++ runner with this backend linked in, on TensorRT 11.1.0.106:ExecutorchBackendConfig(enable_non_cpu_memory_planning=False)produces the same graph —PropagateDevicePassinserts no staging copy and writes the delegate's device straight onto the placeholder's spec — and then plans every tensor into the one host arena regardless of any spec. The buffer's spec is thereforeDeviceType.CUDAunder both configurations, so a spec-device check alone does not discriminate;zero_copy_backend_config()now readsenable_non_cpu_memory_planningoff the config it is building from and hands it to the pass, which refuses a direct argument that is off-CUDA or host-planned, naming the buffer whose update would be lost. The.ptethat motivated this exported clean and passedcheck_zero_copy_kv(), then failed its firstexecute():aliased input 'buf_k_cache' must be device-resident,Error::InvalidArgument.check_zero_copy_kv(), and writes a.ptethat loads, binds both aliases at init, and fails everyexecute()on the argument count (expected at least 7 args, got 6,Error::InvalidArgument).TensorRTBackend.preprocessalready holds both sets ten lines apart; it now compares them and refuses, so the failure lands where the export can be re-run. Teaching the runtime to read the elided names from the compile spec is the other way out and is deliberately not taken: refusing is far smaller, and refused → supported is a compatible progression later.Verified
pytest tests/py/dynamo/executorch/, the same selection those six were counted with), of which 288 pass. Every one of those runs has the same single failure,test_runtime_wheel_pins_cuda_13_native_dependencies, which fails on plainmaintoo since bring back cu126 support #4635 madeTENSORRT_DISTRIBUTIONdynamic without updating it; nothing else fails at any commit.blackat the CI-pinned 26.3.1 andmypyat the pinned 1.15.0: clean at every commit.ruffis clean over the changed files and clean relative toorigin/main.//tests/cpp/executorch:executorch_backend_testsbuilds and passes 4/4 on a real device.bash -nandshellcheck -S warningclean, the workflow parses as YAML, and the argument parser was exercised standalone on three input shapes underset -u. The zero-copy iteration asserts thealiased_iodiscriminator, and was shown to raise on a program finalized withoutzero_copy_backend_config(). The lane itself was then run end to end and exits 0, and the zero-copy no-sync branch was instrumented and observed to execute — the guarded mode takes the skip path on all three of itsexecute()calls, against the unguarded mode in the same binary, which synchronizes every time.Merge note for #4572
#4572 added the coalesced
.ptetoverify-executorch-reference-runner.shas a third positional argument. This stack made the trailing arguments variadic, so it can hand the staged and the zero-copy KV.pteto one run. Those cannot coexist, so the coalesced model is now named —--coalesced=PATH— and the KV models stay variadic. Both sides' behaviour is preserved and only the encoding of #4572's argument changed. The script has two callers, both workflows in this repo. Only the test workflow needed changing, in the same commit as the script; the build workflow passes a single positional model, which the variadic parser takes unchanged.Known gaps
zero_copy_kv=Trueprogram withoutzero_copy_backend_config()does not raise on its own. The engine writes a staging copy that is discarded and the buffer never updates — wrong output, not a crash. The finalized program does carry the evidence, socheck_zero_copy_kv()detects it:save(zero_copy_kv=True)runs that check before writing, and a caller assemblingexportandto_executorchby hand can call it themselves. What remains is that nothing forces them to.insert_write_back_for_buffers_passpairs each mutation spec with a value by walking one counter over thecopy_nodes it created followed by every output it copied nothing for. A rewired mutation is one of the latter, so a method holding a zero-copy cache and an ordinary copy-back buffer comes out ofto_executorch()with the two specs crossed. It is upstream's — stock ExecuTorch reproduces it withrun_reinplace_pass=Trueand no TensorRT involved — and it is going upstream as a bug there. The emitted.pteis byte-identical either way, and every consumer on the emit path tests membership in the map's values rather than reading the pairing, so nothing on this branch is wrong because of it. What is affected is a reader of the finalizedgraph_signature.execute()and needs a real engine; its automated coverage is the CI reference-runner check. The duplicate-aliased_iorejection is done when the blob is parsed, and is unit-tested there._apply_zero_copy_kvloops over methods andcheck_zero_copy_kv()walks every method. The real-engine tests do run the check on a genuinely finalized program, but a single-method one; the multi-method behaviour of both rests on fabricated program shapes, and no two-method.ptewas built.metadata_only=True; folding them into the existing resolved-record handoff is a follow-up, deliberately not bundled here.If
AliasKind.USERever gets a producerAliasKind.USERis a placeholder that nothing emits, so everyaliased_ioentry a.ptecarries today is akv_cache_updateon a lifted buffer. Neither half of zero-copy reads the kind: the partitioner elides an aliased output only when the input it aliases is a buffer placeholder the rewiring marked, and the runtime derives the elided arity from the count of allaliased_ioentries. Those two agree only while every aliased output on an engine is elided together, so a"user"alias on a plain input beside an elided one is a shape neither the binding check norcheck_zero_copy_kv()objects to and the runtime cannot execute.preprocessrefuses that engine outright rather than letting it reach a.pte, which is what a producer of the kind would meet first. Supporting the mix instead means giving the runtime the elided names to read rather than an arity to subtract; refusing now leaves that open.Type of change
Checklist: