[CUDA] Add LinearAttentionGate and GatedRMSNorm contrib operators

docs/ContribOperators.md

@@ -37,6 +37,7 @@ Do not modify directly.*
   * <a href="#com.microsoft.FusedGemm">com.microsoft.FusedGemm</a>
   * <a href="#com.microsoft.FusedMatMul">com.microsoft.FusedMatMul</a>
   * <a href="#com.microsoft.FusedMatMulActivation">com.microsoft.FusedMatMulActivation</a>
+  * <a href="#com.microsoft.GatedRMSNorm">com.microsoft.GatedRMSNorm</a>
   * <a href="#com.microsoft.GatedRelativePositionBias">com.microsoft.GatedRelativePositionBias</a>
   * <a href="#com.microsoft.GatherBlockQuantized">com.microsoft.GatherBlockQuantized</a>
   * <a href="#com.microsoft.GatherND">com.microsoft.GatherND</a>
@@ -51,6 +52,7 @@ Do not modify directly.*
   * <a href="#com.microsoft.Inverse">com.microsoft.Inverse</a>
   * <a href="#com.microsoft.Irfft">com.microsoft.Irfft</a>
   * <a href="#com.microsoft.LinearAttention">com.microsoft.LinearAttention</a>
+  * <a href="#com.microsoft.LinearAttentionGate">com.microsoft.LinearAttentionGate</a>
   * <a href="#com.microsoft.LongformerAttention">com.microsoft.LongformerAttention</a>
   * <a href="#com.microsoft.MatMulBlockQuantizedFp4Weight">com.microsoft.MatMulBlockQuantizedFp4Weight</a>
   * <a href="#com.microsoft.MatMulBlockQuantizedFp8Weight">com.microsoft.MatMulBlockQuantizedFp8Weight</a>
@@ -936,6 +938,8 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dd>Fused activation function. One of: 'silu', 'swish', 'none'. Default is 'none'.</dd>
 <dt><tt>ndim</tt> : int</dt>
 <dd>Spatial dimensionality: 1, 2, or 3. Default is 1.</dd>
+<dt><tt>state_window</tt> : int</dt>
+<dd>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.</dd>
 </dl>
 
 #### Inputs (2 - 4)
@@ -948,7 +952,7 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>bias</tt> (optional) : T</dt>
 <dd>Optional per-channel bias with shape (channels).</dd>
 <dt><tt>past_state</tt> (optional) : T</dt>
-<dd>Carry state from previous step. For ndim=1: (batch_size, channels, k_1 - 1). If not provided, padding is zero.</dd>
+<dd>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.</dd>
 </dl>
 
 #### Outputs
@@ -957,7 +961,7 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>output</tt> : T</dt>
 <dd>Convolution output with same shape as input.</dd>
 <dt><tt>present_state</tt> : T</dt>
-<dd>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.</dd>
+<dd>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).</dd>
 </dl>
 
 #### Type Constraints
@@ -2046,6 +2050,57 @@ This version of the operator has been available since version 1 of the 'com.micr
 </dl>
 
 
+### <a name="com.microsoft.GatedRMSNorm"></a><a name="com.microsoft.gatedrmsnorm">**com.microsoft.GatedRMSNorm**</a>
+
+  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.
+
+#### Version
+
+This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
+
+#### Attributes
+
+<dl>
+<dt><tt>epsilon</tt> : float</dt>
+<dd>Epsilon added to the mean of squares before the reciprocal square root.</dd>
+</dl>
+
+#### Inputs
+
+<dl>
+<dt><tt>X</tt> : T</dt>
+<dd>Input tensor with shape (..., H * C). Normalization is applied over each contiguous group of C elements.</dd>
+<dt><tt>scale</tt> : T</dt>
+<dd>Normalization weight with shape (C).</dd>
+<dt><tt>gate</tt> : T</dt>
+<dd>Gate tensor with the same shape as X.</dd>
+</dl>
+
+#### Outputs
+
+<dl>
+<dt><tt>Y</tt> : T</dt>
+<dd>Output tensor with the same shape as X.</dd>
+</dl>
+
+#### Type Constraints
+
+<dl>
+<dt><tt>T</tt> : tensor(float), tensor(float16), tensor(bfloat16)</dt>
+<dd>Constrain input and output types to float tensors.</dd>
+</dl>
+
+
 ### <a name="com.microsoft.GatedRelativePositionBias"></a><a name="com.microsoft.gatedrelativepositionbias">**com.microsoft.GatedRelativePositionBias**</a>
 
   query_layer = (query_layer + query_bias).reshape(batch_size, seq_len, num_heads, head_size).transpose(1, 2)
@@ -2813,6 +2868,8 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dd>Number of query heads. Always required.</dd>
 <dt><tt>scale</tt> : float</dt>
 <dd>Output scaling factor. When 0.0 (default), derives d_k = query.shape[-1] / q_num_heads and uses 1/sqrt(d_k). Set explicitly to override.</dd>
+<dt><tt>state_window</tt> : int</dt>
+<dd>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.</dd>
 <dt><tt>update_rule</tt> : string</dt>
 <dd>The update rule for the linear attention recurrence. One of: 'linear', 'gated', 'delta', 'gated_delta'. Default is 'gated_delta'.</dd>
 </dl>
@@ -2827,7 +2884,7 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>value</tt> : T</dt>
 <dd>Value vectors with 3D packed shape (B, T, H_kv * d_v).</dd>
 <dt><tt>past_state</tt> (optional) : S</dt>
