Integrate CUDA workspace with activation memory patterns by chilo-ms · Pull Request #32071 · microsoft/onnxruntime

Chi Lo (chilo-ms) · 2026-08-13T20:03:37Z

PR stack

Add workspace estimation accounting and reporting #31962 - Level-1 workspace estimation and reporting
Verify Level-2 workspace reservations #32189 - Level-2 reservation verification
Integrate CUDA workspace with activation memory patterns #32071 - activation-aware CUDA workspace preallocation

This is the final PR. Please review it relative to chilo/level2-workspace-verification; it contains six stack-specific commits.

Summary

register opted-in Level-2 workspace declarations as synthetic entries in the sequential execution plan;
trace workspace allocation/free events during the first compatible run so OrtValuePatternPlanner can pack workspace with non-overlapping activations;
resolve cached workspace pointers from each run's per-device memory-pattern backing buffer through OpKernelContext;
route in-tree CUDA MatMulNBits slot 0 through the planned pointer while retaining GetScratchBuffer() when no usable plan exists;
suppress unused CUTLASS workspace declarations when the cached small-M tactic is GEMV, while retaining conservative declarations on cache misses;
preserve memory-pattern block alignment and use negative synthetic IDs disjoint from OrtValue indices.

Runtime flow

Session finalization stores each declared usable size, aligned allocation size, slot, device, and synthetic pattern ID.
On the first compatible run, the execution frame traces the synthetic workspace lifetime; the kernel still allocates dynamically.
ORT generates and caches one memory pattern containing activation and workspace blocks.
On later compatible runs, the kernel receives backing_buffer + workspace_offset from the execution frame.
Oversized, unresolved, or workspace-free tactic requests continue through the existing dynamic/no-workspace behavior.

Tests

MemPatternPlannerTest.WorkspaceSharesNonOverlappingActivationBlock verifies activation/workspace offset reuse.
MatMulNBitsWorkspace.SequentialChainUsesSharedPlannedWorkspace verifies three distinct synthetic entries, first-run dynamic fallback, cached-run planned workspace, activation overlap, workspace reuse, and output parity.
Tactic-aware tests verify that cached GEMV suppresses the declaration and a cache miss remains conservative.
MatMulNBitsWorkspace.PlannedWorkspaceReducesCudaArenaAllocation verifies that overlapping workspace with dead activation memory reduces second-run BFCArena CUDA allocation.

Validation

onnxruntime_test_all Debug build: passed
MemPatternPlannerTest.*, ResourceAccountantTest.*, RealAccountantTest.*: 20 tests passed

CUDA E2E execution requires CUDA CI or a compatible CUDA developer build. Regenerating the local Windows CUDA test configuration is currently blocked by the known onnxruntime_providers_cuda_ut / onnxruntime_provider_test module dependency cycle.

Current scope

sequential execution only;
one workspace slot per opted-in kernel;
in-tree CUDA MatMulNBits pilot;
Plugin CUDA and parallel/multi-stream lifetime modeling are deferred.

Tracking: #29775

Design document: https://github.com/microsoft/onnxruntime/blob/chilo/workspace-estimation-preallocation-design/docs/annotated_partitioning/workspace_estimation_and_preallocation.md

Add an opt-in sequential static-workspace pilot that reuses one aligned buffer per device and routes MatMulNBits slot zero through it with dynamic fallback. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839

Use the cached fpA/intB tactic for fixed small-M shapes, while retaining conservative workspace declarations for cache misses and CUTLASS GEMM. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839

Compare scratch and planned MatMulNBits execution with a controlled BFCArena strategy and verify that workspace overlap reduces second-run CUDA allocation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839

