[CUDA] Add VarlenCausalConvWithState for continuous batching

docs/ContribOperators.md

@@ -122,6 +122,7 @@ Do not modify directly.*
   * <a href="#com.microsoft.Trilu">com.microsoft.Trilu</a>
   * <a href="#com.microsoft.UnfoldTensor">com.microsoft.UnfoldTensor</a>
   * <a href="#com.microsoft.Unique">com.microsoft.Unique</a>
+  * <a href="#com.microsoft.VarlenCausalConvWithState">com.microsoft.VarlenCausalConvWithState</a>
   * <a href="#com.microsoft.WhisperBeamSearch">com.microsoft.WhisperBeamSearch</a>
   * <a href="#com.microsoft.WordConvEmbedding">com.microsoft.WordConvEmbedding</a>
   * <sub>experimental</sub> <a href="#com.microsoft.IsAllFinite">com.microsoft.IsAllFinite</a>
@@ -7006,6 +7007,90 @@ This version of the operator has been available since version 1 of the 'com.micr
 </dl>
 
 
+### <a name="com.microsoft.VarlenCausalConvWithState"></a><a name="com.microsoft.varlencausalconvwithstate">**com.microsoft.VarlenCausalConvWithState**</a>
+
+  Stateful causal depthwise convolution over a packed, token-major batch of variable-length
+  sequences (CUDA only).
+
+  input and output have shape (total_tokens, channels). cumulative_sequence_length is a
+  device-resident int32 tensor of shape (batch_size + 1); sequence i occupies
+  [cumulative_sequence_length[i], cumulative_sequence_length[i + 1]). Every sequence contributes
+  at least one token. weight has shape (channels, 1, kernel_size), and optional bias has shape
+  (channels). The convolution never reads across a sequence boundary.
+
+  initial_state is required and has shape (batch_size, channels, kernel_size - 1). It contains
+  the committed raw activation samples immediately preceding this call. final_state has the same
+  shape and type and is fully written with the state after each sequence's final token. State
+  uses the activation type because it stores raw samples, not accumulated convolution values.
+  initial_state and final_state may use the same allocation. Such in-place execution is
+  transaction-safe only when the whole operator call is unconditionally committed; a caller that
+  may select a prefix or roll back must preserve initial_state and commit one of the separately
+  produced states instead.
+
+  When requested, prefix_states has shape
+  (max_checkpoints, batch_size, channels, kernel_size - 1). Slot j for request b is the state
+  after local token j when j < min(max_checkpoints, sequence_length[b]). Other slots are
+  unspecified and must not be read. max_checkpoints is static, defaults to zero, and is at most 8.
+  This output lets a transactional caller commit any produced prefix without rerunning the
+  convolution.
+
+  For memory-safety containment, each CUDA work item validates cumulative_sequence_length[0] == 0,
+  cumulative_sequence_length[batch_size] == total_tokens, and its local range
+  0 <= start < end <= total_tokens before accessing input, state, output, or checkpoints.
+  Malformed offsets cause affected work to return without those accesses; outputs are unspecified.
+  This device-side containment is not a synchronous validation or rejection mechanism.
+
+  The optional activation attribute supports none, SiLU, and Swish.
+
+#### Version
+
+This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
+
+#### Attributes
+
+<dl>
+<dt><tt>activation</tt> : string</dt>
+<dd>Fused activation function. One of: 'silu', 'swish', 'none'. Default is 'none'.</dd>
+<dt><tt>max_checkpoints</tt> : int</dt>
+<dd>Static number of per-request prefix states to expose. Checkpoint j is the state after local token j when that token exists. Unwritten slots are unspecified. Valid range is [0, 8].</dd>
+</dl>
+
+#### Inputs
+
+<dl>
+<dt><tt>input</tt> : T</dt>
+<dd>Token-major packed input with shape (total_tokens, channels).</dd>
+<dt><tt>weight</tt> : T</dt>
+<dd>Depthwise convolution kernel with shape (channels, 1, kernel_size).</dd>
+<dt><tt>cumulative_sequence_length</tt> : M</dt>
+<dd>Device tensor with shape (batch_size + 1) giving the half-open packed token range of each sequence.</dd>
+<dt><tt>bias</tt> (optional) : T</dt>
+<dd>Optional per-channel bias with shape (channels). Because the following initial_state input is required, an omitted bias must still occupy this position as an empty input name so initial_state stays at input index 4.</dd>
+<dt><tt>initial_state</tt> : T</dt>
+<dd>Required committed carry state with shape (batch_size, channels, kernel_size - 1).</dd>
+</dl>
+
+#### Outputs (2 - 3)
+
+<dl>
+<dt><tt>output</tt> : T</dt>
+<dd>Token-major convolution output with the same shape as input.</dd>
+<dt><tt>final_state</tt> : T</dt>
+<dd>Fully written state after each sequence's final token, with shape (batch_size, channels, kernel_size - 1).</dd>
+<dt><tt>prefix_states</tt> (optional) : T</dt>
+<dd>Optional prefix checkpoints with shape (max_checkpoints, batch_size, channels, kernel_size - 1). Unwritten slots are unspecified.</dd>
+</dl>
+
+#### Type Constraints
+
+<dl>
+<dt><tt>T</tt> : tensor(float), tensor(float16), tensor(bfloat16)</dt>
+<dd>Constrain input and output types to float tensors.</dd>
+<dt><tt>M</tt> : tensor(int32)</dt>
+<dd>Constrain cumulative_sequence_length to a device int32 tensor.</dd>
+</dl>
+
+
 ### <a name="com.microsoft.WhisperBeamSearch"></a><a name="com.microsoft.whisperbeamsearch">**com.microsoft.WhisperBeamSearch**</a>
 
   Beam Search for whisper model, especially with cross_qk features etc.

