Verify Level-2 workspace reservations

[Chi Lo (chilo-ms)]

Summary

• Preserve accepted per-node workspace reservations after graph partitioning and through post-partition graph mutations.
• Compare final Level-2 workspace declarations with the reservations used during partitioning.
• Report larger, equal, smaller, missing-declaration, missing-reservation, and orphaned-reservation counts and byte totals per graph.
• Warn by default on declaration overruns and support strict constrained-memory validation with session.strict_workspace_verification=1.

Behavior and scope

Verification runs after kernels are created and constant tensors are prepacked, when final Level-2 requirements are available. Strict mode fails session initialization when a declaration exceeds its reservation or when an untracked mutation leaves a nonzero orphaned reservation. Missing partition reservations and zero-byte orphans remain diagnostic-only.

Reservation ownership now follows fused capabilities, selector actions, shared and manual fusions, bottom-up removals, intentional eliminations, and training recomputation clones. This includes MatMulScale and QDQ cleanup, Attention and EmbedLayerNorm cleanup paths, Level-4 initializer fusion, GroupQueryAttention, BitmaskDropout, BiasSoftmaxDropout, SCE loss gradient, Triton, and MemoryOptimizer transformations.

Strict verification is rejected for ORT-format loads because partition-time reservations are not serialized. Reservation storage and verification remain excluded from ordinary minimal builds; extended-minimal ORT-format loads reject strict mode explicitly.

This PR does not repartition the graph or change runtime workspace allocation. #31962 introduced Level-1 workspace estimation and reporting; activation-aware CUDA workspace preallocation remains follow-up work in #32071.

Validation

• Windows Release onnxruntime_test_all build: 123 focused tests passed, 4 hardware-dependent tests skipped.
• Windows training-enabled Release build: all 10 focused mutation-accounting tests passed.
• Windows extended-minimal build: strict ORT-format rejection test passed.
• Windows no-ops MinSizeRel minimal shared-library build passed.
• lintrunner and git diff --check passed.

[Copilot]

🟡 Changes recommended

Fused-node reservations can be omitted or mismatched, and the strict verification paths lack direct tests.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds post-partition verification of Level-2 kernel workspace declarations against accepted reservations, with optional strict validation.

Changes:

• Preserve reservations by graph and node.
• Compare declarations and report discrepancies during session finalization.
• Add strict verification configuration and reservation tests.

File summaries

File	Description
include/onnxruntime/core/framework/resource_accountant.h	Defines reservation maps and retrieval APIs.
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h	Adds the strict verification option.
onnxruntime/core/framework/graph_partitioner.cc	Collects committed reservations after partitioning.
onnxruntime/core/framework/graph_partitioner.h	Exposes reservation collection from partitioning.
onnxruntime/core/framework/resource_accountant.cc	Records reservations by graph and node.
onnxruntime/core/framework/session_state.cc	Performs Level-2 comparison, logging, and strict failure.
onnxruntime/core/framework/session_state.h	Stores session workspace reservations.
onnxruntime/core/session/inference_session.cc	Transfers partition reservations into session state.
onnxruntime/test/framework/resource_accountant_test.cc	Tests reservation retention.

Review details

Suppressed comments (3)

onnxruntime/core/framework/session_state.cc:1874

• Reservations for fused capabilities are keyed to the constituent node indexes, but AccountForAllNodes() commits those indexes before FinalizeFuseSubGraph() removes them. Because this loop visits only surviving nodes, those reservations are never counted as missing declarations, and a declaring fused kernel cannot be compared with them. Remap/aggregate reservations when creating the fused node, or track consumed reservations here and report all unmatched entries after the loop.

    for (const auto& node : graph_viewer_->Nodes()) {

onnxruntime/core/framework/session_state.cc:1958

• This error is also returned when a declaration has no matching reservation (line 1921), so the current text incorrectly tells users that an excess was found. Include both failure conditions in the diagnostic.

    ORT_RETURN_IF(strict_verification_failed,
                  "Level-2 workspace verification failed: one or more declarations exceed "
                  "the workspace reserved during graph partitioning.");

onnxruntime/core/framework/session_state.cc:1915

• This changes the per-node declaration message from VERBOSE to INFO. Models with many workspace-declaring kernels will emit one normal-priority record per node during initialization even though the graph summary already reports the totals; retain VERBOSE here and reserve WARNING for actual excesses.

      LOGS(logger_, INFO) << "Level-2 workspace: node '" << node.Name()

• Files reviewed: 9/9 changed files
• Comments generated: 2
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Chi Lo (chilo-ms)]