## Why `DeclareWorkspaceRequirements()` currently receives a dense `span<const TensorShape>`. That loses positional correspondence when an optional input is omitted and prevents Level 2 workspace estimation for valid signatures with an internal hole, for example: - PackedMultiHeadAttention: missing `bias` with later `token_offset` / `cumulative_sequence_length` - GroupQueryAttention: missing `seqlens_k` with later `total_sequence_length` - MatMulNBits: missing optional `g_idx` with later `bias` The resolver also dropped partial shapes wholesale, conflating “input is missing”, “input is present but shape metadata is unavailable”, and “input is present with unknown dimensions”. ## What changed - Add an owned, positional `WorkspaceInputShape` descriptor with three explicit states: - `Missing` - `PresentWithShape` - `PresentWithoutShape` - Preserve rank-0 tensors, zero extents, partial shapes (`-1` per unknown dimension), and optional holes. - Deep-copy dimensions so descriptors remain valid after graph/shadow-graph teardown. - Replace the unreleased C++ virtual directly; no compatibility overload is retained. - Update MatMulNBits Level 2 to consume the new contract while preserving Level 1/Level 2/runtime workspace parity. - Preserve the adapter-side default no-op. Plugin C-ABI forwarding/invocation remains deferred. ## MatMulNBits parity The shared workspace formula now handles known-empty outputs before architecture-specific arithmetic: - ordinary case: Level 1 = Level 2 = runtime = 1792 bytes - empty output: Level 1 = 0, Level 2 emits no slot, runtime = 0 - known-zero partial shapes such as `[0, unknown, K]` are recognized as empty - invalid negative dimensions and checked-arithmetic overflow remain unavailable/error paths ## microsoft#32071 integration constraints microsoft#32071 uses Level-2 declarations as synthetic negative allocation IDs in the existing activation `MemoryPattern`. The first run for a compatible feed-shape key traces workspace allocation/free while the kernel still uses `GetScratchBuffer()`. Later compatible runs can return an offset in the normal per-run pattern backing buffer; an unavailable pattern or a request larger than the declared capacity falls back to `GetScratchBuffer()`. This is not persistent `Initialize()`-time allocation and does not protect the first run from OOM. The current microsoft#32071 pilot preplans only kernels that explicitly opt in and declare exactly one slot. The PA/PMHA follow-up uses that contract for both operators: - PackedAttention declares one 256-byte-aligned root containing projection and Attention regions. - PackedMultiHeadAttention declares one 256-byte-aligned Attention root. The declarations alone do not opt either kernel into planning. The future integration must atomically add `SupportsPreallocatedWorkspace()`, retrieve slot 0, and split PA's root at its declared aligned Attention offset. Until then, PA retains its two dynamic allocations and PMHA retains its one dynamic allocation. Generic multi-slot framework support remains available, but PA does not require a multi-slot planner extension. ## Deferred Phase-B semantics This PR intentionally does **not** add: - max-shape provenance to the descriptor - a distinct marker for “proven zero” versus “unavailable” - plugin C ABI forwarding - workspace offset planning/allocation - partition/resource-accounting policy Those semantics belong to the planner/accounting integration, where runtime bound checks and fallback behavior can be defined coherently. ## Compatibility with microsoft#32071 This PR should merge first. microsoft#32071 should then rebase and update its Level-2 caller to pass the same positional `WorkspaceInputShape` span while retaining its reservation verification and execution-plan registration. ## Validation - CPU framework/session targeted tests: 29 passed - CUDA MatMulNBits workspace tests: 16 passed - same-session positive-`M` then zero-`M` runtime coverage - SM90 empty-output host regression - core and adapter header-isolation probes - C++ formatting and `git diff --check` Design tracking: microsoft#29775 --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Copilot-Session: c04148cc-7ace-4cf4-b981-0ddc92334e78

