diff --git a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc index 72a97c60df84f..8981f7fc918b2 100644 --- a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "contrib_ops/cpu/bert/causal_conv_with_state.h" +#include "contrib_ops/cpu/bert/causal_conv_with_state_helper.h" #include "core/framework/tensorprotoutils.h" #include "core/common/safeint.h" @@ -47,6 +48,10 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : OpKernel activation_ = info.GetAttrOrDefault("activation", "none"); ORT_ENFORCE(activation_ == "none" || activation_ == "silu" || activation_ == "swish", "activation must be one of: none, silu, swish"); + + ORT_ENFORCE(info.GetAttrOrDefault("state_window", 0) == 0, + "CPU CausalConvWithState does not support state_window > 0 (CUDA EP only)"); + state_window_ = 0; } namespace { @@ -223,7 +228,11 @@ Status CausalConvWithState::Compute(OpKernelContext* context) const { Tensor* output_tensor = context->Output(0, input_shape); float* output_data = output_tensor->MutableData(); - TensorShape state_shape({batch_size, channels, pad}); + // state_window_ is always 0 on CPU, so this is the legacy (B, C, K-1) shape. + TensorShape state_shape; + ORT_RETURN_IF_ERROR(causal_conv_with_state_helper::CheckInputs( + state_window_, static_cast(batch_size), static_cast(channels), + static_cast(pad), past_state_tensor, state_shape)); Tensor* present_state_tensor = context->Output(1, state_shape); float* present_data = present_state_tensor->MutableData(); diff --git a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.h b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.h index e859f69677e80..0e552e7bd27dd 100644 --- a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.h +++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.h @@ -20,6 +20,9 @@ class CausalConvWithState final : public OpKernel { private: int ndim_; std::string activation_; + // Always 0 on CPU (a state window is CUDA-only), but kept so the shared shape helper in + // causal_conv_with_state_helper.h is driven the same way on every EP. + int state_window_; }; } // namespace contrib diff --git a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state_helper.h b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state_helper.h new file mode 100644 index 0000000000000..82fb320324ea2 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state_helper.h @@ -0,0 +1,70 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/framework/op_kernel.h" +#include "core/framework/tensor_shape.h" +#include "core/providers/common.h" + +namespace onnxruntime { +namespace contrib { +namespace causal_conv_with_state_helper { + +// Reads and validates the optional `state_window` attribute. +// +// 0 (the default, i.e. attribute absent) selects the legacy unwindowed state layout. Every model +// exported before the attribute existed lands here, so this must stay the default. +template +Status ParseStateWindow(const TKernelInfo& info, int& state_window) { + const int64_t value = info.template GetAttrOrDefault("state_window", 0); + if (value < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "state_window must be >= 0, got ", value); + } + state_window = static_cast(value); + return Status::OK(); +} + +// Derives the expected past_state / present_state shape and validates past_state against it. +// `state_length` is the carry length along the causal axis, i.e. kernel_size - 1. +// +// state_window == 0 -> (batch_size, channels, state_length). A single state with no window axis: +// the backward-compatible layout that models exported before the attribute existed use. +// +// state_window == W > 0 -> (W, batch_size, channels, state_length). The window holds the carry +// state after each of the last W positions, right-aligned: slot j is the state after position +// (seq_len - W + j), so slot W-1 is the state after the last position and is the only slot read +// back as past_state. W == 1 is therefore the legacy layout with a leading unit axis. +// +// The window axis leads the batch axis so that each slot is one contiguous +// (batch_size, channels, state_length) block. That keeps "fetch/replace the last state" a single +// contiguous range for any batch size, which is what a speculative decoder needs when it crops the +// state back to an accepted prefix. +template +Status CheckInputs(int state_window, + int batch_size, + int channels, + int state_length, + const T* past_state, + TensorShape& state_shape) { + state_shape = state_window > 0 + ? TensorShape({state_window, batch_size, channels, state_length}) + : TensorShape({batch_size, channels, state_length}); + + if (past_state != nullptr && past_state->Shape() != state_shape) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'past_state' is expected to have shape ", state_shape.ToString(), + ", got ", past_state->Shape().ToString(), + ". CausalConvWithState uses (batch_size, channels, kernel_size - 1) when " + "the state_window attribute is absent or 0, and " + "(state_window, batch_size, channels, kernel_size - 1) otherwise."); + } + + return Status::OK(); +} + +} // namespace causal_conv_with_state_helper +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc b/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc index 052e7df8bda14..b08f7b17d4d15 100644 --- a/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/linear_attention.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "contrib_ops/cpu/bert/linear_attention.h" +#include "contrib_ops/cpu/bert/linear_attention_helper.h" #include "core/framework/tensorprotoutils.h" #include "core/common/safeint.h" @@ -60,6 +61,10 @@ LinearAttention::LinearAttention(const OpKernelInfo& info) : OpKernel(info) { int64_t chunk_size = info.GetAttrOrDefault("chunk_size", 64); // chunk_size_ reserved for future chunk-parallel prefill algorithm; not yet used. chunk_size_ = static_cast(chunk_size); + + ORT_ENFORCE(info.GetAttrOrDefault("state_window", 0) == 0, + "CPU LinearAttention does not support state_window > 0 (CUDA EP only)"); + state_window_ = 0; } namespace { @@ -417,21 +422,17 @@ Status LinearAttention::Compute(OpKernelContext* context) const { } // ==== Initialize state: write directly into output present_state ==== - // present_state: (B, H_kv, d_k, d_v) - TensorShape state_shape({batch_size, static_cast(kv_num_heads_), d_k, d_v}); + // state_window_ is always 0 on CPU, so this is the legacy (B, H_kv, d_k, d_v) shape. + TensorShape state_shape; + ORT_RETURN_IF_ERROR(linear_attention_helper::CheckInputs( + state_window_, static_cast(batch_size), kv_num_heads_, + static_cast(d_k), static_cast(d_v), past_state_tensor, state_shape)); Tensor* present_state_tensor = context->Output(1, state_shape); float* state_data = present_state_tensor->MutableData(); int64_t state_per_head = d_k * d_v; int64_t total_state = batch_size * kv_num_heads_ * state_per_head; if (past_state_tensor != nullptr) { - const auto& ps_shape = past_state_tensor->Shape(); - ORT_RETURN_IF_NOT(ps_shape.NumDimensions() == 4 && - ps_shape[0] == batch_size && - ps_shape[1] == kv_num_heads_ && - ps_shape[2] == d_k && - ps_shape[3] == d_v, - "past_state must be (B, H_kv, d_k, d_v)"); const float* ps_data = past_state_tensor->Data(); std::memcpy(state_data, ps_data, static_cast(total_state) * sizeof(float)); } else { diff --git a/onnxruntime/contrib_ops/cpu/bert/linear_attention.h b/onnxruntime/contrib_ops/cpu/bert/linear_attention.h index 9aaa9f80a3369..1fa57028bb95b 100644 --- a/onnxruntime/contrib_ops/cpu/bert/linear_attention.h +++ b/onnxruntime/contrib_ops/cpu/bert/linear_attention.h @@ -23,6 +23,9 @@ class LinearAttention final : public OpKernel { std::string update_rule_; float scale_; int chunk_size_; + // Always 0 on CPU (a state window is CUDA-only), but kept so the shared shape helper in + // linear_attention_helper.h is driven the same way on every EP. + int state_window_; }; } // namespace contrib diff --git a/onnxruntime/contrib_ops/cpu/bert/linear_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/linear_attention_helper.h new file mode 100644 index 0000000000000..8e63de4dac99f --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/bert/linear_attention_helper.h @@ -0,0 +1,70 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/framework/op_kernel.h" +#include "core/framework/tensor_shape.h" +#include "core/providers/common.h" + +namespace onnxruntime { +namespace contrib { +namespace linear_attention_helper { + +// Reads and validates the optional `state_window` attribute. +// +// 0 (the default, i.e. attribute absent) selects the legacy unwindowed state layout. Every model +// exported before the attribute existed lands here, so this must stay the default. +template +Status ParseStateWindow(const TKernelInfo& info, int& state_window) { + const int64_t value = info.template GetAttrOrDefault("state_window", 0); + if (value < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "state_window must be >= 0, got ", value); + } + state_window = static_cast(value); + return Status::OK(); +} + +// Derives the expected past_state / present_state shape and validates past_state against it. +// +// state_window == 0 -> (batch_size, kv_num_heads, d_k, d_v). A single state with no window axis: +// the backward-compatible layout that models exported before the attribute existed use. +// +// state_window == W > 0 -> (W, batch_size, kv_num_heads, d_k, d_v). The window holds the recurrent +// state after each of the last W tokens, right-aligned: slot j is the state after token +// (seq_len - W + j), so slot W-1 is the state after the last token and is the only slot read back +// as past_state. W == 1 is therefore the legacy layout with a leading unit axis. +// +// The window axis leads the batch axis so that each slot is one contiguous +// (batch_size, kv_num_heads, d_k, d_v) block. That keeps "fetch/replace the last state" a single +// contiguous range for any batch size, which is what a speculative decoder needs when it crops +// the state back to an accepted prefix. +template +Status CheckInputs(int state_window, + int batch_size, + int kv_num_heads, + int d_k, + int d_v, + const T* past_state, + TensorShape& state_shape) { + state_shape = state_window > 0 + ? TensorShape({state_window, batch_size, kv_num_heads, d_k, d_v}) + : TensorShape({batch_size, kv_num_heads, d_k, d_v}); + + if (past_state != nullptr && past_state->Shape() != state_shape) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Input 'past_state' is expected to have shape ", state_shape.ToString(), + ", got ", past_state->Shape().ToString(), + ". LinearAttention uses (batch_size, kv_num_heads, d_k, d_v) when the " + "state_window attribute is absent or 0, and " + "(state_window, batch_size, kv_num_heads, d_k, d_v) otherwise."); + } + + return Status::OK(); +} + +} // namespace linear_attention_helper +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc index e60a87afe5f25..b692d051259c4 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc @@ -3,6 +3,7 @@ #include "contrib_ops/cuda/bert/causal_conv_with_state.h" #include "contrib_ops/cuda/bert/causal_conv_with_state_impl.h" +#include "contrib_ops/cpu/bert/causal_conv_with_state_helper.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_type_conversion.h" @@ -35,6 +36,10 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : CudaKern activation_ = info.GetAttrOrDefault("activation", "none"); ORT_ENFORCE(activation_ == "none" || activation_ == "silu" || activation_ == "swish", "activation must be one of: none, silu, swish"); + + // See LinearAttention: only the trailing per-position states are ever consumed, so a window + // caps the allocation and the write traffic for long prompts. 0 keeps the plain single state. + ORT_THROW_IF_ERROR(causal_conv_with_state_helper::ParseStateWindow(info, state_window_)); } template @@ -75,25 +80,29 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { "bias must have shape (", channels, "), got ", bias_shape.ToString()); } - // Validate optional past_state shape - if (past_state_tensor != nullptr) { - const auto& past_shape = past_state_tensor->Shape(); - ORT_RETURN_IF_NOT(past_shape.NumDimensions() == 3, - "past_state must be rank 3 (batch, channels, kernel_size-1), got rank ", past_shape.NumDimensions()); - ORT_RETURN_IF_NOT(past_shape[0] == batch_size && past_shape[1] == channels && past_shape[2] == pad, - "past_state shape mismatch: expected (", batch_size, ", ", channels, ", ", pad, - "), got (", past_shape[0], ", ", past_shape[1], ", ", past_shape[2], ")"); - } + // past_state / present_state are [B, C, K-1], or [W, B, C, K-1] when state_window_ = W > 0. + // Right-aligned: token t lands in slot t + W - L, so slot W-1 always holds the state after the + // last token (and is the slot past_state is read from). + const int state_slots = state_window_ > 0 ? state_window_ : 1; + TensorShape state_shape; + ORT_RETURN_IF_ERROR(causal_conv_with_state_helper::CheckInputs( + state_window_, batch_size, channels, pad, past_state_tensor, state_shape)); // Allocate outputs Tensor* output_tensor = context->Output(0, input_shape); - TensorShape state_shape({batch_size, channels, pad}); Tensor* present_state_tensor = context->Output(1, state_shape); - // Note: no need to zero-initialize present_state — the kernel writes all - // positions unconditionally. When past_state is null, the kernel uses - // zeros for the padding region internally. - // Note: past_state pointer is passed to kernel; kernel reads it directly + // The kernel writes every slot it is responsible for, so present_state normally needs no + // zero-initialization. The exception is a window wider than the sequence: slots below W - L + // belong to positions before this call and are deliberately left alone (that is what bounds the + // per-step work). Zero them when there is no past_state so the output is still fully defined; + // with a past_state the caller owns those slots and is expected to carry them itself. + if (state_window_ > 0 && past_state_tensor == nullptr && state_slots > L) { + CUDA_RETURN_IF_ERROR(cudaMemsetAsync( + present_state_tensor->MutableDataRaw(), 0, + present_state_tensor->SizeInBytes(), + Stream(context))); + } bool apply_silu = (activation_ == "silu" || activation_ == "swish"); @@ -112,7 +121,8 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { L, K, apply_silu, - GetDeviceProp().maxThreadsPerBlock); + GetDeviceProp().maxThreadsPerBlock, + state_slots); } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h index 37a9c29b5e749..f0fb66e8485b9 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h @@ -21,6 +21,8 @@ class CausalConvWithState final : public onnxruntime::cuda::CudaKernel { private: int ndim_; std::string activation_; + // Leading (axis-0) extent of past_state / present_state; 0 means no window axis (single state). + int state_window_; }; } // namespace cuda 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 4f1985f190f93..8c34e672a4797 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 @@ -41,13 +41,14 @@ __global__ void CausalConvDecodeKernel( const T* __restrict__ input, // [B, C, 1] const T* __restrict__ weight, // [C, 1, K] const T* __restrict__ bias, // [C] or nullptr - const T* __restrict__ past_state, // [B, C, K-1] or nullptr + const T* __restrict__ past_state, // [W, B, C, K-1] or nullptr T* __restrict__ output, // [B, C, 1] - T* __restrict__ present_state, // [B, C, K-1] + T* __restrict__ present_state, // [W, B, C, K-1] int batch_channels, // = batch_size * channels (actual element count) int channels, int kernel_size, - bool apply_silu) { + bool apply_silu, + int state_window) { // W: axis-0 extent of past_state / present_state (>= 1) const int bc = blockIdx.x * blockDim.x + threadIdx.x; if (bc >= batch_channels) return; const int b = bc / channels; @@ -58,10 +59,12 @@ __global__ void CausalConvDecodeKernel( // Cache input value in register — avoids redundant global reads const float input_val = to_float(input[(int64_t)b * channels + c]); + // seq_len == 1, so the single position is window slot W-1 for both the read and the write. + // Window-major [W, B, C, K-1]: slot stride is batch_channels*pad and (b, c) flattens to bc. + const int64_t state_offset = (int64_t)(state_window - 1) * batch_channels * pad + (int64_t)bc * pad; + // Cache past_state base pointer for this (b, c) - const T* ps_in = (past_state != nullptr) - ? past_state + (int64_t)b * channels * pad + (int64_t)c * pad - : nullptr; + const T* ps_in = (past_state != nullptr) ? past_state + state_offset : nullptr; // Load weight for this channel: [K] values // weight layout: [C, 1, K], so channel c starts at c * K @@ -82,7 +85,7 @@ __global__ void CausalConvDecodeKernel( output[(int64_t)b * channels + c] = from_float(sum); // Update present_state: shift left by 1, append input - T* ps_out = present_state + (int64_t)b * channels * pad + (int64_t)c * pad; + T* ps_out = present_state + state_offset; for (int k = 0; k < pad - 1; ++k) { ps_out[k] = (ps_in != nullptr) ? ps_in[k + 1] : from_float(0.0f); } @@ -101,7 +104,8 @@ __global__ void CausalConvDecodeKernelFixedK( T* __restrict__ present_state, int batch_channels, int channels, - bool apply_silu) { + bool apply_silu, + int state_window) { // W: axis-0 extent of past_state / present_state (>= 1) const int bc = blockIdx.x * blockDim.x + threadIdx.x; if (bc >= batch_channels) return; @@ -109,11 +113,14 @@ __global__ void CausalConvDecodeKernelFixedK( const int c = bc % channels; constexpr int pad = K - 1; + // seq_len == 1, so the single position is window slot W-1 for both the read and the write. + // Window-major [W, B, C, K-1]: slot stride is batch_channels*pad and (b, c) flattens to bc. + const int64_t state_offset = + static_cast(state_window - 1) * batch_channels * pad + static_cast(bc) * pad; + float sum = (bias != nullptr) ? to_float(bias[c]) : 0.0f; const T* w = weight + static_cast(c) * K; - const T* ps_in = (past_state != nullptr) - ? past_state + static_cast(b) * channels * pad + static_cast(c) * pad - : nullptr; + const T* ps_in = (past_state != nullptr) ? past_state + state_offset : nullptr; if (ps_in != nullptr) { #pragma unroll @@ -128,7 +135,7 @@ __global__ void CausalConvDecodeKernelFixedK( } output[static_cast(b) * channels + c] = from_float(sum); - T* ps_out = present_state + static_cast(b) * channels * pad + static_cast(c) * pad; + T* ps_out = present_state + state_offset; if constexpr (pad > 0) { #pragma unroll for (int k = 0; k < pad - 1; ++k) { @@ -149,13 +156,15 @@ __global__ void CausalConvPrefillKernel( const T* __restrict__ input, // [B, C, L] const T* __restrict__ weight, // [C, 1, K] const T* __restrict__ bias, // [C] or nullptr - const T* __restrict__ past_state, // [B, C, K-1] or nullptr + const T* __restrict__ past_state, // [W, B, C, K-1] or nullptr T* __restrict__ output, // [B, C, L] - T* __restrict__ present_state, // [B, C, K-1] + T* __restrict__ present_state, // [W, B, C, K-1] int seq_len, int channels, int kernel_size, - bool apply_silu) { + bool apply_silu, + int batch_size, + int state_window) { // W: axis-0 extent of past_state / present_state (>= 1) const int b = blockIdx.x; const int c = blockIdx.y; const int tid = threadIdx.x; @@ -163,6 +172,12 @@ __global__ void CausalConvPrefillKernel( const int pad = kernel_size - 1; const int padded_len = pad + seq_len; + // Slot W-1 holds the state after the last token; that is what past_state is read from. + // Window-major, so one slot spans the whole batch. + const int64_t slot_stride = (int64_t)batch_size * channels * pad; + const int64_t last_slot_offset = + (int64_t)(state_window - 1) * slot_stride + ((int64_t)b * channels + c) * pad; + // Shared memory: padded input [pad + L] floats + weight [K] floats extern __shared__ float smem[]; float* s_padded = smem; @@ -172,7 +187,7 @@ __global__ void CausalConvPrefillKernel( // Past state portion: [0..pad-1] for (int i = tid; i < pad; i += blockDim.x) { if (past_state != nullptr) { - s_padded[i] = to_float(past_state[(int64_t)b * channels * pad + (int64_t)c * pad + i]); + s_padded[i] = to_float(past_state[last_slot_offset + i]); } else { s_padded[i] = 0.0f; } @@ -200,11 +215,18 @@ __global__ void CausalConvPrefillKernel( output[((int64_t)b * channels + c) * seq_len + l] = from_float(sum); } - // Save present_state: last K-1 elements of padded input + // Save present_state. The carry state after token t is the pad-length window ending at position + // t in the [past_state, input] stream, i.e. s_padded[t + 1 .. t + pad]; it goes into the + // right-aligned slot t + W - seq_len, and earlier tokens fall outside the window. The last + // token always maps to slot W-1. __syncthreads(); - T* ps = present_state + (int64_t)b * channels * pad + (int64_t)c * pad; - for (int i = tid; i < pad; i += blockDim.x) { - ps[i] = from_float(s_padded[padded_len - pad + i]); + const int first = seq_len > state_window ? seq_len - state_window : 0; + for (int t = first + tid; t < seq_len; t += blockDim.x) { + T* ps = present_state + (int64_t)(t + state_window - seq_len) * slot_stride + + ((int64_t)b * channels + c) * pad; + for (int p = 0; p < pad; ++p) { + ps[p] = from_float(s_padded[t + 1 + p]); + } } } @@ -223,13 +245,15 @@ __global__ void CausalConvPrefillKernelBatched( const T* __restrict__ input, // [B, C, L] const T* __restrict__ weight, // [C, 1, K] const T* __restrict__ bias, // [C] or nullptr - const T* __restrict__ past_state, // [B, C, K-1] or nullptr + const T* __restrict__ past_state, // [W, B, C, K-1] or nullptr T* __restrict__ output, // [B, C, L] - T* __restrict__ present_state, // [B, C, K-1] + T* __restrict__ present_state, // [W, B, C, K-1] int seq_len, int channels, int kernel_size, - bool apply_silu) { + bool apply_silu, + int batch_size, + int state_window) { // W: axis-0 extent of past_state / present_state (>= 1) const int b = blockIdx.x; const int c_base = blockIdx.y * CPB; const int tid = threadIdx.x; @@ -237,6 +261,9 @@ __global__ void CausalConvPrefillKernelBatched( const int pad = kernel_size - 1; const int padded_len = pad + seq_len; + // Window-major [W, B, C, K-1]: one slot spans the whole batch. + const int64_t slot_stride = (int64_t)batch_size * channels * pad; + // Which channel within this block's CPB group does this thread serve? const int threads_per_channel = blockDim.x / CPB; const int local_ch = tid / threads_per_channel; // 0..CPB-1 @@ -250,10 +277,12 @@ __global__ void CausalConvPrefillKernelBatched( float* s_weight = s_padded + padded_len; if (c < channels) { - // Load past state + // Load past state from window slot W-1 (the state after the last token of the previous step) + const int64_t last_slot_offset = + (int64_t)(state_window - 1) * slot_stride + ((int64_t)b * channels + c) * pad; for (int i = local_tid; i < pad; i += threads_per_channel) { if (past_state != nullptr) { - s_padded[i] = to_float(past_state[(int64_t)b * channels * pad + (int64_t)c * pad + i]); + s_padded[i] = to_float(past_state[last_slot_offset + i]); } else { s_padded[i] = 0.0f; } @@ -289,10 +318,15 @@ __global__ void CausalConvPrefillKernelBatched( __syncthreads(); if (c < channels) { - // Save present state - T* ps = present_state + (int64_t)b * channels * pad + (int64_t)c * pad; - for (int i = local_tid; i < pad; i += threads_per_channel) { - ps[i] = from_float(s_padded[padded_len - pad + i]); + // Save the carry state after token t (window s_padded[t+1 .. t+pad]) into the right-aligned + // slot t + W - seq_len; earlier tokens fall outside the window. The last token maps to W-1. + const int first = seq_len > state_window ? seq_len - state_window : 0; + for (int t = first + local_tid; t < seq_len; t += threads_per_channel) { + T* ps = present_state + (int64_t)(t + state_window - seq_len) * slot_stride + + ((int64_t)b * channels + c) * pad; + for (int p = 0; p < pad; ++p) { + ps[p] = from_float(s_padded[t + 1 + p]); + } } } } @@ -313,7 +347,8 @@ Status LaunchCausalConvWithStateKernel( int seq_len, int kernel_size, bool apply_silu, - int max_threads_per_block) { + int max_threads_per_block, + int state_window) { if (seq_len == 1) { // Decode fast-path: one thread per (batch, channel) int total = batch_size * channels; @@ -323,27 +358,27 @@ Status LaunchCausalConvWithStateKernel( case 2: CausalConvDecodeKernelFixedK<<>>( input, weight, bias, past_state, output, present_state, - total, channels, apply_silu); + total, channels, apply_silu, state_window); break; case 3: CausalConvDecodeKernelFixedK<<>>( input, weight, bias, past_state, output, present_state, - total, channels, apply_silu); + total, channels, apply_silu, state_window); break; case 4: CausalConvDecodeKernelFixedK<<>>( input, weight, bias, past_state, output, present_state, - total, channels, apply_silu); + total, channels, apply_silu, state_window); break; case 5: CausalConvDecodeKernelFixedK<<>>( input, weight, bias, past_state, output, present_state, - total, channels, apply_silu); + total, channels, apply_silu, state_window); break; default: CausalConvDecodeKernel<<>>( input, weight, bias, past_state, output, present_state, - total, channels, kernel_size, apply_silu); + total, channels, kernel_size, apply_silu, state_window); break; } } else { @@ -381,7 +416,7 @@ Status LaunchCausalConvWithStateKernel( CausalConvPrefillKernelBatched<<>>( input, weight, bias, past_state, output, present_state, - seq_len, channels, kernel_size, apply_silu); + seq_len, channels, kernel_size, apply_silu, batch_size, state_window); } else { // Original single-channel-per-block path for long sequences const dim3 grid(batch_size, channels, 1); @@ -405,7 +440,7 @@ Status LaunchCausalConvWithStateKernel( CausalConvPrefillKernel<<>>( input, weight, bias, past_state, output, present_state, - seq_len, channels, kernel_size, apply_silu); + seq_len, channels, kernel_size, apply_silu, batch_size, state_window); } } @@ -415,16 +450,16 @@ Status LaunchCausalConvWithStateKernel( // Explicit instantiations template Status LaunchCausalConvWithStateKernel( cudaStream_t, const float*, const float*, const float*, const float*, - float*, float*, int, int, int, int, bool, int); + float*, float*, int, int, int, int, bool, int, int); template Status LaunchCausalConvWithStateKernel( cudaStream_t, const half*, const half*, const half*, const half*, - half*, half*, int, int, int, int, bool, int); + half*, half*, int, int, int, int, bool, int, int); #if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) template Status LaunchCausalConvWithStateKernel<__nv_bfloat16>( cudaStream_t, const __nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*, - __nv_bfloat16*, __nv_bfloat16*, int, int, int, int, bool, int); + __nv_bfloat16*, __nv_bfloat16*, int, int, int, int, bool, int, int); #endif } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h index 4427a1df1fd6d..b07730e6c0d98 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h @@ -20,15 +20,21 @@ Status LaunchCausalConvWithStateKernel( const T* input, // [B, C, L] const T* weight, // [C, 1, K] const T* bias, // [C] or nullptr - const T* past_state, // [B, C, K-1] or nullptr + const T* past_state, // [W, B, C, K-1] or nullptr T* output, // [B, C, L] - T* present_state, // [B, C, K-1] + T* present_state, // [W, B, C, K-1] int batch_size, int channels, int seq_len, int kernel_size, bool apply_silu, - int max_threads_per_block); + int max_threads_per_block, + // Axis-0 extent W of past_state / present_state (>= 1). The window axis leads the batch axis + // so that a slot is one contiguous [B, C, K-1] block. Right-aligned: token t writes slot + // t + W - seq_len and negative slots are skipped, so slot W-1 always holds the state after the + // last token and is the slot past_state is read from. Pass 1 for a plain single-state tensor + // with no window axis. + int state_window = 1); } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 897749a3f6b35..d9b50d318a2b6 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -688,6 +688,8 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons // The fast-decode path lets the flash kernel perform RoPE and KV-append internally, bypassing // PrepareQKV (and therefore the fused QK-Norm prologue). Disable it when q/k norm weights are // present so the regular FlashAttention path (which normalizes via PrepareQKV) is used instead. + // FlashDecoding can handle multi-token decode (sequence_length >= 1): its causal masking and + // split-KV reduction match regular FlashAttention (verified to fp16 tolerance, including MTP-style decode). // It is also disabled for a windowed KV cache: the kernel derives both the absolute RoPE position // and the cache append offset from a single seqlens_k value, which those two no longer share. data.use_flash_attention_fast_decode = use_flash_attention && !disable_flash_decode_ && !parameters.is_first_prompt && parameters.kv_sequence_length > 0 && parameters.past_present_share_buffer && !is_inputs_quantized && !parameters.use_qk_norm && !parameters.is_windowed_kv_cache; diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention.cc b/onnxruntime/contrib_ops/cuda/bert/linear_attention.cc index c8f460b0ca002..fca26a03d36af 100644 --- a/onnxruntime/contrib_ops/cuda/bert/linear_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention.cc @@ -3,6 +3,7 @@ #include "contrib_ops/cuda/bert/linear_attention.h" #include "contrib_ops/cuda/bert/linear_attention_impl.h" +#include "contrib_ops/cpu/bert/linear_attention_helper.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_type_conversion.h" @@ -45,6 +46,11 @@ LinearAttention::LinearAttention(const OpKernelInfo& info) : CudaKernel(info) int64_t chunk_size = info.GetAttrOrDefault("chunk_size", 64); // chunk_size_ reserved for future chunk-parallel prefill algorithm; not yet used. chunk_size_ = static_cast(chunk_size); + + // Only the trailing states are ever consumed (speculative-decoding rollback), while one state + // per token costs d_k*d_v per token per layer -- ~88 GB for a 2.8k prefill on a 30-layer model. + // A window caps both the allocation and the write traffic; 0 keeps the plain 4D single state. + ORT_THROW_IF_ERROR(linear_attention_helper::ParseStateWindow(info, state_window_)); } template @@ -118,31 +124,30 @@ Status LinearAttention::ComputeInternal(OpKernelContext* context) const { TensorShape output_shape({batch_size, seq_len, output_hidden}); Tensor* output_tensor = context->Output(0, output_shape); - TensorShape state_shape({batch_size, kv_num_heads_, d_k, d_v}); + // past_state / present_state are [B, H_kv, d_k, d_v], or [W, B, H_kv, d_k, d_v] when + // state_window_ = W > 0. Right-aligned: token t lands in slot t + W - seq_len, so slot W-1 + // always holds the state after the last token (and is the slot past_state is read from) and, + // when W < seq_len, only the last W states are written. + const int state_slots = state_window_ > 0 ? state_window_ : 1; + TensorShape state_shape; + ORT_RETURN_IF_ERROR(linear_attention_helper::CheckInputs( + state_window_, batch_size, kv_num_heads_, d_k, d_v, past_state_tensor, state_shape)); Tensor* present_state_tensor = context->Output(1, state_shape); - // If past_state is nullptr, zero-init present_state on device + T* present_state_data = present_state_tensor->MutableData(); + const T* initial_state_data = present_state_data; + + // If past_state is nullptr, zero-init the buffer used as the initial state. Only slot W-1 is + // actually read by the kernel, but zeroing the whole thing also defines the slots below + // W - seq_len, which the kernel deliberately leaves alone when the window is wider than the + // sequence (that is what bounds the per-step work). if (past_state_tensor == nullptr) { CUDA_RETURN_IF_ERROR(cudaMemsetAsync( - present_state_tensor->MutableData(), 0, - static_cast(batch_size) * kv_num_heads_ * d_k * d_v * sizeof(T), + present_state_data, 0, + present_state_tensor->SizeInBytes(), Stream(context))); } else { - // Validate past_state shape matches expected (B, H_kv, d_k, d_v) - const auto& past_shape = past_state_tensor->Shape(); - ORT_ENFORCE(past_shape.NumDimensions() == 4, - "past_state must be rank 4 (B, H_kv, d_k, d_v), got rank ", past_shape.NumDimensions()); - ORT_ENFORCE(past_shape[0] == batch_size && past_shape[1] == kv_num_heads_ && - past_shape[2] == d_k && past_shape[3] == d_v, - "past_state shape mismatch: expected (", batch_size, ", ", kv_num_heads_, ", ", d_k, ", ", d_v, - "), got (", past_shape[0], ", ", past_shape[1], ", ", past_shape[2], ", ", past_shape[3], ")"); - // Copy past_state -> present_state (will be updated in-place by kernel) - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync( - present_state_tensor->MutableData(), - past_state_tensor->Data(), - static_cast(batch_size) * kv_num_heads_ * d_k * d_v * sizeof(T), - cudaMemcpyDeviceToDevice, - Stream(context))); + initial_state_data = past_state_tensor->Data(); } typedef typename OrtToCudaType::type CudaT; @@ -155,7 +160,8 @@ Status LinearAttention::ComputeInternal(OpKernelContext* context) const { decay_tensor ? reinterpret_cast(decay_tensor->Data()) : nullptr, beta_tensor ? reinterpret_cast(beta_tensor->Data()) : nullptr, reinterpret_cast(output_tensor->MutableData()), - reinterpret_cast(present_state_tensor->MutableData()), + reinterpret_cast(initial_state_data), + reinterpret_cast(present_state_data), batch_size, seq_len, q_num_heads_, @@ -169,7 +175,8 @@ Status LinearAttention::ComputeInternal(OpKernelContext* context) const { needs_beta, beta_per_head, needs_retrieval, - GetDeviceProp().maxThreadsPerBlock); + GetDeviceProp().maxThreadsPerBlock, + state_slots); } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention.h b/onnxruntime/contrib_ops/cuda/bert/linear_attention.h index ed398218771d0..fe30f14748a17 100644 --- a/onnxruntime/contrib_ops/cuda/bert/linear_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention.h @@ -24,6 +24,8 @@ class LinearAttention final : public onnxruntime::cuda::CudaKernel { std::string update_rule_; float scale_; int chunk_size_; + // Leading (axis-0) extent of past_state / present_state; 0 means no window axis (single state). + int state_window_; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.cc b/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.cc new file mode 100644 index 0000000000000..979120dbc73bd --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.cc @@ -0,0 +1,137 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "contrib_ops/cuda/bert/linear_attention_gates.h" +#include "contrib_ops/cuda/bert/linear_attention_gates_impl.h" +#include "core/providers/cuda/cuda_common.h" +#include "core/providers/cuda/cuda_type_conversion.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; // CudaKernel, OrtToCudaType + +#define REGISTER_KERNEL_TYPED(Op, T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + Op, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TF", DataTypeImpl::GetTensorType()), \ + Op); + +REGISTER_KERNEL_TYPED(LinearAttentionGate, float) +REGISTER_KERNEL_TYPED(LinearAttentionGate, MLFloat16) +REGISTER_KERNEL_TYPED(LinearAttentionGate, BFloat16) + +#undef REGISTER_KERNEL_TYPED + +#define REGISTER_KERNEL_TYPED(Op, T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + Op, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + Op); + +REGISTER_KERNEL_TYPED(GatedRMSNorm, float) +REGISTER_KERNEL_TYPED(GatedRMSNorm, MLFloat16) +REGISTER_KERNEL_TYPED(GatedRMSNorm, BFloat16) + +#undef REGISTER_KERNEL_TYPED + +template +Status LinearAttentionGate::ComputeInternal(OpKernelContext* context) const { + typedef typename OrtToCudaType::type CudaT; + + const Tensor* a = context->Input(0); + const Tensor* dt_bias = context->Input(1); + const Tensor* decay_scale = context->Input(2); + const Tensor* b = context->Input(3); // optional + + const auto& a_shape = a->Shape(); + ORT_RETURN_IF_NOT(a_shape.NumDimensions() >= 1, "a must have rank >= 1"); + const int64_t num_heads = a_shape[a_shape.NumDimensions() - 1]; + ORT_RETURN_IF_NOT(num_heads > 0, "a last dimension must be positive"); + + ORT_RETURN_IF_NOT(dt_bias->Shape().Size() == num_heads, + "dt_bias must have ", num_heads, " elements, got ", dt_bias->Shape().Size()); + ORT_RETURN_IF_NOT(decay_scale->Shape().Size() == num_heads, + "decay_scale must have ", num_heads, " elements, got ", decay_scale->Shape().Size()); + + Tensor* decay = context->Output(0, a_shape); + Tensor* beta = context->Output(1, a_shape); + + if (beta != nullptr) { + ORT_RETURN_IF_NOT(b != nullptr, "the b input is required when the beta output is requested"); + ORT_RETURN_IF_NOT(b->Shape() == a_shape, "b must have the same shape as a"); + } + + const int64_t num_tokens = a_shape.Size() / num_heads; + return LaunchLinearAttentionGateKernel( + Stream(context), + reinterpret_cast(decay->MutableData()), + beta == nullptr ? nullptr : reinterpret_cast(beta->MutableData()), + reinterpret_cast(a->Data()), + b == nullptr ? nullptr : reinterpret_cast(b->Data()), + dt_bias->Data(), + decay_scale->Data(), + num_tokens, + static_cast(num_heads)); +} + +template +GatedRMSNorm::GatedRMSNorm(const OpKernelInfo& info) : CudaKernel(info) { + epsilon_ = info.GetAttrOrDefault("epsilon", 1e-5f); +} + +template +Status GatedRMSNorm::ComputeInternal(OpKernelContext* context) const { + typedef typename OrtToCudaType::type CudaT; + + const Tensor* input = context->Input(0); + const Tensor* scale = context->Input(1); + const Tensor* gate = context->Input(2); + + const auto& shape = input->Shape(); + ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1, "X must have rank >= 1"); + ORT_RETURN_IF_NOT(gate->Shape() == shape, "gate must have the same shape as X"); + + const int64_t norm_size = scale->Shape().Size(); + ORT_RETURN_IF_NOT(norm_size > 0, "scale must not be empty"); + const int64_t last_dim = shape[shape.NumDimensions() - 1]; + ORT_RETURN_IF_NOT(last_dim % norm_size == 0, + "X last dimension (", last_dim, ") must be a multiple of the scale length (", + norm_size, ")"); + + Tensor* output = context->Output(0, shape); + const int64_t num_rows = shape.Size() / norm_size; + + return LaunchGatedRMSNormKernel( + Stream(context), + reinterpret_cast(output->MutableData()), + reinterpret_cast(input->Data()), + reinterpret_cast(scale->Data()), + reinterpret_cast(gate->Data()), + num_rows, + static_cast(norm_size), + epsilon_); +} + +template class LinearAttentionGate; +template class LinearAttentionGate; +template class LinearAttentionGate; +template class GatedRMSNorm; +template class GatedRMSNorm; +template class GatedRMSNorm; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.h b/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.h new file mode 100644 index 0000000000000..6b094b6f8963a --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.h @@ -0,0 +1,34 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/providers/cuda/cuda_kernel.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// decay = decay_scale * Softplus(a + dt_bias), beta = Sigmoid(b). +template +class LinearAttentionGate final : public onnxruntime::cuda::CudaKernel { + public: + explicit LinearAttentionGate(const OpKernelInfo& info) : CudaKernel(info) {} + Status ComputeInternal(OpKernelContext* context) const override; +}; + +// Y = X * rsqrt(mean(X^2) + epsilon) * scale * SiLU(gate). +template +class GatedRMSNorm final : public onnxruntime::cuda::CudaKernel { + public: + explicit GatedRMSNorm(const OpKernelInfo& info); + Status ComputeInternal(OpKernelContext* context) const override; + + private: + float epsilon_; +}; + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.cu b/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.cu new file mode 100644 index 0000000000000..d2b0585948665 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.cu @@ -0,0 +1,176 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Fused gate kernels for the gated-DeltaNet linear-attention layer. +// +// Both fusions exist to remove kernel launches, not to remove FLOPs: the exported graph spends +// eleven launches per layer on tensors of a few thousand elements, purely because the reference +// model computes the gates in float32 while the rest of the graph is float16. Keeping the float32 +// intermediates in registers collapses each chain into a single launch and, under CUDA graphs, +// also returns the per-node replay overhead of the nodes that disappear. + +#include +#include +#include +#include +#include + +#include "contrib_ops/cuda/bert/linear_attention_gates_impl.h" +#include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +namespace { + +// Matches OP_Sigmoid in core/providers/cuda/activation/activations_impl.cu: the branch keeps the +// exponent argument non-positive so large-magnitude inputs cannot overflow. +__device__ __forceinline__ float SigmoidFloat(float x) { + return x > 0.0f ? 1.0f / (1.0f + expf(-x)) : 1.0f - 1.0f / (1.0f + expf(x)); +} + +// Matches OP_Softplus in the same file. +__device__ __forceinline__ float SoftplusFloat(float x) { + return x > 0.0f ? x + logf(expf(-x) + 1.0f) : logf(expf(x) + 1.0f); +} + +template +__global__ void LinearAttentionGateKernel( + T* decay, + T* beta, + const T* a, + const T* b, + const float* dt_bias, + const float* decay_scale, + int64_t count, + int num_heads) { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= count) { + return; + } + + const int h = static_cast(idx % num_heads); + const float biased = to_float(a[idx]) + dt_bias[h]; + decay[idx] = from_float(decay_scale[h] * SoftplusFloat(biased)); + + if (beta != nullptr) { + beta[idx] = from_float(SigmoidFloat(to_float(b[idx]))); + } +} + +// One block per normalization group. The input is read twice (once for the sum of squares, once +// for the output); the group is a few hundred bytes so the second read is an L1 hit. +template +__global__ void GatedRMSNormKernel( + T* output, + const T* input, + const T* scale, + const T* gate, + int norm_size, + float epsilon) { + const int64_t offset = static_cast(blockIdx.x) * norm_size; + const T* x = input + offset; + const T* g = gate + offset; + T* y = output + offset; + + float sum_sq = 0.0f; + for (int i = threadIdx.x; i < norm_size; i += kThreadsPerBlock) { + const float v = to_float(x[i]); + sum_sq += v * v; + } + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + const float total = BlockReduce(temp_storage).Sum(sum_sq); + + __shared__ float shared_inv_rms; + if (threadIdx.x == 0) { + shared_inv_rms = rsqrtf(total / static_cast(norm_size) + epsilon); + } + __syncthreads(); + const float inv_rms = shared_inv_rms; + + for (int i = threadIdx.x; i < norm_size; i += kThreadsPerBlock) { + const float z = to_float(g[i]); + const float normalized = to_float(x[i]) * inv_rms * to_float(scale[i]); + y[i] = from_float(normalized * (z * SigmoidFloat(z))); + } +} + +} // anonymous namespace + +template +Status LaunchLinearAttentionGateKernel( + cudaStream_t stream, + T* decay, + T* beta, + const T* a, + const T* b, + const float* dt_bias, + const float* decay_scale, + int64_t num_tokens, + int num_heads) { + const int64_t count = num_tokens * num_heads; + if (count == 0) { + return Status::OK(); + } + + constexpr int kThreads = 256; + const int64_t blocks = (count + kThreads - 1) / kThreads; + LinearAttentionGateKernel<<(blocks), kThreads, 0, stream>>>( + decay, beta, a, b, dt_bias, decay_scale, count, num_heads); + return CUDA_CALL(cudaGetLastError()); +} + +template +Status LaunchGatedRMSNormKernel( + cudaStream_t stream, + T* output, + const T* input, + const T* scale, + const T* gate, + int64_t num_rows, + int norm_size, + float epsilon) { + if (num_rows == 0) { + return Status::OK(); + } + + const int blocks = static_cast(num_rows); +#define LAUNCH_GATED_RMS_NORM(threads) \ + GatedRMSNormKernel<<>>( \ + output, input, scale, gate, norm_size, epsilon) + + if (norm_size <= 64) { + LAUNCH_GATED_RMS_NORM(64); + } else if (norm_size <= 128) { + LAUNCH_GATED_RMS_NORM(128); + } else if (norm_size <= 256) { + LAUNCH_GATED_RMS_NORM(256); + } else if (norm_size <= 512) { + LAUNCH_GATED_RMS_NORM(512); + } else { + LAUNCH_GATED_RMS_NORM(1024); + } +#undef LAUNCH_GATED_RMS_NORM + + return CUDA_CALL(cudaGetLastError()); +} + +#define INSTANTIATE_LINEAR_ATTENTION_GATES(T) \ + template Status LaunchLinearAttentionGateKernel(cudaStream_t, T*, T*, const T*, const T*, \ + const float*, const float*, int64_t, int); \ + template Status LaunchGatedRMSNormKernel(cudaStream_t, T*, const T*, const T*, const T*, \ + int64_t, int, float); + +INSTANTIATE_LINEAR_ATTENTION_GATES(float) +INSTANTIATE_LINEAR_ATTENTION_GATES(half) +INSTANTIATE_LINEAR_ATTENTION_GATES(__nv_bfloat16) + +#undef INSTANTIATE_LINEAR_ATTENTION_GATES + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.h b/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.h new file mode 100644 index 0000000000000..32b63cc209041 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.h @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include "core/providers/cuda/cuda_common.h" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +// decay = decay_scale * Softplus(a + dt_bias); beta = Sigmoid(b). +// `beta` and `b` may both be nullptr; the two per-head parameter vectors are float32. +template +Status LaunchLinearAttentionGateKernel( + cudaStream_t stream, + T* decay, + T* beta, + const T* a, + const T* b, + const float* dt_bias, + const float* decay_scale, + int64_t num_tokens, + int num_heads); + +// Y = X * rsqrt(mean(X^2) + epsilon) * scale * SiLU(gate), reduced over groups of +// `norm_size` contiguous elements, with all arithmetic in float32. +template +Status LaunchGatedRMSNormKernel( + cudaStream_t stream, + T* output, + const T* input, + const T* scale, + const T* gate, + int64_t num_rows, + int norm_size, + float epsilon); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu index b27641c51c5ad..43b66d34fe85d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.cu @@ -28,6 +28,7 @@ #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" +#include "core/platform/env_var_utils.h" namespace onnxruntime { namespace contrib { @@ -107,7 +108,8 @@ __global__ void LinearAttentionRecurrentKernel( const T* __restrict__ query, // [B, T, H_q * d_k] const T* __restrict__ key, // [B, T, n_k * d_k] const T* __restrict__ value, // [B, T, H_kv * d_v] - T* __restrict__ state, // [B, H_kv, d_k, d_v] — in-place updated + const T* past_state, // [W, B, H_kv, d_k, d_v] -- may alias present_state + T* present_state, // [W, B, H_kv, d_k, d_v] const T* __restrict__ decay, // [B, T, H_kv] or [B, T, H_kv*d_k] or nullptr const T* __restrict__ beta_in, // [B, T, H_kv] or [B, T, 1] or nullptr T* __restrict__ output, // [B, T, max(H_q, H_kv) * d_v] @@ -123,7 +125,9 @@ __global__ void LinearAttentionRecurrentKernel( bool decay_per_key_dim, bool needs_beta, bool beta_per_head, - bool needs_retrieval) { + bool needs_retrieval, + int batch_size, + int state_window) { // W: axis-0 extent of past_state / present_state (>= 1) const int b = blockIdx.x; const int h_kv = blockIdx.y; const int tid = threadIdx.x; @@ -131,8 +135,14 @@ __global__ void LinearAttentionRecurrentKernel( const int kv_per_k = kv_num_heads / n_k_heads; const int h_k = h_kv / kv_per_k; - // Global state pointer for this (batch, head): [d_k, d_v] - T* S_global = state + ((int64_t)b * kv_num_heads + h_kv) * d_k * d_v; + // Global state pointers for this (batch, head): [d_k, d_v] within window slot W-1, the state + // after the last token. They may alias exactly. Window-major layout means one slot spans the + // whole batch, so slot W-1 is a single contiguous [B, H_kv, d_k, d_v] block. + const int64_t slot_stride = (int64_t)kv_num_heads * d_k * d_v; + const int64_t state_offset = ((int64_t)(state_window - 1) * batch_size + b) * slot_stride + + (int64_t)h_kv * d_k * d_v; + const T* S_past = past_state + state_offset; + T* S_present = present_state + state_offset; // Shared memory layout extern __shared__ float smem[]; @@ -142,7 +152,7 @@ __global__ void LinearAttentionRecurrentKernel( // Load state from global memory (type T) into shared memory (fp32) for (int idx = tid; idx < d_k * d_v; idx += num_threads) { - S_smem[idx] = to_float(S_global[idx]); + S_smem[idx] = to_float(S_past[idx]); } __syncthreads(); @@ -276,6 +286,19 @@ __global__ void LinearAttentionRecurrentKernel( } __syncthreads(); + // Emit the recurrent state AFTER processing token t into the right-aligned window slot + // t + W - seq_len (negative => this position falls outside the window and is dropped). + // The last token's slot is always W-1 and is written by the (vectorized) epilogue below. + // Layout [W, B, H_kv, d_k, d_v] row-major; element (i, j) at base_t + i*d_v + j. + const int state_slot = t + state_window - seq_len; + if (state_slot >= 0 && t + 1 < seq_len) { + const int64_t base_t = ((int64_t)state_slot * batch_size + b) * slot_stride + + (int64_t)h_kv * d_k * d_v; + for (int idx = tid; idx < d_k * d_v; idx += num_threads) { + present_state[base_t + idx] = from_float(S_smem[idx]); + } + } + // Step 4: Query readout — output = S^T @ q_t (standard GQA or inverse GQA) if (q_num_heads >= kv_num_heads) { int heads_per_group = q_num_heads / kv_num_heads; @@ -323,7 +346,7 @@ __global__ void LinearAttentionRecurrentKernel( // Write back state from shared memory (fp32) to global memory (type T) for (int idx = tid; idx < d_k * d_v; idx += num_threads) { - S_global[idx] = from_float(S_smem[idx]); + S_present[idx] = from_float(S_smem[idx]); } } @@ -338,7 +361,8 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( const T* __restrict__ query, const T* __restrict__ key, const T* __restrict__ value, - T* __restrict__ state, + const T* past_state, + T* present_state, const T* __restrict__ decay, const T* __restrict__ beta_in, T* __restrict__ output, @@ -352,7 +376,9 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( bool decay_per_key_dim, bool needs_beta, bool beta_per_head, - bool needs_retrieval) { + bool needs_retrieval, + int batch_size, + int state_window) { // W: axis-0 extent of past_state / present_state (>= 1) static_assert(DV % 4 == 0 && DK % 4 == 0, "DK and DV must be multiples of 4 for float4 optimization"); constexpr int DV4 = DV / 4; @@ -362,7 +388,13 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( const int kv_per_k = kv_num_heads / n_k_heads; const int h_k = h_kv / kv_per_k; - T* S_global = state + ((int64_t)b * kv_num_heads + h_kv) * DK * DV; + // Window slot W-1 holds the state after the last token; that is what past_state is read from + // and what the epilogue writes. Window-major, so a slot spans the whole batch. + const int64_t slot_stride = (int64_t)kv_num_heads * DK * DV; + const int64_t state_offset = ((int64_t)(state_window - 1) * batch_size + b) * slot_stride + + (int64_t)h_kv * DK * DV; + const T* S_past = past_state + state_offset; + T* S_present = present_state + state_offset; // Shared memory layout: // S_smem[DK * DV] — recurrent state matrix (fp32) @@ -376,7 +408,7 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( // Load state from global memory (type T) into shared memory (fp32) — vectorized if constexpr (sizeof(T) == 2 && DV % 2 == 0) { // half/bf16: load 2 elements at a time via uint32 - const uint32_t* S_global_u32 = reinterpret_cast(S_global); + const uint32_t* S_global_u32 = reinterpret_cast(S_past); int half_pairs = (DK * DV) / 2; for (int idx = tid; idx < half_pairs; idx += blockDim.x) { uint32_t packed = S_global_u32[idx]; @@ -388,7 +420,7 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( } } else { for (int idx = tid; idx < DK * DV; idx += blockDim.x) { - S_smem[idx] = to_float(S_global[idx]); + S_smem[idx] = to_float(S_past[idx]); } } __syncthreads(); @@ -605,6 +637,19 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( } __syncthreads(); + // Emit the recurrent state AFTER processing token t into the right-aligned window slot + // t + W - seq_len (negative => outside the window, dropped). The last token's slot is always + // W-1 and is written by the vectorized epilogue below. + // Layout [W, B, H_kv, DK, DV] row-major; element (i, j) at base_t + i*DV + j. + const int state_slot = t + state_window - seq_len; + if (state_slot >= 0 && t + 1 < seq_len) { + const int64_t base_t = ((int64_t)state_slot * batch_size + b) * slot_stride + + (int64_t)h_kv * DK * DV; + for (int idx = tid; idx < DK * DV; idx += blockDim.x) { + present_state[base_t + idx] = from_float(S_smem[idx]); + } + } + // ================================================================== // Step 4: Query readout (column dot products — not float4-vectorizable) // ================================================================== @@ -652,7 +697,7 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( // Write back state from shared memory (fp32) to global memory (type T) — vectorized if constexpr (sizeof(T) == 2 && DV % 2 == 0) { - uint32_t* S_global_u32 = reinterpret_cast(S_global); + uint32_t* S_global_u32 = reinterpret_cast(S_present); int half_pairs = (DK * DV) / 2; for (int idx = tid; idx < half_pairs; idx += blockDim.x) { T lo = from_float(S_smem[idx * 2]); @@ -663,7 +708,7 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( S_global_u32[idx] = packed; } } else if constexpr (sizeof(T) == 4 && DV % 4 == 0) { - float4* S_global_f4 = reinterpret_cast(S_global); + float4* S_global_f4 = reinterpret_cast(S_present); int quads = (DK * DV) / 4; for (int idx = tid; idx < quads; idx += blockDim.x) { float4 v; @@ -675,7 +720,7 @@ __global__ void LinearAttentionRecurrentKernelFixedShape( } } else { for (int idx = tid; idx < DK * DV; idx += blockDim.x) { - S_global[idx] = from_float(S_smem[idx]); + S_present[idx] = from_float(S_smem[idx]); } } } @@ -709,7 +754,8 @@ __global__ void LinearAttentionDecodeKernel( const T* __restrict__ query, const T* __restrict__ key, const T* __restrict__ value, - T* __restrict__ state, + const T* past_state, + T* present_state, const T* __restrict__ decay, const T* __restrict__ beta_in, T* __restrict__ output, @@ -724,7 +770,9 @@ __global__ void LinearAttentionDecodeKernel( bool decay_per_key_dim, bool needs_beta, bool beta_per_head, - bool needs_retrieval) { + bool needs_retrieval, + int batch_size, + int state_window) { // W: axis-0 extent of past_state / present_state (>= 1) static_assert(DK % 32 == 0, "DK must be a multiple of warp size (32)"); constexpr int ROWS = DK / 32; @@ -740,11 +788,16 @@ __global__ void LinearAttentionDecodeKernel( 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; + // Window slot W-1 holds the state after the last token; window-major, so a slot spans the batch. + const int64_t slot_stride = (int64_t)kv_num_heads * DK * d_v; + const int64_t state_offset = ((int64_t)(state_window - 1) * batch_size + b) * slot_stride + + (int64_t)h_kv * DK * d_v + col; + const T* S_past_col = past_state + state_offset; + T* S_present_col = present_state + state_offset; 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]); + s_shard[r] = to_float(S_past_col[(int64_t)(r * 32 + lane) * d_v]); } const int k_hidden = n_k_heads * DK; @@ -802,6 +855,20 @@ __global__ void LinearAttentionDecodeKernel( s_shard[r] += k_reg[r] * delta_col; } + // Emit the recurrent state AFTER processing token t into the right-aligned window slot + // t + W - seq_len (negative => outside the window, dropped). The last token's slot is always + // W-1 and is written by the epilogue below. This lane owns rows {r*32 + lane} of column + // `col`. Layout [W, B, H_kv, DK, d_v] row-major; element (row, col) at base_t + row*d_v + col. + const int state_slot = t + state_window - seq_len; + if (state_slot >= 0 && t + 1 < seq_len) { + const int64_t base_t = ((int64_t)state_slot * batch_size + b) * slot_stride + + (int64_t)h_kv * DK * d_v; +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + present_state[base_t + (int64_t)(r * 32 + lane) * d_v + col] = from_float(s_shard[r]); + } + } + // Readout: output = scale * sum_i S[i][col] * q[i]. const int head_count = GetLinearAttentionReadoutHeadCount(q_num_heads, kv_num_heads); for (int group_index = 0; group_index < head_count; ++group_index) { @@ -823,7 +890,7 @@ __global__ void LinearAttentionDecodeKernel( // 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]); + S_present_col[(int64_t)(r * 32 + lane) * d_v] = from_float(s_shard[r]); } } @@ -857,7 +924,8 @@ __global__ void LinearAttentionDecodeColKernel( const T* __restrict__ query, const T* __restrict__ key, const T* __restrict__ value, - T* __restrict__ state, + const T* past_state, + T* present_state, const T* __restrict__ decay, const T* __restrict__ beta_in, T* __restrict__ output, @@ -872,7 +940,10 @@ __global__ void LinearAttentionDecodeColKernel( bool decay_per_key_dim, bool needs_beta, bool beta_per_head, - bool needs_retrieval) { + bool needs_retrieval, + bool force_sequential_state_roundtrip, + int batch_size, + int state_window) { // W: axis-0 extent of past_state / present_state (>= 1) const int b = blockIdx.x; const int h_kv = blockIdx.y; const int tid = threadIdx.x; @@ -883,12 +954,17 @@ __global__ void LinearAttentionDecodeColKernel( 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; + // This thread owns column `col`: S[i][col] lives at i*d_v + col (row-major) within window + // slot W-1, the state after the last token. Window-major, so a slot spans the whole batch. + const int64_t slot_stride = (int64_t)kv_num_heads * DK * d_v; + const int64_t state_offset = ((int64_t)(state_window - 1) * batch_size + b) * slot_stride + + (int64_t)h_kv * DK * d_v + col; + const T* S_past_head = past_state + state_offset; + T* S_present_head = present_state + state_offset; 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]); + s_col[i] = to_float(S_past_head[(int64_t)i * d_v]); } // Per-token broadcasts shared across all columns of this (b, h_kv). @@ -952,6 +1028,20 @@ __global__ void LinearAttentionDecodeColKernel( s_col[i] += k_sh[i] * delta_col; } + // Emit the recurrent state AFTER processing token t into the right-aligned window slot + // t + W - seq_len (negative => outside the window, dropped). The last token's slot is always + // W-1 and is written by the coalesced epilogue below. This thread owns column `col`. + // Layout [W, B, H_kv, DK, d_v] row-major; element (i, col) at base_t + i*d_v + col. + const int state_slot = t + state_window - seq_len; + if (state_slot >= 0 && t + 1 < seq_len) { + const int64_t base_t = ((int64_t)state_slot * batch_size + b) * slot_stride + + (int64_t)h_kv * DK * d_v; +#pragma unroll + for (int i = 0; i < DK; ++i) { + present_state[base_t + (int64_t)i * d_v + col] = from_float(s_col[i]); + } + } + // Readout: output = scale * sum_i S[i][col] * q[i]. const int head_count = GetLinearAttentionReadoutHeadCount(q_num_heads, kv_num_heads); for (int group_index = 0; group_index < head_count; ++group_index) { @@ -969,13 +1059,19 @@ __global__ void LinearAttentionDecodeColKernel( } output[bt * output_hidden + readout_heads.output_head * d_v + col] = from_float(scale * acc); } + if (force_sequential_state_roundtrip && t + 1 < seq_len) { +#pragma unroll + for (int i = 0; i < DK; ++i) { + s_col[i] = to_float(from_float(s_col[i])); + } + } __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]); + S_present_head[(int64_t)i * d_v] = from_float(s_col[i]); } } @@ -990,6 +1086,7 @@ Status LaunchLinearAttentionKernel( const T* decay, const T* beta, T* output, + const T* past_state, T* present_state, int batch_size, int seq_len, @@ -1004,7 +1101,8 @@ Status LaunchLinearAttentionKernel( bool needs_beta, bool beta_per_head, bool needs_retrieval, - int max_threads_per_block) { + int max_threads_per_block, + int state_window) { // Grid: one block per (batch, kv_head) const dim3 grid(batch_size, kv_num_heads, 1); @@ -1028,6 +1126,8 @@ Status LaunchLinearAttentionKernel( // 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 bool force_sequential_state_roundtrip = + ParseEnvironmentVariableWithDefault("ORT_LINEAR_ATTENTION_FORCE_SEQUENTIAL_STATE_ROUNDTRIP", false); const dim3 decode_grid(batch_size, kv_num_heads, (d_v + kColsPerBlock - 1) / kColsPerBlock); const dim3 decode_block(kColsPerBlock, 1, 1); @@ -1035,9 +1135,10 @@ Status LaunchLinearAttentionKernel( auto launch_col = [&](auto dk_tag) -> Status { constexpr int DK = decltype(dk_tag)::value; LinearAttentionDecodeColKernel<<>>( - query, key, value, present_state, decay, beta, output, + query, key, value, past_state, 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); + needs_decay, decay_per_key_dim, needs_beta, beta_per_head, needs_retrieval, + force_sequential_state_roundtrip, batch_size, state_window); return CUDA_CALL(cudaGetLastError()); }; @@ -1055,9 +1156,9 @@ Status LaunchLinearAttentionKernel( auto launch_decode = [&](auto dk_tag) -> Status { constexpr int DK = decltype(dk_tag)::value; LinearAttentionDecodeKernel<<>>( - query, key, value, present_state, decay, beta, output, + query, key, value, past_state, 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); + needs_decay, decay_per_key_dim, needs_beta, beta_per_head, needs_retrieval, batch_size, state_window); return CUDA_CALL(cudaGetLastError()); }; @@ -1089,9 +1190,9 @@ Status LaunchLinearAttentionKernel( } LinearAttentionRecurrentKernelFixedShape<<>>( - query, key, value, present_state, decay, beta, output, + query, key, value, past_state, present_state, decay, beta, output, seq_len, q_num_heads, kv_num_heads, n_k_heads, output_hidden, scale, - needs_decay, decay_per_key_dim, needs_beta, beta_per_head, needs_retrieval); + needs_decay, decay_per_key_dim, needs_beta, beta_per_head, needs_retrieval, batch_size, state_window); return CUDA_CALL(cudaGetLastError()); }; @@ -1136,9 +1237,9 @@ Status LaunchLinearAttentionKernel( } LinearAttentionRecurrentKernel<<>>( - query, key, value, present_state, decay, beta, output, + query, key, value, past_state, present_state, decay, beta, output, seq_len, q_num_heads, kv_num_heads, n_k_heads, d_k, d_v, output_hidden, scale, - needs_decay, decay_per_key_dim, needs_beta, beta_per_head, needs_retrieval); + needs_decay, decay_per_key_dim, needs_beta, beta_per_head, needs_retrieval, batch_size, state_window); return CUDA_CALL(cudaGetLastError()); } @@ -1146,19 +1247,19 @@ Status LaunchLinearAttentionKernel( // Explicit instantiations template Status LaunchLinearAttentionKernel( cudaStream_t, const float*, const float*, const float*, - const float*, const float*, float*, float*, - int, int, int, int, int, int, int, float, bool, bool, bool, bool, bool, int); + const float*, const float*, float*, const float*, float*, + int, int, int, int, int, int, int, float, bool, bool, bool, bool, bool, int, int); template Status LaunchLinearAttentionKernel( cudaStream_t, const half*, const half*, const half*, - const half*, const half*, half*, half*, - int, int, int, int, int, int, int, float, bool, bool, bool, bool, bool, int); + const half*, const half*, half*, const half*, half*, + int, int, int, int, int, int, int, float, bool, bool, bool, bool, bool, int, int); #if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__) template Status LaunchLinearAttentionKernel<__nv_bfloat16>( cudaStream_t, const __nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*, - const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, __nv_bfloat16*, - int, int, int, int, int, int, int, float, bool, bool, bool, bool, bool, int); + const __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, const __nv_bfloat16*, __nv_bfloat16*, + int, int, int, int, int, int, int, float, bool, bool, bool, bool, bool, int, int); #endif } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.h index c30e081e6cd2f..ff609f4d4c64e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.h @@ -16,13 +16,14 @@ namespace cuda { template Status LaunchLinearAttentionKernel( cudaStream_t stream, - const T* query, // [B, T, H_q * d_k] - const T* key, // [B, T, n_k * d_k] - const T* value, // [B, T, H_kv * d_v] - const T* decay, // [B, T, H_kv] or [B, T, H_kv * d_k] or nullptr - const T* beta, // [B, T, H_kv] or [B, T, 1] or nullptr - T* output, // [B, T, max(H_q, H_kv) * d_v] - T* present_state, // [B, H_kv, d_k, d_v] -- in-place (caller pre-fills from past) + const T* query, // [B, T, H_q * d_k] + const T* key, // [B, T, n_k * d_k] + const T* value, // [B, T, H_kv * d_v] + const T* decay, // [B, T, H_kv] or [B, T, H_kv * d_k] or nullptr + const T* beta, // [B, T, H_kv] or [B, T, 1] or nullptr + T* output, // [B, T, max(H_q, H_kv) * d_v] + const T* past_state, // [W, B, H_kv, d_k, d_v] -- may alias present_state + T* present_state, // [W, B, H_kv, d_k, d_v] int batch_size, int seq_len, int q_num_heads, @@ -36,7 +37,13 @@ Status LaunchLinearAttentionKernel( bool needs_beta, bool beta_per_head, bool needs_retrieval, - int max_threads_per_block); + int max_threads_per_block, + // Axis-0 extent W of past_state / present_state (>= 1). The window axis leads the batch axis + // so that a slot is one contiguous [B, H_kv, d_k, d_v] block. Right-aligned: token t writes + // slot t + W - seq_len and slots with a negative index are skipped, so slot W-1 always holds + // the state after the last token and is the slot past_state is read from. Pass 1 for a plain + // single-state tensor with no window axis. + int state_window = 1); } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 1c898ea5912a8..b4669dd71e190 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc @@ -149,6 +149,12 @@ class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, RotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, GemmaRotaryEmbedding); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttention); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttention); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttentionGate); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttentionGate); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, LinearAttentionGate); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, GatedRMSNorm); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, GatedRMSNorm); +class CUDA_MS_OP_TYPED_CLASS_NAME(1, BFloat16, GatedRMSNorm); class CUDA_MS_OP_TYPED_CLASS_NAME(1, float, CausalConvWithState); class CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, CausalConvWithState); #if !defined(DISABLE_GENERATION_OPS) @@ -414,6 +420,12 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, #if !defined(DISABLE_GENERATION_OPS) diff --git a/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc b/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc index 67c1693918267..813d0338019bf 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc @@ -40,6 +40,8 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) std::string activation_str = info.GetAttrOrDefault("activation", "none"); activation_ = ParseCausalConvActivation(activation_str); ORT_ENFORCE(info.GetAttr("ndim", &ndim_).IsOK(), "Attribute 'ndim' is required"); + ORT_ENFORCE(info.GetAttrOrDefault("state_window", 0) == 0, + "WebGPU CausalConvWithState does not support state_window > 0 (CUDA EP only)"); } Status CausalConvWithStateProgram::GenerateShaderCode(ShaderHelper& shader) const { diff --git a/onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc index d71c069d527ed..57c8a56e06f8c 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc @@ -108,6 +108,8 @@ LinearAttention::LinearAttention(const OpKernelInfo& info) scale_ = info.GetAttrOrDefault("scale", 0.0f); q_num_heads_ = static_cast(info.GetAttr("q_num_heads")); kv_num_heads_ = static_cast(info.GetAttr("kv_num_heads")); + ORT_ENFORCE(info.GetAttrOrDefault("state_window", 0) == 0, + "WebGPU LinearAttention does not support state_window > 0 (CUDA EP only)"); } /* diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index f088fe226cfce..97fc6ba0226ff 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -2293,6 +2293,19 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Spatial dimensionality: 1, 2, or 3. Default is 1.", AttributeProto::INT, static_cast(1)) + .Attr("state_window", + "Number of trailing per-position carry states held by past_state and present_state. " + "When 0 (default) the state tensors have no window axis and hold only the state after " + "the last position, i.e. the backward-compatible (batch_size, channels, k_1 - 1). " + "When W > 0 both gain a LEADING axis of extent W, right-aligned: slot j is the state " + "after position (seq_len - W + j), so slot W-1 is always the state after the last " + "position (identical to the W = 0 tensor) and is the slot past_state is read from. " + "The window axis leads the batch axis so that each slot is one contiguous " + "(batch_size, channels, k_1 - 1) block. Slots below max(0, W - seq_len) are not " + "written. A window lets a speculative decoder roll the state back to an accepted " + "prefix without replaying the forward.", + AttributeProto::INT, + static_cast(0)) .Input(0, "input", "Input tensor with shape (batch_size, channels, ...). Channels-first layout. " @@ -2310,8 +2323,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( OpSchema::Optional) .Input(3, "past_state", - "Carry state from previous step. For ndim=1: (batch_size, channels, k_1 - 1). " - "If not provided, padding is zero.", + "Carry state from previous step. For ndim=1: (batch_size, channels, k_1 - 1), or " + "(W, batch_size, channels, k_1 - 1) when state_window = W > 0, in which case only " + "slot W-1 is read. If not provided, padding is zero.", "T", OpSchema::Optional) .Output(0, @@ -2320,8 +2334,10 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "T") .Output(1, "present_state", - "Updated carry state. For ndim=1: (batch_size, channels, k_1 - 1). " - "Contains the last (k-1) values from the virtual input along the causal axis.", + "Updated carry state. For ndim=1: (batch_size, channels, k_1 - 1), or " + "(W, batch_size, channels, k_1 - 1) when state_window = W > 0. Slot W-1 contains " + "the last (k-1) values from the virtual input along the causal axis; slot j contains " + "the same for the prefix ending at position (seq_len - W + j).", "T") .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, @@ -2347,7 +2363,15 @@ ONNX_MS_OPERATOR_SET_SCHEMA( fail_shape_inference("CausalConvWithState: weight must have rank >= 2"); } int64_t ndim = getAttribute(ctx, "ndim", 1); + // state_window = W > 0 prepends a window axis, holding the carry state after each of + // the last W positions (slot W-1 == the W = 0 tensor). The window axis leads the batch + // axis so a slot is one contiguous (batch_size, channels, ...) block. W = 0 keeps the + // legacy (batch_size, channels, ...) shape for backward compatibility. + const int64_t window = getAttribute(ctx, "state_window", 0); TensorShapeProto state_shape; + if (window > 0) { + state_shape.add_dim()->set_dim_value(window); + } *state_shape.add_dim() = input_shape.dim(0); // batch_size *state_shape.add_dim() = input_shape.dim(1); // channels // Copy non-causal spatial dims from input (dims 2 .. 2+ndim-2) @@ -2408,6 +2432,18 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Tuning hint; does not affect output correctness.", AttributeProto::INT, static_cast(64)) + .Attr("state_window", + "Number of trailing per-token recurrent states held by past_state and present_state. " + "When 0 (default) the state tensors are 4D and hold only the state after the last " + "token, i.e. the backward-compatible (B, H_kv, d_k, d_v). When W > 0 both are 5D with " + "a LEADING axis of extent W, right-aligned: slot j is the state after token " + "(T - W + j), so slot W-1 is always the state after the last token (identical to the " + "W = 0 tensor) and is the slot past_state is read from. The window axis leads the " + "batch axis so that each slot is one contiguous (B, H_kv, d_k, d_v) block. Slots " + "below max(0, W - T) are not written. A window lets a speculative decoder roll the " + "state back to an accepted prefix without replaying the forward.", + AttributeProto::INT, + static_cast(0)) .Input(0, "query", "Query vectors with 3D packed shape (B, T, H_q * d_k). " @@ -2424,8 +2460,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "T") .Input(3, "past_state", - "Recurrent state from previous step with shape (B, H_kv, d_k, d_v). " - "Always 4D. If not provided, defaults to zeros.", + "Recurrent state from previous step with shape (B, H_kv, d_k, d_v), or " + "(W, B, H_kv, d_k, d_v) when state_window = W > 0, in which case only slot W-1 is " + "read. If not provided, defaults to zeros.", "S", OpSchema::Optional) .Input(4, @@ -2449,7 +2486,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "T") .Output(1, "present_state", - "Updated recurrent state with shape (B, H_kv, d_k, d_v). Always 4D.", + "Updated recurrent state with shape (B, H_kv, d_k, d_v), or (W, B, H_kv, d_k, d_v) " + "when state_window = W > 0. Slot W-1 is the state after the last token; slot j is " + "the state after token (T - W + j).", "S") .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, @@ -2490,7 +2529,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( updateOutputShape(ctx, 0, output_shape); } - // Output 1: present_state shape (B, H_kv, d_k, d_v) — 4D + // Output 1: present_state shape (B, H_kv, d_k, d_v) — 4D, or (W, B, H_kv, d_k, d_v) — 5D + // when state_window = W > 0. W = 0 keeps the legacy 4D shape for backward compatibility. if (hasInputShape(ctx, 0) && hasInputShape(ctx, 2) && q_num_heads > 0 && kv_num_heads > 0) { auto& query_shape = getInputShape(ctx, 0); auto& value_shape = getInputShape(ctx, 2); @@ -2498,7 +2538,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA( // Already validated in Output 0 block above; skip if shapes are invalid. return; } + const int64_t window = getAttribute(ctx, "state_window", 0); TensorShapeProto state_shape; + if (window > 0) { + state_shape.add_dim()->set_dim_value(window); // W + } *state_shape.add_dim() = query_shape.dim(0); // B state_shape.add_dim()->set_dim_value(kv_num_heads); // H_kv // d_k = query.dim(2) / q_num_heads @@ -2519,5 +2563,121 @@ ONNX_MS_OPERATOR_SET_SCHEMA( } })); +constexpr const char* LinearAttentionGate_ver1_doc = R"DOC( +Fuses the gate projections that feed LinearAttention's gated-delta recurrence: + + decay = decay_scale * Softplus(a + dt_bias) + beta = Sigmoid(b) (only when b is provided) + +Reference implementations compute the decay in float32 because exp(decay) inside the +recurrence exponentially amplifies any precision loss. Exporters therefore emit +Cast -> Add -> Softplus -> Mul -> Cast, which is five kernel launches on a tensor with +only num_heads elements per token. This operator keeps the intermediates in float32 +registers so a single launch replaces the whole chain. + +dt_bias and decay_scale are float32 per-head vectors of length H. decay_scale is the +already-negated -exp(A_log) factor. +)DOC"; + +ONNX_MS_OPERATOR_SET_SCHEMA( + LinearAttentionGate, 1, + OpSchema() + .SetDoc(LinearAttentionGate_ver1_doc) + .Input(0, + "a", + "Decay gate projection with shape (B, T, H).", + "T") + .Input(1, + "dt_bias", + "Per-head float32 bias added to a, with shape (H).", + "TF") + .Input(2, + "decay_scale", + "Per-head float32 multiplier applied to Softplus(a + dt_bias), with shape (H). " + "For gated DeltaNet this is -exp(A_log).", + "TF") + .Input(3, + "b", + "Update-rate projection with shape (B, T, H). Required when the beta output is requested.", + "T", + OpSchema::Optional) + .Output(0, + "decay", + "decay_scale * Softplus(a + dt_bias) with shape (B, T, H).", + "T") + .Output(1, + "beta", + "Sigmoid(b) with shape (B, T, H).", + "T", + OpSchema::Optional) + .TypeConstraint("T", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain gate input and output types to float tensors.") + .TypeConstraint("TF", + {"tensor(float)"}, + "Constrain the per-head parameters to float32.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + propagateShapeFromInputToOutput(ctx, 0, 0); + if (ctx.getNumOutputs() > 1) { + if (ctx.getNumInputs() < 4 || ctx.getInputType(3) == nullptr) { + fail_shape_inference("The b input is required when the beta output is requested."); + } + propagateElemTypeFromInputToOutput(ctx, 0, 1); + if (hasInputShape(ctx, 3)) { + propagateShapeFromInputToOutput(ctx, 3, 1); + } else { + propagateShapeFromInputToOutput(ctx, 0, 1); + } + } + })); + +constexpr const char* GatedRMSNorm_ver1_doc = R"DOC( +Gated RMS normalization as used by Mamba2 / gated DeltaNet attention outputs: + + Y = X * rsqrt(mean(X^2) + epsilon) * scale * SiLU(gate) + +The mean of squares is taken over the trailing `C` elements of each row, where `C` is the +length of `scale`; the input's last dimension must be a multiple of `C`, which lets a +per-head norm run on a packed (B, T, H * C) tensor without any surrounding Reshape. +All arithmetic including SiLU is done in float32 regardless of the tensor type, matching +the reference implementation, so this replaces the exported +SimplifiedLayerNormalization -> Cast -> Sigmoid -> Mul -> Cast -> Mul -> Cast chain with a +single launch. +)DOC"; + +ONNX_MS_OPERATOR_SET_SCHEMA( + GatedRMSNorm, 1, + OpSchema() + .SetDoc(GatedRMSNorm_ver1_doc) + .Attr("epsilon", + "Epsilon added to the mean of squares before the reciprocal square root.", + AttributeProto::FLOAT, + 1e-5f) + .Input(0, + "X", + "Input tensor with shape (..., H * C). Normalization is applied over each " + "contiguous group of C elements.", + "T") + .Input(1, + "scale", + "Normalization weight with shape (C).", + "T") + .Input(2, + "gate", + "Gate tensor with the same shape as X.", + "T") + .Output(0, + "Y", + "Output tensor with the same shape as X.", + "T") + .TypeConstraint("T", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain input and output types to float tensors.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + propagateShapeFromInputToOutput(ctx, 0, 0); + })); + } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/core/graph/contrib_ops/ms_opset.h b/onnxruntime/core/graph/contrib_ops/ms_opset.h index 84ee0acd0bc8e..5cc5bacab0077 100644 --- a/onnxruntime/core/graph/contrib_ops/ms_opset.h +++ b/onnxruntime/core/graph/contrib_ops/ms_opset.h @@ -89,6 +89,8 @@ class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MultiHeadAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GroupQueryAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, PagedAttention); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, LinearAttention); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, LinearAttentionGate); +class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, GatedRMSNorm); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, CausalConvWithState); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, MurmurHash3); class ONNX_OPERATOR_SET_SCHEMA_CLASS_NAME(Microsoft, 1, NGramRepeatBlock); @@ -204,6 +206,8 @@ class OpSet_Microsoft_ver1 { fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); + fn(GetOpSchema()); + fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); fn(GetOpSchema()); diff --git a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc index 2a7837dd1ce73..e7fc44591049a 100644 --- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc +++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc @@ -654,5 +654,144 @@ TEST(CausalConvWithStateTest, LargerDimensions) { batch_size, channels, input_length, kernel_size, "silu"); } +#ifdef USE_CUDA +// The state_window attribute is only implemented by the CUDA kernel. past_state / present_state +// then hold the last W per-position carry states, right-aligned; slot W-1 is the state after the +// last position (what the unwindowed op produces) and is the slot past_state is read from. +// +// The CUDA kernel has four families (fixed-K decode, generic-K decode, batched prefill, +// single-channel prefill) and each computes the slot offsets itself, so every family gets its own +// shape below. `window` > `input_length` is only exercised without a past_state, because that is +// the only case in which the slots below W - L are defined (zeroed rather than left alone). +static void RunCausalConvStateWindowTest(int batch_size, int channels, int input_length, + int kernel_size, int window, bool with_past_state) { + auto ep = DefaultCudaExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "CUDA execution provider not available"; + return; + } + + const int state_length = kernel_size - 1; + const std::string activation = "silu"; + + std::vector input_data(static_cast(batch_size) * channels * input_length); + for (int i = 0; i < static_cast(input_data.size()); ++i) { + input_data[i] = 0.5f * std::sin(static_cast(i) * 0.37f); + } + std::vector weight_data(static_cast(channels) * kernel_size); + for (int i = 0; i < static_cast(weight_data.size()); ++i) { + weight_data[i] = 0.25f * std::cos(static_cast(i) * 0.21f); + } + std::vector bias_data(channels); + for (int i = 0; i < channels; ++i) { + bias_data[i] = 0.01f * static_cast(i); + } + std::vector conv_state_data(static_cast(batch_size) * channels * state_length); + for (int i = 0; i < static_cast(conv_state_data.size()); ++i) { + conv_state_data[i] = with_past_state ? 0.5f * std::sin(static_cast(i) * 0.3f) : 0.0f; + } + const std::vector* past = with_past_state ? &conv_state_data : nullptr; + + // Keep only the first `prefix` positions of a (B, C, L) tensor. + auto slice_prefix = [&](const std::vector& src, int prefix) { + std::vector dst(static_cast(batch_size) * channels * prefix); + for (int bc = 0; bc < batch_size * channels; ++bc) { + for (int l = 0; l < prefix; ++l) { + dst[static_cast(bc) * prefix + l] = src[static_cast(bc) * input_length + l]; + } + } + return dst; + }; + + std::vector expected_output, expected_state; + CausalConvWithStateReference( + input_data, weight_data, &bias_data, past, + expected_output, expected_state, + batch_size, channels, input_length, kernel_size, activation); + + // Slot j holds the state after the first (input_length - window + j + 1) positions; slots for + // non-positive prefixes are never computed by the kernel and stay zero. The window axis leads + // the batch axis, so a slot is exactly one contiguous (B, C, K-1) reference block. + const size_t slot_elems = static_cast(channels) * state_length; + const size_t batch_slot_elems = static_cast(batch_size) * slot_elems; + std::vector expected_state_window(static_cast(window) * batch_slot_elems, 0.0f); + for (int j = 0; j < window; ++j) { + const int prefix = input_length - window + j + 1; + if (prefix <= 0) continue; + std::vector prefix_state; + if (prefix == input_length) { + prefix_state = expected_state; + } else { + std::vector prefix_output; + CausalConvWithStateReference( + slice_prefix(input_data, prefix), weight_data, &bias_data, past, + prefix_output, prefix_state, + batch_size, channels, prefix, kernel_size, activation); + } + std::copy_n(prefix_state.begin(), batch_slot_elems, + expected_state_window.begin() + static_cast(j) * batch_slot_elems); + } + + OpTester test("CausalConvWithState", 1, onnxruntime::kMSDomain); + test.AddAttribute("activation", activation); + test.AddAttribute("state_window", static_cast(window)); + + test.AddInput("input", {batch_size, channels, input_length}, input_data); + test.AddInput("weight", {channels, 1, kernel_size}, weight_data); + test.AddInput("bias", {channels}, bias_data); + if (with_past_state) { + // past_state is windowed too, and only slot W-1 is read. Poison the earlier slots to prove it. + std::vector past_state_window(static_cast(window) * batch_slot_elems, -1e4f); + std::copy_n(conv_state_data.begin(), batch_slot_elems, + past_state_window.begin() + static_cast(window - 1) * batch_slot_elems); + test.AddInput("past_state", {window, batch_size, channels, state_length}, past_state_window); + } else { + test.AddOptionalInputEdge(); + } + + test.AddOutput("output", {batch_size, channels, input_length}, expected_output); + test.AddOutput("present_state", {window, batch_size, channels, state_length}, + expected_state_window); + test.SetOutputAbsErr("output", 0.01f); + test.SetOutputAbsErr("present_state", 0.01f); + + std::vector> execution_providers; + execution_providers.push_back(std::move(ep)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// channels < 4 forces the single-channel prefill kernel. +TEST(CausalConvWithStateTest, StateWindow) { + RunCausalConvStateWindowTest(/*batch_size=*/1, /*channels=*/3, /*input_length=*/5, + /*kernel_size=*/4, /*window=*/3, /*with_past_state=*/true); +} + +// 2 <= L <= 128 with channels >= 4 selects the batched prefill kernel; B = 2 exercises the batch +// stride of the window axis. +TEST(CausalConvWithStateTest, StateWindow_BatchedPrefill) { + RunCausalConvStateWindowTest(/*batch_size=*/2, /*channels=*/8, /*input_length=*/6, + /*kernel_size=*/4, /*window=*/3, /*with_past_state=*/true); +} + +// L > 128 falls back to the single-channel prefill kernel even for wide inputs. +TEST(CausalConvWithStateTest, StateWindow_LongPrefill) { + RunCausalConvStateWindowTest(/*batch_size=*/2, /*channels=*/8, /*input_length=*/140, + /*kernel_size=*/4, /*window=*/5, /*with_past_state=*/true); +} + +// L == 1 with kernel_size in [2, 5] selects the compile-time specialized decode kernel. The window +// is wider than the sequence here, which is the shape genai actually runs during MTP decode. +TEST(CausalConvWithStateTest, StateWindow_DecodeFixedK) { + RunCausalConvStateWindowTest(/*batch_size=*/2, /*channels=*/8, /*input_length=*/1, + /*kernel_size=*/4, /*window=*/3, /*with_past_state=*/false); +} + +// kernel_size > 5 falls back to the generic decode kernel. +TEST(CausalConvWithStateTest, StateWindow_DecodeGenericK) { + RunCausalConvStateWindowTest(/*batch_size=*/2, /*channels=*/8, /*input_length=*/1, + /*kernel_size=*/7, /*window=*/3, /*with_past_state=*/false); +} +#endif // USE_CUDA + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/linear_attention_gates_op_test.cc b/onnxruntime/test/contrib_ops/linear_attention_gates_op_test.cc new file mode 100644 index 0000000000000..41def5b5d06d0 --- /dev/null +++ b/onnxruntime/test/contrib_ops/linear_attention_gates_op_test.cc @@ -0,0 +1,236 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Tests for the two fused linear-attention gate ops: LinearAttentionGate and GatedRMSNorm. +// Both replace float32 elementwise chains that an exporter emits around the LinearAttention op, +// so the references here are the same float32 formulas ORT's Softplus/Sigmoid/RMSNorm kernels use. + +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "test/common/cuda_op_test_utils.h" +#include "test/common/tensor_op_test_utils.h" +#include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" + +namespace onnxruntime { +namespace test { + +namespace { + +float SigmoidRef(float x) { + return x > 0.0f ? 1.0f / (1.0f + std::exp(-x)) : 1.0f - 1.0f / (1.0f + std::exp(x)); +} + +float SoftplusRef(float x) { + return x > 0.0f ? x + std::log(std::exp(-x) + 1.0f) : std::log(std::exp(x) + 1.0f); +} + +std::vector RandomFloats(size_t count, float lo, float hi, uint32_t seed) { + std::mt19937 gen(seed); + std::uniform_real_distribution dist(lo, hi); + std::vector out(count); + for (auto& v : out) { + v = dist(gen); + } + return out; +} + +template +std::vector ToTensorType(const std::vector& data) { + if constexpr (std::is_same_v) { + return ToFloat16(data); + } else if constexpr (std::is_same_v) { + return ToBFloat16(data); + } else { + return data; + } +} + +template +void RunLinearAttentionGateTest(int batch_size, int seq_length, int num_heads, bool with_beta, + float tolerance) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } + + const size_t count = static_cast(batch_size) * seq_length * num_heads; + const auto a = RandomFloats(count, -6.0f, 6.0f, 11); + const auto b = RandomFloats(count, -6.0f, 6.0f, 12); + const auto dt_bias = RandomFloats(static_cast(num_heads), -2.0f, 2.0f, 13); + const auto decay_scale = RandomFloats(static_cast(num_heads), -4.0f, -0.1f, 14); + + std::vector expected_decay(count); + std::vector expected_beta(count); + for (size_t i = 0; i < count; ++i) { + const int h = static_cast(i % num_heads); + expected_decay[i] = decay_scale[h] * SoftplusRef(a[i] + dt_bias[h]); + expected_beta[i] = SigmoidRef(b[i]); + } + + const std::vector dims = {batch_size, seq_length, num_heads}; + const std::vector param_dims = {num_heads}; + + OpTester tester("LinearAttentionGate", 1, onnxruntime::kMSDomain); + tester.AddInput("a", dims, ToTensorType(a)); + tester.AddInput("dt_bias", param_dims, dt_bias); + tester.AddInput("decay_scale", param_dims, decay_scale); + if (with_beta) { + tester.AddInput("b", dims, ToTensorType(b)); + } else { + tester.AddOptionalInputEdge(); + } + tester.AddOutput("decay", dims, ToTensorType(expected_decay), false, tolerance, tolerance); + if (with_beta) { + tester.AddOutput("beta", dims, ToTensorType(expected_beta), false, tolerance, tolerance); + } + + std::vector> execution_providers; + execution_providers.push_back(std::move(cuda_ep)); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +template +void RunGatedRMSNormTest(int batch_size, int seq_length, int num_heads, int head_dim, + float epsilon, float tolerance) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } + + const int hidden = num_heads * head_dim; + const size_t count = static_cast(batch_size) * seq_length * hidden; + const auto x = RandomFloats(count, -3.0f, 3.0f, 21); + const auto gate = RandomFloats(count, -5.0f, 5.0f, 22); + const auto scale = RandomFloats(static_cast(head_dim), 0.2f, 1.8f, 23); + + std::vector expected(count); + const size_t num_rows = count / head_dim; + for (size_t r = 0; r < num_rows; ++r) { + const size_t base = r * head_dim; + float sum_sq = 0.0f; + for (int i = 0; i < head_dim; ++i) { + sum_sq += x[base + i] * x[base + i]; + } + const float inv_rms = 1.0f / std::sqrt(sum_sq / static_cast(head_dim) + epsilon); + for (int i = 0; i < head_dim; ++i) { + const float z = gate[base + i]; + expected[base + i] = x[base + i] * inv_rms * scale[i] * (z * SigmoidRef(z)); + } + } + + const std::vector dims = {batch_size, seq_length, hidden}; + const std::vector scale_dims = {head_dim}; + + OpTester tester("GatedRMSNorm", 1, onnxruntime::kMSDomain); + tester.AddAttribute("epsilon", epsilon); + tester.AddInput("X", dims, ToTensorType(x)); + tester.AddInput("scale", scale_dims, ToTensorType(scale)); + tester.AddInput("gate", dims, ToTensorType(gate)); + tester.AddOutput("Y", dims, ToTensorType(expected), false, tolerance, tolerance); + + std::vector> execution_providers; + execution_providers.push_back(std::move(cuda_ep)); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +} // namespace + +TEST(ContribOpLinearAttentionGateTest, Float_Decode) { + RunLinearAttentionGateTest(1, 1, 32, /*with_beta=*/true, 1e-4f); +} + +TEST(ContribOpLinearAttentionGateTest, Float_SpeculativeDecodeTile) { + RunLinearAttentionGateTest(1, 4, 32, /*with_beta=*/true, 1e-4f); +} + +TEST(ContribOpLinearAttentionGateTest, Float_DecayOnly) { + RunLinearAttentionGateTest(2, 3, 16, /*with_beta=*/false, 1e-4f); +} + +TEST(ContribOpLinearAttentionGateTest, Float16_SpeculativeDecodeTile) { + RunLinearAttentionGateTest(1, 4, 32, /*with_beta=*/true, 2e-3f); +} + +TEST(ContribOpLinearAttentionGateTest, Float16_Prefill) { + RunLinearAttentionGateTest(2, 37, 32, /*with_beta=*/true, 2e-3f); +} + +// num_heads not a multiple of the launch tile, to exercise the bounds check. +TEST(ContribOpLinearAttentionGateTest, Float16_RaggedTail) { + RunLinearAttentionGateTest(1, 5, 7, /*with_beta=*/true, 2e-3f); +} + +TEST(ContribOpLinearAttentionGateTest, BFloat16_SpeculativeDecodeTile) { + if (!CudaHasBF16Support()) { + GTEST_SKIP() << "bfloat16 requires compute capability 8.0 or later"; + } + RunLinearAttentionGateTest(1, 4, 32, /*with_beta=*/true, 2e-2f); +} + +// Requesting beta without b must be rejected by shape inference, not at execution time. +TEST(ContribOpLinearAttentionGateTest, BetaWithoutB_FailsShapeInference) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } + + constexpr int kNumHeads = 8; + const std::vector dims = {1, 2, kNumHeads}; + const std::vector param_dims = {kNumHeads}; + const std::vector values(static_cast(2 * kNumHeads), 0.5f); + const std::vector params(kNumHeads, 0.5f); + + OpTester tester("LinearAttentionGate", 1, onnxruntime::kMSDomain); + tester.AddInput("a", dims, values); + tester.AddInput("dt_bias", param_dims, params); + tester.AddInput("decay_scale", param_dims, params); + tester.AddOptionalInputEdge(); + tester.AddOutput("decay", dims, values); + tester.AddOutput("beta", dims, values); + + std::vector> execution_providers; + execution_providers.push_back(std::move(cuda_ep)); + tester.Run(OpTester::ExpectResult::kExpectFailure, + "The b input is required when the beta output is requested", + {}, nullptr, &execution_providers); +} + +TEST(ContribOpGatedRMSNormTest, Float_PerHead) { + RunGatedRMSNormTest(1, 4, 32, 128, 1e-6f, 1e-4f); +} + +TEST(ContribOpGatedRMSNormTest, Float_SingleGroup) { + RunGatedRMSNormTest(2, 3, 1, 512, 1e-5f, 1e-4f); +} + +TEST(ContribOpGatedRMSNormTest, Float16_PerHead) { + RunGatedRMSNormTest(1, 4, 32, 128, 1e-6f, 2e-3f); +} + +TEST(ContribOpGatedRMSNormTest, Float16_Prefill) { + RunGatedRMSNormTest(2, 17, 32, 128, 1e-6f, 2e-3f); +} + +// norm_size below, at, and above the thread-count dispatch boundaries. +TEST(ContribOpGatedRMSNormTest, Float16_SmallNormSize) { + RunGatedRMSNormTest(1, 4, 8, 48, 1e-6f, 2e-3f); +} + +TEST(ContribOpGatedRMSNormTest, Float16_LargeNormSize) { + RunGatedRMSNormTest(1, 2, 2, 1536, 1e-6f, 4e-3f); +} + +TEST(ContribOpGatedRMSNormTest, BFloat16_PerHead) { + if (!CudaHasBF16Support()) { + GTEST_SKIP() << "bfloat16 requires compute capability 8.0 or later"; + } + RunGatedRMSNormTest(1, 4, 32, 128, 1e-6f, 2e-2f); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc index fd2b648c8badf..f1436810812ad 100644 --- a/onnxruntime/test/contrib_ops/linear_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/linear_attention_op_test.cc @@ -1197,5 +1197,163 @@ TEST(ContribOpLinearAttentionTest, GatedDeltaRule_InverseGQA_LargerDims) { &initial_state, &decay, &beta); } +// state_window: past_state / present_state hold the last W per-token states, right-aligned, so +// slot j is the state after token (T - W + j) and slot W-1 is the state after the last token (the +// tensor the unwindowed op produces). past_state is read from slot W-1. Earlier positions are +// never written -- that is the point of the window (it bounds the arena for long prompts). +#ifdef USE_CUDA +// The state_window attribute is only implemented by the CUDA kernel. The four CUDA kernel families +// (generic recurrent, fixed-shape recurrent, warp-per-column decode, column-per-thread decode) each +// compute the window slot offsets themselves, so every family gets its own shape below. W > T is +// only exercised without a past_state, because that is the only case in which the slots below +// W - T are defined (zeroed rather than left alone). +static void RunLinearAttentionStateWindowTest(int B, int q_H, int kv_H, int n_k, int T, int dk, int dv, int W, + bool with_past_state = true) { + auto ep = DefaultCudaExecutionProvider(); + if (!ep) { + GTEST_SKIP() << "CUDA execution provider not available"; + return; + } + + const float scale = 1.0f / std::sqrt(static_cast(dk)); + + std::vector query(B * q_H * T * dk); + std::vector key(B * n_k * T * dk); + std::vector value(B * kv_H * T * dv); + std::vector decay(B * kv_H * T); + std::vector beta(B * kv_H * T); + for (size_t i = 0; i < query.size(); i++) query[i] = 0.5f * std::sin(static_cast(i) * 0.13f); + for (size_t i = 0; i < key.size(); i++) key[i] = 0.5f * std::cos(static_cast(i) * 0.17f); + for (size_t i = 0; i < value.size(); i++) value[i] = 0.5f * std::sin(static_cast(i) * 0.23f + 0.5f); + for (size_t i = 0; i < decay.size(); i++) { + decay[i] = -0.1f - 0.05f * std::abs(std::sin(static_cast(i) * 0.3f)); + beta[i] = 0.5f + 0.3f * std::sin(static_cast(i) * 0.31f); + } + std::vector initial_state(static_cast(B) * kv_H * dk * dv, 0.0f); + if (with_past_state) { + for (size_t i = 0; i < initial_state.size(); i++) { + initial_state[i] = 0.1f * std::cos(static_cast(i) * 0.07f); + } + } + const std::vector* past = with_past_state ? &initial_state : nullptr; + + // Slice a (B, H, T, D) tensor down to its first `prefix` tokens. + auto slice_prefix = [](const std::vector& src, int seq, int dim, int prefix) { + std::vector dst(src.size() / seq * prefix); + const int bh = static_cast(src.size()) / (seq * dim); + for (int i = 0; i < bh; i++) { + for (int t = 0; t < prefix; t++) { + for (int d = 0; d < dim; d++) { + dst[(i * prefix + t) * dim + d] = src[(i * seq + t) * dim + d]; + } + } + } + return dst; + }; + + std::vector expected_output_4d, expected_state; + LinearAttentionGQAReference("gated_delta", B, q_H, kv_H, n_k, T, dk, dv, scale, + query, key, value, past, &decay, &beta, + expected_output_4d, expected_state); + + // Slot j = the recurrent state after running the first (T - W + j + 1) tokens; slots for + // non-positive prefixes are never computed by the kernel and stay zero. The window axis leads + // the batch axis, so a slot is exactly one contiguous (B, H_kv, d_k, d_v) reference block. + const size_t slot_elems = static_cast(kv_H) * dk * dv; + const size_t batch_slot_elems = static_cast(B) * slot_elems; + std::vector expected_state_window(static_cast(W) * batch_slot_elems, 0.0f); + for (int j = 0; j < W; j++) { + const int prefix = T - W + j + 1; + if (prefix <= 0) continue; + std::vector prefix_state; + if (prefix == T) { + prefix_state = expected_state; + } else { + const std::vector decay_prefix = slice_prefix(decay, T, 1, prefix); + const std::vector beta_prefix = slice_prefix(beta, T, 1, prefix); + std::vector prefix_output; + LinearAttentionGQAReference("gated_delta", B, q_H, kv_H, n_k, prefix, dk, dv, scale, + slice_prefix(query, T, dk, prefix), + slice_prefix(key, T, dk, prefix), + slice_prefix(value, T, dv, prefix), + past, &decay_prefix, &beta_prefix, + prefix_output, prefix_state); + } + std::copy_n(prefix_state.begin(), batch_slot_elems, + expected_state_window.begin() + static_cast(j) * batch_slot_elems); + } + + OpTester tester("LinearAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("update_rule", "gated_delta"); + tester.AddAttribute("scale", scale); + tester.AddAttribute("q_num_heads", static_cast(q_H)); + tester.AddAttribute("kv_num_heads", static_cast(kv_H)); + tester.AddAttribute("state_window", static_cast(W)); + + tester.AddInput("query", {B, T, q_H * dk}, PackBHTD_to_BTHD(query, B, q_H, T, dk)); + tester.AddInput("key", {B, T, n_k * dk}, PackBHTD_to_BTHD(key, B, n_k, T, dk)); + tester.AddInput("value", {B, T, kv_H * dv}, PackBHTD_to_BTHD(value, B, kv_H, T, dv)); + if (with_past_state) { + // past_state is windowed too, and only slot W-1 is read. Poison the earlier slots to prove it. + std::vector past_state_window(static_cast(W) * batch_slot_elems, -1e4f); + std::copy_n(initial_state.begin(), batch_slot_elems, + past_state_window.begin() + static_cast(W - 1) * batch_slot_elems); + tester.AddInput("past_state", {W, B, kv_H, dk, dv}, past_state_window); + } else { + tester.AddOptionalInputEdge(); + } + tester.AddInput("decay", {B, T, kv_H}, TransposeBHT_to_BTH(decay, B, kv_H, T)); + tester.AddInput("beta", {B, T, kv_H}, TransposeBHT_to_BTH(beta, B, kv_H, T)); + + tester.AddOutput("output", {B, T, kv_H * dv}, + PackBHTD_to_BTHD(expected_output_4d, B, kv_H, T, dv), + false, 0.005f, 0.005f); + tester.AddOutput("present_state", {W, B, kv_H, dk, dv}, expected_state_window, + false, 0.005f, 0.005f); + + std::vector> execution_providers; + execution_providers.push_back(std::move(ep)); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// d_k = 4 is not a decode fast-path shape, so this lands on the generic recurrent kernel. +TEST(ContribOpLinearAttentionTest, GatedDeltaRule_StateWindow) { + RunLinearAttentionStateWindowTest(/*B=*/1, /*q_H=*/2, /*kv_H=*/2, /*n_k=*/2, /*T=*/5, + /*dk=*/4, /*dv=*/4, /*W=*/3); +} + +// B > 1 exercises the batch stride of the window axis on the generic recurrent kernel. +TEST(ContribOpLinearAttentionTest, GatedDeltaRule_StateWindow_MultiBatch) { + RunLinearAttentionStateWindowTest(/*B=*/2, /*q_H=*/2, /*kv_H=*/2, /*n_k=*/2, /*T=*/5, + /*dk=*/4, /*dv=*/4, /*W=*/3); +} + +// T <= 16 with d_k = 128 and d_v % 32 == 0 selects the column-per-thread decode kernel. +TEST(ContribOpLinearAttentionTest, GatedDeltaRule_StateWindow_DecodeColKernel) { + RunLinearAttentionStateWindowTest(/*B=*/2, /*q_H=*/2, /*kv_H=*/2, /*n_k=*/1, /*T=*/4, + /*dk=*/128, /*dv=*/128, /*W=*/3); +} + +// d_k = 256 falls back to the warp-per-column decode kernel. +TEST(ContribOpLinearAttentionTest, GatedDeltaRule_StateWindow_DecodeWarpKernel) { + RunLinearAttentionStateWindowTest(/*B=*/1, /*q_H=*/2, /*kv_H=*/2, /*n_k=*/1, /*T=*/4, + /*dk=*/256, /*dv=*/64, /*W=*/3); +} + +// T > 16 with a (d_k, d_v) fast-path pair selects the compile-time specialized recurrent kernel, +// whose final state write is a vectorized epilogue into slot W-1. +TEST(ContribOpLinearAttentionTest, GatedDeltaRule_StateWindow_FixedShapeKernel) { + RunLinearAttentionStateWindowTest(/*B=*/2, /*q_H=*/2, /*kv_H=*/2, /*n_k=*/1, /*T=*/24, + /*dk=*/128, /*dv=*/128, /*W=*/4); +} + +// W > T is the shape genai actually runs during MTP decode: the leading W - T slots belong to +// positions before this call and are left alone (zeroed here because there is no past_state). +TEST(ContribOpLinearAttentionTest, GatedDeltaRule_StateWindow_WiderThanSequence) { + RunLinearAttentionStateWindowTest(/*B=*/2, /*q_H=*/2, /*kv_H=*/2, /*n_k=*/1, /*T=*/2, + /*dk=*/128, /*dv=*/128, /*W=*/5, /*with_past_state=*/false); +} +#endif // USE_CUDA + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py index df0209962cd4a..0dc71fc0f4b80 100644 --- a/onnxruntime/test/python/transformers/test_gqa.py +++ b/onnxruntime/test/python/transformers/test_gqa.py @@ -512,7 +512,9 @@ def create_group_query_attention_graph_prompt(config: GQAConfig, ort_type, share config, ort_type, share_buffer, is_past=False ) graph = helper.make_graph([node], "GroupQueryAttention_Graph", graph_input, graph_output, initializer=initializers) - model = helper.make_model(graph) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 26), helper.make_opsetid("com.microsoft", 1)] + ) return model.SerializeToString() @@ -521,7 +523,9 @@ def create_group_query_attention_graph_past(config: GQAConfig, ort_type, share_b config, ort_type, share_buffer, is_past=True, head_sink_values=head_sink_values ) graph = helper.make_graph([node], "GroupQueryAttention_Graph", graph_input, graph_output, initializer=initializers) - model = helper.make_model(graph) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 26), helper.make_opsetid("com.microsoft", 1)] + ) return model.SerializeToString() @@ -2814,6 +2818,127 @@ def fused_kernel_test_cases(): yield f"fused_config_{i}", config +@unittest.skipIf(not has_flash_attention(), "Flash Attention is not available, skipping tests.") +class TestFlashDecodeMultiTokenParity(unittest.TestCase): + def test_shared_buffer_multitoken_decode_matches_flash_attention(self): + device = "cuda" + torch_type = torch.float16 + ort_type = TensorProto.FLOAT16 + + config = GQAConfig( + batch_size=1, + q_sequence_length=4, + kv_sequence_length=4, + num_heads=16, + kv_num_heads=2, + head_size=128, + past_kv_sequence_length=64, + buffer_sequence_length=128, + rotary=False, + packed=False, + share_buffer=True, + softcap=0.0, + ) + + torch.manual_seed(123) + std = 0.1 + q = ( + torch.randn( + config.batch_size, + config.q_sequence_length, + config.num_heads, + config.head_size, + device=device, + dtype=torch_type, + ) + * std + ) + k = ( + torch.randn( + config.batch_size, + config.kv_num_heads, + config.buffer_sequence_length, + config.head_size, + device=device, + dtype=torch_type, + ) + * std + ) + v = torch.randn_like(k) * std + + past_len = config.past_kv_sequence_length + k[:, :, past_len:, :] = 0 + v[:, :, past_len:, :] = 0 + + new_k = ( + torch.randn( + config.batch_size, + config.kv_sequence_length, + config.kv_num_heads, + config.head_size, + device=device, + dtype=torch_type, + ) + * std + ) + new_v = torch.randn_like(new_k) * std + + seqlens_k = torch.tensor([past_len + config.kv_sequence_length - 1], dtype=torch.int32, device=device) + + def run_once(disable_flash_decode: bool): + with ( + scoped_env_var("ORT_ENABLE_XQA", "0"), + scoped_env_var("ORT_ENABLE_CUDNN_FLASH_ATTENTION", "0"), + scoped_env_var("ORT_DISABLE_FLASH_ATTENTION", "0"), + scoped_env_var("ORT_DISABLE_FLASH_DECODE", "1" if disable_flash_decode else "0"), + ): + out, present_k, present_v = gqa_past_func( + q=q, + k=k.clone(), + v=v.clone(), + config=config, + new_k=new_k, + new_v=new_v, + cos=None, + sin=None, + seqlens_k=seqlens_k, + position_ids=None, + attention_bias=None, + head_sink=None, + k_scale=None, + v_scale=None, + ep="CUDAExecutionProvider", + device=device, + share_buffer=True, + ort_type=ort_type, + ) + return out, present_k, present_v + + out_fast, pk_fast, pv_fast = run_once(disable_flash_decode=False) + out_ref, pk_ref, pv_ref = run_once(disable_flash_decode=True) + + numpy.testing.assert_allclose( + out_fast.to(torch.float32).detach().cpu().numpy(), + out_ref.to(torch.float32).detach().cpu().numpy(), + rtol=2e-3, + atol=2e-3, + ) + + valid_len = past_len + config.kv_sequence_length + numpy.testing.assert_allclose( + pk_fast[:, :, :valid_len, :].to(torch.float32).detach().cpu().numpy(), + pk_ref[:, :, :valid_len, :].to(torch.float32).detach().cpu().numpy(), + rtol=2e-3, + atol=2e-3, + ) + numpy.testing.assert_allclose( + pv_fast[:, :, :valid_len, :].to(torch.float32).detach().cpu().numpy(), + pv_ref[:, :, :valid_len, :].to(torch.float32).detach().cpu().numpy(), + rtol=2e-3, + atol=2e-3, + ) + + def gqa_xqa_test_cases(): # Decoding config (seq_len=1, share_buffer=True) # Testing different group sizes and query types