docs/OperatorKernels.md

@@ -1141,6 +1141,7 @@ The **OpSet Version** column uses the following notation:
 |TransposeMatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
 |Trilu|*in* X:**T**<br> *in* k:**tensor(int64)**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
 |UnfoldTensor|*in* input:**T**<br> *out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
+|VarlenCausalConvWithState|*in* input:**T**<br> *in* weight:**T**<br> *in* cumulative_sequence_length:**M**<br> *in* bias:**T**<br> *in* initial_state:**T**<br> *out* output:**T**<br> *out* final_state:**T**<br> *out* prefix_states:**T**|1+|**M** = tensor(int32)<br/> **T** = tensor(bfloat16), tensor(float), tensor(float16)|
 |WhisperBeamSearch|*in* input_ids:**F**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* num_beams:**I**<br> *in* num_return_sequences:**I**<br> *in* length_penalty:**T**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**M**<br> *in* prefix_vocab_mask:**M**<br> *in* attention_mask:**I**<br> *in* decoder_input_ids:**I**<br> *in* logits_processor:**I**<br> *in* cross_qk_layer_head:**I**<br> *in* extra_decoding_ids:**I**<br> *in* temperature:**T**<br> *out* sequences:**I**<br> *out* sequences_scores:**T**<br> *out* scores:**T**<br> *out* cross_qk:**V**<br> *out* non_speech_probs:**T**|1+|**T** = tensor(float), tensor(float16)|
 | |
 | |

onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc

@@ -232,7 +232,7 @@ Status CausalConvWithState<T>::Compute(OpKernelContext* context) const {
     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));
+        static_cast<int>(pad), past_state_tensor, state_shape, "CausalConvWithState"));
     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_helper.h

@@ -3,6 +3,8 @@
 
 #pragma once
 
+#include <string_view>
+
 #include "contrib_ops/cpu/bert/attention_common.h"
 #include "core/common/common.h"
 #include "core/framework/op_kernel.h"
@@ -49,7 +51,8 @@ Status CheckInputs(int state_window,
                    int channels,
                    int state_length,
                    const T* past_state,
-                   TensorShape& state_shape) {
+                   TensorShape& state_shape,
+                   std::string_view op_name) {
   state_shape = state_window > 0
                     ? TensorShape({state_window, batch_size, channels, state_length})
                     : TensorShape({batch_size, channels, state_length});
@@ -58,7 +61,8 @@ Status CheckInputs(int state_window,
     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 "
+                           ". ", op_name,
+                           " 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.");
   }

onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc

@@ -89,7 +89,7 @@ Status CausalConvWithState<T>::ComputeInternal(OpKernelContext* context) const {
   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));
+      state_window_, batch_size, channels, pad, past_state_tensor, state_shape, "CausalConvWithState"));
 
   // Allocate outputs
   Tensor* output_tensor = context->Output(0, input_shape);

onnxruntime/contrib_ops/cuda/bert/varlen_causal_conv_with_state.cc

