Add workspace estimation accounting and reporting

[Chi Lo (chilo-ms)]

Description

Adds workspace-memory accounting and reporting for resource-constrained graph partitioning.

• tracks pending workspace during GetCapability and commits it only for accepted nodes
• preserves correct accounting across the two-pass NHWC layout transformation flow
• combines workspace sources using the policy: profile only, estimator only, max(profile, estimator) when both exist, and fallback when neither is available
• applies the Level-1 MatMulNBits workspace estimator before CUDA budget decisions
• reports non-workspace, workspace, total estimated memory, and accepted-node source counts after partitioning
• reports accepted-node profile-versus-Level-1-estimator comparisons, including which source is larger and aggregate bytes from each source
• reports Level-2 per-node declarations and aggregate declared workspace bytes during session initialization; the aggregate is not a peak-live or preallocated-memory value
• renames the internal implementation to SizeBasedResourceAccountant to reflect that it handles profile, fallback, and operator-estimated costs

Related 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.

[Copilot]

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.

[Ti-Tai Wang (titaiwangms)]

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.

Major

1. No-profile path drops the safety floor entirely instead of taking max()
onnxruntime/core/framework/resource_accountant.cc (SizeBasedResourceAccountant::ComputeResourceCount, ad-hoc/no-profile branch): selected_workspace = level1_workspace_estimate.value_or(fallback_workspace) unconditionally substitutes the Level-1 estimate for the 1.5x fallback margin. The has-profile branch correctly does max(stats.total_temp_allocations, *level1_workspace_estimate), but this branch doesn't apply the same floor.

EstimateMatMulNBitsWorkspace only covers the fpA/intB GEMM kernel's own scratch buffer — it doesn't account for the node's prepack/persistent weight buffers (GetTransientScratchBuffer + fpA_intB_weight_buffer_/fpA_intB_scale_buffer_/fpA_intB_zero_buffer_ in matmul_nbits.cc). For a representative 4-bit 4096×4096 layer, budget drops ~33% (≈14.16MB → ≈9.44MB) vs. old fallback, while actual device usage for that node is closer to ~2x the weight blob. Net effect: less CUDA-OOM protection for exactly the models this feature targets.

Suggested fix: max(fallback_workspace, level1_workspace_estimate) here too, until the estimator covers all of a node's allocations (or explicitly document/gate the narrower scope with a comment + issue reference).

2. Stats-miss nodes silently go from cost=0 to fully-charged, and double-count shared initializers across the two accounting models
Previously, a node absent from the stats file returned 0 (free, always accepted). This PR falls through to the ad-hoc estimation path instead, which is a real behavior change for existing users of ep.cuda.resource_partitioning_settings with partial-coverage stats files — previously-accepted nodes can now be rejected/pushed to CPU on an unchanged model + unchanged stats file.

Second-order issue: the stats path doesn't touch pending_weights_/committed_weights_, while the ad-hoc fallthrough path does — so an initializer shared between a stats-known node and a now-fallthrough stats-unknown node gets double-charged, breaking the documented "never under-count, may over-count" dedup invariant across the mixed accounting session.

Suggested fix: either keep return 0 on stats-miss (staying in scope for this PR), or explicitly call out the behavior change in the PR description and reconcile the initializer bookkeeping between the two paths.

Minor