Follow-up fixes for the Copilot review are in 7d18338: fused capabilities now consolidate constituent workspace reservations onto the surviving fused node, per-node declaration totals are VERBOSE rather than INFO, and default warning-mode overrun behavior has direct session-level coverage. The focused Release suite passed 35/35 tests before merging the latest unrelated main commits.

[Copilot]

🟡 Changes recommended

Post-partition node replacements can invalidate index-keyed reservations and bypass strict overrun detection.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 11/11 changed files
• Comments generated: 1
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

[Ti-Tai Wang (titaiwangms)]

Full review-team verdict: changes requested.

Major findings

1. Post-partition node replacement bypasses strict verification. InferenceSession::TransformGraph snapshots and installs reservations immediately after GraphPartitioner::Partition (inference_session.cc:1844-1849), but Level 2-4 transforms run afterward (:1933-1968). Replacement actions create a new node/index (selectors_actions/actions.cc:82-117) while the reservation remains keyed to the removed partition-time index. During finalization, the replacement kernel falls into missing reservation (session_state.cc:1919-1921), which is diagnostic-only even when session.strict_workspace_verification=1; the orphaned old reservation is not included in any summary bucket. This affects paths such as post-partition QDQ -> MatMulNBits and other ReplaceWithNew/fusion transforms. Please propagate/consolidate reservation ownership through post-partition graph mutations (or reconcile/recompute it after the final transform loop) and add an end-to-end strict-mode test covering partition -> replacement -> final Level-2 declaration.

2. Strict verification is silently ineffective for ORT-format loads. The ORT-format path calls GraphPartitioner::Partition(..., Mode::kOrtFormatLoad) without the new reservation output and never installs a map on SessionState (inference_session.cc:2509-2517). Consequently all Level-2 declarations lack reservations and strict mode cannot reject an overrun, although the public option does not state a model-format limitation. Please either collect/install comparable reservations for this path or explicitly reject/document strict mode as unsupported for ORT-format sessions, with a focused test.

Review notes

The new partition-fusion consolidation is sound for the paths it covers: constituent byte fields are added with checked arithmetic, initialization scratch uses peak/max semantics, and graph identity correctly disambiguates subgraph node indices. The default warning behavior and static matched-node tests also look good.

Not checked: runtime execution or performance; the review was static against head e70994f55493ee0a9f8f2c50548d9df323771cae.

[Chi Lo (chilo-ms)]

Addressed the latest review in 393dee7. Post-partition selector-action replacements and graph fusions now transfer/aggregate reservation ownership onto the surviving node; strict mode rejects any remaining declaration-plus-orphan mismatch rather than silently bypassing verification. ORT-format loads now explicitly reject session.strict_workspace_verification=1, with public documentation and a focused test, because partition-time reservations are not serialized in ORT format. Windows Release onnxruntime_test_all built successfully and all 39 focused tests passed.

[Copilot]

🟢 Approval recommended

The reservation lifecycle, strict-mode behavior, graph mutation handling, and targeted regression coverage are coherent with no unresolved correctness issues found.

Review details

• Files reviewed: 16/16 changed files
• Comments generated: 0 new
• Review effort level: Balanced (auto)

Note

Copilot is running an experiment and ran this review at Balanced.

[Ti-Tai Wang (titaiwangms)]

Updated full-team review of 393dee76d1: changes still requested. The new callback correctly fixes selector-action ReplaceWithNew, Graph::FinalizeFuseSubGraph, callback lifetime, and the ORT-format strict-mode gap. Two Major issues remain.

Major findings

1. The replacement callback still does not cover the production post-partition mutation surface. Notifications currently come only from Graph::FinalizeFuseSubGraph and selector-action ReplaceWithNew. Existing Level-2 paths still preserve EP assignment while directly adding/removing nodes without notifying:

   • graph_utils::FinalizeNodeFusion removes fused nodes directly (graph_utils.cc:1045-1062) and is used by multiple attention/layer-norm fusions.
   • GroupQueryAttentionPreNormFusion removes the original GQA and related nodes, creates a new EP-assigned GQA, and never transfers reservations (group_query_attention_pre_norm_fusion.cc:343-373).
   • GroupQueryAttentionFusion creates an EP-assigned MatMul/MatMulNBits, removes Q/K/V/Rotary nodes, and mutates the surviving GQA without notification (group_query_attention_fusion.cc:374-399,603-690).

   These paths still leave reservations on removed indices and/or leave the replacement kernel unreserved. Please notify at the shared mutation seams (FinalizeNodeFusion) and explicitly handle manual/many-to-many fusions such as the GQA paths, with a real Level-2 integration test—not only a selector-action test.

2. missing_reservation > 0 && orphaned_reservations > 0 is not an ownership check. session_state.cc:1939-1985 correlates two graph-wide counts without establishing lineage. An unrelated removed reservation plus an independently inserted workspace-declaring node can falsely fail strict initialization; a missing reservation passes whenever no orphan exists; and an orphan passes whenever no new declaration exists. Track the disposition/transfer of each reservation during a mutation, or fail directly on the precise unsupported condition rather than pairing unrelated aggregate counts.

Minor

The public option comment still says “Missing reservations remain diagnostic only,” but the new orphan+missing condition now fails strict initialization. Please update the documented contract to match the implemented behavior.

Verified improvements: reservation consolidation uses checked arithmetic and deduplicates source indices; the root/subgraph callback routing and gsl::finally cleanup are sound; ORT-format strict mode now fails closed with a focused test.

Not checked: runtime performance or the full test matrix; this was a static review of head 393dee76d1780367e97ff61cb1c5104ccf75292e.

[Chi Lo (chilo-ms)]

Addressed in 88c98a54a3:

• graph_utils::FinalizeNodeFusion now transfers reservations in both the survivor and replacement-node overloads.
• GroupQueryAttentionPreNormFusion transfers all seven removed-node reservations to the new GQA node.
• GroupQueryAttentionFusion transfers the Q/K/V projection and rotary reservations to the generated MatMul/MatMulNBits node.
• Consolidation now preserves and aggregates an existing destination reservation instead of overwriting it.
• Strict verification now fails directly on any orphaned reservation; it no longer correlates unrelated graph-wide missing/orphan counts. Nodes that never had a partition-time reservation remain diagnostic-only, and the public option documentation now states this contract.

Added production-path tests for both FinalizeNodeFusion overloads and both GQA transformers, plus a strict orphan test with no missing reservation. Windows Release onnxruntime_test_all rebuilt successfully; the expanded focused suite passed all 44 tests.

[Chi Lo (chilo-ms)]

CI triage and fix:

• AndroidBinarySizeCheckJob_MinimalBaseline was a real size regression: section size was 1,592,830 bytes, exceeding the 1,590,272-byte limit by 2,558 bytes. Commit 254ddd9839 excludes Level-2 reservation storage/verification from baseline minimal builds, which do not run graph partitioning and therefore cannot produce partition-time reservations.
• build_x86_release compiled successfully and all native CTest tests passed; its only failure was the unrelated C# test CompileApiTests.GetInitializerLocationDelegateThatReusesExternalInitializers reporting protobuf parsing failure.

Validated the fix with a full Release build, all 44 focused workspace/accounting tests, and a no-ops MinSizeRel minimal shared-library build with exceptions disabled.

