[CUDA] Add LinearAttentionGate and GatedRMSNorm contrib operators

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<T>::CausalConvWithState(const OpKernelInfo& info) : OpKernel
   activation_ = info.GetAttrOrDefault<std::string>("activation", "none");
   ORT_ENFORCE(activation_ == "none" || activation_ == "silu" || activation_ == "swish",
               "activation must be one of: none, silu, swish");
+
+  ORT_ENFORCE(info.GetAttrOrDefault<int64_t>("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<T>::Compute(OpKernelContext* context) const {
     Tensor* output_tensor = context->Output(0, input_shape);
     float* output_data = output_tensor->MutableData<float>();
 
-    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<int>(batch_size), static_cast<int>(channels),
+        static_cast<int>(pad), past_state_tensor, state_shape));
     Tensor* present_state_tensor = context->Output(1, state_shape);
     float* present_data = present_state_tensor->MutableData<float>();
 

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

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 <typename TKernelInfo>
+Status ParseStateWindow(const TKernelInfo& info, int& state_window) {
+  const int64_t value = info.template GetAttrOrDefault<int64_t>("state_window", 0);
+  if (value < 0) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "state_window must be >= 0, got ", value);
+  }
+  state_window = static_cast<int>(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 <typename T>
+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

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<T>::LinearAttention(const OpKernelInfo& info) : OpKernel(info) {
   int64_t chunk_size = info.GetAttrOrDefault<int64_t>("chunk_size", 64);
   // chunk_size_ reserved for future chunk-parallel prefill algorithm; not yet used.
   chunk_size_ = static_cast<int>(chunk_size);
+
+  ORT_ENFORCE(info.GetAttrOrDefault<int64_t>("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<T>::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<int64_t>(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<int>(batch_size), kv_num_heads_,
+      static_cast<int>(d_k), static_cast<int>(d_v), past_state_tensor, state_shape));
   Tensor* present_state_tensor = context->Output(1, state_shape);
   float* state_data = present_state_tensor->MutableData<float>();
   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<float>();
     std::memcpy(state_data, ps_data, static_cast<size_t>(total_state) * sizeof(float));
   } else {

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

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 <typename TKernelInfo>
+Status ParseStateWindow(const TKernelInfo& info, int& state_window) {
+  const int64_t value = info.template GetAttrOrDefault<int64_t>("state_window", 0);
+  if (value < 0) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "state_window must be >= 0, got ", value);
+  }
+  state_window = static_cast<int>(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 <typename T>
+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

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<T>::CausalConvWithState(const OpKernelInfo& info) : CudaKern
   activation_ = info.GetAttrOrDefault<std::string>("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 <typename T>
@@ -75,25 +80,29 @@ Status CausalConvWithState<T>::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<T>::ComputeInternal(OpKernelContext* context) const {
       L,
       K,
       apply_silu,
-      GetDeviceProp().maxThreadsPerBlock);
+      GetDeviceProp().maxThreadsPerBlock,
+      state_slots);
 }
 
 }  // namespace cuda

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