• Two-pass NHWC workspace replay path has no dedicated regression test — existing test only asserts total cost via GetConsumedAmount(), not GetCommittedWorkspaceEstimate()/GetWorkspaceEstimateSourceCounts().
• internal_testing_partitioning_tests.cc's AccountingNhwcTestExecutionProvider still calls the 1-arg ComputeResourceCount(), so it no longer mirrors the real CUDA EP's 2-arg call and won't exercise the new estimator-aware path.
• Plugin EP host bridge (ep_plugin_provider_interfaces.cc) still calls the 1-arg form, so plugin CUDA EP gets no Level-1 estimate at all — permanently diverging from in-tree CUDA behavior. The design doc (future_directions_constrained_env.md) describes this differently: a non-member helper as "the single decision point for both in-tree and plugin host bridge," rather than a virtual signature change. Worth confirming this divergence is the intended incremental step (per the issue [Tracking] Constrained-environment memory and layer-offload roadmap #29775 roadmap) vs. an oversight.
• AddConsumedAmount is noexcept and wraps on overflow, while the new CommitWorkspaceEstimate uses SafeInt and throws — inconsistent overflow policy between the two accumulators feeding the same "non-workspace = total - workspace" log computation; under a wrap the clamp-to-zero would mask real corruption instead of surfacing it.
• WorkspaceEstimateSelection only populates 2 of 4 members for kEstimator/kFallback sources (level1_estimated_bytes/profiled_bytes left at 0), contradicting the documented meaning of those fields for future Level-2-comparison consumers.
• Per-node INFO-level logging in session_state.cc scales linearly with node count — consider VERBOSE for per-node detail, INFO for the aggregate summary only.
• New enum/struct types in resource_accountant.h (WorkspaceEstimateSource, WorkspaceEstimateSelection, WorkspaceEstimateSourceCounts) have no doc comments despite living in a semi-public header; bytes field name is ambiguous next to its profiled_bytes/level1_estimated_bytes siblings.
• Stale /*workspace_estimate=*/ parameter-name comments in tests — actual param is level1_workspace_estimate.
• The dropped TODO in session_state.cc ("For now, log the declarations so we can verify the wiring works") documented that Level-2 data isn't persisted anywhere yet; removing it without replacement makes the surrounding code look more wired-up than it is.

What's solid

• Two-pass NHWC capture-before-reset ordering is correct (verified independently by two reviewers) — this was the highest-risk part of the diff to get subtly wrong.
• Fallback-path refactor from multiplier to additive form is bit-exact with prior behavior (x + ⌊x/2⌋ = ⌊3x/2⌋, verified exhaustively).
• SafeInt correctly introduced on paths that previously used raw arithmetic over model-derived values.

Design-doc drift (multiple stale references to renamed methods/signatures in docs/annotated_partitioning/future_directions_constrained_env.md) is out of this diff's scope but worth a follow-up pass once this PR's approach is settled.

[Ti-Tai Wang (titaiwangms)]

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 addressed

Previous Major #1 (no floor against unknown allocations) — now addressed via a structural fix, not just a floor.
Instead of maxing the estimator against the old 1.5x fallback, this PR introduces Level1MemoryEstimate{runtime_workspace_bytes, persistent_prepack_bytes, temporary_prepack_bytes} and now charges persistent_prepack_bytes/temporary_prepack_bytes cumulatively and unconditionally on top of the selected workspace, in both the profile and no-profile paths (resource_accountant.cc). This directly covers the previously-missing prepack/persistent buffers (fpA_intB_weight_buffer_, _scale_buffer_, _zero_buffer_) that the estimator didn't account for — good fix, and cuda_kernel_workspace_inventory.md / the design doc were updated to document the three-way split and that "the current byte-count accountant conservatively charges all three." New unit tests (Level1MemoryEstimateReplacesFallbackAndAddsPrepackMemory, PrepackMemorySeparatesPersistentAndTemporaryBytes, PrepackMemoryRejectsInvalidMetadata) cover this well.

One residual nuance: runtime_workspace_bytes in the no-profile path is still value_or(fallback_workspace) (substitution, not max()) — but since the CUTLASS workspace formula itself is described as "verified against the runtime value" (exact, not a lower bound), and the real risk (prepack buffers) is now separately and conservatively charged, this is much lower-risk than before. Not blocking.

Still open

1. [Major] Stats-file-miss node still goes from cost=0 to fully-charged fallback, with an unreconciled double-count risk
resource_accountant.cc — a node absent from the stats file still falls through to the full ad-hoc/fallback path (now explicitly intentional, per the updated comment "A stats file may have incomplete coverage. Unknown nodes use the ad-hoc fallback" and the renamed test StatsPath_UnknownNodeUsesFallback). This is a real behavior change for existing ep.cuda.resource_partitioning_settings users with partial-coverage stats files (previously-free/always-accepted nodes can now be rejected on an unchanged model+stats file) — worth calling out explicitly in the PR description if not already, since it's user-visible.

The double-count concern also still stands: the stats-hit path doesn't touch pending_weights_/committed_weights_, while the fallthrough (stats-miss) path does — so an initializer shared between a stats-known node and a stats-unknown node in the same session is still charged twice. Direction is conservative (won't under-count budget), but it means the documented "shared initializer may be re-counted but never under-counted" invariant is now also true within a single mixed-accounting pass, not just across partitioning iterations — might be worth a comment noting this explicitly since it's a second reason (beyond cross-iteration re-counting) that total_estimate can grow.

