Optimize CUDA DynamicSparseAttention for large-context inference

[Copilot]

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]

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

[Copilot]

🟡 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.