[CUDA] Add decode-optimized LinearAttention (GatedDeltaNet) kernels

[tianleiwu]

Description

Adds a decode-optimized CUDA path for the LinearAttention contrib op (the
gated-delta / linear-attention recurrence used by hybrid models such as
Qwen3-Next / Qwen3.6). The existing recurrent kernels are tuned for prefill; at
decode (seq_len == 1) they leave the GPU mostly idle. This PR adds two
decode-specific kernels that saturate the GPU and access the recurrent state in
a coalesced pattern, without changing the op's present_state layout or
numerics.

Motivation

LinearAttentionRecurrentKernelFixedShape launches one block per
(batch, kv_head) and keeps the full d_k x d_v state in shared memory across
the token loop, with block-wide __syncthreads at every step. That design
amortizes state I/O during prefill, but at decode it:

• launches only kv_num_heads blocks (e.g. 32) — a fraction of the SMs, and
• gets no amortization from the shared-memory state cache (one token per launch),

so the op is latency-bound. On an H200 profile of Qwen3.6-35B-A3B it was the
single most expensive decode kernel after the dense/MoE GEMMs (~0.69 ms/token).

Key Changes

All in onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu:

Addition	Description
warp_reduce_sum	__shfl_xor_sync full-warp sum helper.
LinearAttentionDecodeKernel<T, DK>	Warp-per-column decode kernel: grid (kv_num_heads, batch, ceil(d_v/4)), each warp owns one output column with the state column sharded into registers; reductions via warp shuffles, no shared memory, no block barriers. Handles any d_v.
LinearAttentionDecodeColKernel<T, DK>	Column-per-thread decode kernel (default): one thread owns a full state column in registers. For a fixed row i, consecutive threads read consecutive addresses i*d_v + col, so state load/store is fully coalesced with no transpose — the row-major [d_k, d_v] present_state layout is unchanged. Used when d_v % 32 == 0; otherwise falls back to the warp kernel.
Dispatch in LaunchLinearAttentionKernel	Routes seq_len <= 16 and d_k in {64,128,256} to the decode kernels; all other shapes fall through to the existing recurrent kernels, so the prefill path is unchanged.

Both kernels cover the full op semantics: linear / gated / delta /
gated_delta update rules, scalar and per-key-dim decay, per-head and scalar
beta, standard GQA and inverse GQA, and n_k_heads K-sharing.

Performance

H200, Qwen3.6-35B-A3B (INT4), single-sequence decode, CUDA graph on. Kernel time
measured with Nsight Systems (steady-state, warmup excluded):

Kernel	Time / token
LinearAttentionRecurrentKernelFixedShape (existing)	693 µs
LinearAttentionDecodeKernel (warp-per-column)	346 µs (2.0x)
LinearAttentionDecodeColKernel (column-per-thread)	202 µs (3.4x)

End-to-end decode throughput improved measurably (the kernel is ~half its prior
cost), with no change to prefill.

Testing

• All 26 ContribOpLinearAttentionTest parity tests pass (the decode kernels are
  exercised by the single-token, inverse-GQA, KGQA, and Qwen3.5-like cases):

  ./onnxruntime_provider_test --gtest_filter='*LinearAttention*'
  

• No public API or present_state layout change; the decode path is opt-in by
  shape and falls back to the existing kernels for unsupported d_k / d_v.

Motivation and Context

Decode-throughput optimization for hybrid linear-attention + MoE models. No
breaking changes; numerics and output layout are preserved.

[Copilot]

Pull request overview

This PR adds a decode-optimized CUDA fast path for the LinearAttention contrib op to improve latency when seq_len is small (notably decode), while preserving the existing present_state layout and numerics and leaving the prefill-tuned recurrent kernels as the fallback.

Changes:

• Added a full-warp sum helper (warp_reduce_sum) for shuffle-based reductions.
• Introduced two decode-focused kernels (LinearAttentionDecodeKernel and LinearAttentionDecodeColKernel) optimized for better GPU occupancy and more coalesced state access.
• Updated LaunchLinearAttentionKernel dispatch to route common decode shapes (seq_len <= 16, d_k ∈ {64,128,256}) to the new kernels and otherwise keep existing behavior.

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.