Skip to content

fix(executorch): support KV-cache aliased I/O in the TensorRT delegate - #4445

Open
Conarnar wants to merge 1 commit into
pytorch:mainfrom
Conarnar:fix/executorch-kv-alias-bindings
Open

fix(executorch): support KV-cache aliased I/O in the TensorRT delegate#4445
Conarnar wants to merge 1 commit into
pytorch:mainfrom
Conarnar:fix/executorch-kv-alias-bindings

Conversation

@Conarnar

Copy link
Copy Markdown
Contributor

Description

The dynamo TensorRT converter emits in-place KV-cache updates as aliased engine outputs (an output binding sharing device memory with an input binding). The ExecuTorch delegate path did not support this: the blob carried no aliasing info and the C++ backend expected a separate delegate arg for every engine input/output, so a model with KV-cache aliasing failed at execute() with an arg-count mismatch (0x12).

Thread the aliasing through the blob and let the delegate own the cache:

  • serialization.py / backend.py: carry aliased_io (output -> input, kind) in the TR01 blob (list form so the small C++ parser can walk it).
  • TensorRTBlobHeader.{h,cpp}: parse the optional aliased_io array (absent in older blobs -> empty, backward compatible).
  • TensorRTBackend.{h,cpp}: at init, resolve each aliased output to its input's binding index. The aliased KV buffers are not threaded as delegate args (with the graph-level copy_ gone, ExecuTorch sees a non-mutated buffer and freezes it as a constant), so the delegate self-owns them: it allocates a persistent device buffer per aliased input, binds both the input and its aliased output in-place to that buffer, and reuses it across execute() calls so the cache accumulates over the decode loop. The arg-count is (num_inputs - num_self_owned) + (num_outputs - num_aliased_outputs).

Tests

  • test_serialization.py: aliased_io survives a serialize/deserialize round trip; a blob without the key parses to an empty mapping.
  • test_executorch_blob_header.cpp: parse aliased_io (incl. kind), plus the missing/empty cases stay backward compatible.

@meta-cla meta-cla Bot added the cla signed label Jul 29, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests component: api [Python] Issues re: Python API component: api [C++] Issues re: C++ API labels Jul 29, 2026
@github-actions
github-actions Bot requested a review from narendasan July 29, 2026 19:08
The dynamo TensorRT converter emits in-place KV-cache updates as aliased
engine outputs (an output binding sharing device memory with an input
binding). The ExecuTorch delegate path did not support this: the blob carried
no aliasing info and the C++ backend expected a separate delegate arg for
every engine input/output, so a model with KV-cache aliasing failed at
execute() with an arg-count mismatch (0x12).

Thread the aliasing through the blob and let the delegate own the cache:

- serialization.py / backend.py: carry aliased_io (output -> input, kind) in
  the TR01 blob (list form so the small C++ parser can walk it).
- TensorRTBlobHeader.{h,cpp}: parse the optional aliased_io array (absent in
  older blobs -> empty, backward compatible).
- TensorRTBackend.{h,cpp}: at init, resolve each aliased output to its input's
  binding index. The aliased KV buffers are not threaded as delegate args
  (with the graph-level copy_ gone, ExecuTorch sees a non-mutated buffer and
  freezes it as a constant), so the delegate self-owns them: it allocates a
  persistent device buffer per aliased input, binds both the input and its
  aliased output in-place to that buffer, and reuses it across execute() calls
  so the cache accumulates over the decode loop. The arg-count is
  (num_inputs - num_self_owned) + (num_outputs - num_aliased_outputs).

Tests:
- test_serialization.py: aliased_io survives a serialize/deserialize round
  trip; a blob without the key parses to an empty mapping.
- test_executorch_blob_header.cpp: parse aliased_io (incl. kind), plus the
  missing/empty cases stay backward compatible.
@Conarnar
Conarnar force-pushed the fix/executorch-kv-alias-bindings branch from 52d316b to c5ab1e4 Compare July 29, 2026 21:13
@shoumikhin

Copy link
Copy Markdown
Contributor

Thanks for adding aliased-I/O metadata to the ExecuTorch blob. This fixes a real missing capability, and the serialization changes look reasonable.

I found a blocking issue that applies even when only one method is exported. The patch removes every aliased input from the delegate arguments and replaces it with a private, zero-initialized buffer. If the caller supplies the cache tensor, its existing contents are ignored and the caller cannot observe the update.

For example, the expected behavior is:

caller cache    -> TensorRT input
TensorRT output -> same caller cache

The new behavior is:

caller cache        -> ignored
private empty cache -> TensorRT input and output

Checking only kind == "kv_cache_update" is not enough, because a compiler-generated KV update can still use caller-owned storage. Alias kind describes who enforces the alias, not who owns the buffer, so I think the metadata needs to describe storage ownership separately from alias kind.

Runtime-owned storage also needs a defined lifetime. The current buffer is initialized once, has no reset operation, and is shared by every call using that loaded method. Two conversations, or concurrent requests, would therefore use the same cache.

Could caller-owned aliases stay explicit delegate inputs, with the output bound to the same pointer? If runtime-owned storage is genuinely needed, please add explicit ownership and a stable state identity, plus reset or session-selection behavior.

On testing, it would help to pin the runtime behavior directly rather than through serialization round-trips: caller-visible mutation, repeated calls on one loaded method, a fresh load starting from a known state, and sequence isolation.

For the multi-method case: multi-method Torch-TensorRT ExecuTorch export is still in flight (#4440), so a natural next step is to build a two-method prefill/decode test on top of it and assert that decode observes cache state written by prefill. I want to flag the expected outcome up front: with the current design I believe that test fails, because each method loads as its own delegate with its own private cache, so there is no shared object for the two methods to write through. That is really why I think the cache has to be owned above the delegate and bound into both methods; the shared-cache test is the acceptance criterion for that ownership change rather than something stacking alone will make pass.

For mixed TensorRT and CUDA execution there are two independent requirements worth testing separately: (1) both methods bind the same KV-cache storage, and (2) dependent GPU work from the two delegates is ordered. Ordering needs a shared caller stream (there is separate in-flight work for that, #4421), so a mixed test should run on top of it, but note the shared stream only provides (2). Property (1), shared storage across the TensorRT and CUDA delegates, cannot come from a delegate-private buffer, so it also depends on moving cache ownership above the delegate.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [C++] Issues re: C++ API component: api [Python] Issues re: Python API component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants