feat(executorch): let a TensorRT engine write the caller's KV cache in place - #4667
feat(executorch): let a TensorRT engine write the caller's KV cache in place#4667Conarnar wants to merge 15 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
left a comment
shoumikhin
left a comment
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 |
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.
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 |
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, |
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.
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 |
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.
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.
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, to keep the inline review comments attached. `_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.
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.
2 participants
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: