[Chi Lo (chilo-ms)]

Description

Adds a design document for policy-based, route-aware Level-1 (L1) workspace estimation, the foundation for capacity-aware CPU offloading on VRAM-constrained machines.

Problem

ORT's partition-time workspace estimate can over-reserve VRAM by including routes the committed runtime policy will not execute, such as GQA's quadratic unfused path. This needlessly pushes heavy nodes off GPU. An estimate is also unsound if runtime remains free to escape to an omitted route.

Proposal

Bind a committed dispatch policy: (admitted workload envelope, phase-specific permitted routes, workspace bound). Runtime is contractually constrained to that policy, and L1 budgets every permitted route for each admitted phase, then takes the maximum across prefill, decode, and chunked prefill while phases remain mutually exclusive per Run.

The document covers:

• the policy contract and soundness invariant;
• route descriptors as workload/device-dependent functions;
• query-length and KV-capacity envelopes;
• per-node/per-phase commitment for prefill, decode, and chunked prefill;
• maximum-over-phases L1 accounting and Level-2 lifetime implications;
• explicit runtime enforcement and fallback closure;
• portable auto / latency / memory / safe intent semantics;
• empirical profiling as validation rather than proof;
• llama.cpp, vLLM, and SGLang offloading comparisons;
• a worked GQA policy and staged implementation roadmap;
• framework-owned allocation identity/lifetime accounting for Initialize and Run;
• physical backing, allocator retention, external KV/I/O ownership, and reserve composition;
• pre-allocation budget planning, immutable contract bindings, and L2 reconciliation;
• deterministic budget-constrained contract/CPU-region selection, with latency calibration later.

Implementation alignment

#33316 implements the GQA-first committed-policy slice: L1 and L2 resolve the same phase-specific masks, runtime dispatch is gated by the current phase's mask, and memory retains selected general routes as bounded decode fallbacks. Whole-session lifetime/physical-backing accounting and budget-constrained contract selection are the next framework follow-ups. Optional ORT-side workload-envelope enforcement and additional bounded fallback coverage are explicit deferred TODOs.

Scope

• Docs only. No code or behavior change in this PR.
• Design foundation for the workspace-estimation and policy-enforcement PR series.

Files

• docs/annotated_partitioning/policy_based_workspace_estimation.md — the design document.
• docs/annotated_partitioning/future_directions_constrained_env.md — cross-reference paragraph.

[Copilot]

Copilot review overview

🟡 Changes recommended

The soundness model relies on an unenforced shape hint and contains factual inconsistencies with current GQA behavior and configuration.

Review effort: Balanced
Findings: 1 · 2 · 2

Open (5)

• Max shape override does not enforce the runtime workload envelope · New
• Absolute sequence length is conflated with resident cache capacity · New
• Route enforcement prerequisite contradicts documented case handling · New
• L1 estimate is not equal to concrete runtime allocation · New
• Unimplemented workspace option is presented as an existing setting · New

What changed in this PR

Documents a route-aware workspace-estimation design for safer, less conservative GPU partitioning.

Changes:

• Defines enforceable dispatch policies and workload envelopes.
• Covers GQA routing, configuration, validation, and rollout.
• Cross-references the design from constrained-environment guidance.

File	Description
policy_based_workspace_estimation.md	Adds the workspace-policy design.
future_directions_constrained_env.md	Links the new design from existing guidance.

---

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

[Ti-Tai Wang (titaiwangms)]

So we basically add "dispatch order" to set up the correct priority as opposed to choosing the fallback (conservative) choice. The only risk is that we don't know for sure whether it will be taken during runtime? We only know it's accessible.

[Chi Lo (chilo-ms)]

Exactly right. L1 never predicts the taken route — it enumerates the reachable set and takes the max over it. "Reachable but maybe not taken" is fine for an upper bound: a seq-free-ineligible route is excluded; a reachable one is included in the max even if a given Run picks something cheaper. Dispatch order matters only to prove the unfused terminal unreachable (its guard requires every higher flag false). The corollary — which your next comment correctly presses — is that every reachable route must be sized, not just Flash/XQA.

[Ti-Tai Wang (titaiwangms)]

