Skip to content

Integrate CUDA workspace with activation memory patterns - #32071

Draft
Chi Lo (chilo-ms) wants to merge 4 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 4 commits into
chilo/level2-workspace-verificationfrom
chilo/static-workspace-preallocation

Conversation

@chilo-ms

Copy link
Copy Markdown
Contributor

Summary

This is a stacked draft PR on top of chilo/level2-workspace-verification.

  • 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 the request exceeds the declaration or no usable pattern exists;
  • remove the earlier session-persistent buffer and its workspace-specific concurrent-Run/synchronization restrictions;
  • preserve memory-pattern block alignment and use negative synthetic IDs disjoint from OrtValue indices.

Runtime flow

  1. Session finalization stores the declared usable size, aligned allocation size, slot, device, and synthetic pattern ID.
  2. On the first 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 or unresolved requests continue through the existing dynamic allocation path.

Tests

  • Added MemPatternPlannerTest.WorkspaceSharesNonOverlappingActivationBlock.
  • Extended the CUDA E2E test to verify first-run dynamic fallback and second-run planned workspace use.
  • Added a three-node sequential MatMulNBits chain that verifies:
    • one distinct synthetic entry per node;
    • all three non-overlapping workspaces reuse the same cached offset;
    • all kernels use dynamic workspace during tracing and planned workspace on the cached run;
    • planned-workspace output exactly matches the nonzero dynamic-reference output.

Local validation:

onnxruntime_test_all build: passed
MemPatternPlannerTest.*: 2 tests passed

The available Windows build is CPU-only (onnxruntime_USE_CUDA=OFF), so the CUDA E2E tests require CUDA CI or a CUDA-enabled developer build.

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
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
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