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.