Make workspace input shapes optional-aware

[Ti-Tai Wang (titaiwangms)]

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

[Copilot]

Pull request overview

Introduces optional-aware, positional workspace shape metadata for Level-2 estimation and updates CUDA MatMulNBits workspace handling.

Changes:

• Adds three-state WorkspaceInputShape resolution while preserving optional-input holes.
• Handles zero-sized MatMulNBits outputs consistently across estimation and runtime.
• Expands framework/CUDA tests and updates design documentation.

No actionable issues identified.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
include/onnxruntime/core/framework/workspace_input_shape.h	Defines presence-aware shape metadata.
include/onnxruntime/core/framework/op_kernel.h	Updates the Level-2 virtual contract.
include/onnxruntime/ep/adapter/op_kernel.h	Mirrors the contract for plugin builds.
onnxruntime/core/framework/node_shape_resolver.h	Preserves positional input metadata.
onnxruntime/core/framework/session_state.cc	Passes resolved entries to kernels.
onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h	Updates MatMulNBits declarations.
onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc	Handles partial and empty shapes.
onnxruntime/contrib_ops/cuda/quantization/matmul_nbits_workspace_estimate.h	Exposes shared dimension-product logic.
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemm/fpA_intB_gemm.h	Returns zero workspace for empty outputs.
onnxruntime/test/framework/max_shape_override_test.cc	Tests all shape states and optional holes.
onnxruntime/test/framework/session_state_test.cc	Tests the default declaration behavior.
onnxruntime/test/providers/cuda/test_cases/matmul_nbits_workspace_test.cc	Tests formula edge cases.
onnxruntime/test/providers/cuda/test_cases/matmul_nbits_e2e_workspace_test.cc	Tests positional and runtime parity.
docs/annotated_partitioning/future_directions_constrained_env.md	Documents the revised contract and scope.

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Ti-Tai Wang (titaiwangms)]

Chi Lo (@chilo-ms) Can you take a look and check if the PR understands #32071 correctly? We will need multi-slots support as Attention is much more complicated.

[Copilot]

🔵 Needs a closer look

The broad framework API and CUDA workspace changes require final human review.

Review details

• Files reviewed: 15/15 changed files
• Comments generated: 0 new
• Review effort level: Balanced