2. [Major] Plugin EP host bridge still uses the 1-arg ComputeResourceCount() — permanent policy divergence from in-tree CUDA
ep_plugin_provider_interfaces.cc is not touched in this push. Plugin EPs still get zero Level-1 estimate and always fall back to the 1.5x/profile-only path, while in-tree CUDA now gets the full three-component estimate. The design doc's own architecture note ("Multiplier handling — non-member helper approach... Rather than changing ComputeResourceCount()'s signature... single decision point for both in-tree and plugin host bridge") is unchanged in this push and still describes a different approach than what's implemented (a virtual signature change, policy inside the accountant) — worth either reconciling the doc with the shipped design, or confirming this divergence is the deliberately-deferred next step per the #29775 roadmap (chilo-ms's step 3, "plugin-compatible Level-1 estimator ABI").

Minor (unchanged from before)

• No dedicated test for the two-pass NHWC workspace replay specifically (GetCommittedWorkspaceEstimate()/GetWorkspaceEstimateSourceCounts() after pass-2 survival) — internal_testing_partitioning_tests.cc's NHWC test still only asserts total GetConsumedAmount(), and the AccountingNhwcTestExecutionProvider still calls the 1-arg form (comment renamed, behavior unchanged).
• AddConsumedAmount (noexcept, wraps on overflow) vs. CommitWorkspaceEstimate (SafeInt, throws) — still an inconsistent overflow policy between the two accumulators feeding the non_workspace_estimate log computation. (The underflow-clamp guard itself is now more robust — it correctly nets out workspace + persistent_prepack + temporary_prepack as categorized_estimate before subtracting, which is a nice improvement over the previous version.)

New surface added since last review (not previously covered) — looks solid

• ComputeMatMulNBitsPrepackMemoryEstimate correctly validates metadata (rejects n=0, invalid nbits) and its test cases (PrepackMemorySeparatesPersistentAndTemporaryBytes) hand-verify the byte arithmetic for both prepacked and non-prepacked weight layouts, including the zero-point case doubling the scale-buffer charge — good coverage.
• WorkspaceEstimateComparisonSummary (profile vs. estimator larger/equal counts) is a nice addition for the eventual Level-1-vs-Level-2 reconciliation work.

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.

[Tianlei Wu (tianleiwu)]

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.

[Ti-Tai Wang (titaiwangms)]

Review with build + test verification

I 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, CMAKE_CUDA_ARCHITECTURES=80, PR commit 2c33db5.

Build and test status: green

• Build passes, 0 errors. The new tests compile cleanly — no signature/include drift.
• 29/29 relevant tests pass:
  • ResourceAccountantTest.* — 11/11
  • InternalTestingEP.NhwcTwoPassAccountingCommitsOnlySurvivors — PASS
  • InternalTestingEP.NhwcSecondPassDropFallsBackFromCpuKernelNode — PASS
  • MatMulNBitsWorkspace.* — 16/16
• The only failure in the CUDA suite was AllocatorTest.CUDAAllocatorFallbackTest, which is unrelated to this PR (no allocator file is touched) and appears environment-specific.
• GetCapabilityBudgetUsesLevel1Estimate is not vacuously passing: it ran for 1426 ms, was not skipped, and both the 600 KiB and 500 KiB sessions initialize successfully (so the CPU fallback path is genuinely exercised).

Note for anyone reproducing: the CUDA test cases need -Donnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON; they build into libonnxruntime_providers_cuda_ut.so and are driven by CUDA_EP_Unittest.All.

---

Critical — the ad-hoc path double-counts prepacked weights, and the headline test encodes that bug

In the ad-hoc branch of resource_accountant.cc, initializer bytes are accumulated into total_size (line 147), folded into estimated (line 191), and then persistent_prepack_bytes is added on top (lines 217-218). But persistent_prepack_bytes describes the packed copy of those same tensors.

Per session_state.cc:678-686, once PrePack sets is_packed and the constant's use count reaches zero, ORT erases the original initializer:

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 cuda_kernel_workspace_inventory.md now states. Separately, temporary_prepack_bytes is scratch from GetTransientScratchBuffer that is released before the next node prepacks, yet it is summed across nodes rather than tracked as a peak.

Using the exact constants from the new e2e test (N=256, K=1024, M=256, block_size=32, bits=4):

packed weight W       = 131,072
scales        S       =  16,384
initializers (W+S)    = 147,456   <- counted in total_size
output                = 131,072
persistent_prepack    = 147,456   <- byte-for-byte identical, counted again
temporary_prepack     = 131,200   <- released after prepack, summed anyway

