[CUDA] Add state_window to LinearAttention and CausalConvWithState to support MTP

docs/ContribOperators.md

@@ -936,6 +936,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) hold no position from this call and are filled with zeros. A window lets a speculative decoder roll the state back to an accepted prefix without replaying the forward. Valid range is [0, 8].</dd>
 </dl>
 
 #### Inputs (2 - 4)
@@ -948,7 +950,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 +959,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
@@ -2813,6 +2815,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) hold no token from this call and are filled with zeros. A window lets a speculative decoder roll the state back to an accepted prefix without replaying the forward. Valid range is [0, 8].</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 +2831,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 +2844,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

onnxruntime/contrib_ops/cpu/bert/attention_common.h

@@ -88,6 +88,11 @@ inline KVQuantizationType StringToKVQuantizationType(std::string s) {
 constexpr bool LAYOUT_BSNH = false;
 constexpr bool LAYOUT_BNSH = true;
 
+// Upper bound on the `state_window` attribute of LinearAttention and CausalConvWithState. The
+// window only has to cover the tokens a multi-token predictor can propose, and the state tensors
+// grow linearly with it, so cap it rather than let a model request an unbounded allocation.
+constexpr int64_t kMaxStateWindow = 8;
+
 namespace sparse_attention {
 // Environment variable to enable or disable sparse attention v1 kernel. Default is 0 (enabled).
 constexpr const char* kDisableSparseAttentionV1 = "ORT_DISABLE_SPARSE_ATTENTION_V1";

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,71 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "contrib_ops/cpu/bert/attention_common.h"
+#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 TOpKernelInfo>
+inline Status ParseStateWindow(const TOpKernelInfo& info, int& state_window) {
+  const int64_t value = info.template GetAttrOrDefault<int64_t>("state_window", 0);
+  if (value < 0 || value > kMaxStateWindow) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "state_window must be in [0, ", kMaxStateWindow, "], got ", value);
+  }
+  state_window = static_cast<int>(value);
+  return Status::OK();
+}
+
+// Derives the expected past_state / present_state shape and validates past_state against it.
+// `state_length` is the carry length along the causal axis, i.e. kernel_size - 1.
+//
+// state_window == 0 -> (batch_size, channels, state_length). A single state with no window axis:
+// the backward-compatible layout that models exported before the attribute existed use.
+//
+// state_window == W > 0 -> (W, batch_size, channels, state_length). The window holds the carry
+// state after each of the last W positions, right-aligned: slot j is the state after position
+// (seq_len - W + j), so slot W-1 is the state after the last position and is the only slot read
+// back as past_state. W == 1 is therefore the legacy layout with a leading unit axis.
+//
+// The window axis leads the batch axis so that each slot is one contiguous
+// (batch_size, channels, state_length) block. That keeps "fetch/replace the last state" a single
+// contiguous range for any batch size, which is what a speculative decoder needs when it crops the
+// state back to an accepted prefix.
+template <typename T>
+Status CheckInputs(int state_window,
+                   int batch_size,
+                   int channels,
+                   int state_length,
+                   const T* past_state,
+                   TensorShape& state_shape) {
+  state_shape = state_window > 0
+                    ? TensorShape({state_window, batch_size, channels, state_length})
+                    : TensorShape({batch_size, channels, state_length});
+
+  if (past_state != nullptr && past_state->Shape() != state_shape) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "Input 'past_state' is expected to have shape ", state_shape.ToString(),
+                           ", got ", past_state->Shape().ToString(),
+                           ". CausalConvWithState uses (batch_size, channels, kernel_size - 1) when "
+                           "the state_window attribute is absent or 0, and "
+                           "(state_window, batch_size, channels, kernel_size - 1) otherwise.");
+  }
+
+  return Status::OK();
+}
+
+}  // namespace causal_conv_with_state_helper
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cpu/bert/linear_attention.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "contrib_ops/cpu/bert/linear_attention.h"
+#include "contrib_ops/cpu/bert/linear_attention_helper.h"
 
 #include "core/framework/tensorprotoutils.h"
 #include "core/common/safeint.h"
@@ -60,6 +61,10 @@ LinearAttention<T>::LinearAttention(const OpKernelInfo& info) : OpKernel(info) {
   int64_t chunk_size = info.GetAttrOrDefault<int64_t>("chunk_size", 64);
   // chunk_size_ reserved for future chunk-parallel prefill algorithm; not yet used.
   chunk_size_ = static_cast<int>(chunk_size);
+
+  ORT_ENFORCE(info.GetAttrOrDefault<int64_t>("state_window", 0) == 0,
+              "CPU LinearAttention does not support state_window > 0 (CUDA EP only)");
+  state_window_ = 0;
 }
 
 namespace {
@@ -417,21 +422,17 @@ Status LinearAttention<T>::Compute(OpKernelContext* context) const {
   }
 
   // ==== Initialize state: write directly into output present_state ====
-  // present_state: (B, H_kv, d_k, d_v)
-  TensorShape state_shape({batch_size, static_cast<int64_t>(kv_num_heads_), d_k, d_v});
+  // state_window_ is always 0 on CPU, so this is the legacy (B, H_kv, d_k, d_v) shape.
+  TensorShape state_shape;
+  ORT_RETURN_IF_ERROR(linear_attention_helper::CheckInputs(
+      state_window_, static_cast<int>(batch_size), kv_num_heads_,
+      static_cast<int>(d_k), static_cast<int>(d_v), past_state_tensor, state_shape));
   Tensor* present_state_tensor = context->Output(1, state_shape);
   float* state_data = present_state_tensor->MutableData<float>();
   int64_t state_per_head = d_k * d_v;
   int64_t total_state = batch_size * kv_num_heads_ * state_per_head;
 
   if (past_state_tensor != nullptr) {
-    const auto& ps_shape = past_state_tensor->Shape();
-    ORT_RETURN_IF_NOT(ps_shape.NumDimensions() == 4 &&
-                          ps_shape[0] == batch_size &&
-                          ps_shape[1] == kv_num_heads_ &&
-                          ps_shape[2] == d_k &&
-                          ps_shape[3] == d_v,
-                      "past_state must be (B, H_kv, d_k, d_v)");
     const float* ps_data = past_state_tensor->Data<float>();
     std::memcpy(state_data, ps_data, static_cast<size_t>(total_state) * sizeof(float));
   } else {

onnxruntime/contrib_ops/cpu/bert/linear_attention.h

@@ -23,6 +23,9 @@ class LinearAttention final : public OpKernel {
   std::string update_rule_;
   float scale_;
   int chunk_size_;
+  // Always 0 on CPU (a state window is CUDA-only), but kept so the shared shape helper in
+  // linear_attention_helper.h is driven the same way on every EP.
+  int state_window_;
 };
 
 }  // namespace contrib

onnxruntime/contrib_ops/cpu/bert/linear_attention_helper.h

@@ -0,0 +1,71 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "contrib_ops/cpu/bert/attention_common.h"
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+#include "core/framework/tensor_shape.h"
+#include "core/providers/common.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace linear_attention_helper {
+
+// Reads and validates the optional `state_window` attribute.
+//
+// 0 (the default, i.e. attribute absent) selects the legacy unwindowed state layout. Every model
+// exported before the attribute existed lands here, so this must stay the default.
+template <typename TOpKernelInfo>
+inline Status ParseStateWindow(const TOpKernelInfo& info, int& state_window) {
+  const int64_t value = info.template GetAttrOrDefault<int64_t>("state_window", 0);
+  if (value < 0 || value > kMaxStateWindow) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "state_window must be in [0, ", kMaxStateWindow, "], got ", value);
+  }
+  state_window = static_cast<int>(value);
+  return Status::OK();
+}
+
+// Derives the expected past_state / present_state shape and validates past_state against it.
+//
+// state_window == 0 -> (batch_size, kv_num_heads, d_k, d_v). A single state with no window axis:
+// the backward-compatible layout that models exported before the attribute existed use.
+//
+// state_window == W > 0 -> (W, batch_size, kv_num_heads, d_k, d_v). The window holds the recurrent
+// state after each of the last W tokens, right-aligned: slot j is the state after token
+// (seq_len - W + j), so slot W-1 is the state after the last token and is the only slot read back
+// as past_state. W == 1 is therefore the legacy layout with a leading unit axis.
+//
+// The window axis leads the batch axis so that each slot is one contiguous
+// (batch_size, kv_num_heads, d_k, d_v) block. That keeps "fetch/replace the last state" a single
+// contiguous range for any batch size, which is what a speculative decoder needs when it crops
+// the state back to an accepted prefix.
+template <typename T>
+Status CheckInputs(int state_window,
+                   int batch_size,
+                   int kv_num_heads,
+                   int d_k,
+                   int d_v,
+                   const T* past_state,
+                   TensorShape& state_shape) {
+  state_shape = state_window > 0
+                    ? TensorShape({state_window, batch_size, kv_num_heads, d_k, d_v})
+                    : TensorShape({batch_size, kv_num_heads, d_k, d_v});
+
+  if (past_state != nullptr && past_state->Shape() != state_shape) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "Input 'past_state' is expected to have shape ", state_shape.ToString(),
+                           ", got ", past_state->Shape().ToString(),
+                           ". LinearAttention uses (batch_size, kv_num_heads, d_k, d_v) when the "
+                           "state_window attribute is absent or 0, and "
+                           "(state_window, batch_size, kv_num_heads, d_k, d_v) otherwise.");
+  }
+
+  return Status::OK();
+}
+
+}  // namespace linear_attention_helper
+}  // namespace contrib
+}  // namespace onnxruntime