## Summary Add MatMulNBits-equivalent operator-side workspace estimation for CUDA `PackedAttention` and `PackedMultiHeadAttention`: - Level 1 derives a conservative estimate from the node, resolved input shapes, CUDA device properties, and the EP's resolved attention options. - Level 2 declares the same estimate from positional `WorkspaceInputShape` metadata and constructed kernel state. - Existing graph-free runtime recipes remain the single source of truth for workspace bytes and layouts. - PackedAttention declares one 256-byte-aligned root in slot 0, with internal projection and attention regions. - PackedMultiHeadAttention declares one 256-byte-aligned attention root in slot 0. Level 1 is log-only, matching the current MatMulNBits pilot. This PR does not add #32071-specific planner APIs or change runtime `GetScratchBuffer()` behavior. ## Route aggregation Runtime routes are mutually exclusive, so the estimate uses: ```text PackedAttention: align_up(projection_bytes, 256) + max(feasible TRT, MEA, unfused recipes) PackedMultiHeadAttention: max(feasible Flash, TRT, MEA, unfused recipes) ``` Route reachability is evaluated conservatively for every runtime shape up to the supplied maximum geometry. This is necessary because Flash/MEA thresholds and attention-bias alignment gates are not monotonic when moving from a maximum shape to a smaller runtime shape. Unfused fallback is always retained, and a failure to size any included route makes the estimate unavailable rather than silently underestimating. ## Shape and zero semantics - Missing mandatory inputs, shapeless required inputs, unknown dimensions, malformed geometry, and checked-arithmetic overflow produce no estimate. - `WorkspaceInputShape` does not carry max-shape provenance, so zero-shaped framework hints are conservatively treated as unavailable. - Exact zero behavior remains supported by the graph-free runtime recipes. - At the current Level-2 boundary, both unavailable and zero are represented by an empty requirements list. ## Planner integration - Both operators fit #32071's current one-slot pilot. Generic framework multi-slot support remains unchanged. - A declaration alone is not planner opt-in. `SupportsPreallocatedWorkspace()`, slot-0 retrieval, and PA root slicing must land atomically in the planner integration. - Until then, PA retains its two dynamic allocations and PMHA retains its one dynamic allocation. ## Build boundaries The framework adapters and kernel overrides are excluded from: - CUDA minimal builds - `DISABLE_CONTRIB_OPS` builds - CUDA plugin EP builds The graph-free workspace recipes remain available to the shared BERT attention infrastructure where required. ## Validation - 18/18 PA/PMHA workspace estimator tests - includes direct production-kernel Level-2 declaration tests - route-threshold, max-not-sum, aligned-root padding/no-padding, optional-hole, zero, overflow, and malformed geometry coverage - 23/23 existing packed-attention workspace recipe tests - 20/20 existing hand-calculated runtime parity cases - 26/26 PackedAttention/PackedMultiHeadAttention runtime operator tests - CUDA provider-test build - 145 CUDA internal tests executed: 143 passed, 2 unrelated LeanAttention skips - `DISABLE_CONTRIB_OPS` and CUDA-minimal compile-guard probes - C++ formatting and diff checks ## Dependency This is a stacked follow-up to #32312. The base should change to `main` after #32312 merges. Tracking: #29775 --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Copilot-Session: c04148cc-7ace-4cf4-b981-0ddc92334e78

## Summary Add the graph-free workspace preparation foundation for CUDA `GroupQueryAttention`, following the checked-recipe architecture established by #32321. - model windowed KV-cache staging and compaction without changing runtime allocations; - model 3-vector/6-vector sequence metadata and Flash fast-decode suppression; - model route-selected QKV preprocessing requirements; - use checked arithmetic and one internally 256-byte-aligned preparation layout; - distinguish transient staging/compaction from output KV cache and persistent/prepacked state. ## Scope This is the first GQA workspace-estimation PR in the #29775 rollout. It intentionally does **not** add backend-internal workspace recipes, route aggregation, Level-1 estimation, Level-2 declaration, cuDNN workspace queries, or #32071 planner integration. Existing `GetScratchBuffer()` allocation count, size, lifetime, and pointer layout are unchanged. Follow-up PRs will add: 1. XQA and Flash regular/fast-decode workspace recipes. 2. CUTLASS MEA and unfused fallback recipes plus the complete route aggregate. 3. Optional-aware Level-1/Level-2 adapters with one 256-byte-aligned operator root. 4. cuDNN support only if an exact runtime-parity workspace query can be established. ## Validation - 30 GQA workspace recipe and validation tests passed. - 73 related GroupQueryAttention runtime tests passed; 12 WebGPU-only tests skipped because WebGPU EP was unavailable. - CUDA provider test target built successfully. Tracking: #29775 Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b24d04bb-12ab-41a0-8941-76eb9769005d

