fix(executorch): support KV-cache aliased I/O in the TensorRT delegate - #4445
fix(executorch): support KV-cache aliased I/O in the TensorRT delegate#4445Conarnar wants to merge 1 commit into
Conversation
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.
52d316b to
c5ab1e4
Compare
|
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: The new behavior is: Checking only 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. |
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:
Tests