PR charge  = 278,528 + 147,456 + 131,200 = 557,184  (544.1 KiB)
steady state =                              278,528  (272.0 KiB)
over-charge  = 2.00x

persistent_prepack_bytes coming out exactly equal to the initializer bytes is the double count made visible.

The part I'd most like you to look at: this directly determines the outcome of GetCapabilityBudgetUsesLevel1Estimate.

legacy 1.5x fallback = 408.0 KiB  -> below the 500 KiB budget, node accepted
PR structured total  = 544.1 KiB  -> above 500 KiB, node rejected

The test's 500/600 KiB straddle works because of the double count. Remove it (charge max(initializer_bytes, persistent_prepack_bytes) for prepack-consumed inputs, and track temporary_prepack_bytes as a running max instead of a sum) and the node drops to roughly 272 KiB + runtime workspace, so it would be accepted at 500 KiB and the test would fail.

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, total_temp_allocations is collected from kernel-context allocations at Compute time and never includes PrePack allocations, so persistent_prepack_bytes genuinely adds missing information there.

---

Major — Level 1 omits the largest init-time allocation for this exact kernel

matmul_nbits.h calls RunGemmProfile(has_fpA_intB_gemv_, 1, max_m) from the kernel constructor, and gemm_profiler.h:430 allocates computeTmpSize(...) on the device allocator. Per fpA_intB_gemm_profiler.cc:80-92 that includes a k * n * sizeof(half) term — for 4-bit weights that alone is 2NK, i.e. 4× the packed weight and ~8× the temporary_prepack_bytes this PR does model.

While EstimateWorkspace was log-only this was harmless. Now that Level 1 drives accept/reject, and since §20 of cuda_kernel_workspace_inventory.md (edited here) is meant to be the inventory of workspace-allocating kernels, this omission is material. Either add the row plus a profiler_scratch_bytes field, or state explicitly that construction-time tactic profiling is out of scope for Level 1 and why.

Major — partitioning now depends on how fpA_intB was enabled

The estimator reads only the ORT_FPA_INTB_GEMM environment variable, while the kernel resolves eligibility through ResolveFpAIntBConfigOrEnv, which also honors the ep.cuda.fpa_intb_gemm session config.

This PR removes the justification that made that divergence safe:

"Because Level 1 is currently log-only and does not change the partition budget, this divergence is safe."

but keeps describing the nullopt return as "conservative". It no longer is. With the option set via session config and the env var unset, the same node is charged the 1.5x fallback (408.0 KiB in the numbers above) instead of the structured estimate (544.1 KiB), so it is accepted under a 500 KiB budget even though the kernel that actually runs is the fpA_intB one. That fails open toward OOM rather than conservatively.

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 sizing

The new if (graph_for_shapes != nullptr && !max_shapes.Empty()) block changes dynamic outputs from contributing 0 bytes to contributing max_shape->Size() * elem_size, with the 1.5x heuristic then applied to the larger base. This changes which nodes are accepted for any session using session.max_shape_override, and it is not in the PR description or in any of the three docs updated here.

It also runs against max_shape_inference.h:22-25, which cautions that inferred shapes are "estimation hints, not proven upper bounds". Using them for a hard budget deserves at least a documented caveat.

Two small things in the same block:

• utils::GetSizeInBytesFromTensorTypeProto<0>(type_proto->tensor_type(), &size).IsOK(); discards the Status as a bare statement — prefer if (!....IsOK()) { size = 0; }.
• graph_for_shapes re-calls node.GetContainingGraph() a few lines below the existing graph; they are the same pointer.

---

Minor