[Ti-Tai Wang (titaiwangms)]

Updated review of current head 254ddd9839 (the findings are in code inherited from 88c98a54a3; the latest minimal-build size fix does not affect them). The previously reported FinalizeNodeFusion and GQA fusion gaps are fixed, destination reservations are preserved correctly, and the public option text now distinguishes orphaned from missing reservations. Two Major issues remain before approval.

1. Other normal Level-2 replacements still orphan reservations

Reservation transfer remains opt-in at individual transformer call sites. MatMulScaleFusion is registered for CUDA at Level 2 (graph_transformer_utils.cc:455), creates an EP-assigned FusedMatMul, and directly removes the original MatMul and merged scale nodes (matmul_scale_fusion.cc:268-302) without calling NotifyNodeReplacement. Those removed nodes retain reservation entries under dead indices.

Because strict mode now fails on any orphaned reservation, a valid MatMul+Scale optimization can fail session initialization even when no Level-2 workspace declaration exceeds its reservation. The same structural risk remains in other direct AddNode/RemoveNode fusion paths that do not use the newly instrumented helpers.

Please either instrument all post-partition replacement paths or move reservation tracking into a graph-mutation mechanism where a transformer cannot accidentally omit it. A focused strict-mode test using MatMulScaleFusion would demonstrate the remaining issue and protect the fix.

2. Intentional deletion has no reservation disposition

QDQFinalCleanupTransformer, registered at Level 2 (graph_transformer_utils.cc:500-502), legitimately removes Q/DQ nodes without necessarily creating a replacement (qdq_final_cleanup.cc:85-135). The current callback only represents sources -> destination; it cannot represent “these nodes were intentionally eliminated.” Their reservations therefore become orphans, and strict mode rejects the otherwise-valid optimized graph.

This is especially broad because the resource accountant records a reservation selection for every committed accepted node, including zero-byte workspace selections. Strictly failing on every orphan means even removal of a node with no workspace can block initialization.

Please add an explicit intentional-removal disposition, such as NotifyNodesRemoved(indices), or make the replacement event accept an optional destination. Strict verification can then distinguish an expected elimination from an untracked mutation. At minimum, zero-byte orphan reservations should not cause a constrained-memory failure.

Verified improvements

• Both graph_utils::FinalizeNodeFusion overloads now transfer reservations.
• GroupQueryAttentionFusion and GroupQueryAttentionPreNormFusion transfer the intended removed-node reservations.
• Existing destination reservations are merged rather than overwritten.
• ORT-format strict mode fails closed with a clear unsupported message.
• Strict orphan handling no longer relies on the unrelated missing && orphan heuristic.

Verdict: changes requested. Static review of current head 254ddd983977d2b445e0a9175a5b9c6d5c49ae99; runtime/performance testing was not repeated.

[Chi Lo (chilo-ms)]

Addressed the latest review and completed a broader final mutation audit in b9660e51a8 and e54601200f:

• MatMulScaleFusion transfers every removed MatMul/scale reservation to FusedMatMul.
• QDQFinalCleanupTransformer reports only nodes it actually eliminates; intentional removals erase reservations instead of leaving false orphans.
• Bottom-up cleanup now reports the indexes actually removed, so Attention and EmbedLayerNorm transfer only real removals while shared survivors keep their reservations.
• Remaining post-partition Level-4 and training paths now have explicit dispositions: initializer elimination, BitmaskDropout, BiasSoftmaxDropout, SCE gradient, Triton, and MemoryOptimizer recomputation clones.
• Extended-minimal ORT-format loads now reject unsupported strict verification, fixing the failing OrtModelOnlyTests.RejectsStrictWorkspaceVerification CI test.

Added production-path reservation tests for these ownership shapes. The final Windows Release suite passed 123 tests with 4 expected hardware skips; the training-enabled suite passed all 10 focused tests; extended-minimal and ordinary minimal builds also passed.