microsoft / microsoft/onnxruntime

[Feature Request] Create a shared non-CPU allocator that outlives inference sessions

Open
#25,785 2 comments 0 reactions 0 assignees View on GitHub
ep:WebGPU feature request WebNN
Dominant language
C++
Stars
21.9k
Forks
4.2k
Avg merge
4d 11h
Merged PRs (30d)
184

Description

### Describe the feature request

When I tried to create on-device tensors for WebGPU EP or OV EP, I found ORT currently doesn't support creating a shared non-CPU allocator which is required for creating the tensors.
1. [Environment::CreateAndRegisterAllocator](https://github.com/microsoft/onnxruntime/blob/7d719758f34ef99123cc26aa18bbfda508108c57/onnxruntime/core/session/environment.cc#L127) can only create CPU allocators.
2. [Environment::CreateAndRegisterAllocatorV2](https://github.com/microsoft/onnxruntime/blob/7d719758f34ef99123cc26aa18bbfda508108c57/onnxruntime/core/session/environment.cc#L348) can only create allocators for `kCudaExecutionProvider` and `kCpuExecutionProvider`.
3. [Environment::CreateSharedAllocator](https://github.com/microsoft/onnxruntime/blob/cb4d4af2ce20d9ff45ce04394fcf6c623a5b64a7/onnxruntime/core/session/environment.cc#L638) with a given `ep_device` reports error `"Invalid memory type for OrtEpDevice."` because the `ep_device.device_memory_info` and `ep_device.host_accessible_memory_info` are both nullptr for some EPs that we are focusing on: WebGPU, OV and even default CPU EP.

As the first step, we plan to create a WebNN prototype based on the WebGPU EP since the support for the OV EP is still far away. Do you have any suggestions and plans?

### Describe scenario use case

For Whisper and other language models, WebNN needs to keep the KV cache tensors on device that avoids host-and-device data copies during the decoder inference iterations. WebNN needs these pre-allocated on-device tensors shared among different ORT sessions.

@huningxin @fdwr @skottmckay

Contributor guide

Open the contributing guide

Research direction

Start with the linked Environment::CreateAndRegisterAllocator, CreateAndRegisterAllocatorV2, and CreateSharedAllocator entry points in onnxruntime/core/session/environment.cc, then inspect how WebGPU, OV, and CPU execution providers describe device and host-accessible memory. Done means defining and validating a shared non-CPU allocator that can provide pre-allocated device tensors across inference sessions.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
ai-infra-agents
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.