• matmul_nbits.cc: on SafeInt<int> narrowing failure the catch returns std::nullopt for the whole estimate, discarding already-computed prepack bytes. Consider nulling only runtime_workspace_bytes and still returning the prepack fields.
• ComputeMatMulNBitsPrepackMemoryEstimate hard-codes sizeof(uint16_t) with no dtype parameter. Correct today only because CheckFpAIntBEligibility rejects non-FP16/BF16, but the helper is now public and directly unit-tested, so a future FP8/FP32 caller would silently get wrong bytes. Add an element_size parameter or document the precondition.
• WorkspaceEstimateSelection is built with positional aggregate initialization at both production call sites, and has six same-typed size_t fields — transposing profiled_bytes and level1_estimated_bytes would compile silently. Designated initializers would remove that class of bug.
• Level-2 aggregate log line reports Σ nodes Σ slots. The PR description disclaims that this is not a peak-live value, but that disclaimer does not appear in the emitted log, which is all the user sees. Suggest "sum of per-node declared workspace (upper bound; not peak-live): N bytes; max single-node: M bytes".
• Skill compliance: .agents/skills/workspace-estimation-shared-header asks that reusable estimator math stay separable from graph parsing. The helper is correctly graph-type-free, but lives in matmul_nbits.h, which pulls in cuda_kernel.h / fpA_intB_gemm_profiler.h, so a plugin TU cannot include it. This also conflicts with this PR's own new doc bullet: "keep reusable estimator math graph-type-free, with separate in-tree and plugin parsing wrappers". Consider a slim header alongside ComputeFpAIntBGemmWorkspaceSize.
• ABI note: ComputeResourceCount's virtual signature changes and three virtuals are inserted mid-class in a header included by core/providers/shared_library/provider_interfaces.h. Fine for the version-locked in-tree build, but worth confirming no packaging scenario mixes a provider DLL and core from different commits.

Pre-existing, not introduced here

In GetCapabilityForEP, survivor costs are restored via AddConsumedAmount only after pass 2 has made its budget decisions, so pass-2 nodes are checked against a deflated consumed. I initially flagged this, then confirmed against the diff that the AddConsumedAmount call is pre-existing context — this PR only adds the parallel AddCommittedWorkspaceEstimate. Flagging it as a possible follow-up, since larger per-node costs widen the window. The comment claiming the budget "can never be exceeded" is also pre-existing.

---

What I think is genuinely well done

• The fallback refactor is exactly rounding-preserving: floor(3x/2) == x + floor(x/2) for all integers (checked by parity), and the (150 - 100) form keeps the constant's provenance readable while raising the overflow ceiling.
• CommitWorkspaceEstimate computes all three SafeInt sums into locals before mutating any member, so an overflow leaves the accountant untouched — and CommittedWorkspaceRejectsOverflow asserts exactly that post-throw invariant. That is deliberate strong-exception-safety, and it passes.
• The pass-1 capture is placed before ResetForNewPass(), which is the only correct position and the easiest thing here to get subtly wrong. The two-pass test confirms survivors commit exactly once.
• The prepack formulas mirror the kernel expression-for-expression, including the (8 / nbits) divisor and the ceil form, so future divergence shows up as a visible textual diff.
• level1_memory_estimate.h correctly avoids the known Node forward-declaration ODR trap for shared-provider headers.
• The SizeBasedStatsAccountant → SizeBasedResourceAccountant and CommitWeightsForNode → CommitResourcesForNode renames are complete across the repo, including tests and comments.
• Documenting that "a node missing from a supplied stats file has zero accounting cost" captures genuinely surprising pre-existing behavior that was previously only discoverable by reading the code.

Happy to share the exact build commands or rerun anything against a revision.

[Chi Lo (chilo-ms)]

Critical — the ad-hoc path double-counts prepacked weights, and the headline test encodes that bug

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

[Tianlei Wu (tianleiwu)]

Both of my earlier blocking threads are verified fixed on this head and resolved: pass-2 admission now sees the pass-1 survivor budget, and the lazy tactic-profiler scratch is reserved at max(configured bucket, RoundUpProfileM(M)) (or kMaxProfileM for dynamic M), each with a matching regression.

Two issues remain in the code that landed with those fixes.

1. The provisional reservation can leak permanently. pass1_node_costs is built from every node of every accounting-enabled pass-1 capability, but the rollback loop iterates nodes_temporarily_assigned_to_ep, which TryAssignNodes only populates for nodes it newly tagged. It returns false up front (assigning nothing) when any node in a capability is nullptr or already carries a different EP type, and it skips nodes that already carry this EP type. Those nodes keep their AddConsumedAmount forever — phantom budget, which is exactly what the previous deferred-commit design avoided. Latent for the in-tree CUDA EP today, but GetCapabilityForEP is generic and nothing enforces the invariant. Deriving both sets from nodes_temporarily_assigned_to_ep is a one-line fix.

2. Initialization scratch is summed per node although it is a shared, one-at-a-time allocation. Now that Level 1 gates the hard budget rather than only logging, charging the full fpA_intB tactic-profiler scratch to every MatMulNBits node is a very large over-estimate: profileTactics frees workspace_tmp before returning, and s_profilerManager's mMNKProfileMap is shared by all kernels so same-shape nodes allocate nothing at all. The direction is safe, but the magnitude can push entire models to CPU once a budget is configured.

