[Selena Yang (selenayang888)]

Summary

Add an EP device API for querying total memory capacity and currently available memory, with support for both in-tree and plugin CUDA execution providers. This lets consumers such as Foundry Local and ONNX Runtime GenAI query capacity without provider-specific code.

Changes

• Add OrtApi::EpDevice_GetMemoryCapacity and C++ wrappers:
  • GetMemoryCapacity() returns {total, available} in bytes.
  • GetTotalMemoryCapacity() queries total capacity only.
• Add the optional OrtEpFactory::GetDeviceMemoryCapacity callback with version/null gating and internal factory/provider-bridge forwarding.
• Allow either output to be omitted, requiring at least one non-null output.
• Return ORT_NOT_IMPLEMENTED for unsupported providers or memory types, and preserve output values on failure.
• Implement CUDA queries using cuDeviceTotalMem for device-global capacity and cuMemGetInfo for available memory.
• Avoid CUDA context initialization for total-only queries.
• Preserve the caller’s CUDA context on successful queries and attempt recovery before releasing the primary context after a failed pop. Document unrecoverable driver-error limitations.
• Include CUDA operation, error code, name, and description in failures using fixed-size error buffers.
• Document memory-type selection, partition-aware capacity, free-memory snapshot semantics, and context-initialization costs.
• Add CUDA driver linkage for minimal CUDA builds and GPU tests.
• Update plugin compatibility documentation without raising its compatibility floor.
• Adjust the Android minimal binary-size threshold by 1 KiB.

API Semantics

Total memory is the device-global capacity exposed for the selected EP device and memory type, not a per-process allocation budget. Partitioned or virtual devices report their exposed capacity, not necessarily the entire physical board.

Available memory is an advisory free-memory snapshot, not a reservation or allocation guarantee.

For CUDA, total-only queries do not initialize a context. Available-memory queries may create and destroy a primary context when nothing else retains it, adding latency and transient memory usage. Context initialization can fail under severe memory pressure.

[Copilot]

🟡 Changes recommended

The minimal CUDA build lacks the required driver-library dependency, and the new mock capacities overflow on 32-bit targets.

2 open issues

• Link CUDA minimal builds against the CUDA driver library · New
• Use portable mock capacities for 32-bit AutoEP tests · New

What changed in this PR

Adds a provider-independent API for querying device memory capacity, giving downstream consumers access to total and currently available memory through ONNX Runtime.

Changes:

• Adds C and C++ query APIs with an optional execution-provider factory callback.
• Implements CUDA queries through shared logic that manages context switching and preserves outputs on failure.
• Adds example-plugin coverage and deterministic regression tests.

File	Description
onnxruntime/​test/​autoep/​test_registration.cc	Adds API, context, and failure-path tests.
onnxruntime/​test/​autoep/​library/​example_plugin_ep/​ep_factory.h	Declares the example callback.
onnxruntime/​test/​autoep/​library/​example_plugin_ep/​ep_factory.cc	Provides synthetic memory capacities.
onnxruntime/​core/​session/​plugin_ep/​forward_to_factory_impl.h	Forwards queries to internal factories.
onnxruntime/​core/​session/​plugin_ep/​ep_factory_provider_bridge.h	Adds version-checked provider forwarding.
onnxruntime/​core/​session/​plugin_ep/​ep_factory_internal.h	Exposes internal query forwarding.
onnxruntime/​core/​session/​plugin_ep/​ep_factory_internal.cc	Registers the forwarding callback.
onnxruntime/​core/​session/​plugin_ep/​ep_factory_internal_impl.h	Supplies the unsupported-query default.
onnxruntime/​core/​session/​ort_apis.h	Declares the C API implementation.
onnxruntime/​core/​session/​onnxruntime_c_api.cc	Validates, dispatches, and registers queries.
onnxruntime/​core/​providers/​cuda/​shared_inc/​device_memory_capacity.h	Adds testable query and cleanup logic.
onnxruntime/​core/​providers/​cuda/​shared_inc/​cuda_memory_capacity.h	Wraps CUDA Driver API calls.
onnxruntime/​core/​providers/​cuda/​plugin/​cuda_ep_factory.h	Declares the CUDA plugin callback.
onnxruntime/​core/​providers/​cuda/​plugin/​cuda_ep_factory.cc	Implements CUDA plugin queries.
onnxruntime/​core/​providers/​cuda/​cuda_provider_factory.cc	Implements shared-provider CUDA queries.
include/​onnxruntime/​core/​session/​onnxruntime_ep_c_api.h	Appends and documents the optional callback.
include/​onnxruntime/​core/​session/​onnxruntime_cxx_inline.h	Implements the C++ wrapper.
include/​onnxruntime/​core/​session/​onnxruntime_cxx_api.h	Declares the C++ wrapper.
include/​onnxruntime/​core/​session/​onnxruntime_c_api.h	Appends and documents the public query API.