@@ -0,0 +1,147 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/cuda/bert/varlen_causal_conv_with_state.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"
+
+#include <limits>
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+using namespace onnxruntime::cuda;  // CudaKernel, Stream, GetDeviceProp, ToCudaType
+
+#define REGISTER_KERNEL_TYPED(T)                                        \
+  ONNX_OPERATOR_TYPED_KERNEL_EX(                                        \
+      VarlenCausalConvWithState,                                        \
+      kMSDomain,                                                        \
+      1,                                                                \
+      T,                                                                \
+      kCudaExecutionProvider,                                           \
+      (*KernelDefBuilder::Create())                                     \
+          .MayInplace(4, 1)                                             \
+          .TypeConstraint("T", DataTypeImpl::GetTensorType<T>())        \
+          .TypeConstraint("M", DataTypeImpl::GetTensorType<int32_t>()), \
+      VarlenCausalConvWithState<T>);
+
+REGISTER_KERNEL_TYPED(float)
+REGISTER_KERNEL_TYPED(MLFloat16)
+REGISTER_KERNEL_TYPED(BFloat16)
+
+template <typename T>
+VarlenCausalConvWithState<T>::VarlenCausalConvWithState(const OpKernelInfo& info) : CudaKernel(info) {
+  activation_ = info.GetAttrOrDefault<std::string>("activation", "none");
+  ORT_ENFORCE(activation_ == "none" || activation_ == "silu" || activation_ == "swish",
+              "activation must be one of: none, silu, swish");
+
+  const int64_t max_checkpoints = info.GetAttrOrDefault<int64_t>("max_checkpoints", 0);
+  ORT_ENFORCE(max_checkpoints >= 0 && max_checkpoints <= kMaxStateWindow,
+              "max_checkpoints must be in [0, ", kMaxStateWindow, "]");
+  max_checkpoints_ = static_cast<int>(max_checkpoints);
+}
+
+template <typename T>
+Status VarlenCausalConvWithState<T>::ComputeInternal(OpKernelContext* context) const {
+  const Tensor* input_tensor = context->Input<Tensor>(0);
+  const Tensor* weight_tensor = context->Input<Tensor>(1);
+  const Tensor* cu_seqlens_tensor = context->Input<Tensor>(2);
+  const Tensor* bias_tensor = context->Input<Tensor>(3);  // optional
+  const Tensor* initial_state_tensor = context->Input<Tensor>(4);
+
+  ORT_RETURN_IF_NOT(input_tensor != nullptr, "input is required");
+  ORT_RETURN_IF_NOT(weight_tensor != nullptr, "weight is required");
+  ORT_RETURN_IF_NOT(cu_seqlens_tensor != nullptr, "cumulative_sequence_length input is required");
+  ORT_RETURN_IF_NOT(initial_state_tensor != nullptr, "initial_state input is required");
+
+  const auto& input_shape = input_tensor->Shape();
+  const auto& weight_shape = weight_tensor->Shape();
+
+  ORT_RETURN_IF_NOT(input_shape.NumDimensions() == 2,
+                    "input must be rank 2 (total_tokens, channels), got rank ", input_shape.NumDimensions());
+  ORT_RETURN_IF_NOT(weight_shape.NumDimensions() == 3,
+                    "weight must be rank 3 (channels, 1, kernel_size), got rank ", weight_shape.NumDimensions());
+
+  const auto& cu_seqlens_shape = cu_seqlens_tensor->Shape();
+  ORT_RETURN_IF_NOT(cu_seqlens_shape.NumDimensions() == 1,
+                    "cumulative_sequence_length must be rank 1 (batch_size + 1), got rank ",
+                    cu_seqlens_shape.NumDimensions());
+  ORT_RETURN_IF_NOT(cu_seqlens_shape[0] >= 2,
+                    "cumulative_sequence_length must have at least 2 elements (batch_size >= 1), got ",
+                    cu_seqlens_shape[0]);
+  // batch_size = cu_seqlens.Shape()[0] - 1, never derived from total_tokens: a ragged batch has
+  // no fixed relationship between the two beyond total_tokens >= batch_size.
+  const int64_t batch_size_64 = cu_seqlens_shape[0] - 1;
+  ORT_RETURN_IF_NOT(batch_size_64 <= std::numeric_limits<int>::max(),
+                    "batch size is too large for the CUDA kernel");
+  const int batch_size = static_cast<int>(batch_size_64);
+
+  const int64_t total_tokens_64 = input_shape[0];
+  const int64_t channels_64 = input_shape[1];
+  ORT_RETURN_IF_NOT(total_tokens_64 <= std::numeric_limits<int>::max() &&
+                        channels_64 <= std::numeric_limits<int>::max(),
+                    "input dimensions are too large for the CUDA kernel");
+  ORT_RETURN_IF_NOT(total_tokens_64 >= batch_size_64,
+                    "total_tokens must be at least batch_size because every sequence must contain a token");
+  ORT_RETURN_IF_NOT(channels_64 > 0, "input channel dimension must be positive");
+  const int channels = static_cast<int>(channels_64);
+
+  ORT_RETURN_IF_NOT(weight_shape[0] == channels_64,
+                    "weight[0] (", weight_shape[0], ") must match input channels (", channels_64, ")");
+  ORT_RETURN_IF_NOT(weight_shape[1] == 1,
+                    "weight[1] must be 1 for depthwise convolution, got ", weight_shape[1]);
+  const int64_t kernel_size_64 = weight_shape[2];
+  ORT_RETURN_IF_NOT(kernel_size_64 >= 1 && kernel_size_64 <= std::numeric_limits<int>::max(),
+                    "weight last dim (kernel_size) must be positive, got ", kernel_size_64);
+  const int kernel_size = static_cast<int>(kernel_size_64);
+  const int pad = kernel_size - 1;
+
+  if (bias_tensor != nullptr) {
+    const auto& bias_shape = bias_tensor->Shape();
+    ORT_RETURN_IF_NOT(bias_shape.NumDimensions() == 1 && bias_shape[0] == channels_64,
+                      "bias must have shape (", channels_64, "), got ", bias_shape.ToString());
+  }
+
+  const TensorShape state_shape({batch_size_64, channels_64, pad});
+  ORT_RETURN_IF_NOT(initial_state_tensor->Shape() == state_shape,
+                    "initial_state must have shape ", state_shape.ToString(), ", got ",
+                    initial_state_tensor->Shape().ToString());
+
+  Tensor* output_tensor = context->Output(0, input_shape);
+  Tensor* final_state_tensor = context->Output(1, state_shape);
+  const TensorShape checkpoint_shape({max_checkpoints_, batch_size_64, channels_64, pad});
+  Tensor* prefix_states_tensor = context->Output(2, checkpoint_shape);
+
+  bool apply_silu = (activation_ == "silu" || activation_ == "swish");
+
+  // total_tokens == batch_size selects a candidate decode path. That path still reads and
+  // validates each exact [b, b + 1] offset interval before accessing any data.
+  const bool all_ones = (total_tokens_64 == batch_size_64);
+
+  typedef typename OrtToCudaType<T>::type CudaT;
+
+  return LaunchVarlenCausalConvWithStateKernel<CudaT>(
+      Stream(context),
+      reinterpret_cast<const CudaT*>(input_tensor->Data<T>()),
+      reinterpret_cast<const CudaT*>(weight_tensor->Data<T>()),
+      bias_tensor ? reinterpret_cast<const CudaT*>(bias_tensor->Data<T>()) : nullptr,
+      reinterpret_cast<const CudaT*>(initial_state_tensor->Data<T>()),
+      reinterpret_cast<CudaT*>(output_tensor->MutableData<T>()),
+      reinterpret_cast<CudaT*>(final_state_tensor->MutableData<T>()),
+      prefix_states_tensor ? reinterpret_cast<CudaT*>(prefix_states_tensor->MutableData<T>()) : nullptr,
+      cu_seqlens_tensor->Data<int32_t>(),
+      batch_size,
+      static_cast<int>(total_tokens_64),
+      all_ones,
+      channels,
+      kernel_size,
+      apply_silu,
+      GetDeviceProp().maxThreadsPerBlock,
+      max_checkpoints_);
+}
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/bert/varlen_causal_conv_with_state.h

