microsoft / microsoft/onnxruntime

Register Level-1 workspace estimators with kernel metadata

Open
#32,341 0 comments 0 reactions 1 assignee Claimed by @titaiwangms View on GitHub
Dominant language
C++
Stars
21.9k
Forks
4.2k
Avg merge
4d 11h
Merged PRs (30d)
184

Description

## Problem

CUDA workspace Level-1 pilots currently dispatch estimators explicitly from
`CUDAExecutionProvider::GetCapability()` by checking operator domain and type.
This is acceptable for the MatMulNBits and PA/PMHA pilots, but extending the
same conditional block for every future kernel will create a second,
CUDA-specific registry that can drift from kernel registration.

An `op_type -> estimator` map inside `GetCapability()` would shorten the
conditional chain, but estimator inputs and build boundaries differ between
kernels. The longer-term design should use registration-based discovery rather
than another ad hoc map.

## Goal

Define a generic way for Level-1 workspace estimators to be registered and
discovered alongside kernel metadata, such as through `KernelCreateInfo` or an
equivalent framework-owned registry.

## Requirements

- The selected kernel registration identifies its Level-1 workspace estimator.
- `GetCapability()` does not require operator-specific string dispatch for each
newly supported kernel.
- The contract can consume the node, resolved/max input-shape metadata, and
provider/device information needed by operator-specific adapters.
- Unavailable estimates remain distinguishable from a valid byte count.
- Build configurations without a kernel or estimator do not create unresolved
references.
- In-tree and future plugin/C-ABI boundaries are explicit; no in-tree graph
types leak into reusable graph-free workspace formulas.
- The result can feed the resource-accounting/planner work tracked by #29775
and #32071 rather than remaining log-only.
- MatMulNBits and PA/PMHA migrate from the temporary explicit CUDA
`GetCapability()` blocks.

## Context

- #29811 introduced the MatMulNBits two-level workspace pilot.
- #32312 adds positional optional-aware Level-2 input shapes.
- #32321 adds PA/PMHA Level-1 and Level-2 adapters.
- #32071 integrates declared workspace with activation memory patterns.

This issue tracks dispatch/registration architecture only. Operator-specific
workspace formulas and planned-root runtime integration remain separate work.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.