Design review (documentation only): I agree with the central distinction: we need not predict a single route, but must prove an upper bound for every route and input the runtime may execute. The Case A Flash/MEA argument can rule out unfused when its complete static eligibility predicate holds for every admitted Run; that alone does not make the proposed lower L1 bound sound. Two additional, actionable gaps are noted inline. I also agree with the existing comments that session.max_shape_override does not enforce a runtime envelope and that absolute total length differs from windowed-cache capacity. Please gate Stage 2's smaller estimate on a hard, cache-mode-aware envelope and complete sizing of every reachable route; otherwise preserve the conservative bound or return estimate unavailable. Profiling in Stage 4 can validate an upper bound but should not replace it with a sampled peak.

[Ti-Tai Wang (titaiwangms)]

Case A proves at most that unfused is unreachable; it does not prove every higher-priority route has a workspace upper bound. group_query_attention.cc:741-771 may choose cuDNN before Flash (including auto-selection on SM >= 90), while group_query_attention_workspace_bounds.cc:241-244 explicitly declines to bound a reachable cuDNN route without a graph-free workspace oracle. Flash being statically eligible is therefore insufficient to budget only Flash/XQA. Please require a sound cuDNN bound or prohibit cuDNN at runtime before enabling a smaller Stage-2 estimate; otherwise retain the conservative/unavailable result.

[Chi Lo (chilo-ms)]

Agreed that Case A proves only unfused unreachable, not that Flash/XQA bound the higher routes. This is already enforced in the estimator: BuildBounds flags a reachable cuDNN route (group_query_attention_workspace_estimate.cc:383-390, gate UseCudnnFlashAttention() || (AllowCudnnFlashAttentionAuto() && device_prop.major >= 9)) and GetGQAWorkspaceAggregateForBounds returns Unavailable for it (group_query_attention_workspace_bounds.cc:241-244) → the estimate is nullopt and the conservative reservation is retained. The smaller Stage-2 estimate is therefore never emitted when cuDNN is reachable. The real gap is the doc: the Case A paragraph's "max(prefill, decode, chunk) is sound as-is" silently assumes cuDNN is unreachable. I''ll add the cuDNN-reachability caveat so the prose matches the guard. Thanks for catching it.

[Ti-Tai Wang (titaiwangms)]

This corner need not bound Flash workspace even when the route does not switch. flash_attention/flash_api.cc:199-269 chooses num_splits by an occupancy heuristic: a large S_q=P can use one split (no accumulators), while a shorter admitted chunk at the same KV extent can use multiple splits and allocate substantial lse_accum and out_accum buffers. group_query_attention_workspace_bounds.cc:104-148 already bounds split counts rather than evaluating only the endpoint. Please bound all heuristic regimes (including preparation buffers); a single Stage-4 profile at the envelope corner cannot replace that proof.

[Chi Lo (chilo-ms)]

Agreed the worked-example corner (S_q=P, S_kv=C) labelled "[linear]" over-simplifies the Flash term. The estimator does not evaluate the heuristic at the corner: FlashBackendEnvelope (the group_query_attention_workspace_bounds.cc:104-148 you cite) bounds the split count by its structural maximum max_splits = min(128, multi_processor_count, ceil((kv + block_n − 1) / block_n)) and sizes lse_accum/out_accum with it — the code comment notes it "remains sound across its discontinuities," i.e. the one-split and many-split regimes are both covered; preparation/QKV buffers come from GQABufferRequirements. Stage-4 profiling validates this bound rather than replacing it with a sampled peak. I''ll rewrite the worked example to reference the max_splits bound instead of implying a single endpoint is evaluated.

[Chi Lo (chilo-ms)]

Thanks Ti-Tai Wang (@titaiwangms) — this is the right frame: prove an upper bound for every reachable route, don't predict one. Three doc fixes incoming (behavior-level checks already exist in the estimator; these align the prose with the code):

1. Gate Stage 2's smaller estimate explicitly on a hard, cache-mode-aware envelope and cuDNN-unreachability, else return Unavailable. The estimator already enforces the cuDNN case — BuildBounds marks a reachable cuDNN route and GetGQAWorkspaceAggregateForBounds returns Unavailable (group_query_attention_workspace_bounds.cc:241-244) — so the smaller estimate is never emitted when cuDNN is reachable; the Case A prose will say so.
2. Correct the Flash worked example to cite FlashBackendEnvelope's structural max_splits = min(128, SM_count, ceil(kv/block_n)) bound (bounds.cc:104-148) rather than implying a single envelope corner is evaluated. Stage-4 profiling validates this bound, it does not replace it with a sampled peak.
3. Relabel the knobs: session.max_shape_override as an estimation hint (not a runtime envelope enforcer) and ep.cuda.gqa_workspace_max_total_sequence_length as a proposed option — per the bot comments I'm also addressing — and note a hard runtime envelope check (query + KV) is required for the no-OOM claim.

