Fix LinearAttention on GPUs with limited shared memory

[Copilot]

Description

• Detect when the recurrent CUDA kernel exceeds the device’s opt-in shared-memory limit.
• Route supported shapes to the column-parallel kernel instead.
• Return an actionable error when no compatible fallback exists.

Motivation and Context

The fixed-shape LinearAttention kernel requested 66,560 bytes of shared memory, exceeding the 65,536-byte limit on some GPUs. cudaFuncSetAttribute consequently failed with cudaErrorInvalidValue.

[azure-pipelines]

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

[Copilot]

Pull request overview

This PR improves the CUDA LinearAttention kernel dispatch to avoid failing on GPUs whose opt-in per-block shared-memory limit is below the recurrent kernel’s shared-memory requirement. It does so by detecting the over-limit case up front, routing eligible shapes to the column-parallel kernel, and producing a clear failure when no compatible fallback exists.

Changes:

• Pass the device opt-in shared-memory limit (sharedMemPerBlockOptin) into the CUDA LinearAttention kernel launcher.
• Detect when the recurrent kernel’s dynamic shared-memory requirement exceeds the device limit and preferentially select the column-parallel kernel for supported shapes.
• Update test commentary to reflect the new “recurrent-or-column-fallback” behavior for the fixed-shape prefill case.

Reviewed changes

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

File	Description
onnxruntime/test/contrib_ops/linear_attention_op_test.cc	Updates test comment to reflect shared-memory-based fallback behavior.
onnxruntime/contrib_ops/cuda/bert/linear_attention.cc	Plumbs the device opt-in shared-memory limit into the kernel dispatcher.
onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.h	Extends LaunchLinearAttentionKernel signature with max_shared_memory_per_block.
onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu	Adds shared-memory limit checks and routes to column kernels or returns an error when needed.

---

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