_{Review effort: Balanced}

[Justin Chu (justinchuby)]

Does it say what device this info is about? What if an ep manages multiple different memory?

[Tianlei Wu (tianleiwu)]

Re-review at bb453a3e1d (head unchanged since 07:13 UTC; previous review at 7b7fce9e65).

API/ABI looks good. EpDevice_GetMemoryCapacity is appended at the end of OrtApi, GetDeviceMemoryCapacity is the last OrtEpFactory member, both \since 1.32 with correct version/null gating; unsupported EPs return ORT_NOT_IMPLEMENTED with the outputs untouched; the memory_type selector answers the earlier "which memory?" question; my wrapper-placement thread is addressed.

Blocking: the new real-CUDA test (ep_device_memory_capacity_test.cc).

1. cuCtxCreate_v2 (L63) is not declared by the CUDA 13 headers, so onnxruntime_test_all does not build on CUDA 13. Linux CUDA x64 Release and its no-cuDNN variant already fail on it at this head (four jobs at 02b760c76a).
2. Where it compiles, registering and unregistering the in-tree CUDA provider library (L36/L50) shuts down the process-wide CUDA kernel registry. At 02b760c76a the test passes, then the same nine unrelated tests fail in three GPU jobs (Windows CUDA, Linux and Windows TensorRT); they passed before the test existed.

Inline suggestions: keep the CUresult in errors (a nearly full GPU makes the query fail with an out-of-memory error), document the per-call context-creation cost, update the plugin design doc. The PR description is stale (still says no GPU validation).

Validation: head CI at 08:03 UTC is 51 success / 2 failed (the compile error above) / 34 pending; the Windows CUDA, plugin and TensorRT test jobs have not reported for this head, so (2) rests on the identical registration code at 02b760c76a plus code reading. Local (A100, CUDA 13.0): reproduced (1), verified a guarded fix on the 12.6/12.8/12.9/13.0 headers, and ran a standalone harness for the query; no full ORT build or onnxruntime_test_all run.

[Tianlei Wu (tianleiwu)]

Update at bb453a3e1d: the CI jobs that were still pending when I posted my review have now finished, and they confirm both blocking points.

1. cuCtxCreate_v2 (ep_device_memory_capacity_test.cc:63) fails in four CUDA 13.0 build jobs at this head: Build Linux CUDA x64 Release, Build Linux CUDA x64 Release without cuDNN link, Build Linux CUDA Plugin EP x64 Release ('cuCtxCreate_v2' was not declared in this scope; did you mean 'cuCtxCreate_v4'?), and Windows CUDA Plugin EP Build without cuDNN (error C3861: 'cuCtxCreate_v2': identifier not found).
2. The nine unrelated failures are now reproduced at this head instead of being inferred from the earlier one. In Test Linux TensorRT x64 Release, Windows GPU TensorRT CI Pipeline Test Job and Windows GPU CUDA CI Pipeline Test Job, EpDeviceMemoryCapacityTest.CudaCapacityPreservesContext passes and the same run then reports [ FAILED ] 9 tests: InferenceSessionTests.TestArenaShrinkageAfterRun; MoeExpertCountingTest.CudaMoE, CudaQMoE, CudaQMoEPackedIntGemv, CudaQMoESingleToken, CudaQMoETiled; and SessionStateTest.TestLayeringPartitioning, TestResourceAwarePartitioning_LargeLimit, TestResourceAwarePartitioning_NoLimit. Windows CUDA Plugin EP Test succeeds.

wasm_Debug / build-wasm also fails (external-data loading in the pattern-rewriter tests: Module.MountedFiles is not available), which this PR does not touch.

[github-actions]

You can commit the suggested changes from lintrunner.

[Baiju Meswani (baijumeswani)]