I''ll push these as a follow-up commit on this branch for your re-review.

[Chi Lo (chilo-ms)]

Updated in 324f7c46c7 to make the proof boundary explicit: analytical upper bounds, empirical peaks, and the generic 50% heuristic are distinct; unavailable GQA proof bounds fail closed for capacity-aware CUDA placement; profiling validates/calibrates rather than replaces an upper bound; and whole-graph peak/lifetime accounting is separate from per-node route bounds. The caller-declared query/KV envelope remains explicitly non-enforced and is not presented as a no-OOM guarantee.

[Ti-Tai Wang (titaiwangms)]

Follow-up design review at 324f7c46c77f919f8d56d0fd9d3ade2812cf75a3, informed by #29775 and the related implementation PRs.

The main earlier concerns are addressed: profiling now validates rather than replaces the analytical proof; declared estimates are distinguished from runtime-enforced envelopes; and a whole-session no-OOM claim depends on global lifetime accounting and controlled allocator capacity/reserve, not just a node's scratch bound.

I have four remaining documentation corrections, detailed inline:

1. Label fail-closed placement on an unavailable GQA estimate as proposed, rather than current behavior.
2. Separate embedding gather from the dense vocabulary projection performed by lm_head.
3. Specify the execution/copy mechanism for MoE offload; redirecting an initializer to a CPU allocator alone is not CPU execution or GPU weight streaming.
4. Separate measured concrete-allocation savings from analytical admission-bound savings in the 38x worked example.

These are documentation/design corrections, not claims that this docs-only PR changes runtime behavior. The current two failed CI jobs do not establish a documentation regression: the lint job failed resolving a Docker Hub image with HTTP429, and the ASan job received a runner shutdown during Python setup before build/test. No local build, performance measurement, or external-framework offload experiment was run.

[Ti-Tai Wang (titaiwangms)]

Clarify current behavior versus the proposed fail-closed contract. At this pinned revision, CUDAExecutionProvider::GetCapability still calls ComputeResourceCount after a GQA estimate returns unavailable; the accountant can use generic 50% fallback or profile data. There is no GQA-specific unavailable-to-decline branch. Please label analytical-bound-only placement here and in the earlier cuDNN discussion as a future requirement, and explicitly state today's fallback behavior. Otherwise the descriptive section suggests that enforcement is already implemented.

[Ti-Tai Wang (titaiwangms)]

Separate embedding from lm_head in this cost model. Embedding is a gather, but lm_head normally computes a dense hidden-to-vocabulary projection. Its arithmetic and result-transfer costs are not the same as transferring a few embedding rows. Please split the row and evaluate the projection's offload tradeoff separately rather than assigning both near-zero gather compute and the same recommendation.

[Ti-Tai Wang (titaiwangms)]

Initializer placement alone does not implement the proposed offload. The current CUDA MoE/QMoE registrations expect device inputs for expert weights; placing the source initializer on CPU does not make the fused node execute on CPU or make its CUDA kernel stream host-resident weights. Please distinguish whole-node CPU assignment, with supported dtype/feature checks and activation-copy boundaries, from a new GPU weight-staging/streaming mechanism. State which mechanism the proposed option selects and account for any retained device/prepacked copies before claiming the residency savings.

[Ti-Tai Wang (titaiwangms)]

The structural envelope does not establish a 38x admission reduction. The earlier 123 MB versus approximately 3.2 MB comparison is a concrete allocation observation. The structural split ceiling below can combine query and split extrema that do not occur together, so its analytical bound can be much larger than that concrete Flash allocation. Please label the measured comparison separately and avoid carrying its ratio into W_L1 without a bound calculation under the same geometry and assumptions. The displayed CompleteWorkspace expression also needs to include preparation/QKV terms, or be explicitly labeled backend-only.

[Tianlei Wu (tianleiwu)]

Well-structured design document covering policy-based, route-aware L1/L2 workspace estimation and whole-session GPU memory planning. Left two inline comments on GPU-resident total_sequence_length phase classification and on how current post-partition reservation callbacks / empty L2 declarations interact with the proposed verification invariants.

[Tianlei Wu (tianleiwu)]

Clarify GPU-resident total_sequence_length phase classification for single-token first prompts and the no past input ⇒ prefill-only table row.

