Repository navigation
Add workspace estimation accounting and reporting - #31962
Conversation
Move workspace reporting out of GetCapabilityForEP and into the completed GraphPartitioner::Partition boundary. Track workspace estimates as pending during capability probing and commit them only for nodes that ORT actually assigns. Replace the generic safety-margin workspace with the MatMulNBits Level-1 estimate before CUDA's budget decision. Preserve workspace estimates across the layout-transformation deferred-commit path, and report non-workspace, workspace, and total estimated memory once per resource-accounted EP. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839
There was a problem hiding this comment.
Pull request overview
Adds workspace-memory estimation, accounting, and reporting for resource-aware CUDA graph partitioning.
Changes:
- Combines profile, estimator, and fallback workspace costs.
- Preserves workspace accounting across two-pass NHWC partitioning.
- Adds Level-1/Level-2 reporting and accountant tests.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
onnxruntime/core/framework/graph_partitioner.cc |
Commits and reports accepted workspace estimates. |
onnxruntime/core/framework/resource_accountant.cc |
Implements workspace selection and accounting. |
onnxruntime/core/framework/session_state.cc |
Reports Level-2 workspace declarations. |
onnxruntime/core/providers/cuda/cuda_execution_provider.cc |
Applies MatMulNBits estimates to budgets. |
include/onnxruntime/core/framework/resource_accountant.h |
Defines workspace accounting interfaces. |
include/onnxruntime/core/graph/indexed_sub_graph.h |
Commits complete per-node resources. |
onnxruntime/test/framework/resource_accountant_test.cc |
Tests workspace accounting policies. |
onnxruntime/test/internal_testing_ep/internal_testing_partitioning_tests.cc |
Updates renamed accountant reference. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Review summary (multi-agent pass: readability, correctness, adversarial, spec/math, cross-module)Reviewed the full diff. Two findings look like real semantic risk worth resolving before merge; the rest are minor/nits. Major1. No-profile path drops the safety floor entirely instead of taking
Suggested fix: 2. Stats-miss nodes silently go from cost=0 to fully-charged, and double-count shared initializers across the two accounting models Second-order issue: the stats path doesn't touch Suggested fix: either keep Minor
What's solid
Design-doc drift (multiple stale references to renamed methods/signatures in |
Separate runtime workspace, persistent prepack buffers, and temporary prepack scratch in Level-1 estimates, and conservatively charge each component during CUDA partitioning. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Re-review of latest push (draft → ready, 8→17 files)Re-ran the same review pass against the updated diff. Good progress — one of the two Major findings from the previous review is substantially addressed; the other two remain. Resolved / substantially addressedPrevious Major #1 (no floor against unknown allocations) — now addressed via a structural fix, not just a floor. One residual nuance: Still open1. [Major] Stats-file-miss node still goes from cost=0 to fully-charged fallback, with an unreconciled double-count risk The double-count concern also still stands: the stats-hit path doesn't touch 2. [Major] Plugin EP host bridge still uses the 1-arg Minor (unchanged from before)
New surface added since last review (not previously covered) — looks solid
Overall: solid iteration, the prepack-memory gap is closed. The stats-miss double-count and plugin-bridge divergence are the two things I'd still want addressed or explicitly deferred-with-tracking-issue before merge. |
Keep nodes missing from a supplied stats file at zero cost so profile-based sessions do not fall through to ad-hoc initializer and workspace bookkeeping. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add estimator-aware NHWC and CUDA budget coverage, make resource arithmetic checked, clarify per-graph Level-2 reporting, and document deferred plugin estimator parity. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1faf9517-5062-4655-a967-3834143334d0
Extract Level-1 runtime workspace into explicitly initialized scalar state so GCC 14 reduced builds do not diagnose nested optional storage as maybe uninitialized. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1faf9517-5062-4655-a967-3834143334d0
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
One blocking accounting issue remains. The structured lifetime split, checked byte arithmetic, dynamic-shape fallback, accepted-node commit semantics, and diagnostics look sound, and the earlier test/reporting requests are addressed on this head. However, the offline-prepacked GPU path is charged for a duplicate packed-weight allocation that PrePack_B does not make.
Review with build + test verificationI reviewed this PR and, since the description notes "Tests were not run for this update", I built the branch and ran the new tests on real hardware. Environment: 4× A100-SXM4-80GB (SM80), CUDA 12.9, cuDNN 9.10.2, Release, Build and test status: green
Note for anyone reproducing: the CUDA test cases need Critical — the ad-hoc path double-counts prepacked weights, and the headline test encodes that bugIn the ad-hoc branch of Per if (is_packed) {
++number_of_prepacks_counter_;
if (constant_initializers_use_count.count(input_name) && --constant_initializers_use_count[input_name] == 0) {
st->initialized_tensors_.erase(ort_value_idx);
constant_initialized_tensors.erase(ort_value_idx);
}
}So the original copy does not exist in steady state, and charging both is a genuine double count — not merely conservative, as Using the exact constants from the new e2e test (
The part I'd most like you to look at: this directly determines the outcome of The test's 500/600 KiB straddle works because of the double count. Remove it (charge That means the budget thresholds in the test were effectively back-derived from the current behavior. When fixing the accounting, please recompute them deliberately — otherwise the correction will look like a regression. One asymmetry worth capturing in the docs: this is only a double count on the ad-hoc path. On the profile path, Major — Level 1 omits the largest init-time allocation for this exact kernel
While Major — partitioning now depends on how fpA_intB was enabledThe estimator reads only the This PR removes the justification that made that divergence safe:
but keeps describing the Worth noting: my test run used the env-var path. No test covers the session-config path, so this gap is currently invisible to CI. Suggest resolving the effective option once at EP/session level, feeding both capability estimation and kernel creation, and adding a config-only budget regression test. Major — undeclared behavior change: max-shape inference now feeds ad-hoc output sizingThe new It also runs against Two small things in the same block:
Minor
Pre-existing, not introduced hereIn What I think is genuinely well done
Happy to share the exact build commands or rerun anything against a revision. |
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1faf9517-5062-4655-a967-3834143334d0
This is a valid observation about precision, but I want to clarify the CUDA initialization sequence. Before PrePack() runs, ORT’s allocation planner assigns each initializer to its first consumer’s device. SaveInitializedTensors() then allocates and copies all CUDA initializers to that device. Only afterward does ORT create kernels and invoke PrePack() sequentially node by node. Therefore, although the model’s external data may originally be mmap-backed on the host, the source initializers relevant to CUDA prepacking already consume CUDA device memory when PrePack() begins. For runtime prepacking, the source CUDA initializer and newly allocated packed destination coexist until the source’s final successful prepack consumer releases it. Packed destinations remain live for the session. Level-1 consequently reports a conservative initialization-time upper bound; it does not yet model the sequential release of source initializers or calculate the maximum live set across the prepack order. Exact peak accounting would require explicit source and destination lifetimes, initializer consumer counts, and prepack ordering. The offline-prepacked case is different. MatMulNBits::PrePack_B() reuses an already GPU-resident offline-prepacked B initializer in place and does not allocate another B destination. The base accountant still charges the retained initializer, while persistent_prepack_bytes now excludes the nonexistent duplicate B allocation. Newly allocated scale and zero-point prepack destinations remain included. A CUDA integration test verifies that the offline-prepacked weight is not charged twice. I’ve clarified this contract in Level1MemoryEstimate : it describes additional persistent prepack destinations and initialization scratch, while original CUDA initializer storage remains part of the accountant’s base initializer cost |
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1faf9517-5062-4655-a967-3834143334d0
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1faf9517-5062-4655-a967-3834143334d0
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1faf9517-5062-4655-a967-3834143334d0
Document that Level-1 prepack estimates include only allocations additional to already-accounted initializers, and distinguish sequential MatMulNBits tactic profiling from PrePack scratch. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1faf9517-5062-4655-a967-3834143334d0
Description
Adds workspace-memory accounting and reporting for resource-constrained graph partitioning.
GetCapabilityand commits it only for accepted nodesmax(profile, estimator)when both exist, and fallback when neither is availableMatMulNBitsworkspace estimator before CUDA budget decisionsSizeBasedResourceAccountantto reflect that it handles profile, fallback, and operator-estimated costsRelated to #29775.
Testing
Added resource-accountant and two-pass partitioning unit coverage, including profile-versus-estimator comparison reporting. Tests were not run for this update.