Big-picture assessment from the Foundry Local use case: this is an appropriate, useful foundation for excluding model variants that require more VRAM than the target device provides. Foundry's catalog vramsize metadata supplies the model side of that comparison, so I do not think this PR needs to introduce a model-memory-requirement API or expand into a general model evaluator.

We should keep three decisions distinct:

• Hardware memory fit: compare the catalog's required runtime VRAM with the capacity of an actual selectable EP device, accounting separately for configured allocation limits. The metadata needs defined units and workload assumptions.
• Availability now: current free memory is an advisory readiness signal, not an allocation guarantee. For a hardware-compatible catalog, temporary memory pressure should not permanently hide an otherwise supported model; if we promise 'loadable now', represent temporary unavailability separately and recheck at load time.
• EP/model compatibility: driver, architecture, and artifact compatibility remain separate EP-specific checks. A successful capacity query alone does not establish them.

Foundry should retain individual device identities through filtering and placement, rather than summing unrelated GPUs' VRAM. Unsupported queries should mean unknown, not zero capacity; coverage beyond CUDA depends on those EPs implementing the optional callback.

The append-only API placement, version-guarded forwarding, failure-output preservation, and normal-path CUDA context handling look sound. I have left two inline comments: a correctness concern in the failed-context-pop cleanup path, and a recommendation to allow total-capacity queries without paying for a dynamic free-memory query. The latter is particularly relevant to lightweight catalog discovery now that vramsize supplies the requirement.

Validation performed for this review: the PR's three deterministic helper tests passed unchanged in a standalone build. Additional probes exercised the shared helper on eight real CUDA GPUs, confirming normal-path context preservation and allocation-sensitive free-memory reporting, and reproduced the failed-pop state using the PR's mock. Cold queries took approximately 480-520 ms per device on this machine. This was not a full ORT build or public-API/plugin integration test.

[Tianlei Wu (tianleiwu)]

Re-review at 8e49c3619c (head unchanged for the whole review). Approving the code, with one lint fix to land first and CUDA CI still to confirm.

Previous blockers are fixed. cuCtxCreate is guarded for CUDA 13 (#if CUDA_VERSION >= 13000), and the real-GPU test no longer registers or unregisters the in-tree CUDA provider. baijumeswani's two open threads are addressed at head as well (I replied on both with the evidence).

Before merging

1. "Python format" is red because of this PR: clang-format wants two more spaces on device_memory_capacity.h:19 (one-click suggestion inline; lintrunner -a does the same). It is the only failure so far (39 passed, 47 pending, 1 failed at 21:07Z).
2. CUDA CI has not finished (Linux CUDA, CUDA Plugin EP, Windows CUDA, TensorRT), so the CUDA 13 compile fix is not confirmed by CI yet. I checked that branch against the CUDA 13.0 cuda.h locally and it compiles. Please also confirm EpDeviceMemoryCapacityTest.CudaCapacityPreservesContext shows [ OK ] and not [ SKIPPED ] in the Plugin EP jobs.

Non-blocking (inline): license headers on the three new files; the in-tree (provider-bridge) CUDA path has no test that runs, with a small test suggested; a doc note for the C++ wrapper. The PR description is also stale: it still says the lint checks passed and that CUDA compilation and real-GPU validation are outstanding. Please refresh it after the lint fix.

What I validated on a real GPU. I compiled the PR's device_memory_capacity.h and cuda_memory_capacity.h unchanged against CUDA 13.0 headers (-Wall -Wextra, no warnings) into a standalone program and ran it on 8x A100-80GB (driver 580.159.04):

• total-only query: about 0.01 ms (driver already initialized), creates no context
• cold available query: about 230-330 ms (creates and destroys the primary context); 30 repeated calls leaked nothing, and the caller's current context and stack depth were preserved
• caller already holding the device's primary context: about 0.03 ms, and available matched a direct cuMemGetInfo exactly with a live 1 GiB allocation
• caller with another device's context current: preserved, and the queried device stays inactive
• invalid ordinal: outputs untouched, error text cuDeviceGet failed: CUDA error 101 (CUDA_ERROR_INVALID_DEVICE): invalid device ordinal
• 4 threads x 200 queries while another thread uses the same device, and parallel cold queries on 4 devices: no errors, no leaks

Not covered: the failed-pop recovery path (it cannot be induced on a real driver; the mock tests cover it) and the ORT plumbing, since I did not build ORT with CUDA locally.