Add packed attention workspace estimates

[Ti-Tai Wang (titaiwangms)]

Summary

Add MatMulNBits-equivalent operator-side workspace estimation for CUDA
PackedAttention and PackedMultiHeadAttention:

• Level 1 derives a conservative estimate from the node, resolved input shapes,
  CUDA device properties, and the EP's resolved attention options.
• Level 2 declares the same estimate from positional WorkspaceInputShape
  metadata and constructed kernel state.
• Existing graph-free runtime recipes remain the single source of truth for
  workspace bytes and layouts.
• PackedAttention declares one 256-byte-aligned root in slot 0, with internal
  projection and attention regions.
• PackedMultiHeadAttention declares one 256-byte-aligned attention root in slot 0.

Level 1 is log-only, matching the current MatMulNBits pilot. This PR does not
add #32071-specific planner APIs or change runtime GetScratchBuffer() behavior.

Route aggregation

Runtime routes are mutually exclusive, so the estimate uses:

PackedAttention:
  align_up(projection_bytes, 256) + max(feasible TRT, MEA, unfused recipes)

PackedMultiHeadAttention:
  max(feasible Flash, TRT, MEA, unfused recipes)

Route reachability is evaluated conservatively for every runtime shape up to
the supplied maximum geometry. This is necessary because Flash/MEA thresholds
and attention-bias alignment gates are not monotonic when moving from a maximum
shape to a smaller runtime shape. Unfused fallback is always retained, and a
failure to size any included route makes the estimate unavailable rather than
silently underestimating.

Shape and zero semantics

• Missing mandatory inputs, shapeless required inputs, unknown dimensions,
  malformed geometry, and checked-arithmetic overflow produce no estimate.
• WorkspaceInputShape does not carry max-shape provenance, so zero-shaped
  framework hints are conservatively treated as unavailable.
• Exact zero behavior remains supported by the graph-free runtime recipes.
• At the current Level-2 boundary, both unavailable and zero are represented by
  an empty requirements list.

Planner integration

• Both operators fit Integrate CUDA workspace with activation memory patterns #32071's current one-slot pilot. Generic framework
  multi-slot support remains unchanged.
• A declaration alone is not planner opt-in. SupportsPreallocatedWorkspace(),
  slot-0 retrieval, and PA root slicing must land atomically in the planner
  integration.
• Until then, PA retains its two dynamic allocations and PMHA retains its one
  dynamic allocation.

Build boundaries

The framework adapters and kernel overrides are excluded from:

• CUDA minimal builds
• DISABLE_CONTRIB_OPS builds
• CUDA plugin EP builds

The graph-free workspace recipes remain available to the shared BERT attention
infrastructure where required.

Validation

• 18/18 PA/PMHA workspace estimator tests
  • includes direct production-kernel Level-2 declaration tests
  • route-threshold, max-not-sum, aligned-root padding/no-padding,
    optional-hole, zero, overflow, and malformed geometry coverage
• 23/23 existing packed-attention workspace recipe tests
• 20/20 existing hand-calculated runtime parity cases
• 26/26 PackedAttention/PackedMultiHeadAttention runtime operator tests
• CUDA provider-test build
• 145 CUDA internal tests executed: 143 passed, 2 unrelated LeanAttention skips
• DISABLE_CONTRIB_OPS and CUDA-minimal compile-guard probes
• C++ formatting and diff checks

Dependency

This is a stacked follow-up to #32312. The base should change to main after
#32312 merges.

Tracking: #29775

[Xavier Dupré (xadupre)]

Is it possible to add a link to the main document designing the feature, it is difficult to understand how this bricks is connected with the others.

[Copilot]

Pull request overview

Adds conservative CUDA workspace estimates and Level-2 declarations for packed attention operators. No actionable issues found.

Changes:

• Aggregates mutually exclusive backend workspace recipes safely.
• Adds graph/kernel adapters and stable workspace slots.
• Adds extensive estimator, boundary, and integration tests.

Reviewed changes

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

Show a summary per file

File	Description
cuda_execution_provider.cc	Logs Level-1 estimates.
packed_attention.h	Declares Level-2 override.
packed_attention.cc	Implements PA declaration.
packed_multihead_attention.h	Declares Level-2 override.
packed_multihead_attention.cc	Implements PMHA declaration.
packed_attention_workspace.h	Defines masks and aggregates.
packed_attention_workspace.cc	Aggregates route recipes.
packed_attention_workspace_estimate.h	Defines estimator API.
packed_attention_workspace_estimate.cc	Implements shape and route estimation.
packed_attention_workspace_estimate_test.cc	Adds estimator coverage.
cuda_kernel_workspace_inventory.md	Documents workspace slots.
attention_workspace_estimation.md	Documents framework integration.

---

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

[Ti-Tai Wang (titaiwangms)]

Thanks Xavier Dupré (@xadupre). The Related design and rollout section at the top of the Attention workspace roadmap now links #29775, the generic framework design, #32312, #32071, and this PA/PMHA adapter PR. The latest update also documents the one-aligned-root contract and explicit planner opt-in so the dependency chain and runtime transition are visible from the document.

[Ti-Tai Wang (titaiwangms)]

Follow-up after the planner discussion with Chi Lo (@chilo-ms): PA no longer declares two framework slots. The two slots mirrored its current projection and Attention GetScratchBuffer() allocations, but they are not an inherent planner requirement. PA now declares one 256-byte-aligned slot-0 root with an aligned Attention subregion; PMHA declares one aligned Attention root. Generic framework multi-slot support remains unchanged. Neither kernel opts into planning yet: SupportsPreallocatedWorkspace(), root retrieval, and PA slicing must land atomically with the #32071 integration, while the current dynamic allocations remain unchanged.

[Copilot]

🟡 Changes recommended

Dynamic head-size bounds can omit a reachable MEA route and produce an undersized workspace declaration.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (1)

onnxruntime/contrib_ops/cuda/bert/packed_attention_workspace_estimate.cc:407

• PMHA has the same max-head-only reachability gap. A packed FP16 maximum shape [1, 1, 3, 1032] excludes MEA here and declares only the 6,704-byte unfused workspace, while the valid smaller runtime shape [1, 1, 3, 1024] can select MEA and requires 10,240 bytes. Enumerate backend support boundaries for variable head dimensions, not just sequence-length thresholds, so every runtime geometry up to the hint is covered.

      has_memory_efficient_attention(
          sm, problem.element_size == 2, false,
          problem.qk_head_size, problem.v_head_size);

• Files reviewed: 12/12 changed files
• Comments generated: 1
• Review effort level: Balanced