1. In group_query_attention_helper.h:412-427, when total_sequence_length resides on the GPU (!is_total_seqlen_on_cpu), CheckInputs skips the prompt-type block entirely and leaves is_first_prompt = false and is_subsequent_prompt = false for all steps:
   • A single-token first prompt (sequence_length == 1, total_sequence_length == 1 on GPU) has !is_first_prompt && sequence_length == 1, so runtime classifies it as decode rather than prefill.
   • The batch_size == 1 guard for sequence_length > 1 with past context (group_query_attention_helper.h:417-420, cited on line 313 as the guard protecting max across phases) is also bypassed when total_sequence_length is on GPU.
2. Line 302's table states no past input ⇒ prefill-only, whereas BuildBounds (group_query_attention_workspace_estimate.cc:244) currently requires past_key/past_value (if (!Present(shapes, kPastKey)) return std::nullopt;). Moreover, if a future prefill-only mode omitted decode and chunk from AdmittedPhases, any run with GPU-resident total_sequence_length would set is_first_prompt = false and classify sequence_length == 1 as decode and sequence_length > 1 as chunk, escaping a prefill-only phase mask.

[Chi Lo (chilo-ms)]

Thanks, agreed. Updated the document in 1934035 to make both GPU-resident cases explicit: a single-token first prompt uses the decode mask, and a multi-token first prompt uses the chunked-prefill mask. These are runtime dispatch classes, not necessarily the semantic prompt stage.

The document now states that BuildBounds currently requires both past K/V shapes and declines no-past configurations; absent past inputs do not establish prefill-only coverage. It also removes the CPU-only batch guard as a universal justification for max across phases: alternatives are mutually exclusive within one GQA invocation, but that does not prove phase lockstep across an entire Run or exclude overlapping allocations.

I also opened #33420, stacked on #33316, to tie phase admission explicitly to the shared runtime classifier and query-length bound, with regression coverage for CPU/GPU total-length metadata, missing past inputs, and phase-mask coverage. #33316 already admits the conservative phase set; this follow-up protects against future narrowing rather than claiming an existing escape. It introduces no device-to-host scalar read, runtime routing change, or new GPU-path batch restriction.

[Tianlei Wu (tianleiwu)]

Document how current post-partition callbacks and empty L2 declarations differ from the target verification invariants.

1. Lines 815–818 state that orphaned bindings fail verification and A fused region requires its own contract, not the sum of constituent scratch peaks. Today, InferenceSession::TransformGraph (inference_session.cc:1882-1907) registers SetNodeReplacementCallback → ConsolidateWorkspaceReservations (resource_accountant.cc:26-46), which sums reservation.bytes across source_node_indices, and SetNodeRemovalCallback, which erases removed node_index entries from workspace_reservations before FinalizeSessionStateImpl's orphaned_reservations check (session_state.cc:2227-2236) runs.
2. Lines 823–825 state that a warning followed by dynamic growth cannot satisfy hard-budget/strict verification. In SessionState::FinalizeSessionStateImpl (session_state.cc:2191-2196), if kernel->DeclareWorkspaceRequirements(...) returns Status::OK() with requirements.empty() for a node that has a nonzero L1 reservation, it increments missing_declaration and continues without setting strict_verification_failed.

Calling out both gaps in the integration table below (lines 860–869) would make the required Stage 4/5 framework changes explicit.

[Chi Lo (chilo-ms)]

Thanks, agreed. Addressed in 1934035 with a new "Current reconciliation gaps, not implemented guarantees" discussion, corresponding integration-table entries, and explicit Stage 4/5 follow-ups:

• Replacement callbacks currently aggregate source reservations; that is bookkeeping, not a target implementation estimate or fused-region contract. The target design requires the replacement's own requirements and binding.
• Removal callbacks erase entries before orphan checking, so the final reservation map cannot reveal an obligation already lost. The proposed verification must track explicit retirement, validated transfer, or invalidation/re-estimation provenance.
• A successful empty L2 declaration with a nonzero L1 reservation currently increments missing_declaration without setting strict_verification_failed. The proposed opt-in strict mode must reject unresolved obligations or require evidence of a valid zero-workspace replacement; an empty declaration alone is not that evidence.

These are now clearly identified as required framework work, not guarantees provided by the current callbacks/finalization path. The update also calls for focused coverage of each gap and preserving legacy behavior outside the proposed opt-in contract/budget mode. This PR remains documentation-only; it does not claim to implement those framework fixes.