From 1ca8d3dbedcfacc4839a4e1d99c49eeab3af82b9 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 10 Jun 2026 09:01:44 -0700 Subject: [PATCH 1/6] Improve LinarAttention decode --- .../cuda/bert/linear_attention_impl.cu | 407 ++++++++++++++++++ 1 file changed, 407 insertions(+) diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu index 9137fcf9b25b9..2456fcf1233b6 100644 --- a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu @@ -63,6 +63,15 @@ template <> __device__ __forceinline__ __nv_bfloat16 from_float(float val) { return __float2bfloat16(val); } #endif +// Full-warp sum reduction (result broadcast to all lanes). +__device__ __forceinline__ float warp_reduce_sum(float v) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + v += __shfl_xor_sync(0xffffffffu, v, offset); + } + return v; +} + // ============================================================================= // Fused recurrent linear attention kernel // @@ -654,6 +663,346 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( } } +// ============================================================================= +// Decode-optimized kernel (small seq_len). +// +// Unlike the recurrent kernels above (one block per head, state cached in shared +// memory across the token loop, block-wide __syncthreads barriers), this kernel +// parallelizes the d_v output columns across many warps/blocks so the whole GPU +// is saturated at decode (batch=1) and there are no block-wide barriers — only +// warp shuffles. This mirrors the flash-linear-attention / llama.cpp decode +// design and is far better for the latency-bound seq_len<=few case where the +// shared-memory state caching of the recurrent kernels yields no amortization. +// +// Grid: (kv_num_heads, batch_size, ceil(d_v / kWarpsPerBlock)) +// Block: (32, kWarpsPerBlock) +// Each warp owns exactly one output column `col`. The recurrent state column +// S[:, col] is sharded into registers across the warp's 32 lanes (DK/32 rows per +// lane). Matrix-vector reductions (S^T@k, S^T@q) use warp_reduce_sum. +// +// Requires DK % 32 == 0. State global layout is unchanged (row-major [DK, d_v]), +// so the present_state is bit-compatible with the recurrent kernels above and +// can be freely interchanged between prefill and decode steps. +// ============================================================================= +constexpr int kWarpsPerBlock = 4; + +template +__global__ void LinearAttentionDecodeKernel( + const T* __restrict__ query, + const T* __restrict__ key, + const T* __restrict__ value, + T* __restrict__ state, + const T* __restrict__ decay, + const T* __restrict__ beta_in, + T* __restrict__ output, + int seq_len, + int q_num_heads, + int kv_num_heads, + int n_k_heads, + int d_v, + int output_hidden, + float scale, + bool needs_decay, + bool decay_per_key_dim, + bool needs_beta, + bool beta_per_head, + bool needs_retrieval) { + static_assert(DK % 32 == 0, "DK must be a multiple of warp size (32)"); + constexpr int ROWS = DK / 32; + + const int b = blockIdx.x; + const int h_kv = blockIdx.y; + const int lane = threadIdx.x; + const int col = blockIdx.z * kWarpsPerBlock + threadIdx.y; + if (col >= d_v) { + return; + } + + const int kv_per_k = kv_num_heads / n_k_heads; + const int h_k = h_kv / kv_per_k; + + // State column S[:, col] sharded into registers: lane holds rows {r*32 + lane}. + T* S_col = state + ((int64_t)b * kv_num_heads + h_kv) * DK * d_v + col; + float s_shard[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + s_shard[r] = to_float(S_col[(int64_t)(r * 32 + lane) * d_v]); + } + + const int k_hidden = n_k_heads * DK; + const int q_hidden = q_num_heads * DK; + const int v_hidden = kv_num_heads * d_v; + + for (int t = 0; t < seq_len; ++t) { + const int64_t bt = (int64_t)b * seq_len + t; + + // Load this lane's key shard. + float k_reg[ROWS]; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + k_reg[r] = to_float(key[bt * k_hidden + h_k * DK + (r * 32 + lane)]); + } + + // Per-row decay multipliers (scalar broadcast or per-key-dim). + float exp_g[ROWS]; + if (needs_decay) { + if (decay_per_key_dim) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + exp_g[r] = expf(to_float(decay[bt * (kv_num_heads * DK) + h_kv * DK + (r * 32 + lane)])); + } + } else { + float g = expf(to_float(decay[bt * kv_num_heads + h_kv])); +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + exp_g[r] = g; + } + } +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + s_shard[r] *= exp_g[r]; + } + } + + // Retrieval: r_col = sum_i (decayed S[i][col]) * k[i]. + float r_col = 0.0f; + if (needs_retrieval) { + float partial = 0.0f; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + partial += s_shard[r] * k_reg[r]; + } + r_col = warp_reduce_sum(partial); + } + + // delta_col = beta * (v[col] - r_col), or just v[col] when no beta. + float v_col = to_float(value[bt * v_hidden + h_kv * d_v + col]); + float delta_col; + if (needs_beta) { + float beta_v = beta_per_head ? to_float(beta_in[bt * kv_num_heads + h_kv]) + : to_float(beta_in[bt]); + delta_col = beta_v * (v_col - r_col); + } else { + delta_col = v_col; + } + + // State update: S[i][col] = decayed S[i][col] + k[i] * delta_col. +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + s_shard[r] += k_reg[r] * delta_col; + } + + // Readout: output = scale * sum_i S[i][col] * q[i]. + if (q_num_heads >= kv_num_heads) { + const int heads_per_group = q_num_heads / kv_num_heads; + for (int g = 0; g < heads_per_group; ++g) { + const int h_q = h_kv * heads_per_group + g; + float partial = 0.0f; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + float q_reg = to_float(query[bt * q_hidden + h_q * DK + (r * 32 + lane)]); + partial += s_shard[r] * q_reg; + } + float acc = warp_reduce_sum(partial); + if (lane == 0) { + output[bt * output_hidden + h_q * d_v + col] = from_float(scale * acc); + } + } + } else { + const int h_q = h_kv * q_num_heads / kv_num_heads; + float partial = 0.0f; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + float q_reg = to_float(query[bt * q_hidden + h_q * DK + (r * 32 + lane)]); + partial += s_shard[r] * q_reg; + } + float acc = warp_reduce_sum(partial); + if (lane == 0) { + output[bt * output_hidden + h_kv * d_v + col] = from_float(scale * acc); + } + } + } + + // Write the updated state column back (row-major layout, strided). +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + S_col[(int64_t)(r * 32 + lane) * d_v] = from_float(s_shard[r]); + } +} + +// ============================================================================= +// Decode-optimized kernel v2 — column-per-thread, coalesced row-major state. +// +// The v1 warp-per-column kernel above shards the state column across a warp's +// lanes (DK/32 rows each). In the row-major [DK, d_v] state layout a column is +// strided by d_v, so the per-token state load/store is fully uncoalesced (32 +// lanes hit 32 separate sectors). llama.cpp avoids this by storing the state +// transposed, but that would change this op's present_state output layout. +// +// This v2 keeps the state layout unchanged (row-major [DK, d_v], contract and +// parity preserved) and instead maps ONE THREAD per output column. Thread `col` +// owns the whole column S[:, col] in registers (DK values). For a fixed row i, +// consecutive threads read consecutive addresses (i*d_v + col), so the state +// load/store is fully COALESCED — no transpose required. Per-column reductions +// (S^T@k, S^T@q) are sequential within the thread, so there are no cross-thread +// reductions; the only shared data are the per-token k/q/decay broadcasts. +// +// Grid: (batch_size, kv_num_heads, ceil(d_v / kColsPerBlock)) +// Block: (kColsPerBlock) +// Requires d_v % kColsPerBlock == 0 (so every thread maps to a valid column and +// all threads participate in the cooperative broadcast loads / __syncthreads). +// ============================================================================= +constexpr int kColsPerBlock = 32; + +template +__global__ void LinearAttentionDecodeColKernel( + const T* __restrict__ query, + const T* __restrict__ key, + const T* __restrict__ value, + T* __restrict__ state, + const T* __restrict__ decay, + const T* __restrict__ beta_in, + T* __restrict__ output, + int seq_len, + int q_num_heads, + int kv_num_heads, + int n_k_heads, + int d_v, + int output_hidden, + float scale, + bool needs_decay, + bool decay_per_key_dim, + bool needs_beta, + bool beta_per_head, + bool needs_retrieval) { + const int b = blockIdx.x; + const int h_kv = blockIdx.y; + const int tid = threadIdx.x; + const int col = blockIdx.z * kColsPerBlock + tid; + // d_v is required to be a multiple of kColsPerBlock by the dispatcher, so all + // threads have a valid column and may participate in the barriers below. + + const int kv_per_k = kv_num_heads / n_k_heads; + const int h_k = h_kv / kv_per_k; + + // This thread owns column `col`: S[i][col] lives at i*d_v + col (row-major). + T* S_head = state + ((int64_t)b * kv_num_heads + h_kv) * DK * d_v + col; + float s_col[DK]; +#pragma unroll + for (int i = 0; i < DK; ++i) { + s_col[i] = to_float(S_head[(int64_t)i * d_v]); + } + + // Per-token broadcasts shared across all columns of this (b, h_kv). + __shared__ float k_sh[DK]; + __shared__ float q_sh[DK]; + __shared__ float g_sh[DK]; + __shared__ float scalar_g; + + const int k_hidden = n_k_heads * DK; + const int q_hidden = q_num_heads * DK; + const int v_hidden = kv_num_heads * d_v; + + for (int t = 0; t < seq_len; ++t) { + const int64_t bt = (int64_t)b * seq_len + t; + + for (int i = tid; i < DK; i += kColsPerBlock) { + k_sh[i] = to_float(key[bt * k_hidden + h_k * DK + i]); + } + if (needs_decay) { + if (decay_per_key_dim) { + for (int i = tid; i < DK; i += kColsPerBlock) { + g_sh[i] = expf(to_float(decay[bt * (kv_num_heads * DK) + h_kv * DK + i])); + } + } else if (tid == 0) { + scalar_g = expf(to_float(decay[bt * kv_num_heads + h_kv])); + } + } + __syncthreads(); + + // Decay. + if (needs_decay) { + if (decay_per_key_dim) { +#pragma unroll + for (int i = 0; i < DK; ++i) { + s_col[i] *= g_sh[i]; + } + } else { + const float g = scalar_g; +#pragma unroll + for (int i = 0; i < DK; ++i) { + s_col[i] *= g; + } + } + } + + // Retrieval: r_col = sum_i (decayed S[i][col]) * k[i]. + float r_col = 0.0f; + if (needs_retrieval) { +#pragma unroll + for (int i = 0; i < DK; ++i) { + r_col += s_col[i] * k_sh[i]; + } + } + + // delta_col = beta * (v[col] - r_col), or just v[col] when no beta. + const float v_col = to_float(value[bt * v_hidden + h_kv * d_v + col]); + float delta_col; + if (needs_beta) { + const float beta_v = beta_per_head ? to_float(beta_in[bt * kv_num_heads + h_kv]) + : to_float(beta_in[bt]); + delta_col = beta_v * (v_col - r_col); + } else { + delta_col = v_col; + } + + // State update: S[i][col] = decayed S[i][col] + k[i] * delta_col. +#pragma unroll + for (int i = 0; i < DK; ++i) { + s_col[i] += k_sh[i] * delta_col; + } + + // Readout: output = scale * sum_i S[i][col] * q[i]. + if (q_num_heads >= kv_num_heads) { + const int heads_per_group = q_num_heads / kv_num_heads; + for (int g = 0; g < heads_per_group; ++g) { + const int h_q = h_kv * heads_per_group + g; + __syncthreads(); + for (int i = tid; i < DK; i += kColsPerBlock) { + q_sh[i] = to_float(query[bt * q_hidden + h_q * DK + i]); + } + __syncthreads(); + float acc = 0.0f; +#pragma unroll + for (int i = 0; i < DK; ++i) { + acc += s_col[i] * q_sh[i]; + } + output[bt * output_hidden + h_q * d_v + col] = from_float(scale * acc); + } + } else { + const int h_q = h_kv * q_num_heads / kv_num_heads; + __syncthreads(); + for (int i = tid; i < DK; i += kColsPerBlock) { + q_sh[i] = to_float(query[bt * q_hidden + h_q * DK + i]); + } + __syncthreads(); + float acc = 0.0f; +#pragma unroll + for (int i = 0; i < DK; ++i) { + acc += s_col[i] * q_sh[i]; + } + output[bt * output_hidden + h_kv * d_v + col] = from_float(scale * acc); + } + __syncthreads(); // before next token overwrites k_sh/g_sh + } + + // Store the updated column back (coalesced). +#pragma unroll + for (int i = 0; i < DK; ++i) { + S_head[(int64_t)i * d_v] = from_float(s_col[i]); + } +} + } // anonymous namespace template @@ -685,6 +1034,64 @@ Status LaunchLinearAttentionKernel( int output_hidden = std::max(q_num_heads, kv_num_heads) * d_v; + // --------------------------------------------------------------------------- + // Decode fast path: small seq_len is latency-bound. The recurrent kernels keep + // state in shared memory across the token loop (great for prefill) but at + // decode that caching yields no amortization while only using kv_num_heads + // blocks. The column-parallel decode kernel saturates the GPU and avoids + // block barriers. Engaged for the common decode shapes (DK in {64,128,256}); + // everything else falls through to the recurrent kernels below. + // --------------------------------------------------------------------------- + constexpr int kDecodeSeqThreshold = 16; + if (seq_len <= kDecodeSeqThreshold && (d_k == 64 || d_k == 128 || d_k == 256)) { + // v2 (column-per-thread, coalesced row-major state) is the default. It + // requires d_v % kColsPerBlock == 0; otherwise fall back to the v1 + // warp-per-column kernel (which handles any d_v). + if (d_v % kColsPerBlock == 0) { + const dim3 decode_grid(batch_size, kv_num_heads, + (d_v + kColsPerBlock - 1) / kColsPerBlock); + const dim3 decode_block(kColsPerBlock, 1, 1); + + auto launch_col = [&](auto dk_tag) -> Status { + constexpr int DK = decltype(dk_tag)::value; + LinearAttentionDecodeColKernel<<>>( + query, key, value, present_state, decay, beta, output, + seq_len, q_num_heads, kv_num_heads, n_k_heads, d_v, output_hidden, scale, + needs_decay, decay_per_key_dim, needs_beta, beta_per_head, needs_retrieval); + return CUDA_CALL(cudaGetLastError()); + }; + + if (d_k == 128) { + return launch_col(std::integral_constant{}); + } else if (d_k == 64) { + return launch_col(std::integral_constant{}); + } else { // d_k == 256 + return launch_col(std::integral_constant{}); + } + } + + const dim3 decode_grid(batch_size, kv_num_heads, + (d_v + kWarpsPerBlock - 1) / kWarpsPerBlock); + const dim3 decode_block(32, kWarpsPerBlock, 1); + + auto launch_decode = [&](auto dk_tag) -> Status { + constexpr int DK = decltype(dk_tag)::value; + LinearAttentionDecodeKernel<<>>( + query, key, value, present_state, decay, beta, output, + seq_len, q_num_heads, kv_num_heads, n_k_heads, d_v, output_hidden, scale, + needs_decay, decay_per_key_dim, needs_beta, beta_per_head, needs_retrieval); + return CUDA_CALL(cudaGetLastError()); + }; + + if (d_k == 128) { + return launch_decode(std::integral_constant{}); + } else if (d_k == 64) { + return launch_decode(std::integral_constant{}); + } else { // d_k == 256 + return launch_decode(std::integral_constant{}); + } + } + auto launch_fixed = [&](auto dk_tag, auto dv_tag) -> Status { constexpr int DK = decltype(dk_tag)::value; constexpr int DV = decltype(dv_tag)::value; From ebcd859666b0533f8724f2c705d1e4c7d65d2340 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 11 Jun 2026 17:38:56 -0700 Subject: [PATCH 2/6] address feedbacks --- .../cuda/bert/linear_attention_impl.cu | 140 +++++++++--------- 1 file changed, 73 insertions(+), 67 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu index 2456fcf1233b6..e3a00209d6cba 100644 --- a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu @@ -26,6 +26,7 @@ #include #include #include "contrib_ops/cuda/bert/linear_attention_impl.h" +#include "core/providers/cuda/cu_inc/common.cuh" namespace onnxruntime { namespace contrib { @@ -67,11 +68,57 @@ __device__ __forceinline__ __nv_bfloat16 from_float(float val) { return __float2 __device__ __forceinline__ float warp_reduce_sum(float v) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { - v += __shfl_xor_sync(0xffffffffu, v, offset); + v += onnxruntime::cuda::WARP_SHFL_XOR(v, offset); } return v; } +template +__device__ __forceinline__ float ComputeLinearAttentionDeltaColumn( + const T* value, + const T* beta_in, + int64_t batch_token_offset, + int value_hidden, + int kv_head, + int d_v, + int col, + int kv_num_heads, + bool needs_beta, + bool beta_per_head, + float retrieval_col) { + const float value_col = to_float(value[batch_token_offset * value_hidden + kv_head * d_v + col]); + if (!needs_beta) { + return value_col; + } + + const float beta_value = beta_per_head ? to_float(beta_in[batch_token_offset * kv_num_heads + kv_head]) + : to_float(beta_in[batch_token_offset]); + return beta_value * (value_col - retrieval_col); +} + +struct LinearAttentionReadoutHeads { + int query_head; + int output_head; +}; + +__device__ __forceinline__ int GetLinearAttentionReadoutHeadCount(int q_num_heads, int kv_num_heads) { + return q_num_heads >= kv_num_heads ? q_num_heads / kv_num_heads : 1; +} + +__device__ __forceinline__ LinearAttentionReadoutHeads GetLinearAttentionReadoutHeads( + int kv_head, + int q_num_heads, + int kv_num_heads, + int group_index) { + if (q_num_heads >= kv_num_heads) { + const int heads_per_group = q_num_heads / kv_num_heads; + const int query_head = kv_head * heads_per_group + group_index; + return {query_head, query_head}; + } + + return {kv_head * q_num_heads / kv_num_heads, kv_head}; +} + // ============================================================================= // Fused recurrent linear attention kernel // @@ -674,7 +721,7 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( // design and is far better for the latency-bound seq_len<=few case where the // shared-memory state caching of the recurrent kernels yields no amortization. // -// Grid: (kv_num_heads, batch_size, ceil(d_v / kWarpsPerBlock)) +// Grid: (batch_size, kv_num_heads, ceil(d_v / kWarpsPerBlock)) // Block: (32, kWarpsPerBlock) // Each warp owns exactly one output column `col`. The recurrent state column // S[:, col] is sharded into registers across the warp's 32 lanes (DK/32 rows per @@ -775,16 +822,8 @@ __global__ void LinearAttentionDecodeKernel( r_col = warp_reduce_sum(partial); } - // delta_col = beta * (v[col] - r_col), or just v[col] when no beta. - float v_col = to_float(value[bt * v_hidden + h_kv * d_v + col]); - float delta_col; - if (needs_beta) { - float beta_v = beta_per_head ? to_float(beta_in[bt * kv_num_heads + h_kv]) - : to_float(beta_in[bt]); - delta_col = beta_v * (v_col - r_col); - } else { - delta_col = v_col; - } + const float delta_col = ComputeLinearAttentionDeltaColumn( + value, beta_in, bt, v_hidden, h_kv, d_v, col, kv_num_heads, needs_beta, beta_per_head, r_col); // State update: S[i][col] = decayed S[i][col] + k[i] * delta_col. #pragma unroll @@ -793,32 +832,19 @@ __global__ void LinearAttentionDecodeKernel( } // Readout: output = scale * sum_i S[i][col] * q[i]. - if (q_num_heads >= kv_num_heads) { - const int heads_per_group = q_num_heads / kv_num_heads; - for (int g = 0; g < heads_per_group; ++g) { - const int h_q = h_kv * heads_per_group + g; - float partial = 0.0f; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - float q_reg = to_float(query[bt * q_hidden + h_q * DK + (r * 32 + lane)]); - partial += s_shard[r] * q_reg; - } - float acc = warp_reduce_sum(partial); - if (lane == 0) { - output[bt * output_hidden + h_q * d_v + col] = from_float(scale * acc); - } - } - } else { - const int h_q = h_kv * q_num_heads / kv_num_heads; + const int head_count = GetLinearAttentionReadoutHeadCount(q_num_heads, kv_num_heads); + for (int group_index = 0; group_index < head_count; ++group_index) { + const LinearAttentionReadoutHeads readout_heads = + GetLinearAttentionReadoutHeads(h_kv, q_num_heads, kv_num_heads, group_index); float partial = 0.0f; #pragma unroll for (int r = 0; r < ROWS; ++r) { - float q_reg = to_float(query[bt * q_hidden + h_q * DK + (r * 32 + lane)]); + float q_reg = to_float(query[bt * q_hidden + readout_heads.query_head * DK + (r * 32 + lane)]); partial += s_shard[r] * q_reg; } float acc = warp_reduce_sum(partial); if (lane == 0) { - output[bt * output_hidden + h_kv * d_v + col] = from_float(scale * acc); + output[bt * output_hidden + readout_heads.output_head * d_v + col] = from_float(scale * acc); } } } @@ -945,16 +971,8 @@ __global__ void LinearAttentionDecodeColKernel( } } - // delta_col = beta * (v[col] - r_col), or just v[col] when no beta. - const float v_col = to_float(value[bt * v_hidden + h_kv * d_v + col]); - float delta_col; - if (needs_beta) { - const float beta_v = beta_per_head ? to_float(beta_in[bt * kv_num_heads + h_kv]) - : to_float(beta_in[bt]); - delta_col = beta_v * (v_col - r_col); - } else { - delta_col = v_col; - } + const float delta_col = ComputeLinearAttentionDeltaColumn( + value, beta_in, bt, v_hidden, h_kv, d_v, col, kv_num_heads, needs_beta, beta_per_head, r_col); // State update: S[i][col] = decayed S[i][col] + k[i] * delta_col. #pragma unroll @@ -963,27 +981,13 @@ __global__ void LinearAttentionDecodeColKernel( } // Readout: output = scale * sum_i S[i][col] * q[i]. - if (q_num_heads >= kv_num_heads) { - const int heads_per_group = q_num_heads / kv_num_heads; - for (int g = 0; g < heads_per_group; ++g) { - const int h_q = h_kv * heads_per_group + g; - __syncthreads(); - for (int i = tid; i < DK; i += kColsPerBlock) { - q_sh[i] = to_float(query[bt * q_hidden + h_q * DK + i]); - } - __syncthreads(); - float acc = 0.0f; -#pragma unroll - for (int i = 0; i < DK; ++i) { - acc += s_col[i] * q_sh[i]; - } - output[bt * output_hidden + h_q * d_v + col] = from_float(scale * acc); - } - } else { - const int h_q = h_kv * q_num_heads / kv_num_heads; + const int head_count = GetLinearAttentionReadoutHeadCount(q_num_heads, kv_num_heads); + for (int group_index = 0; group_index < head_count; ++group_index) { + const LinearAttentionReadoutHeads readout_heads = + GetLinearAttentionReadoutHeads(h_kv, q_num_heads, kv_num_heads, group_index); __syncthreads(); for (int i = tid; i < DK; i += kColsPerBlock) { - q_sh[i] = to_float(query[bt * q_hidden + h_q * DK + i]); + q_sh[i] = to_float(query[bt * q_hidden + readout_heads.query_head * DK + i]); } __syncthreads(); float acc = 0.0f; @@ -991,7 +995,7 @@ __global__ void LinearAttentionDecodeColKernel( for (int i = 0; i < DK; ++i) { acc += s_col[i] * q_sh[i]; } - output[bt * output_hidden + h_kv * d_v + col] = from_float(scale * acc); + output[bt * output_hidden + readout_heads.output_head * d_v + col] = from_float(scale * acc); } __syncthreads(); // before next token overwrites k_sh/g_sh } @@ -1042,6 +1046,8 @@ Status LaunchLinearAttentionKernel( // block barriers. Engaged for the common decode shapes (DK in {64,128,256}); // everything else falls through to the recurrent kernels below. // --------------------------------------------------------------------------- + // This cutoff keeps short decode/continuation runs on the column-parallel kernels; + // longer prefill-like runs benefit from amortizing state in shared memory below. constexpr int kDecodeSeqThreshold = 16; if (seq_len <= kDecodeSeqThreshold && (d_k == 64 || d_k == 128 || d_k == 256)) { // v2 (column-per-thread, coalesced row-major state) is the default. It @@ -1061,10 +1067,10 @@ Status LaunchLinearAttentionKernel( return CUDA_CALL(cudaGetLastError()); }; - if (d_k == 128) { - return launch_col(std::integral_constant{}); - } else if (d_k == 64) { + if (d_k == 64) { return launch_col(std::integral_constant{}); + } else if (d_k == 128) { + return launch_col(std::integral_constant{}); } else { // d_k == 256 return launch_col(std::integral_constant{}); } @@ -1083,10 +1089,10 @@ Status LaunchLinearAttentionKernel( return CUDA_CALL(cudaGetLastError()); }; - if (d_k == 128) { - return launch_decode(std::integral_constant{}); - } else if (d_k == 64) { + if (d_k == 64) { return launch_decode(std::integral_constant{}); + } else if (d_k == 128) { + return launch_decode(std::integral_constant{}); } else { // d_k == 256 return launch_decode(std::integral_constant{}); } From 74945fd0adc157a6b1addc5812d0ce9c32141d6e Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 11 Jun 2026 18:20:00 -0700 Subject: [PATCH 3/6] refactoring --- .../cuda/bert/causal_conv_with_state_impl.cu | 29 +------------ .../cuda/bert/linear_attention_impl.cu | 32 +-------------- .../cuda/cu_inc/cuda_type_helper.cuh | 41 +++++++++++++++++++ 3 files changed, 43 insertions(+), 59 deletions(-) create mode 100644 onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu index 87774f7205bc9..4f1985f190f93 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu @@ -18,6 +18,7 @@ #include #include #include "contrib_ops/cuda/bert/causal_conv_with_state_impl.h" +#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" namespace onnxruntime { namespace contrib { @@ -25,34 +26,6 @@ namespace cuda { namespace { -template -__device__ __forceinline__ float to_float(T val); - -template <> -__device__ __forceinline__ float to_float(float val) { return val; } - -template <> -__device__ __forceinline__ float to_float(half val) { return __half2float(val); } - -#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) -template <> -__device__ __forceinline__ float to_float(__nv_bfloat16 val) { return __bfloat162float(val); } -#endif - -template -__device__ __forceinline__ T from_float(float val); - -template <> -__device__ __forceinline__ float from_float(float val) { return val; } - -template <> -__device__ __forceinline__ half from_float(float val) { return __float2half(val); } - -#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) -template <> -__device__ __forceinline__ __nv_bfloat16 from_float(float val) { return __float2bfloat16(val); } -#endif - __device__ __forceinline__ float silu_fn(float x) { return x / (1.0f + expf(-x)); } diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu index e3a00209d6cba..b12b5ab7ba872 100644 --- a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu @@ -27,43 +27,13 @@ #include #include "contrib_ops/cuda/bert/linear_attention_impl.h" #include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" namespace onnxruntime { namespace contrib { namespace cuda { namespace { - -// Convert half/bfloat16 to float -template -__device__ __forceinline__ float to_float(T val); - -template <> -__device__ __forceinline__ float to_float(float val) { return val; } - -template <> -__device__ __forceinline__ float to_float(half val) { return __half2float(val); } - -#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) -template <> -__device__ __forceinline__ float to_float(__nv_bfloat16 val) { return __bfloat162float(val); } -#endif - -// Convert float to half/bfloat16/float -template -__device__ __forceinline__ T from_float(float val); - -template <> -__device__ __forceinline__ float from_float(float val) { return val; } - -template <> -__device__ __forceinline__ half from_float(float val) { return __float2half(val); } - -#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) -template <> -__device__ __forceinline__ __nv_bfloat16 from_float(float val) { return __float2bfloat16(val); } -#endif - // Full-warp sum reduction (result broadcast to all lanes). __device__ __forceinline__ float warp_reduce_sum(float v) { #pragma unroll diff --git a/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh b/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh new file mode 100644 index 0000000000000..7289864797bb5 --- /dev/null +++ b/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh @@ -0,0 +1,41 @@ +namespace onnxruntime { +namespace contrib { +namespace cuda { + +namespace { + +// Convert half/bfloat16 to float +template +__device__ __forceinline__ float to_float(T val); + +template <> +__device__ __forceinline__ float to_float(float val) { return val; } + +template <> +__device__ __forceinline__ float to_float(half val) { return __half2float(val); } + +#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) +template <> +__device__ __forceinline__ float to_float(__nv_bfloat16 val) { return __bfloat162float(val); } +#endif + +// Convert float to half/bfloat16/float +template +__device__ __forceinline__ T from_float(float val); + +template <> +__device__ __forceinline__ float from_float(float val) { return val; } + +template <> +__device__ __forceinline__ half from_float(float val) { return __float2half(val); } + +#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) +template <> +__device__ __forceinline__ __nv_bfloat16 from_float(float val) { return __float2bfloat16(val); } +#endif + +} // namespace + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime From d01b89557be092092b7bba84be3bde62c4c46887 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 11 Jun 2026 18:25:08 -0700 Subject: [PATCH 4/6] lintrunner -a --- .../core/providers/cuda/cu_inc/cuda_type_helper.cuh | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh b/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh index 7289864797bb5..21295611ff03c 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh @@ -34,8 +34,8 @@ template <> __device__ __forceinline__ __nv_bfloat16 from_float(float val) { return __float2bfloat16(val); } #endif -} // namespace +} // namespace -} // namespace cuda -} // namespace contrib -} // namespace onnxruntime +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime From 6ee345a925d2d9e88d58d8634fd8e4ac139a9719 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 12 Jun 2026 22:26:57 +0000 Subject: [PATCH 5/6] Address review feedback in CUDA helper and decode comments --- .../contrib_ops/cuda/bert/linear_attention_impl.cu | 5 +++-- .../core/providers/cuda/cu_inc/cuda_type_helper.cuh | 8 ++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu index b12b5ab7ba872..1f099319d1fce 100644 --- a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu @@ -1012,8 +1012,9 @@ Status LaunchLinearAttentionKernel( // Decode fast path: small seq_len is latency-bound. The recurrent kernels keep // state in shared memory across the token loop (great for prefill) but at // decode that caching yields no amortization while only using kv_num_heads - // blocks. The column-parallel decode kernel saturates the GPU and avoids - // block barriers. Engaged for the common decode shapes (DK in {64,128,256}); + // blocks. The column-parallel decode kernel increases occupancy with + // finer-grained work distribution. Engaged for the common decode shapes + // (DK in {64,128,256}); // everything else falls through to the recurrent kernels below. // --------------------------------------------------------------------------- // This cutoff keeps short decode/continuation runs on the column-parallel kernels; diff --git a/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh b/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh index 21295611ff03c..895e331a4cd02 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh @@ -1,3 +1,11 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include + namespace onnxruntime { namespace contrib { namespace cuda { From cdcea5e12785a3a2e5f5fb1c225b7fe6a14f67bb Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 18 Jun 2026 06:27:29 +0000 Subject: [PATCH 6/6] address more feedbacks --- .../cuda/bert/linear_attention_impl.cu | 23 ++++++++++--------- .../cuda/cu_inc/cuda_type_helper.cuh | 4 ++-- ...est_parity_linear_attention_causal_conv.py | 6 +++++ 3 files changed, 20 insertions(+), 13 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu index 1f099319d1fce..b27641c51c5ad 100644 --- a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu @@ -686,10 +686,11 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( // Unlike the recurrent kernels above (one block per head, state cached in shared // memory across the token loop, block-wide __syncthreads barriers), this kernel // parallelizes the d_v output columns across many warps/blocks so the whole GPU -// is saturated at decode (batch=1) and there are no block-wide barriers — only -// warp shuffles. This mirrors the flash-linear-attention / llama.cpp decode -// design and is far better for the latency-bound seq_len<=few case where the -// shared-memory state caching of the recurrent kernels yields no amortization. +// is saturated at decode (batch=1). This v1 warp-per-column kernel has no +// block-wide barriers, only warp shuffles. It mirrors the flash-linear-attention +// / llama.cpp decode design and is far better for the latency-bound +// seq_len<=few case where the shared-memory state caching of the recurrent +// kernels yields no amortization. // // Grid: (batch_size, kv_num_heads, ceil(d_v / kWarpsPerBlock)) // Block: (32, kWarpsPerBlock) @@ -841,7 +842,8 @@ __global__ void LinearAttentionDecodeKernel( // consecutive threads read consecutive addresses (i*d_v + col), so the state // load/store is fully COALESCED — no transpose required. Per-column reductions // (S^T@k, S^T@q) are sequential within the thread, so there are no cross-thread -// reductions; the only shared data are the per-token k/q/decay broadcasts. +// reductions; the only shared data are the per-token k/q/decay broadcasts, with +// block-wide barriers around those cooperative loads. // // Grid: (batch_size, kv_num_heads, ceil(d_v / kColsPerBlock)) // Block: (kColsPerBlock) @@ -1021,10 +1023,11 @@ Status LaunchLinearAttentionKernel( // longer prefill-like runs benefit from amortizing state in shared memory below. constexpr int kDecodeSeqThreshold = 16; if (seq_len <= kDecodeSeqThreshold && (d_k == 64 || d_k == 128 || d_k == 256)) { - // v2 (column-per-thread, coalesced row-major state) is the default. It - // requires d_v % kColsPerBlock == 0; otherwise fall back to the v1 - // warp-per-column kernel (which handles any d_v). - if (d_v % kColsPerBlock == 0) { + // v2 (column-per-thread, coalesced row-major state) is the default for + // DK <= 128. It requires d_v % kColsPerBlock == 0; otherwise fall back + // to the v1 warp-per-column kernel (which handles any d_v). DK=256 also + // uses v1 to avoid the high per-thread register footprint of s_col[256]. + if (d_k <= 128 && d_v % kColsPerBlock == 0) { const dim3 decode_grid(batch_size, kv_num_heads, (d_v + kColsPerBlock - 1) / kColsPerBlock); const dim3 decode_block(kColsPerBlock, 1, 1); @@ -1042,8 +1045,6 @@ Status LaunchLinearAttentionKernel( return launch_col(std::integral_constant{}); } else if (d_k == 128) { return launch_col(std::integral_constant{}); - } else { // d_k == 256 - return launch_col(std::integral_constant{}); } } diff --git a/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh b/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh index 895e331a4cd02..8d53891709ca1 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh @@ -14,7 +14,7 @@ namespace { // Convert half/bfloat16 to float template -__device__ __forceinline__ float to_float(T val); +__device__ __forceinline__ float to_float(T val) = delete; template <> __device__ __forceinline__ float to_float(float val) { return val; } @@ -29,7 +29,7 @@ __device__ __forceinline__ float to_float(__nv_bfloat16 val) { return __bfloat16 // Convert float to half/bfloat16/float template -__device__ __forceinline__ T from_float(float val); +__device__ __forceinline__ T from_float(float val) = delete; template <> __device__ __forceinline__ float from_float(float val) { return val; } diff --git a/onnxruntime/test/python/transformers/test_parity_linear_attention_causal_conv.py b/onnxruntime/test/python/transformers/test_parity_linear_attention_causal_conv.py index d3b053f1a094a..c234bd908b99f 100644 --- a/onnxruntime/test/python/transformers/test_parity_linear_attention_causal_conv.py +++ b/onnxruntime/test/python/transformers/test_parity_linear_attention_causal_conv.py @@ -337,6 +337,12 @@ def test_linear_attention_d128(self): "gated_delta", q_num_heads=4, kv_num_heads=2, d_k=128, d_v=128, batch=2, seq_lens=(1, 5) ) + def test_linear_attention_decode_v1_fallback(self): + """d_v not divisible by 32 — exercises the warp-per-column decode fallback.""" + self._run_linear_attention_test( + "gated_delta", q_num_heads=4, kv_num_heads=2, d_k=64, d_v=40, batch=2, seq_lens=(1, 5) + ) + def test_linear_attention_per_key_dim_decay(self): """Per-key-dim decay: decay shape [B, T, H_kv * d_k] — exercises per-row expf path.""" self._run_linear_attention_test(