Repository navigation
[Framework] Add opt-in fixed-placement session memory diagnostics - #33657
Draft
Chi Lo (chilo-ms) wants to merge 8 commits into
Draft
Chi Lo (chilo-ms) wants to merge 8 commits into
Chi Lo (chilo-ms) wants to merge 8 commits into
Conversation
…oute-aware L1 estimate) Make the GQA Level-1 workspace estimator produce a sound bound for the non-windowed KV-cache case instead of declining. Previously BuildBounds returned nullopt for every non-windowed node because total_sequence_length is a runtime scalar not recoverable from shapes. It now uses the declared KV-length envelope (config.max_total_sequence_length, from session option ep.cuda.gqa_workspace_max_total_sequence_length, plumbed in the prior PR) as present_kv_cache_capacity_bound and evaluates the existing route-aware, phase-cornered aggregate over it. Opt-in and default-off: when the envelope is unset (== 0) the non-windowed path still returns nullopt, so behavior is byte-identical to today. The windowed path is untouched -- every change is guarded by the windowed flag. Soundness: - Backend workspace is non-decreasing in KV length and the runtime sizes non-windowed scratch from total_sequence_length <= envelope (GetGQAEffectiveWorkspaceKvLength), so capacity_bound = envelope is an upper bound. - Non-windowed partial aliasing (exactly one past/present K/V pair shared) copies the full past cache into preservation scratch that coexists with the backend workspace. This is now charged on top of the per-route maximum via GQAWorkspaceBounds::account_partial_alias_preservation, bounded by the envelope and head_size_bound. Windowed GQA rejects non-shared buffers, so it never incurs this copy. - cuDNN (when reachable) and attention bias conservatively decline (nullopt), keeping the estimate an upper bound. The small seq_lens buffer remains an unmodeled pre-existing gap shared with the windowed path. Tests (onnxruntime_provider_test): non-windowed without an envelope still declines; with an envelope it is bounded and monotonic in the envelope; the preservation copy is additive on top of the selected route. Part of the policy-based workspace estimation effort; see docs/annotated_partitioning/policy_based_workspace_estimation.md and the PR chain (#32944, #32945, #33001, #33014). Base retargets to main after #33014 merges. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fa98d07-aec3-41b7-92b8-b0ec317a95bf
Use the larger of allocated past capacity and the total-KV envelope to cover MEA expansion and full partial-alias preservation. Keep the query-length envelope check independent and add static-cache and runtime-ABI regressions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 7fa98d07-aec3-41b7-92b8-b0ec317a95bf
Reject out-of-ABI envelopes during initialization, require a proven route-aware GQA estimate for capacity-aware CUDA placement, and keep the caller-declared envelope contract explicit. Share the validated parser for Level-2 reuse and cover invalid values plus fail-closed placement. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Reuse the validated workspace-envelope parser during CUDA GQA kernel construction and pass the stored value into the shared estimator for workspace declaration. Add direct L1/L2 agreement and strict-verification coverage for bounded non-windowed GQA. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Retain the envelope PR integration regressions while stacking on the shared policy and Level-2 plumbing in #33316. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6d876416-81ec-4982-a9b0-38b8b28b5e4a
Require an available kernel declaration for nonzero workspace reservations in strict mode, including replacement kernels. Preserve explicit zero workspace and reuse existing reservation transfer and retirement callbacks. Match constructed legacy and fpA-intB MatMulNBits chunking envelopes, retain GQA cross-level agreement, and cover unavailable, zero, and dynamic-shape cases with focused regressions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6d876416-81ec-4982-a9b0-38b8b28b5e4a
Capture local-graph initializer snapshots and alias-aware tensor/workspace lifetimes. Keep allocator-wide observations separate from modeled payloads and physical peaks unavailable until uncovered backing terms are resolved. Add opt-in logging, documentation, and focused CPU regressions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6d876416-81ec-4982-a9b0-38b8b28b5e4a
Base automatically changed from
chilo/gqa-level2-envelope
to
chilo/gqa-committed-policy
November 6, 2026 18:14
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Add opt-in fixed-placement memory diagnostics with
session.enable_memory_backing_report=1(default off).This is the first diagnostic adapter, not complete session backing accounting or a no-OOM guarantee. It changes no allocation, placement, contract-selection, or budget policy and adds no device allocation or explicit synchronization. Control-flow/subgraph reports are not composed into a session peak.
Validation
Windows CPU Release, Ninja/MSVC 2022, using the existing isolated incremental build with the Windows memory leak checker disabled.
From the configured CPU Release build directory:
Passed.
Passed: 22 tests, no skips. Eight new cases cover sequential/parallel scratch, input/initializer aliases, source release and packed destinations, dynamic shapes, known-zero versus unavailable workspace, reporting disabled, and a real report-enabled CPU Run with and without an arena. The runtime case verifies allocator-retained bytes are not added to the modeled payload peak and unsupported allocator statistics stay unavailable.
From the source worktree:
Both passed. No new CUDA, training-enabled, minimal-build, or cross-stream runtime validation was performed for this slice.
Motivation and Context
Implements the first fixed-placement reporting slice from #32944, stacked on #33240's L1/L2 reconciliation. The stack is #33316 -> #33240 -> this PR.
The report deliberately separates conditional logical lifetime estimates from physical backing observations. Retained pools, alignment/fragmentation, packed ownership/sharing, constructor/preparation scratch, caller I/O/KV capacities, private library/EP allocations, capture, retained outputs, subgraph residency, and concurrency still prevent a complete physical peak.
The remaining sequence is to complete those backing adapters, introduce enforceable pre-allocation session budgets, then add generic contract selection and CPU-region placement. CUDA workspace/activation-pattern integration remains in #32071, WebGPU region integration in #32720, and the placement design in #33372; this PR does not duplicate them.