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