Resolve workspace-planning conflicts against current main while preserving the activation-memory integration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Include the latest reservation verification and fused-node accounting changes from PR 32189. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> # Conflicts: # include/onnxruntime/core/framework/resource_accountant.h # onnxruntime/core/framework/graph_partitioner.cc # onnxruntime/core/framework/resource_accountant.cc # onnxruntime/core/framework/session_state.cc # onnxruntime/test/framework/resource_accountant_test.cc # onnxruntime/test/providers/cuda/test_cases/matmul_nbits_e2e_workspace_test.cc

Preserve Level-2 workspace reservations through partitioning and update CUDA workspace tests for the WorkspaceInputShape API. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

## Summary Stacked on #32446, add graph-free checked workspace recipes for the CUDA GroupQueryAttention decode backends: - reproduce XQA semaphore and multi-block scratch sizing from plain device/shape scalars; - retain XQA RoPE Q/K and dynamic FP32 head-sink conversion allocation parity; - reproduce Flash regular and fast-decode LSE/split-accumulator sizing; - preserve the GQA fast-decode distinction between KV heads used for split selection and query heads used for accumulator allocation; - fail closed for unsupported XQA cache/storage configurations. ## Flash bounded-estimation constraint Flash workspace is not monotonic in KV length. For `B=1`, `S_q=1`, two heads, head size 64, and 108 SMs, increasing KV length from 13,824 to 13,825 changes the selected split count from 54 to 28 and reduces workspace. Future Level-1/Level-2 aggregation must compute a conservative envelope or report unavailable rather than evaluating only the maximum shape. The split heuristic matches runtime double-threshold semantics, including a regression boundary where using `0.85f` would select 17 splits instead of the runtime 20. ## Scope This PR adds concrete selected-backend recipes only. It does not add MEA/unfused recipes, complete-route composition, dynamic-bound route reachability, L1/L2 adapters, cuDNN workspace queries, runtime allocation changes, or #32071 planner integration. ## Validation - CUDA provider test target built successfully. - 201 CUDA internal tests passed; 2 unrelated Lean Attention tests skipped. - 44 GQA preparation/XQA/Flash tests passed, including runtime-helper parity and Flash discontinuity regressions. Tracking: #29775 Dependency: #32446 > This draft temporarily targets `main` because GitHub cannot use a fork-only branch as the base of an upstream PR. Its diff will reduce to this commit after #32446 merges. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: b24d04bb-12ab-41a0-8941-76eb9769005d

Include reservation tracking across graph replacements and Level-2 mutations from PR 32189. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bf9b9b9a-ce06-44c9-94b2-784fda3c6c39

Copilot

🟡 Changes recommended

Training-enabled builds bypass workspace tracing, and the new usage probe can retain stale state.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Integrates CUDA kernel workspace allocations into activation memory-pattern planning, piloted with MatMulNBits.

Changes:

Adds synthetic workspace lifetimes and cached pointer resolution.
Uses planned workspace in CUDA MatMulNBits, with dynamic fallback.
Adds tactic-aware declarations and CUDA/framework tests.

File summaries