@@ -0,0 +1,56 @@
+// 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"
+
+#include <cstdint>
+#include <string>
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+template <typename T>
+class VarlenCausalConvWithState final : public onnxruntime::cuda::CudaKernel {
+ public:
+  VarlenCausalConvWithState(const OpKernelInfo& info);
+  Status ComputeInternal(OpKernelContext* context) const override;
+
+ private:
+  std::string activation_;
+  int max_checkpoints_;
+};
+
+// Launches the packed varlen causal-conv recurrence.
+//
+// input/output hold every sequence's tokens back to back along axis 0. cu_seqlens is a device
+// int32 tensor of length (batch_size + 1): sequence r occupies the half-open token range
+// [cu_seqlens[r], cu_seqlens[r + 1]). Offset *values* are a trusted producer precondition -- the
+// device kernel validates the global endpoints and its local interval before any data access.
+// Invalid offsets are contained by returning without data accesses; outputs are unspecified.
+template <typename T>
+Status LaunchVarlenCausalConvWithStateKernel(
+    cudaStream_t stream,
+    const T* input,             // [total_tokens, channels]
+    const T* weight,            // [channels, 1, kernel_size]
+    const T* bias,              // [channels] or nullptr
+    const T* initial_state,     // [batch_size, channels, kernel_size - 1], required
+    T* output,                  // [total_tokens, channels]
+    T* final_state,             // [batch_size, channels, kernel_size - 1]
+    T* prefix_states,           // [max_checkpoints, batch_size, channels, kernel_size - 1] or nullptr
+    const int32_t* cu_seqlens,  // [batch_size + 1], device-resident
+    int batch_size,
+    int total_tokens,
+    bool all_ones,
+    int channels,
+    int kernel_size,
+    bool apply_silu,
+    int max_threads_per_block,
+    int max_checkpoints);
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime