Skip to content

Add packed attention workspace estimates - #32321

Draft
Ti-Tai Wang (titaiwangms) wants to merge 1 commit into
titaiwangms/workspace-optional-input-shapesfrom
titaiwangms/packed-attention-workspace-l1-l2
Draft

Add packed attention workspace estimates#32321
Ti-Tai Wang (titaiwangms) wants to merge 1 commit into
titaiwangms/workspace-optional-input-shapesfrom
titaiwangms/packed-attention-workspace-l1-l2

Conversation

@titaiwangms

Copy link
Copy Markdown
Contributor

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 two simultaneously live slots:
    • slot 0: projection workspace
    • slot 1: attention workspace
  • PackedMultiHeadAttention declares one attention workspace 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:
  projection_bytes + 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

  • PMHA's one-slot declaration can use Integrate CUDA workspace with activation memory patterns #32071's current pilot after its explicit
    preallocation opt-in is added.
  • PA preserves its two-slot runtime allocation topology and therefore requires
    chilo-ms's multi-slot planner extension before planned workspace can be used.
  • Until then, both operators retain their existing dynamic allocation fallback.

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

  • 17/17 PA/PMHA workspace estimator tests
    • includes direct production-kernel Level-2 declaration tests
    • route-threshold, max-not-sum, 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
  • 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

Add route-aware Level-1 estimates and Level-2 workspace declarations for PackedAttention and PackedMultiHeadAttention while preserving the shared runtime recipes and dynamic allocation fallback.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: c04148cc-7ace-4cf4-b981-0ddc92334e78
注册 for free to join this conversation on GitHub. Already have an account? 登录 to comment

标签

None yet

项目

None yet

Development

Successfully merging this pull request may close these issues.

2 participants