microsoft / microsoft/onnxruntime
[Feature Request] Create a shared non-CPU allocator that outlives inference sessions
- 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
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