[Tianlei Wu (tianleiwu)]

Description

Stacked on #33224 (base kvaishnavi/index-share; review only the top commit). Adds an optional state_status output to PackedSparseAttentionIndexer that reports, per request, whether this call's state update was applied or rejected, and why.

Today a rejected step (inconsistent cursor or state_capacity overflow) is a safe no-op on state with an empty selection, but the caller cannot tell it happened: the rejection flags live in a scratch buffer and indexshare_status is only produced on the IndexShare interface. A CUDA-graph-replayed decode/verify loop cannot synchronize inside the operator, so it needs a device tensor it can read on the host after replay.

Summary of Changes

Schema and docs

File	Change
bert_defs.cc	New optional output 8 state_status (int32, shape (batch_size)); declared output range is now 6-9; type/shape inference; one doc sentence
packed_sparse_attention_indexer_common.h	kStateStatus = 8, kOutputCount = 9, StateStatusCode {0 ok, 1 invalid device metadata, 2 capacity overflow}
docs/ContribOperators.md, docs/contrib_ops/packed_sparse_attention_indexer.md	Output and status contract

CUDA / WebGPU

File	Change
packed_sparse_attention_indexer_impl.cu/.h	The QSA and CSA update kernels write the code (1 or 2) instead of the boolean 1 into the existing overflow_flags; all readers still test != 0, so selection behaviour is unchanged
packed_sparse_attention_indexer.cc	When output 8 is requested it is used as the overflow_flags buffer (no extra kernel, no copy, no host sync). On the IndexShare reuse path that holds state (QSA, CSA hold) it is zero-filled with cudaMemsetAsync
webgpu/.../packed_sparse_attention_indexer.cc	Fails if state_status is requested (CUDA only, like IndexShare)

Semantics

• A non-zero request leaves its state unchanged and emits an empty selection (selected_counts == 0, indices -1), exactly as before.
• Codes follow PackedSparseAttentionIndexerMerge (1 invalid device metadata, 2 capacity overflow).
• state_status is independent of indexshare_status (per-query selection status, forwarded on reuse and zero on refresh). On CSA reuse with advance, both are set when a row is rejected.
• Default behaviour is unchanged: the output is optional and absent in existing models.

Testing

Built with CUDA 13.0, sm_90, in-tree CUDA EP, and run on one H200:

CUDA_VISIBLE_DEVICES=5 ./onnxruntime_provider_test \
  --gtest_filter='*PackedSparseAttentionIndexer*:*DynamicSparseAttention*:*SparseAttentionIndexer*:*SparsePagedAttention*'

177 passed (17 WebGPU tests skipped, no WebGPU in this build); PackedSparseAttentionIndexerTest.* repeated 20 times, no flake.

New tests in packed_sparse_attention_indexer_op_test.cc (Qsa/CsaStateStatus*, IndexShareCsaAdvanceReportsStateStatus, QsaTwoRowStepEqualsChainedSingleRowSteps):

• status values for a healthy batch, capacity overflow, inconsistent cursor, malformed state lengths, and mixed batches {ok, invalid, overflow}, for QSA and CSA, checked against the host reference (which now also computes the codes);
• requesting state_status leaves selections and state bit-identical to not requesting it;
• IndexShare CSA reuse with state overflow reports indexshare_status = 1 and state_status = 2;
• a two-row QSA step (one request, two tokens, compress_ratio 2 and 4, all three cursor phases) is bit-identical in selections, block state and pending entries to two chained one-row steps. This passes without any kernel change, i.e. the multi-row verify step is already supported by main; the test pins that contract.

Mutation checks: with the output plumbing removed 7 new tests fail; with code 2 reported as 1, 4 fail; existing tests pass in both.

Not covered: WebGPU compile (not in this build), plugin-EP build, an end-to-end CUDA-graph capture/replay (the kernels are launch-only and the reuse-path zero fill is cudaMemsetAsync, so nothing in the path synchronizes).

Motivation and Context

Needed by decode/verify loops that capture the indexer in a CUDA graph and drive it with device-resident lengths (as in the multi-token verify path of #33328). Reuses the status convention introduced in #33224 so callers see one scheme. Touches different hunks from #33328 (scoring kernels); a trivial merge is expected for the test file.

Checklist

• Tests added/updated
• No breaking changes (new optional output; the declared output range grows from 6-8 to 6-9)
• Documentation updated

[Copilot]

🟡 Changes recommended

Empty-batch CSA reuse can exceed the status output’s extent, and the rejection contract incorrectly promises unchanged state lengths.

2 open issues

• Skip overflow_flags memset for empty CSA advance batches · New
• Clarify rejected updates clamp invalid state lengths · New

What changed in this PR

Adds optional per-request state_status reporting to PackedSparseAttentionIndexer, letting CUDA graph callers detect rejected state updates after replay.

Changes:

• Extends the schema and documentation with status codes and output shape.
• Reuses CUDA rejection flags as the output; rejects unsupported WebGPU requests.
• Adds status-reporting and multi-token QSA regression tests.

File	Description
onnxruntime/​test/​contrib_ops/​packed_sparse_attention_indexer_op_test.cc	Adds status and chained-step tests.
onnxruntime/​core/​graph/​contrib_ops/​bert_defs.cc	Defines and infers the optional output.
onnxruntime/​contrib_ops/​webgpu/​bert/​packed_sparse_attention_indexer.cc	Rejects requests for the CUDA-only output.
onnxruntime/​contrib_ops/​cuda/​sparse/​packed_sparse_attention_indexer.cc	Connects status output to rejection flags.
onnxruntime/​contrib_ops/​cuda/​sparse/​packed_sparse_attention_indexer_impl.h	Documents the status buffer contract.
onnxruntime/​contrib_ops/​cuda/​sparse/​packed_sparse_attention_indexer_impl.cu	Writes distinct rejection codes.
onnxruntime/​contrib_ops/​cpu/​sparse/​packed_sparse_attention_indexer_common.h	Adds output index and status constants.
docs/​OperatorKernels.md	Lists the new output.
docs/​ContribOperators.md	Updates the generated operator contract.
docs/​contrib_ops/​packed_sparse_attention_indexer.md	Explains status semantics and provider support.

_{Review effort: Balanced}

[Akshay Sonawane (apsonawane)]

lgtm

[Copilot]

🟢 Approval recommended

The optional-output changes have focused regression coverage, and both previously reported issues are resolved.

No open issues 🎉

2 resolved since last review

• Clarify rejected updates clamp invalid state lengths
• Skip overflow_flags memset for empty CSA advance batches

_{Review effort: Balanced}