-<dd>Recurrent state from previous step with shape (B, H_kv, d_k, d_v). Always 4D. If not provided, defaults to zeros.</dd>
+<dd>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.</dd>
 <dt><tt>decay</tt> (optional) : T</dt>
 <dd>Exponential decay gate in log-space. 3D packed shape: (B, T, H_kv * d_k) for per-key-dimension decay (GLA/RWKV-6), or (B, T, H_kv) for per-head scalar decay (DeltaNet/RetNet). Required for 'gated' and 'gated_delta' modes.</dd>
 <dt><tt>beta</tt> (optional) : T</dt>
@@ -2840,7 +2897,7 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>output</tt> : T</dt>
 <dd>Attention output with 3D packed shape (B, T, H_q * d_v).</dd>
 <dt><tt>present_state</tt> : S</dt>
-<dd>Updated recurrent state with shape (B, H_kv, d_k, d_v). Always 4D.</dd>
+<dd>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).</dd>
 </dl>
 
 #### Type Constraints
@@ -2853,6 +2910,58 @@ This version of the operator has been available since version 1 of the 'com.micr
 </dl>
 
 
+### <a name="com.microsoft.LinearAttentionGate"></a><a name="com.microsoft.linearattentiongate">**com.microsoft.LinearAttentionGate**</a>
+
+  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.
+
+#### Version
+
+This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
+
+#### Inputs (3 - 4)
+
+<dl>
+<dt><tt>a</tt> : T</dt>
+<dd>Decay gate projection with shape (B, T, H).</dd>
+<dt><tt>dt_bias</tt> : TF</dt>
+<dd>Per-head float32 bias added to a, with shape (H).</dd>
+<dt><tt>decay_scale</tt> : TF</dt>
+<dd>Per-head float32 multiplier applied to Softplus(a + dt_bias), with shape (H). For gated DeltaNet this is -exp(A_log).</dd>
+<dt><tt>b</tt> (optional) : T</dt>
+<dd>Update-rate projection with shape (B, T, H). Required when the beta output is requested.</dd>
+</dl>
+
+#### Outputs (1 - 2)
+
+<dl>
+<dt><tt>decay</tt> : T</dt>
+<dd>decay_scale * Softplus(a + dt_bias) with shape (B, T, H).</dd>
+<dt><tt>beta</tt> (optional) : T</dt>
+<dd>Sigmoid(b) with shape (B, T, H).</dd>
+</dl>
+
+#### Type Constraints
+
+<dl>
+<dt><tt>T</tt> : tensor(float), tensor(float16), tensor(bfloat16)</dt>
+<dd>Constrain gate input and output types to float tensors.</dd>
+<dt><tt>TF</tt> : tensor(float)</dt>
+<dd>Constrain the per-head parameters to float32.</dd>
+</dl>
+
+
 ### <a name="com.microsoft.LongformerAttention"></a><a name="com.microsoft.longformerattention">**com.microsoft.LongformerAttention**</a>
 
   Longformer Self Attention with a local context and a global context. Tokens attend locally: Each token

docs/OperatorKernels.md

@@ -1086,6 +1086,7 @@ The **OpSet Version** column uses the following notation:
 |FastGelu|*in* X:**T**<br> *in* bias:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
 |FusedConv|*in* X:**T**<br> *in* W:**T**<br> *in* B:**T**<br> *in* Z:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
 |FusedMatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
+|GatedRMSNorm|*in* X:**T**<br> *in* scale:**T**<br> *in* gate:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
 |GatedRelativePositionBias|*in* query_layer:**T**<br> *in* query_bias:**T**<br> *in* rel_pos:**T**<br> *in* weight:**T**<br> *in* bias:**T**<br> *in* eco_a:**T**<br> *in* token_offset:**M**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
 |GatherBlockQuantized|*in* data:**T1**<br> *in* indices:**Tind**<br> *in* scales:**T2**<br> *in* zero_points:**T1**<br> *out* output:**T2**|1+|**T1** = tensor(int4), tensor(uint4), tensor(uint8)<br/> **T2** = tensor(bfloat16), tensor(float), tensor(float16)<br/> **Tind** = tensor(int32), tensor(int64)|
 |Gelu|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
@@ -1098,6 +1099,7 @@ The **OpSet Version** column uses the following notation:
 |Inverse|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
 |Irfft|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
 |LinearAttention|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* past_state:**S**<br> *in* decay:**T**<br> *in* beta:**T**<br> *out* output:**T**<br> *out* present_state:**S**|1+|**T** = tensor(float), tensor(float16)|
+|LinearAttentionGate|*in* a:**T**<br> *in* dt_bias:**TF**<br> *in* decay_scale:**TF**<br> *in* b:**T**<br> *out* decay:**T**<br> *out* beta:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)<br/> **TF** = tensor(float)|
 |LongformerAttention|*in* input:**T**<br> *in* weight:**T**<br> *in* bias:**T**<br> *in* mask:**T**<br> *in* global_weight:**T**<br> *in* global_bias:**T**<br> *in* global:**G**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
 |MatMulBlockQuantizedFp4Weight|*in* A:**T**<br> *in* B:**T1**<br> *in* weight_scale:**T2**<br> *in* weight_scale_2:**T3**<br> *in* input_scale:**T3**<br> *in* bias:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float16)<br/> **T1** = tensor(uint8)<br/> **T2** = tensor(uint8)<br/> **T3** = tensor(float)|
 |MatMulBlockQuantizedFp8Weight|*in* A:**T**<br> *in* B:**T1**<br> *in* b_scale:**T2**<br> *in* a_scale:**T2**<br> *in* bias:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float16)<br/> **T1** = tensor(float8e4m3fn)<br/> **T2** = tensor(float)|

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