Skip to content

Integrate CUDA workspace with activation memory patterns - #32071

Draft
Chi Lo (chilo-ms) wants to merge 6 commits into
chilo/level2-workspace-verificationfrom
chilo/static-workspace-preallocation
Draft

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

Conversation

@chilo-ms

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

Copy link
Copy Markdown
Contributor

PR stack

  1. Add workspace estimation accounting and reporting #31962 - Level-1 workspace estimation and reporting
  2. Verify Level-2 workspace reservations #32189 - Level-2 reservation verification
  3. 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

  1. Session finalization stores each declared usable size, aligned allocation size, slot, device, and synthetic pattern ID.
  2. On the first compatible run, the execution frame traces the synthetic workspace lifetime; the kernel still allocates dynamically.
  3. ORT generates and caches one memory pattern containing activation and workspace blocks.
  4. On later compatible runs, the kernel receives backing_buffer + workspace_offset from the execution frame.
  5. 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

@chilo-ms
Chi Lo (chilo-ms) force-pushed the chilo/level2-workspace-verification branch from d05286b to 1948ce2 Compare August 20, 2026 17:47
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
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant