diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index 2331780e8e932..96cfc7863f731 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -37,6 +37,7 @@ Do not modify directly.*
* com.microsoft.FusedGemm
* com.microsoft.FusedMatMul
* com.microsoft.FusedMatMulActivation
+ * com.microsoft.GatedRMSNorm
* com.microsoft.GatedRelativePositionBias
* com.microsoft.GatherBlockQuantized
* com.microsoft.GatherND
@@ -51,6 +52,7 @@ Do not modify directly.*
* com.microsoft.Inverse
* com.microsoft.Irfft
* com.microsoft.LinearAttention
+ * com.microsoft.LinearAttentionGate
* com.microsoft.LongformerAttention
* com.microsoft.MatMulBlockQuantizedFp4Weight
* com.microsoft.MatMulBlockQuantizedFp8Weight
@@ -936,6 +938,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Fused activation function. One of: 'silu', 'swish', 'none'. Default is 'none'.
ndim : int
Spatial dimensionality: 1, 2, or 3. Default is 1.
+state_window : int
+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.
#### Inputs (2 - 4)
@@ -948,7 +952,7 @@ This version of the operator has been available since version 1 of the 'com.micr
bias (optional) : T
Optional per-channel bias with shape (channels).
past_state (optional) : T
-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.
#### Outputs
@@ -957,7 +961,7 @@ This version of the operator has been available since version 1 of the 'com.micr
output : T
Convolution output with same shape as input.
present_state : T
-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).
#### Type Constraints
@@ -2046,6 +2050,57 @@ This version of the operator has been available since version 1 of the 'com.micr
+### **com.microsoft.GatedRMSNorm**
+
+ 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
+
+
+- epsilon : float
+- Epsilon added to the mean of squares before the reciprocal square root.
+
+
+#### Inputs
+
+
+- X : T
+- Input tensor with shape (..., H * C). Normalization is applied over each contiguous group of C elements.
+- scale : T
+- Normalization weight with shape (C).
+- gate : T
+- Gate tensor with the same shape as X.
+
+
+#### Outputs
+
+
+- Y : T
+- Output tensor with the same shape as X.
+
+
+#### Type Constraints
+
+
+- T : tensor(float), tensor(float16), tensor(bfloat16)
+- Constrain input and output types to float tensors.
+
+
+
### **com.microsoft.GatedRelativePositionBias**
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
Number of query heads. Always required.
scale : float
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.
+state_window : int
+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.
update_rule : string
The update rule for the linear attention recurrence. One of: 'linear', 'gated', 'delta', 'gated_delta'. Default is 'gated_delta'.
@@ -2827,7 +2884,7 @@ This version of the operator has been available since version 1 of the 'com.micr
value : T
Value vectors with 3D packed shape (B, T, H_kv * d_v).
past_state (optional) : S
-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.
decay (optional) : T
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.
beta (optional) : T
@@ -2840,7 +2897,7 @@ This version of the operator has been available since version 1 of the 'com.micr
output : T
Attention output with 3D packed shape (B, T, H_q * d_v).
present_state : S
-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).
#### Type Constraints
@@ -2853,6 +2910,58 @@ This version of the operator has been available since version 1 of the 'com.micr
+### **com.microsoft.LinearAttentionGate**
+
+ 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)
+
+
+- a : T
+- Decay gate projection with shape (B, T, H).
+- dt_bias : TF
+- Per-head float32 bias added to a, with shape (H).
+- decay_scale : TF
+- Per-head float32 multiplier applied to Softplus(a + dt_bias), with shape (H). For gated DeltaNet this is -exp(A_log).
+- b (optional) : T
+- Update-rate projection with shape (B, T, H). Required when the beta output is requested.
+
+
+#### Outputs (1 - 2)
+
+
+- decay : T
+- decay_scale * Softplus(a + dt_bias) with shape (B, T, H).
+- beta (optional) : T
+- Sigmoid(b) with shape (B, T, H).
+
+
+#### Type Constraints
+
+
+- T : tensor(float), tensor(float16), tensor(bfloat16)
+- Constrain gate input and output types to float tensors.
+- TF : tensor(float)
+- Constrain the per-head parameters to float32.
+
+
+
### **com.microsoft.LongformerAttention**
Longformer Self Attention with a local context and a global context. Tokens attend locally: Each token
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index d359aa0e2b523..6b0a514423b67 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -1086,6 +1086,7 @@ The **OpSet Version** column uses the following notation:
|FastGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
+|GatedRMSNorm|*in* X:**T**
*in* scale:**T**
*in* gate:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|GatedRelativePositionBias|*in* query_layer:**T**
*in* query_bias:**T**
*in* rel_pos:**T**
*in* weight:**T**
*in* bias:**T**
*in* eco_a:**T**
*in* token_offset:**M**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|GatherBlockQuantized|*in* data:**T1**
*in* indices:**Tind**
*in* scales:**T2**
*in* zero_points:**T1**
*out* output:**T2**|1+|**T1** = tensor(int4), tensor(uint4), tensor(uint8)
**T2** = tensor(bfloat16), tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)|
|Gelu|*in* X:**T**
*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**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|Irfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|LinearAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_state:**S**
*in* decay:**T**
*in* beta:**T**
*out* output:**T**
*out* present_state:**S**|1+|**T** = tensor(float), tensor(float16)|
+|LinearAttentionGate|*in* a:**T**
*in* dt_bias:**TF**
*in* decay_scale:**TF**
*in* b:**T**
*out* decay:**T**
*out* beta:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TF** = tensor(float)|
|LongformerAttention|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* mask:**T**
*in* global_weight:**T**
*in* global_bias:**T**
*in* global:**G**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|MatMulBlockQuantizedFp4Weight|*in* A:**T**
*in* B:**T1**
*in* weight_scale:**T2**
*in* weight_scale_2:**T3**
*in* input_scale:**T3**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(uint8)
**T2** = tensor(uint8)
**T3** = tensor(float)|
|MatMulBlockQuantizedFp8Weight|*in* A:**T**
*in* B:**T1**
*in* b_scale:**T2**
*in* a_scale:**T2**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(float8e4m3fn)
**T2** = tensor(float)|
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..861416b18b288 100644
--- a/onnxruntime/contrib_ops/cuda/bert/linear_attention.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention.cc
@@ -3,9 +3,12 @@
#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"
+#include
+
namespace onnxruntime {
namespace contrib {
namespace cuda {
@@ -45,6 +48,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
@@ -61,19 +69,40 @@ Status LinearAttention::ComputeInternal(OpKernelContext* context) const {
const auto& query_shape = query_tensor->Shape();
ORT_RETURN_IF_NOT(query_shape.NumDimensions() == 3, "query must be 3D");
- const int batch_size = static_cast(query_shape[0]);
- const int seq_len = static_cast(query_shape[1]);
- const int query_hidden = static_cast(query_shape[2]);
+ const int64_t batch_size_64 = query_shape[0];
+ const int64_t seq_len_64 = query_shape[1];
+ const int64_t query_hidden_64 = query_shape[2];
+ ORT_RETURN_IF_NOT(batch_size_64 <= std::numeric_limits::max() &&
+ seq_len_64 <= std::numeric_limits::max() &&
+ query_hidden_64 <= std::numeric_limits::max(),
+ "query dimensions are too large for the CUDA kernel");
+ const int batch_size = static_cast(batch_size_64);
+ const int seq_len = static_cast(seq_len_64);
ORT_RETURN_IF_NOT(key_tensor != nullptr && value_tensor != nullptr, "key and value inputs are required");
const auto& key_shape = key_tensor->Shape();
const auto& value_shape = value_tensor->Shape();
- int d_k = query_hidden / q_num_heads_;
+ ORT_RETURN_IF_NOT(key_shape.NumDimensions() == 3 && value_shape.NumDimensions() == 3,
+ "key and value must be 3D");
+ ORT_RETURN_IF_NOT(key_shape[2] <= std::numeric_limits::max() &&
+ value_shape[2] <= std::numeric_limits::max(),
+ "key and value dimensions are too large for the CUDA kernel");
+ ORT_RETURN_IF_NOT(key_shape[0] == query_shape[0] && value_shape[0] == query_shape[0],
+ "key and value batch dimensions must match query");
+ ORT_RETURN_IF_NOT(key_shape[1] == query_shape[1] && value_shape[1] == query_shape[1],
+ "key and value sequence dimensions must match query");
+ ORT_RETURN_IF_NOT(query_hidden_64 > 0 && query_hidden_64 % q_num_heads_ == 0,
+ "query last dim (", query_hidden_64, ") must be positive and divisible by q_num_heads (",
+ q_num_heads_, ")");
+ ORT_RETURN_IF_NOT(value_shape[2] > 0 && value_shape[2] % kv_num_heads_ == 0,
+ "value last dim (", value_shape[2], ") must be positive and divisible by kv_num_heads (",
+ kv_num_heads_, ")");
+ const int d_k = static_cast(query_hidden_64 / q_num_heads_);
int d_v = static_cast(value_shape[2]) / kv_num_heads_;
- ORT_ENFORCE(static_cast(key_shape[2]) % d_k == 0,
- "key last dim (", key_shape[2], ") must be divisible by d_k (", d_k, ")");
+ ORT_RETURN_IF_NOT(key_shape[2] > 0 && key_shape[2] % d_k == 0,
+ "key last dim (", key_shape[2], ") must be divisible by d_k (", d_k, ")");
int n_k_heads = static_cast(key_shape[2]) / d_k;
// GQA head mapping validations
@@ -103,46 +132,67 @@ Status LinearAttention::ComputeInternal(OpKernelContext* context) const {
bool decay_per_key_dim = false;
if (decay_tensor != nullptr) {
- int64_t decay_last = decay_tensor->Shape()[2];
- decay_per_key_dim = (decay_last == kv_num_heads_ * d_k);
+ const auto& decay_shape = decay_tensor->Shape();
+ ORT_RETURN_IF_NOT(decay_shape.NumDimensions() == 3,
+ "decay must be rank 3 (B, T, ...), got rank ", decay_shape.NumDimensions());
+ ORT_RETURN_IF_NOT(decay_shape[0] == batch_size && decay_shape[1] == seq_len,
+ "decay batch/sequence dimensions must match query");
+ const int64_t decay_last = decay_shape[2];
+ if (decay_last == static_cast(kv_num_heads_) * d_k) {
+ decay_per_key_dim = true;
+ } else {
+ ORT_RETURN_IF_NOT(decay_last == kv_num_heads_,
+ "decay last dim must be H_kv or H_kv*d_k");
+ }
}
bool beta_per_head = false;
if (beta_tensor != nullptr) {
- int64_t beta_last = beta_tensor->Shape()[2];
- beta_per_head = (beta_last == kv_num_heads_);
+ const auto& beta_shape = beta_tensor->Shape();
+ ORT_RETURN_IF_NOT(beta_shape.NumDimensions() == 3,
+ "beta must be rank 3 (B, T, ...), got rank ", beta_shape.NumDimensions());
+ ORT_RETURN_IF_NOT(beta_shape[0] == batch_size && beta_shape[1] == seq_len,
+ "beta batch/sequence dimensions must match query");
+ const int64_t beta_last = beta_shape[2];
+ if (beta_last == kv_num_heads_) {
+ beta_per_head = true;
+ } else {
+ ORT_RETURN_IF_NOT(beta_last == 1, "beta last dim must be H_kv or 1");
+ }
}
// Allocate outputs
- int output_hidden = std::max(q_num_heads_, kv_num_heads_) * d_v;
+ const int64_t output_hidden_64 = static_cast(std::max(q_num_heads_, kv_num_heads_)) * d_v;
+ ORT_RETURN_IF_NOT(output_hidden_64 <= std::numeric_limits::max(),
+ "output hidden dimension is too large for the CUDA kernel");
+ int output_hidden = static_cast(output_hidden_64);
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 +205,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 +220,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..a65b8c53750a6
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.cc
@@ -0,0 +1,143 @@
+// 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"
+
+#include
+
+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(num_heads <= std::numeric_limits::max(),
+ "a last dimension is too large for the CUDA kernel");
+
+ 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");
+ ORT_RETURN_IF_NOT(norm_size <= std::numeric_limits::max(),
+ "scale is too large for the CUDA kernel");
+ 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..1388f3072380a
--- /dev/null
+++ b/onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.cu
@@ -0,0 +1,181 @@
+// 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
+
+#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 - 1) / kThreads + 1;
+ ORT_RETURN_IF_NOT(blocks <= std::numeric_limits::max(),
+ "LinearAttentionGate launch requires too many blocks");
+ 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();
+ }
+
+ ORT_RETURN_IF_NOT(num_rows <= std::numeric_limits::max(),
+ "GatedRMSNorm launch requires too many blocks");
+ 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 446d71f457052..efd3cffaffe36 100644
--- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
+++ b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc
@@ -155,6 +155,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)
@@ -426,6 +432,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 bc8bd383383de..570effdfeeb8f 100644
--- a/onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc
+++ b/onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc
@@ -110,6 +110,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