Optimize QSA sparse attention indexing for long-context decode by Copilot · Pull Request #32656 · microsoft/onnxruntime

Copilot · 2026-09-16T20:20:54Z

Description

Extend SparseAttentionIndexer to support query RMS normalization, optional packed QK input, INT64 padding masks with internal causal filtering, and shared rank-2 or batched rank-3 rotary caches across CUDA and WebGPU. Regenerate the operator documentation for the updated contract.
Optimize dense CUDA QSA decode by preparing normalized/rotated queries once per query head, avoiding visible-index materialization for causal-prefix masks, skipping redundant past-key copies when state aliases in place, and replacing full-context sorting/repeated scans with deterministic bounded TopK selection for the common block_topk <= 512 case.
Optimize packed CUDA QSA by hoisting query preparation out of block scoring and using deterministic single-read block TopK selection for up to 32 selected blocks, with the correctness fallback retained for larger values.
Preserve score-descending/index-ascending tie ordering, exact handling of padding-mask holes, cache state semantics, and CUDA graph compatibility.
Add CUDA and WebGPU schema/numerical coverage, shared-rotary and packed-QK tests, long-context parity at 8K/32K/64K/128K/256K for compress_ratio=4 and token_budget=2048, CUDA graph kernel-topology and latency regression checks, and a Python decode benchmark sweep.

Motivation and Context

A 65K-context CUDA profile showed QSA index selection dominating decode time, with additional context-scaled work in visibility compaction, query preparation, cache copying, and temporary-buffer traffic. This change removes repeated TopK scans and full-row sorting, bounds selection storage and writes by the configured token budget, avoids redundant work on the common causal-prefix/in-place decode path, and keeps exact QSA output ordering and fallback behavior.

Exact query-specific scoring and padding-mask inspection still scale with the visible context because each query can change, but the selection algorithm and its intermediate traffic no longer scale with the entire candidate row. The schema updates also match exporter/model layouts without requiring expanded per-batch rotary tables or materialized per-query causal masks.

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

azure-pipelines · 2026-09-16T20:21:03Z

Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run.

Copilot

🟡 Changes recommended

The dense scoring kernel can exceed CUDA shared-memory limits, and TopK boundary coverage is incomplete.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Optimizes CUDA QSA scoring and deterministic TopK selection while reducing workspace usage.

Changes:

Fuses query rotation into dense and packed scoring kernels.
Adds single-pass TopK for up to 32 blocks.
Updates CUDA tests and performance documentation.

File summaries

File	Description
`sparse_attention_indexer_op_test.cc`	Adds focused TopK coverage.
`sparse_attention_indexer_impl.cu`	Optimizes dense CUDA QSA execution.
`sparse_attention_indexer_device_math.cuh`	Adds deterministic block TopK.
`packed_sparse_attention_indexer_impl.cu`	Applies optimizations to packed QSA.
`packed_sparse_attention_indexer.md`	Documents packed CUDA performance.
`cuda/sparse_attention_indexer.md`	Updates optimization limitations.

Review details

Files reviewed: 6/6 changed files
Comments generated: 2
Review effort level: Balanced

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

…sparse-attention-indexer' into copilot/copilotcopilotimplement-packed-sparse-attention-in Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

…sparse-attention-indexer' into copilot/copilotcopilotimplement-packed-sparse-attention-in Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

github-actions

You can commit the suggested changes from lintrunner.

github-actions · 2026-09-17T16:52:27Z

+                                                     const SparseAttentionIndexerParams&,                      \
+                                                     const T*, const T*, const T*, const T*, const T*,         \
+                                                     const bool*, const T*, int32_t*, T*, float*, int32_t*);   \
  template Status LaunchCsaSparseAttentionIndexer<T>(                                                        \
      cudaStream_t, const SparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, const T*,   \


Suggested change

const SparseAttentionIndexerParams&, \

const T*, const T*, const T*, const T*, const T*, \

const bool*, const T*, int32_t*, T*, float*, int32_t*); \

template Status LaunchCsaSparseAttentionIndexer<T>( \

cudaStream_t, const SparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, const T*, \

const SparseAttentionIndexerParams&, \

const T*, const T*, const T*, const T*, const T*, \

const bool*, const T*, int32_t*, T*, float*, int32_t*); \

