Integrate CUDA workspace with activation memory patterns

[Chi Lo (chilo-ms)]

Summary

Integrates opted-in in-tree CUDA kernel workspace with ORT's run-scoped activation memory patterns. Review against main; this PR includes the implementation, tests, diagnostics, and benchmark updates from chilo/static-workspace-preallocation-test.

Builds on Level-1 workspace estimation/reporting (#31962) and Level-2 reservation verification (#32189). Tracking: #29775.

• Register declared workspace as negative synthetic pattern IDs, preserving block alignment and separation from OrtValue indices.
• Trace the first compatible run, then resolve workspace from each run's activation-pattern backing buffer. Training-enabled builds also trace workspace instead of caching an activation-only static pattern.
• Support CUDA MatMulNBits CUTLASS and legacy dequantize-plus-GEMM workspace, including chunked dequantization. Cached workspace-free GEMV tactics omit unused declarations.
• Support CUDA GroupQueryAttention workspace, including windowed caches and non-windowed execution with an explicit total-sequence-length bound. Retain dynamic scratch whenever a request or suballocation does not fit.
• Add opt-in allocator diagnostics and local CUDA/ORT GenAI benchmark tooling and recorded results.

Configuration and runtime flow

Set session.enable_static_workspace_preallocation=1. Sequential execution and active memory-pattern optimization are required. For non-windowed CUDA GQA, additionally set the positive bound ep.cuda.gqa_workspace_max_total_sequence_length.

1. Session finalization records the declared usable size, padded allocation size, slot, device, and synthetic pattern ID.
2. The first compatible run traces workspace lifetimes while kernels use dynamic scratch.
3. ORT caches a combined activation/workspace memory pattern for the input shapes.
4. Later compatible runs use the current frame's backing buffer at the planned offset.
5. Missing patterns, oversized requests, and out-of-bound GQA requests continue through dynamic allocation. Workspace-free tactics remain workspace-free.

Review follow-up

• Address all three substantive review comments: training cache behavior, stale MatMulNBits usage probes, and missing oversized-request coverage.
• Keep synthetic workspace entries out of tensor-only memory-profiler metadata.
• Preserve CUDA plugin compilation, CUDA builds without fpA/intB GEMM, and Linux internal-test symbol exports.
• Exclude unavailable workspace planning and arena diagnostics from minimal builds rather than raising the Android binary-size limit.
• Preserve diagnostic C API out-parameters on failure and document registration/synchronization requirements.
• Apply the benchmark formatting fix separately from functional changes.

Regression coverage

• ExecutionFrameTest.WorkspacePatternCachesAndRejectsOversizedRequests: first-run tracing, cached reuse, a 129-byte request against a 128-byte declaration despite larger alignment padding, dynamic fallback, and subsequent reuse.
• MemPatternPlannerTest.WorkspaceSharesNonOverlappingActivationBlock: activation/workspace offset reuse.
• CUDA MatMulNBits tests: cross-level agreement, sequential-chain reuse, output parity, tactic-aware declarations, empty-output probe reset, legacy/chunked workspace, and arena savings.
• GQA tests: explicit non-windowed bounds, partial-alias preservation scratch, and checked backend workspace bounds.
• Experimental C API tests: invalid-argument output preservation and minimal-build unsupported behavior.

See the Checks tab for current platform validation. Historical benchmark measurements and their limitations are recorded in onnxruntime/test/providers/cuda/workspace_preallocation_benchmark_results.md and onnxruntime/test/providers/cuda/ort_genai_workspace_benchmark_results.md; they are not blanket performance guarantees.

Scope

• Sequential, memory-pattern-compatible execution; one workspace slot per opted-in kernel.
• In-tree CUDA MatMulNBits and GroupQueryAttention integration.
• Plugin CUDA planned-workspace consumption and parallel/multi-stream lifetime modeling remain deferred.
• Minimal builds retain dynamic scratch; the opt-in GPU-arena diagnostic hook returns not-implemented there.

[Chi Lo (chilo-ms)]

Final review follow-up is pushed through e964fed. All four review threads have been answered and resolved, and GitHub reports the branch as mergeable (remaining checks still gate merging).

The changes address training workspace tracing, stale MatMulNBits usage probes, oversized-request fallback coverage, memory-profiler compatibility, CUDA plugin/no-fpA build issues, Linux internal-test exports, diagnostic output preservation, and benchmark formatting. Minimal builds retain dynamic scratch without unavailable workspace machinery. Allocator statistics retain all 11 fields, including ReservedBytes, using a shared serialization loop to avoid repeated code.

Validation on the final source:

• Fresh Windows Release build with training and memory profiling: 17 targeted framework/arena/workspace tests passed.
• Fresh shared-library build: 12 experimental C API and allocator tests passed, including diagnostic failure-output preservation and the complete allocator-statistics output.
• Android minimal binary-size gate passed: 1,594,797 bytes against the unchanged 1,595,392-byte limit (595 bytes below the limit).
• Python format and lintrunner checks passed.

These are targeted local results, not a full-suite or CUDA E2E result. The remaining platform jobs are still running. The separate Suggest fixes job failed during Setup Rust before linting; I requested a retry of that infrastructure failure. The PR description now reflects review against main and the complete test-branch implementation, diagnostics, and benchmarks.

[Copilot]

Copilot review overview

🟡 Changes recommended

Unresolved ABI, Windows test-target, API-versioning, and GQA cache issues can break compatibility or workspace reuse.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 3 · 3

Open (6)

• Preserve CUDA internal module exports and Windows linkage · New
• Register new experimental API as SinceV31 · New
• Append ProviderHost method to preserve vtable compatibility · New
• Preserve unique CUDA test coverage on Windows · New
• Keep CUDA unit test module dependency on all platforms · New
• Cache workspace traces for out-of-bound runs · New

Resolved since last review (3)

• In ENABLE_TRAINING builds, SessionState::GetMemoryPatternGroup() eagerly creates and caches an…
• The new oversized-request fallback is not exercised by the added tests, although it is the guard…
• This flag is only updated in the CUTLASS GEMM branch. After a planned GEMM stores true, a later…

[Chi Lo (chilo-ms)]

Addressed the six new review threads in 02f4fad and 327f615:

• Correct experimental arena diagnostics to SinceV31 (including documentation and lookup regressions).
• Append the workspace ProviderHost method without shifting pre-existing vtable entries.
• Trace declared GQA workspace on an out-of-bound first run, while retaining scratch fallback whenever the runtime scalar bound is exceeded.
• Restore the complete Windows CUDA session-test coverage using a deferred test-only DLL interface, retain the module dependency on all platforms, and keep both chain regressions with distinct names.

Validation: the rebuilt shared library passed all 10 ExperimentalCApiTest cases. The affected Windows CUDA test/support objects and new GQA regression compiled. A synthetic-DLL smoke check covered deferred loading and selected forwarding paths; a COFF check covered all 26 provider-internal references. No-FPA, no-contrib, and CUDA-minimal support-file compilation checks passed. All 75 moved test declarations are retained (one disambiguated name), and all 266 generated provider-test sources were scanned for duplicate literal test registrations. Formatting and whitespace checks passed.

Limitations: these are focused checks, not a full CUDA DLL/executable link or GPU test run. The new GQA runtime regression and restored CUDA session suites still need GPU execution; fresh CI is pending. Each review thread has a specific reply and is resolved.

[Ti-Tai Wang (titaiwangms)]

I think the GQA results point to an important distinction between the admission bound and the runtime allocation policy.

Today the integration effectively uses the same value for three different quantities:

• W: the Level-1/Level-2 worst-case bound across all reachable backends and admitted shapes;
• C: the workspace capacity recorded in a cached memory pattern;
• R: the complete workspace root actually required by the backend selected for the current invocation.

GroupQueryAttention::DeclareWorkspaceRequirements() stores the conservative declaration in declared_workspace_bytes_, and ComputeInternal() requests that full value from GetPreallocatedWorkspace(). The framework then traces/looks up the plan's fixed allocation_bytes. Consequently, the memory-pattern cache is being used for placement and lifetime overlap, but not for pattern-specific GQA sizing: every participating pattern receives the same W-sized slot.

That appears consistent with the benchmark data:

• Qwen 2.5 1.5B: 123,273,472-byte root versus 56,064 bytes consumed during decode;
• Qwen 2.5 7B: 287,441,152-byte root versus 130,304 bytes consumed during decode.

The conservative maximum still looks appropriate for partition admission. The concern is using it as the allocation size for every cached pattern, especially when the unfused fallback is reachable and dominates the envelope.

A relatively bounded follow-up could separate the quantities without requiring mutable cache replacement:

1. Keep W unchanged as the immutable admission/domain ceiling.
2. Refactor GQA to select the route and build one complete, all-or-nothing workspace recipe before any workspace-dependent copies or launches.
3. Let the recording run allocate only R, while tracing a conservative route/domain capacity C, with R <= C <= W.
4. On a cache hit, accept the planned root whenever its usable capacity covers the newly computed R. Do not require route identity; the root is untyped scratch.
5. If R > C, strict mode should fail before any additional scratch allocation. A non-strict mode may use one complete dynamic R root, but should report that the already-materialized C plus R no longer preserves the W-sized peak claim.
6. Avoid the current partial behavior where some subregions come from the planned root and later regions independently fall back to GetScratchBuffer().

C cannot simply be the first observed R. The existing memory-pattern key does not include scalar contents or binding/alias facts, so identical feed shapes can see different total_sequence_length values, prompt/decode classification, past/present aliasing, or backend selection. C therefore needs to be a checked envelope for the selected route/domain over the admitted T range and relevant alias uncertainty, and every invocation still needs R <= C and R <= W validation.

This would preserve the value of the Level-1 worst case for deterministic partitioning while preventing the normal fixed-capacity decode path from materializing the global unfused envelope. It would also keep the current cache entries immutable. More general guaranteed-success behavior across route changes would require a later pre-materialization variant/key design; it does not seem necessary for the first memory-efficiency improvement.

Given the current benchmark results, I would describe GQA's present behavior as conservative workspace placement/accounting rather than an effective per-pattern preallocation optimization. It may be worth either separating W/C/R before enabling GQA by default in this pilot, or explicitly scoping the current GQA path as accounting-only until that follow-up lands.

[Copilot]

Copilot review overview

🟡 Changes recommended

Benchmark device selection can measure the wrong GPU, no-fpA legacy coverage is missing, and one analysis document contains an unedited transcript.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 2 · 1

Open (3)

• CUDA device ordinal may select the wrong GPU in nvidia-smi · New
• Legacy workspace tests excluded from non-FPA GEMM builds · New
• Remove raw Copilot transcript from checked-in analysis · New

Resolved since last review (6)

• Append ProviderHost method to preserve vtable compatibility
• Register new experimental API as SinceV31
• Preserve CUDA internal module exports and Windows linkage
• Cache workspace traces for out-of-bound runs
• Keep CUDA unit test module dependency on all platforms
• Preserve unique CUDA test coverage on Windows

[Ti-Tai Wang (titaiwangms)]

Reviewed at 832850b9508fc458ad7ae72c8012a1aca4af5213, concentrating on production planning/allocation contracts and directly related tests.

Please address the three blocking items below before approval:

1. The refactored single-node MatMulNBits builder supplies three inputs, but retained agreement assertions require six positional inputs.
2. The GQA regression verifies outputs and a recorded slot, but does not detect unconditional dynamic fallback instead of planned-workspace consumption.
3. The unconditional GQA constructor parser rejects explicit "0", contrary to the already-merged envelope option's unspecified-value contract.

Each is detailed inline. The test-fixture defect and explicit-zero rejection are established from source; I am not attributing the existing CI failures to them. The GQA item is a central-behavior coverage gap, not a claim of incorrect current numerical output.

Non-blocking follow-ups:

• Document first-Run tracing and subsequent run-scoped pattern reuse, the one-slot opt-in restriction, and non-owning workspace lifetime; this is not Initialize-time persistent allocation.
• Consolidate generic framework tracing and use the repository logger rather than duplicated kernel-specific environment switches and once-only stderr output.
• GQA currently uses the worst-case declaration W as cached capacity C; an out-of-envelope fallback can coexist with that materialized block. This is a structural allocation-policy observation, not a measured OOM or performance finding.
• Document or handle partial-training support: bypassing activation-only static caching does not cause PartialExecuteThePlan to publish a combined pattern, so opted-in partial execution does not get the same reuse as the full executor.
• Reconcile the older-base option declaration, estimator eligibility and layout with the merged option and the #33016/#33118 changes before stacking.

Concerns considered but not reported as blockers: No shipped planner-device/alignment mismatch was demonstrated. The public experimental arena API preserves output parameters on failure through local temporaries. The proposed source-extension configure failure was refuted on an inert CMake 4.1.2/Linux/Ninja glob containing .cc/.md/.py: configuration succeeded and only .cc had a compilation rule; the ORT executable wrapper does not override source language. This does not certify a full ORT configure or other generators.

Validation/exclusions: No local ORT build, CUDA test, sanitizer or benchmark was run. Older failing CI logs were unavailable, so their causes remain unattributed. Combined-tree behavior, non-FPA runtime coverage, plugin compatibility, asynchronous lifetimes and arena-destruction concurrency remain unverified rather than implied safe.

[Ti-Tai Wang (titaiwangms)]

Major: the refactored single-node fixture no longer matches its retained six-input assertions. This helper appends only A, B and scales, and BuildMatMulNBitsModelBytes no longer adds the two optional holes plus bias-or-empty. ResolveNodeInputShapes emits one entry per InputDefs(); graph resolution normalizes argument counts, not the input-definition vector. The assertions at lines 832 and 953 therefore see three entries instead of six, preventing the agreement test from reaching its planned-workspace checks. Please restore [A, B, scales, "", "", bias-or-empty] and the conditional bias initializer in the single-node builder; the chain helper can remain three-input. This is a source-established mismatch in eligible FPA-enabled configurations, not an observed CI-root-cause claim.

[Ti-Tai Wang (titaiwangms)]

Major coverage gap: this does not prove GQA consumes the planned block. The kernel traces the declared slot before deciding whether to discard its pointer and use dynamic scratch. If it unconditionally discarded the pointer, these block-size and numerical-output assertions would still pass. Please add test-only allocation-source observability for the existing {128,64,128,64} sequence: recording/out-of-envelope runs should use dynamic fallback, while cached in-envelope runs must consume planned regions. Retain the numerical oracle. MatMulNBits already has analogous planned-use checks.

[Ti-Tai Wang (titaiwangms)]

Major: explicit zero is valid under the merged option contract but fails here. #33014 defines ep.cuda.gqa_workspace_max_total_sequence_length="0" as unspecified, equivalent to unset. This parser accepts unset but requires parsed > 0 for an explicit value. Because the GQA constructor calls it unconditionally, a valid zero setting now fails Initialize even without workspace-preallocation opt-in. Please reuse the validated nonnegative configuration, or accept parsed >= 0 with zero meaning unspecified, and add an initialization regression covering explicit zero. Also reconcile the older-base duplicate key/documentation on rebase.

[Chi Lo (chilo-ms)]

Closing this broad integration PR for now to avoid coupling review of the staged GQA sizing and estimation work to an outdated combined tree.

The unique activation-memory-pattern integration is still valuable and the branch is retained as source material. It will be rebased or extracted into focused follow-up work after #33118, #33016, and the Level-2 declaration plumbing settle. The three current blocking findings are acknowledged and will be addressed in that reconstructed follow-up rather than patched onto this conflicted 64-file/87-commit branch.