onnxruntime
f98e809f - Make workspace input shapes optional-aware (#32312)

Commit
50 days ago
Make workspace input shapes optional-aware (#32312) ## 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 ## #32071 integration constraints #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 #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 #32071 This PR should merge first. #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: #29775 --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Copilot-Session: c04148cc-7ace-4cf4-b981-0ddc92334e78
Author
Parents
Loading