Skip to content

[CUDA] Optimize DynamicSparseAttention for large-context inference - #32671

Merged
kunal-vaishnavi merged 12 commits into
mainfrom
copilot/optimize-dynamic-sparse-attention
Oct 3, 2026
Merged

kunal-vaishnavi merged 12 commits into
mainfrom
copilot/optimize-dynamic-sparse-attention

Conversation

Copilot AI commented Sep 17, 2026

Copy link
Copy Markdown
Contributor

Description

  • Parallelize selected-token attention across fixed 256-candidate tiles.
  • Use online softmax and stable cross-tile reduction.
  • Bound split count to control workspace growth for large batches and 256K contexts.
  • Bypass intermediate workspace and merge kernels when one split suffices.
  • Replace quadratic duplicate validation with a GPU bitmap.
  • Replace cache initialization kernels with asynchronous copies and clears.
  • Preserve CUDA graph capture by avoiding synchronous validation during capture.
  • Add FP16, multi-batch, prefill, ragged-selection, sink, and split-boundary coverage.

Motivation and Context

CUDA DynamicSparseAttention serialized work across selected tokens, causing decode latency to scale poorly with retained context. The optimized paths increase GPU occupancy, eliminate per-candidate block synchronization and large-context fallback cliffs, and bound temporary memory for production-scale workloads.

@azure-pipelines

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

@kunal-vaishnavi
kunal-vaishnavi changed the base branch from main to copilot/copilotadd-cuda-dynamicsparseattention September 17, 2026 09:48
@kunal-vaishnavi
kunal-vaishnavi force-pushed the copilot/optimize-dynamic-sparse-attention branch from f9bc5f9 to e9c1dd6 Compare September 17, 2026 09:57

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟑 Changes recommended

The newly introduced CUDA graph capture path lacks regression coverage.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Optimizes CUDA DynamicSparseAttention for large-context inference.

Changes:

  • Adds tiled split attention with stable online softmax reduction.
  • Introduces bitmap validation and asynchronous cache initialization.
  • Expands CUDA coverage for large, ragged, FP16, and prefill cases.
File summaries
File Description
dynamic_sparse_attention.cc Allocates workspaces and supports capture-time execution.
dynamic_sparse_attention_impl.h Extends CUDA implementation interfaces.
dynamic_sparse_attention_impl.cu Implements tiled kernels, reductions, validation, and cache copies.
dynamic_sparse_attention_op_test.cc Adds large-context and split-path tests.
Review details
  • Files reviewed: 4/4 changed files
  • Comments generated: 1
  • Review effort level: Balanced

πŸ’‘ Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

@kunal-vaishnavi
kunal-vaishnavi marked this pull request as ready for review October 1, 2026 06:07
@kunal-vaishnavi kunal-vaishnavi changed the title Optimize CUDA DynamicSparseAttention for large-context inference [CUDA] Optimize DynamicSparseAttention for large-context inference Oct 1, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🟑 Changes recommended

Both new attention paths incorrectly treat the valid finite logit -FLT_MAX as an invalid candidate.

Review effort: Balanced
Findings: 1 High severity

Open (1)
Resolved since last review (1)

Base automatically changed from copilot/copilotadd-cuda-dynamicsparseattention to main October 3, 2026 02:11
Copilot AI and others added 12 commits October 2, 2026 19:11
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
@kunal-vaishnavi
kunal-vaishnavi force-pushed the copilot/optimize-dynamic-sparse-attention branch from 2888589 to fdd0ad3 Compare October 3, 2026 02:11
@kunal-vaishnavi
kunal-vaishnavi merged commit 8478fe9 into main Oct 3, 2026
@kunal-vaishnavi
kunal-vaishnavi deleted the copilot/optimize-dynamic-sparse-attention branch October 3, 2026 18:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants