Add GetUserProvidedOutput to KernelContext. - #32089
Open
Eric Crawford (ericcraw) wants to merge 1 commit into
Open
Add GetUserProvidedOutput to KernelContext.#32089Eric Crawford (ericcraw) wants to merge 1 commit into
Eric Crawford (ericcraw) wants to merge 1 commit into
Conversation
Provide a method to query preallocated output tensors supplied by the caller that are available to a kernel. This allows an execution provider to inspect and write directly to user-provided output buffers whose shapes may be unknown to the EP before inference but are known by the caller. This is particularly useful for dynamically shaped transformer models with large KV caches. An EP can reuse preallocated output tensors without requiring shape inference before execution or allocating an intermediate buffer and copying the result afterward. The EP remains responsible for validating the output tensor's shape and element type before writing to it.
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
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
1 participant
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
Add
KernelContext::GetUserProvidedOutputand the corresponding C API,KernelContext_GetUserProvidedOutput.The API allows an execution provider to query whether the caller supplied a
preallocated output tensor for the current kernel output and retrieve it as a
borrowed
OrtValue. It does not allocate, resize, or replace output tensors.The C++ wrapper and runtime execution-frame tracking are included, along with
tests covering
Run,IoBinding, missing user outputs, invalid indices, anddirect writes to the caller-provided buffer.
Motivation and Context
Execution providers need to efficiently handle dynamically shaped models where
the caller knows the expected output shape but the EP cannot quickly or easily
determine it before inference execution. This is particularly important for
transformer models with large KV caches.
The new API allows an EP to reuse a caller-provided output buffer directly,
avoiding an intermediate allocation and subsequent copy. Since the API returns
a borrowed buffer without performing dynamic output-shape validation, the EP
must validate the tensor's element type and shape before writing to it.