Smaller items are inline. Also, per AGENTS.md, please build and run the tests before merge — the CUDA integration tests encode exact byte budgets that depend on the new profiler-scratch formula, and the PR description states tests were not run.

Positives worth calling out: replacing the collapsed workspace scalar with a structured Level1MemoryEstimate is the right model; RoundUpProfileM fixes a real signed-overflow in std::min(nextPowerOfTwo(v), getMaxProfileM()); bundling bytes and source in one WorkspaceEstimateSelection prevents the reported source from diverging from the charged size; and CommitWorkspaceEstimate computing all three totals into locals before assigning keeps state consistent when SafeInt throws.

[Chi Lo (chilo-ms)]

Addressed the latest review round in 539a10d:

• made NHWC pass-1 accounting capture/reservation symmetric with actual tentative assignments;
• treated sequential initialization scratch as a reported peak rather than additive per-node budget;
• narrowed retained estimator configuration;
• aligned MatMulNBits profile sizing/rounding and restored requested error handling/cleanup;
• added regression coverage for rejected capabilities and scratch-peak semantics.

The normal Release CUDA targets build, and all 22 targeted framework/accounting tests pass. The CUDA internal MatMulNBits test source also compiles; running that internal-test module on this Windows build remains blocked by the existing onnxruntime_provider_test / onnxruntime_providers_cuda_ut dependency and missing import-library arrangement.

[Tianlei Wu (tianleiwu)]

The earlier two-pass accounting, accepted-node commit, overflow handling, configuration, and profiler-sizing concerns are fixed on this head, and I resolved those addressed threads. One blocking lifetime mismatch remains in the existing MatMulNBits profiler-scratch thread: scratch for an unseen bucket can be allocated during Run(), but the new Level-1 estimate classifies all profiler scratch as initialization-only and excludes it from the hard partitioning budget. The same expression also sums non-overlapping constructor and PrePack scratch instead of reporting their peak. Details and the requested regression coverage are in the continued inline thread.

[Tianlei Wu (tianleiwu)]

One blocking budget issue and one smaller reporting inconsistency remain. The structured lifetime split, checked arithmetic, two-pass NHWC reservation/rollback, shared estimator math, and the executed static lazy-profile regression are strong improvements. The earlier lazy-profile lifetime concern is fixed on this head; however, the estimator still treats a propagated max-shape hint like an exact runtime M, so sparse initial profile buckets can leave a valid smaller runtime bucket's lazy scratch unreserved.

[Tianlei Wu (tianleiwu)]

The bounded-shape and stats-attribution fixes on this head are correct, and I resolved those prior threads. Three issues remain: two newly added MatMulNBits assertions conflict with the lazy-profile accounting and fail current CUDA/TensorRT CI; NHWC pass-two admission can permanently reject a candidate based on a provisional cost that is immediately rolled back; and the roadmap still documents the opposite initialization-scratch policy. Details are inline.

[Copilot]

🟡 Changes recommended

The NHWC retry path can remove a surviving node’s reservation and subsequently admit it at zero cost, violating the configured memory budget.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 23/23 changed files
• Comments generated: 2
• Review effort level: Balanced

[Tianlei Wu (tianleiwu)]

Summary

This update completes workspace-memory accounting and reporting for resource-constrained graph partitioning, specifically integrating Level-1 and Level-2 memory estimation across the two-pass NHWC layout transformation flow.

Key Highlights

• Survivor Discovery & Reconciliation: Pass-2 survivor discovery is executed without budget gating, ensuring later survivors are not prematurely dropped due to an early budget cutoff. Reservations for all confirmed survivors are rebuilt before evaluating pass-2 candidate admission.
• Compiling Capabilities: Compiling/fused survivor capabilities are reconciled iteratively against accountant-aware groupings to guarantee stable convergence.
• Shared Initializers: Shared initializers are charged exactly once to the finalized survivor set, avoiding double-counting or phantom reservations.
• Documentation & Inventory: Diagnostic, non-additive initialization scratch is clearly separated from budgeted runtime workspace and persistent prepack memory.
• Test Coverage: Added NhwcTwoPassAccountingPreservesLaterSurvivorReservation covering the Relu -> Conv -> LogSoftmax 1000 KiB scenario along with comprehensive tests for fused, overlapping, and mixed survivor capabilities. Test node lifetime has been hardened to use stable node indices.

All prior review comments have been verified and resolved, and all CI pipelines are passing.