[Tracking] Constrained-environment memory and layer-offload roadmap

[titaiwangms]

Background

This continues the roadmap documented alongside the annotation-based partitioning
work (#27595, #28028, #28623, #28903). See docs/annotated_partitioning/
future_directions_constrained_env.md for the full design notes.

Goal: make ORT as easy to use as ollama/llama.cpp for running large models on
memory-constrained devices — automatic layer offloading without requiring exact
node annotations, and eventually static/pre-planned memory usage close to
llama.cpp's model.

Status

Near-term, low effort, high value:

1. Name-based substring matching (session.name_based_layer_assignment,
   SubstringMatcher) — DONE, shipped in Implement Name Based partitioning and update documents #28903.
2. Precise per-node memory/workspace estimation — NOT STARTED, next
   actionable step.
   • Add OrtKernelWorkspaceEstimateFunc (Level 1 — at partitioning time,
     before a kernel instance exists; used both in-tree and via the plugin
     host bridge, i.e. the layer that lets out-of-tree/plugin EPs call into
     the same estimation logic) and DeclareWorkspaceRequirements()
     (Level 2 — after kernel creation).
   • Resolve the chicken-and-egg problem: workspace size is only known
     after a kernel exists, but EP assignment happens before kernel
     creation (see doc §"The Chicken-and-Egg Problem" for resolution
     options).
   • Bypass the 1.5x safety multiplier in IResourceAccountant (ORT's
     per-node memory-budget tracker used during partitioning) per-kernel
     once an estimate is available; retain it as fallback for
     un-instrumented kernels and dynamic/unknown shapes.
   • ⚠️ Touches the plugin EP C ABI (OrtEp / PluginExecutionProvider,
     see ep_plugin_provider_interfaces.cc) — a sensitive, versioned
     boundary; coordinate with plugin EP owners.

Mid-term, medium effort — NOT STARTED:
3. Auto-partitioning with memory budget only — requires new design work
(layer discovery/boundary detection is NOT solved by the existing
per-node IResourceAccountant; needs explicit decisions on ordering
and atomicity before implementation). Can proceed independently of #2.
4. Static allocation mode (doc's Phase A: workspace pre-declaration +
Phase B: arena bypass for static-shape models). Depends on #2's
Level-2 workspace declaration.

• ⚠️ Requires an upfront workspace-declaration and planning phase
  during FinalizeSessionState() / allocation planning (allocation_planner.cc),
  plus execution-plan/frame state and kernel workspace retrieval changes.
  Existing activations already bypass the arena via BFCArena::Reserve();
  undeclared/dynamic workspace must retain the current arena fallback.

Long-term, high effort, ollama-parity — NOT STARTED (new work, not
scoped by the original design doc's two Directions):
5. Layer prefetch pipeline (stream weights CPU↔GPU during execution).
6. GenAI integration (KV-cache-aware planning + continuous batching).

• ⚠️ KV-cache/continuous-batching logic lives in the separate
  microsoft/onnxruntime-genai repo — this item likely needs
  coordinated PRs across both repos (author annotation, not from the
  source doc).

How to pick this up

Sub-issues will be linked here as each item above is picked up — comment
on this issue or self-assign to claim one.

Reference

Full design notes: docs/annotated_partitioning/future_directions_constrained_env.md

[titaiwangms]

Follow-up risk found while reviewing #28903: name-based matching fails silently, not just "unstable"

While walking through the #28903 implementation in detail, found a sharper version of the
"Name stability" risk already noted in the design doc (future_directions_constrained_env.md
§"Risks / Open Questions": "Node names aren't guaranteed stable across exports").

Root cause: NodeProto.name is an optional ONNX field with no format/uniqueness
guarantee (onnx.proto: "An optional identifier for this node in a graph. This field MAY be
absent in this version of the IR."). The layers.N/ style paths that session.name_based_layer_assignment
relies on are purely an export-tool convention (PyTorch torch.onnx.export, HF Optimum), not
an ONNX or ORT contract — and could change across exporter versions without any ORT-side
warning.

Confirmed by code inspection (onnxruntime/core/framework/layering_annotations.cc):

• LayeringIndex::Create() only validates that each rule maps to an available EP
  (matched_rule_count) — it does not track whether a rule's pattern ever matched any
  node in the graph.
• ProcessGraph() calls substring_matcher_->Match(node.Name()) per node with no aggregate
  match-coverage bookkeeping, and no LOGS(logger, WARNING/ERROR) anywhere in the file for
  "pattern X matched 0 nodes."
• A node that isn't matched simply keeps its default EP assignment — no error surfaced.

Practical impact: If an exporter changes its node-naming format (e.g. /model/layers.5/...
→ model.layers.5.... or a different path scheme), a previously-working
session.name_based_layer_assignment config silently stops matching all nodes. The
session still initializes successfully; the user only notices via unexpected performance/OOM
behavior (nodes not landing on the intended device), with no log signal pointing at the real
cause. This is a silent functional regression, distinct from (and arguably worse than) a
traditional API/ABI break.

Suggested new near-term item: match-coverage diagnostics for name-based layer assignment

• After ProcessGraph() completes (or during LayeringIndex::Create()), compute per-rule
  match counts for name-based mode and LOGS(logger, WARNING) for any pattern that matched
  zero nodes.
• Consider an opt-in strict mode that turns "0 nodes matched for pattern X" into a hard
  Status error at Initialize() time, so exporter-format drift is caught immediately rather
  than discovered as a silent regression.
• Small, self-contained, and can be scoped as its own sub-issue independent of the
  Direction 2 items above.

[titaiwangms]

Follow-up clarification: node-name origin is NOT uniformly "external/uncontrollable"

Investigated where the layers.N/...-style node names actually come from, since this
determines how much control ORT has over the risk noted in the previous comment. There
are two distinct production paths with very different risk profiles — sub-issue authors
should evaluate them separately rather than treating "exporter changed the naming format" as
one uniform risk.

Path A: Models built via mobius + onnxscript.nn (internal, ORT-controlled)

mobius (github.com/onnxruntime/mobius) does not trace/export PyTorch models. Per its
own README: "Rather than tracing or exporting PyTorch models, it constructs the ONNX graph
declaratively" using onnxscript.nn.Module, then loads pretrained HF weights only
(graph structure/naming is independent of the HF/PyTorch model code).

Traced the exact naming logic in onnxscript/_internal/builder.py:

def push_module(self, module: str, class_name: str = "") -> None:
    """module: The attribute name of the module (e.g. "layers.0")."""
    self._scope_stack.append((module, class_name))

def _qualify_node_name(self, name: str) -> str:
    """Qualify a node name with the current scope using '/' separator."""
    parts = self._scope_name_parts()
    if parts:
        return "/".join(parts) + "/" + name
    return name

and the layers.0, layers.1 indexing in onnxscript/nn/_module_list.py (a ModuleList-like
container: child._set_name(f"{name}.{key}")).

Implication: for mobius-built models, the /-separated, layers.N-indexed node-name
convention is entirely defined by onnxscript's builder.py — a dependency inside the same
Microsoft/ORT ecosystem. If this naming ever changes, it's a coordinable, internal
change (onnxscript + mobius + ORT's name_based_layer_assignment users can be updated
together), not a silent external break.

Path B: Models from torch.onnx.export / HF Optimum (external, uncontrolled)

For models produced by the traditional PyTorch/HF export path (not mobius), the node-naming
convention genuinely is external and uncontrolled by ORT — this is the scenario the
previous comment's risk (silent match-coverage failure) is actually about. ORT has no
visibility into, or influence over, if/when PyTorch's exporter or HF Optimum change their
internal naming logic.

Recommendation for sub-issue scoping

When scoping the "match-coverage diagnostics" work (previous comment) or any future
name-based-matching hardening, treat these as separate risk tiers:

• Path A (mobius/onnxscript): lower priority — coordinate directly with the onnxscript
  maintainers if/when naming changes; consider a shared contract or version pin.
• Path B (external exporters): higher priority for the diagnostics work — this is the
  case with no coordination channel, where silent failure is most likely and most costly to
  debug for end users.

[titaiwangms]

Correction & additional context: existing prior art for "Direction 2 / auto-partitioning"

Follow-up investigation surfaced two existing mechanisms that partially address what this issue's Direction 2 discussion characterized as unstarted work. This comment corrects that framing and proposes splitting Direction 2 into two separate ledgers, since the two goals it bundles together are not equally covered by existing prior art.

1. ORT already ships a profile-guided, budget-based partitioning path (CUDA Execution Provider, EP, only)

@yuslepukhin built a working, narrower-scope version of automatic memory-budget partitioning before this issue was opened — the "nobody has picked this up" framing above is not accurate for this specific capability:

• NodeStatsRecorder (include/onnxruntime/core/framework/resource_accountant.h) hooks into SequentialExecutor::Execute (onnxruntime/core/framework/sequential_executor.cc:589-643), right after each kernel's real Compute() call, and records per-node stats (input_sizes, initializers_sizes, total_dynamic_sizes, total_temp_allocations) via AllocatorStats/GetAllocatorStats. This sidesteps the "chicken-and-egg" estimation problem discussed elsewhere in this issue by measuring empirically rather than requiring a static per-kernel declaration API. Caveat: initializers_sizes is always recorded as 0 in the current implementation (constant/initializer inputs are folded into input_sizes instead) — the four fields are not independently meaningful; only their sum matters for budgeting.
• SizeBasedStatsAccountant + the kOrtSessionOptionsResourceCudaPartitioningSettings config option (format: "<cuda_memory_limit_KB>,<node_stats_file>") feed those recorded stats plus a byte budget into CUDAExecutionProvider::GetCapability() (onnxruntime/core/providers/cuda/cuda_execution_provider.cc:3418-3444) for per-node accept/reject against the budget. When a stats file is loaded, the budget check uses the measured bytes directly with no safety multiplier — the kAdHocSafetyMultiplierPercent = 150 multiplier (resource_accountant.cc:117) only applies in the no-stats estimation branch, not on top of measured stats.
• Precise trigger condition (corrected): if the partitioning option is enabled but its budget field is left empty, the same code path calls cudaMemGetInfo() to probe actual free VRAM at runtime as the threshold — this is opt-in (the config key must be present), not an automatic default when nothing is configured. In that no-stats-file case, per-node cost falls back to known initializer bytes + statically-known output bytes × 1.5 — dynamic outputs and workspace are not counted, so this fallback is closer to a "known-weight-size, contiguous-node-prefix" heuristic than a full hardware-adaptive equivalent of llama.cpp's layer offload.
• Shipping history predates this issue: PR Capacity aware partitioning #22766 (2025-02, "Capacity aware partitioning"), Annotation based partitioning along with resource accounting #27595 (2026-03), [CUDA EP Plugin] ResourceAcountant integration #28028 (2026-04) — all by the same author as PR Implement Name Based partitioning and update documents #28903 (Direction 1).
• The roadmap doc (future_directions_constrained_env.md) never mentions NodeStatsRecorder/the measured-stats path anywhere; its "IResourceAccountant Precision" section (~lines 276-296) only describes the no-stats estimation branch and its 1.5x multiplier.

What this does NOT solve, even setting aside the CUDA-only scope: the accept/reject cutoff is node-level, not layer-level (the doc itself notes cutoffs tend to land near layer boundaries but that isn't guaranteed); it requires a representative profiling run first, which itself requires the model to successfully execute once on the target EP — a model that can't fit/run at all up front gets no help from this path; and node coverage is limited to what the profiling run actually executed (an unexecuted If/Loop branch has no recorded stats — the same class of silent-gap risk raised in the match-coverage comment on this issue).

Important distinction — this does not reduce the need for Direction 2's static-workspace-declaration goal (item 2). The stats this mechanism records are a single aggregated per-node byte total, not a per-kernel workspace size usable by OrtValuePatternPlanner for static buffer layout. It answers "should this node/layer be accepted into this EP's budget" (roadmap item 3, auto-partitioning) — it does not and structurally cannot answer "what is the exact workspace this kernel needs so we can pre-allocate everything once and guarantee Run() cannot OOM" (roadmap item 2 / Phase A-B, the actual arena-elimination goal of Direction 2). These are two separate ledgers:

• Auto-partitioning (item 3): has real, shipped, CUDA-only prior art (this section).
• Static pre-allocation / workspace declaration (item 2, Phase A-B): still has no prior art; this profiling mechanism cannot substitute for it.

2. Olive already ships an offline, weight-size-based splitting pipeline (used in QNN recipes)

Olive's CaptureSplitInfo (olive/passes/pytorch/capture_split_info.py) + SplitModel (olive/passes/onnx/split.py) passes split a transformer model into N physically separate .onnx sub-models at export time. These are consumed by onnxruntime-genai's pipeline execution (src/models/decoder_only_pipeline.h/.cpp), which creates one OrtSession per configured pipeline entry (not necessarily 1:1 with SplitModel's raw output — Olive's ComposeOnnxModels pass (olive/passes/onnx/compose.py) can recombine context/iterator splits before the final pipeline is assembled). All CaptureSplitInfo-using configs found in the olive-recipes repo for Llama-3.1, Phi-3/3.5/4, and Qwen2.5 (QNN targets) hardcode num_splits (commonly 4); none reference the pass's own cost_model-driven auto-budget path, even though that code exists. This is suggestive, not conclusive, evidence about practitioner preference for manual tuning over automation — the recipe files alone don't establish why cost_model went unused.

This Olive path is inherently offline/ahead-of-time — the target device is known at export time — which is a different deployment scenario than the in-session CUDA EP path above (which can adapt at load time to whatever hardware is actually present), and different again from the "one model file, arbitrary/unknown end-user hardware" experience that motivated this roadmap (matching ollama's -ngl flag — number of layers offloaded to GPU — style UX). Olive's approach is a complete answer for the known-target-device case; it cannot address the unknown-until-runtime-hardware case.

Architectural note for future sub-issue scoping: Olive's split produces physically separate .onnx files run as independent sessions; ORT's LayeringIndex/NodeStatsRecorder mechanisms operate on node assignment within a single graph in one session. These are structurally incompatible today — any future work assuming the two can be directly unified should treat that as an open design question, not a given.

Suggested reframing for sub-issue scoping

• The auto-partitioning sub-goal (roadmap item 3) has real shipped prior art (the CUDA-only profile path above). The most concrete near-term extension is generalizing it to other EPs — but note this is not purely a config change: resource_accountant.h's own header comments state that budget enforcement for bridge-based/plugin EPs needs to move into the graph partitioner as a separate follow-up, and while IExecutionProvider::GetCapability()'s signature already accepts an IResourceAccountant* across all EPs, only the CUDA EP's implementation currently consults it — the other EPs' parameters are unused placeholders. Scope this as "EP-specific enforcement logic + partitioner-level enforcement for bridge/plugin EPs," not a small task.
• Adding a layer-boundary-aware accept/reject decision on top of the existing per-node accountant (closing the gap between "node-level budget cutoff" and the "how many layers fit" semantics users actually want — see the -ngl-analogy discussion earlier in this thread) is independent of the above and can be scoped separately.
• The static pre-allocation / workspace declaration sub-goal (roadmap item 2, this issue's original item 2) remains unaddressed by anything found so far and should stay a distinct, undiminished priority — it is not substitutable by either the CUDA profiling path or Olive's offline split, since neither produces per-kernel workspace sizes suitable for static buffer planning.

Filing this as context ahead of breaking out sub-issues, so scoping reflects what already exists rather than re-deriving it from scratch, and so auto-partitioning and static-pre-allocation are tracked as separate efforts with different levels of existing coverage.

[titaiwangms]

Opened pilot sub-issue for MatMulNBits (CUDA EP) two-level workspace estimation: #29810

[chilo-ms]

Proposed incremental workspace-estimation work

The current recommendation is to evolve workspace estimation in reviewable increments rather than combine host accounting, a plugin ABI change, and preallocation in one PR.

1. In-tree CUDA pilot and source-resolution policy — draft PR #31962

Draft PR #31962 uses MatMulNBits as the Level-1 pilot and adds framework accounting/reporting:

• retain workspace as pending during GetCapability and commit it only for accepted nodes, including the two-pass NHWC flow;
• resolve workspace per node as estimator-only, profile-only, max(profile, estimator) when both exist, or the existing fallback when neither exists;
• apply the selected workspace before the CUDA budget decision;
• report non-workspace memory, workspace memory, total memory, and source coverage;
• report Level-2 declarations after kernel creation as an initial step toward post-assignment verification.

This PR intentionally keeps the pilot in the in-tree CUDA path. It validates the estimator contract and resolution policy without also introducing a plugin ABI.

2. Centralize generic budget enforcement in the ORT host

Move source selection, cost calculation, threshold checking, accepted-capability commitment, and reporting into framework-owned partitioning code so in-tree and plugin EPs follow the same policy.

The EP should continue to own support discovery, device/kernel-specific estimates, and an explicit capacity-exceeded policy (for example, CUDA's current stop-on-first-overflow behavior, which avoids disconnected CUDA partitions). The host should own generic arithmetic and enforcement.

Plugin CUDA already has partial prior art: GetAvailableResource() supplies device capacity and PluginExecutionProvider performs host-side accounting. This step would centralize that enforcement rather than leave separate implementations in in-tree CUDA and the plugin host adapter.

3. Add a plugin-compatible Level-1 estimator contract

Define an ABI-safe mechanism for an EP to return an optional per-node workspace estimate using planning/max-shape context. ORT consumes that estimate through the same resolver introduced by the pilot.

Keep this separate because it changes the versioned Plugin EP boundary and requires coordination with plugin EP owners.

4. Implement CUDA Plugin estimator parity

Expose the MatMulNBits estimator through the new contract and add parity coverage between in-tree CUDA and CUDA Plugin. Additional high-workspace operators can then be migrated incrementally.

5. Level-2 verification and static preallocation

After Level-1 coverage is established:

• compare post-kernel Level-2 declarations with the reservation used during partitioning;
• define explicit behavior when Level 2 exceeds Level 1;
• add workspace slots/lifetimes to memory planning;
• route declared workspace to planned buffers while retaining arena fallback for unresolved or dynamic requirements.

Proposed design document

A more complete proposed policy and preallocation design is available at:

Workspace Estimation and Preallocation

The document covers profile compatibility, max(profile, estimator) resolution, dynamic shapes, Level-2 verification, workspace lifetimes, profiling coverage gaps, rollout phases, and open decisions. It is a design proposal, not documentation of fully implemented behavior, and can be moved into the main documentation tree once the policy is reviewed.