File	Description
`onnxruntime/test/providers/cuda/test_cases/matmul_nbits_workspace_test.cc`	Exposes workspace-use instrumentation.
`onnxruntime/test/providers/cuda/test_cases/matmul_nbits_workspace_test_probe.h`	Declares the new test probe.
`onnxruntime/test/providers/cuda/test_cases/matmul_nbits_e2e_workspace_test.cc`	Adds tactic, reuse, correctness, and arena tests.
`onnxruntime/test/framework/mem_pattern_planner_test.cc`	Tests activation/workspace offset reuse.
`onnxruntime/core/session/provider_bridge_ort.cc`	Bridges workspace retrieval to providers.
`onnxruntime/core/session/inference_session.cc`	Validates sequential mode and reuses synchronization configuration.
`onnxruntime/core/providers/shared_library/provider_wrappedtypes.h`	Wraps workspace retrieval for shared providers.
`onnxruntime/core/providers/shared_library/provider_interfaces.h`	Extends the provider-host interface.
`onnxruntime/core/framework/session_state.cc`	Registers declared workspace plans.
`onnxruntime/core/framework/sequential_execution_plan.h`	Stores synthetic workspace allocation metadata.
`onnxruntime/core/framework/ort_value_pattern_planner.h`	Adds device-specific workspace tracing APIs.
`onnxruntime/core/framework/ort_value_pattern_planner.cc`	Implements workspace tracing.
`onnxruntime/core/framework/op_kernel_context_internal.h`	Resolves and releases planned workspace.
`onnxruntime/core/framework/execution_frame.h`	Adds execution-frame workspace APIs.
`onnxruntime/core/framework/execution_frame.cc`	Traces or resolves workspace blocks.
`onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h`	Opts `MatMulNBits` into preallocation.
`onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc`	Uses planned workspace with scratch fallback.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.h`	Defines the GEMV tactic boundary.
`onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm_profiler.cc`	Applies the shared GEMV boundary.
`include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h`	Documents the opt-in setting.
`include/onnxruntime/core/framework/op_kernel.h`	Adds kernel preallocation capability opt-in.
`include/onnxruntime/core/framework/op_kernel_context.h`	Adds planned-workspace retrieval.

Review details

Files reviewed: 22/22 changed files
Comments generated: 3
Review effort level: Balanced

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

+  const auto* pattern = mem_patterns_->GetPatterns(location);
+  const auto* block = pattern == nullptr ? nullptr : pattern->GetBlock(pattern_id);
+  auto buffer_it = buffers_.find(location);
+  if (block == nullptr || block->size_ != allocation_bytes || buffer_it == buffers_.end()) {
+    return Status::OK();


+        const bool use_preallocated_workspace = workspace != nullptr;
+        last_compute_used_preallocated_workspace_.store(
+            use_preallocated_workspace, std::memory_order_relaxed);


+    for (const auto& workspace_plan : node_it->second) {
+      if (workspace_plan.slot_id != slot_id || requested_bytes > workspace_plan.size_bytes) {
+        continue;


Chi Lo (chilo-ms) added 6 commits August 20, 2026 10:48

Plan CUDA workspace with activations

f867bdd

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839

Test planned workspace across MatMulNBits chain

84409b9

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839

Bridge planned workspace access to CUDA

01165aa

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: cdd38fe6-fcf3-45c2-acea-b8e6206a7839

Chi Lo (chilo-ms) force-pushed the chilo/static-workspace-preallocation branch from 33bca06 to f4e45cb Compare August 20, 2026 19:38

Chi Lo (chilo-ms) mentioned this pull request Aug 20, 2026

Verify Level-2 workspace reservations #32189

Open

This was referenced Sep 4, 2026

Add GQA workspace preparation recipes #32446

Merged

Add GQA XQA and Flash workspace recipes #32453

Merged

Add complete GQA workspace route recipes #32454

Open

Ti-Tai Wang (titaiwangms) added this pull request to stack #32601 September 14, 2026 23:10

Chi Lo (chilo-ms) and others added 3 commits September 14, 2026 16:15

Merge current main into static workspace preallocation

28e723b

Resolve workspace-planning conflicts against current main while preserving the activation-memory integration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Fix workspace integration after branch merge

19c8b90

Preserve Level-2 workspace reservations through partitioning and update CUDA workspace tests for the WorkspaceInputShape API. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Merge latest Level-2 workspace verification updates

3258464

Include reservation tracking across graph replacements and Level-2 mutations from PR 32189. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bf9b9b9a-ce06-44c9-94b2-784fda3c6c39

Chi Lo (chilo-ms) marked this pull request as ready for review September 15, 2026 23:28

Copilot AI balanced review requested due to automatic review settings September 15, 2026 23:28

Copilot started reviewing on behalf of Chi Lo (chilo-ms) September 15, 2026 23:29 View session

Copilot AI reviewed Sep 15, 2026

View reviewed changes

Navigation Menu

Uh oh!

Integrate CUDA workspace with activation memory patterns - #32071

Integrate CUDA workspace with activation memory patterns#32071
Chi Lo (chilo-ms) wants to merge 10 commits into
chilo/level2-workspace-verificationfrom
chilo/static-workspace-preallocation

Chi Lo (chilo-ms) commented Aug 13, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

Chi Lo (chilo-ms) commented Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

PR stack

Summary

Runtime flow

Tests

Validation

Current scope

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Chi Lo (chilo-ms) commented Aug 13, 2026 •

edited

Loading