Compute SparseAttention CUDA buffer sizes and offsets in size_t

[Akshay Sonawane (apsonawane)]

The scratch allocation sizes in SparseAttention::ComputeInternal and the Q/K/V and rotary offsets in QkvToContext were evaluated as products of int shape fields and only widened afterwards, so a large batch_size or sequence_length could wrap the product before it reached GetScratchBuffer or the pointer arithmetic. Use SafeInt<size_t> so the products are computed at full width and overflow throws instead of wrapping.

The Triton kernel parameter structs take 32-bit strides, so also bound the corresponding element counts in CheckInputs to keep those strides representable.

[Copilot]

Pull request overview

This PR hardens SparseAttention’s CUDA path against integer overflow by computing scratch-buffer sizes and Q/K/V/rotary offsets in size_t (via SafeInt<size_t>), and adds input validation to ensure Triton kernel stride-related products remain representable in 32-bit integers.

Changes:

• Use SafeInt<size_t> for scratch allocation size computations in SparseAttention::ComputeInternal to avoid pre-widening overflow.
• Use SafeInt<size_t> for Q/K/V element-count offsets and rotary buffer offset math in QkvToContext.
• Add input-shape product bounds in CheckInputs to keep Triton stride-related quantities within int32_t limits.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc	Uses SafeInt<size_t> to compute scratch buffer sizes without intermediate overflow.
onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu	Uses SafeInt<size_t> for Q/K/V and rotary offset computations used in pointer arithmetic.
onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h	Adds shape-product bounds checks intended to keep Triton stride-related quantities within int32_t.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.