template Status LaunchCsaSparseAttentionIndexer<T>( \

cudaStream_t, const SparseAttentionIndexerParams&, const T*, const T*, const T*, const T*, const T*, \

github-actions · 2026-09-17T16:52:27Z

-  template Status LaunchQsaSparseAttentionIndexer<T>(cudaStream_t, const SparseAttentionIndexerParams&,      \
-                                                     const T*, const T*, const T*, const T*, const T*,       \
-                                                     const bool*, const T*, int32_t*, T*, float*, int32_t*); \
+#define INSTANTIATE_SPARSE_ATTENTION_INDEXER(T)                                                                \


Suggested change

#define INSTANTIATE_SPARSE_ATTENTION_INDEXER(T) \

#define INSTANTIATE_SPARSE_ATTENTION_INDEXER(T) \

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

@@ -0,0 +1,114 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.


+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import argparse


+import time
+
+import numpy as np
+import onnx


+import time
+
+import numpy as np
+import onnx


+import time
+
+import numpy as np
+import onnx


github-actions

You can commit the suggested changes from lintrunner.

github-actions · 2026-09-17T23:41:09Z

+import onnx
+import onnxruntime as ort


Suggested change

import onnx

import onnxruntime as ort

import onnxruntime as ort

# Conflicts: # docs/contrib_ops/cuda/sparse_attention_indexer.md # onnxruntime/contrib_ops/cuda/sparse/packed_sparse_attention_indexer_impl.cu # onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu # onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.cc # onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc

# Conflicts: # docs/ContribOperators.md # docs/contrib_ops/cuda/sparse_attention_indexer.md # docs/contrib_ops/webgpu/sparse_attention_indexer.md # onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer.cc # onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu # onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.h # onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.cc # onnxruntime/core/graph/contrib_ops/bert_defs.cc # onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc

# Conflicts: # onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu

# Conflicts: # onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.h # onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc

github-actions

You can commit the suggested changes from lintrunner.

github-actions · 2026-09-20T18:13:54Z

@@ -54,6 +54,7 @@ class SparseAttentionIndexerQsaSelectProgram final
    {"query_row_stride", ProgramUniformVariableDataType::Uint32},


Suggested change

{"query_row_stride", ProgramUniformVariableDataType::Uint32},

{"query_row_stride", ProgramUniformVariableDataType::Uint32},

Copilot AI and others added 2 commits September 16, 2026 20:15

Optimize CUDA QSA scoring and selection

ebc6bdf

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Add single-pass TopK coverage

ecabf39

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Copilot AI assigned Copilot and kunal-vaishnavi Sep 16, 2026

Copilot created this pull request from a session on behalf of kunal-vaishnavi September 16, 2026 20:20 View session

kunal-vaishnavi changed the base branch from main to copilot/copilotimplement-packed-sparse-attention-indexer September 16, 2026 20:21

kunal-vaishnavi added this pull request to stack #32530 September 16, 2026 20:21

kunal-vaishnavi marked this pull request as ready for review September 16, 2026 20:21

Copilot AI balanced review requested due to automatic review settings September 16, 2026 20:21

Copilot started reviewing on behalf of kunal-vaishnavi September 16, 2026 20:22 View session

Copilot AI reviewed Sep 16, 2026

View reviewed changes

Comment thread onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu Outdated

Comment thread onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc Outdated

Copilot started work on behalf of kunal-vaishnavi September 16, 2026 21:31 View session

Copilot AI and others added 2 commits September 16, 2026 21:33

Merge remote-tracking branch 'origin/copilot/copilotimplement-packed-…

49ade0c

…sparse-attention-indexer' into copilot/copilotcopilotimplement-packed-sparse-attention-in Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Address CUDA QSA review feedback

8f506aa

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Copilot finished work on behalf of kunal-vaishnavi September 16, 2026 21:35

Copilot AI requested a review from kunal-vaishnavi September 16, 2026 21:35

Copilot started work on behalf of kunal-vaishnavi September 17, 2026 01:36 View session

Merge remote-tracking branch 'origin/copilot/copilotimplement-packed-…

51b6403

…sparse-attention-indexer' into copilot/copilotcopilotimplement-packed-sparse-attention-in Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Copilot finished work on behalf of kunal-vaishnavi September 17, 2026 01:39

Copilot started work on behalf of kunal-vaishnavi September 17, 2026 01:56 View session

Copilot finished work on behalf of kunal-vaishnavi September 17, 2026 02:38

Copilot stopped work on behalf of kunal-vaishnavi due to an error September 17, 2026 02:38
The session was cancelled by the user.

Copilot started work on behalf of kunal-vaishnavi September 17, 2026 02:39 View session

Copilot finished work on behalf of kunal-vaishnavi September 17, 2026 02:42

kunal-vaishnavi marked this pull request as draft September 17, 2026 09:57

Copilot started work on behalf of kunal-vaishnavi September 17, 2026 10:00 View session

Copilot AI and others added 3 commits September 17, 2026 10:08

Implement profile-driven QSA selection optimizations

24bccfd

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Fix and format QSA radix selection

1593b64

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Include CUDA kernel scratch API

5b77441

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Copilot finished work on behalf of kunal-vaishnavi September 17, 2026 10:16

Fix sparse attention indexer CUDA kernel type

1f12360

github-actions Bot reviewed Sep 17, 2026

View reviewed changes

Copilot started work on behalf of kunal-vaishnavi September 17, 2026 21:24 View session

Copilot AI and others added 2 commits September 17, 2026 21:33

Implement bounded exact QSA selection

d2de258

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Finalize bounded QSA decode path

a6c7873

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

Copilot finished work on behalf of kunal-vaishnavi September 17, 2026 21:48

github-advanced-security Bot found potential problems Sep 17, 2026

View reviewed changes

Comment thread onnxruntime/python/tools/microbench/sparse_attention_indexer.py

import time

import numpy as np

import onnx

Comment thread onnxruntime/python/tools/microbench/sparse_attention_indexer.py

import time

import numpy as np

import onnx

github-actions Bot reviewed Sep 17, 2026

View reviewed changes

kunal-vaishnavi added 10 commits September 18, 2026 10:07

Support shared rotary caches in sparse indexer

15d73ba

Accept padding masks in sparse attention indexer

3d6bf1a

Merge branch 'stack-32618' into stack-32656

61e9e7a

Merge branch 'stack-32618' into stack-32656

240f37a

# Conflicts: # onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.cu

Hoist packed QSA query preparation out of scoring

04c7a92

Guard QSA decode performance in CUDA tests

cdbf1fe

Enforce QSA decode latency budgets

7b9d8bc

Merge branch 'stack-32618' into stack-32656

31dd38b

# Conflicts: # onnxruntime/contrib_ops/cuda/sparse/sparse_attention_indexer_impl.h # onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc

github-actions Bot reviewed Sep 20, 2026

View reviewed changes

Copilot started work on behalf of kunal-vaishnavi September 20, 2026 18:14 View session

kunal-vaishnavi changed the title ~~Optimize CUDA QSA sparse attention indexing~~ Optimize QSA sparse attention indexing for long-context decode Sep 20, 2026

Copilot finished work on behalf of kunal-vaishnavi September 20, 2026 18:17

kunal-vaishnavi added 3 commits September 21, 2026 00:04

Optimize prefix-causal sparse attention indexing

59253d6

Optimize QSA TopK and block scoring

548ff66

Test multi-row QSA partial TopK

5a01fad

Navigation Menu

Optimize QSA sparse attention indexing for long-context decode - #32656

Optimize QSA sparse attention indexing for long-context decode#32656
kunal-vaishnavi with Copilot wants to merge 24 commits into
copilot/copilotimplement-packed-sparse-attention-indexerfrom
copilot/copilotcopilotimplement-packed-sparse-attention-in

Copilot AI commented Sep 16, 2026 •

edited by kunal-vaishnavi

Loading

Uh oh!

azure-pipelines Bot commented Sep 16, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

github-actions Bot left a comment

Uh oh!

github-actions Bot Sep 17, 2026

Uh oh!

github-actions Bot Sep 17, 2026

Uh oh!

github-actions Bot left a comment

Uh oh!

github-actions Bot Sep 17, 2026

Uh oh!

github-actions Bot left a comment

Uh oh!

github-actions Bot Sep 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	#define INSTANTIATE_SPARSE_ATTENTION_INDEXER(T) \
	#define INSTANTIATE_SPARSE_ATTENTION_INDEXER(T) \

		@@ -0,0 +1,114 @@
		# Copyright (c) Microsoft Corporation. All rights reserved.

		@@ -54,6 +54,7 @@ class SparseAttentionIndexerQsaSelectProgram final
		{"query_row_stride", ProgramUniformVariableDataType::Uint32},

Conversation

Copilot AI commented Sep 16, 2026 • edited by kunal-vaishnavi Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

Uh oh!

azure-pipelines Bot commented Sep 16, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Sep 20, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Copilot AI commented Sep 16, 2026 •

edited by kunal-vaishnavi

Loading