Extend GDN and causal convolution fused state handling

onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc

@@ -41,6 +41,12 @@ REGISTER_KERNEL_TYPED(float)
 
 template <typename T>
 CausalConvWithState<T>::CausalConvWithState(const OpKernelInfo& info) : OpKernel(info) {
+  ORT_ENFORCE(info.GetAttrOrDefault<int64_t>("input_channels", 0) == 0 &&
+                  info.GetAttrOrDefault<int64_t>("input_channel_offset", 0) == 0 &&
+                  info.GetAttrOrDefault<int64_t>("state_length", 0) == 0 &&
+                  info.GetAttrOrDefault<int64_t>("state_channels_last", -1) == -1 &&
+                  info.node().OutputDefs().size() <= 2,
+              "CPU CausalConvWithState does not support extended channel/state options (CUDA only)");
   int64_t ndim = info.GetAttrOrDefault<int64_t>("ndim", 1);
   ORT_ENFORCE(ndim == 1, "CPU CausalConvWithState only supports ndim=1");
   ndim_ = static_cast<int>(ndim);

onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc

@@ -6,6 +6,7 @@
 #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 "core/common/safeint.h"
 
 namespace onnxruntime {
 namespace contrib {
@@ -45,6 +46,13 @@ CausalConvWithState<T>::CausalConvWithState(const OpKernelInfo& info) : CudaKern
   // 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_));
+  input_channels_ = SafeInt<int>(info.GetAttrOrDefault<int64_t>("input_channels", 0));
+  input_channel_offset_ = SafeInt<int>(info.GetAttrOrDefault<int64_t>("input_channel_offset", 0));
+  state_length_ = SafeInt<int>(info.GetAttrOrDefault<int64_t>("state_length", 0));
+  state_channels_last_ = SafeInt<int>(info.GetAttrOrDefault<int64_t>("state_channels_last", -1));
+  ORT_ENFORCE(input_channels_ >= 0 && input_channel_offset_ >= 0 && state_length_ >= 0 &&
+                  state_channels_last_ >= -1 && state_channels_last_ <= 1,
+              "invalid packed-channel or state attributes");
 }
 
 template <typename T>
@@ -60,6 +68,47 @@ Status CausalConvWithState<T>::ComputeInternal(OpKernelContext* context) const {
   const auto& input_shape = input_tensor->Shape();
   const auto& weight_shape = weight_tensor->Shape();
 
+  const bool extended = input_channels_ != 0 || input_channel_offset_ != 0 || state_length_ != 0 ||
+                        state_channels_last_ != -1 || context->OutputCount() > 2;
+  if (extended) {
+    ORT_RETURN_IF_NOT(channels_last_ && state_window_ == 0 && input_shape.NumDimensions() == 3 &&
+                          weight_shape.NumDimensions() == 3 && weight_shape[1] == 1,
+                      "extended CausalConvWithState requires rank-3 channels-last input, rank-3 weight and state_window=0");
+    const int batch = SafeInt<int>(input_shape[0]);
+    const int length = SafeInt<int>(input_shape[1]);
+    const int width = SafeInt<int>(input_shape[2]);
+    const int channels = input_channels_ == 0 ? width : input_channels_;
+    const int kernel = SafeInt<int>(weight_shape[2]);
+    const int pad = SafeInt<int>(kernel - 1) * dilation_;
+    const int state_length = state_length_ == 0 ? pad : state_length_;
+    const bool state_last = state_channels_last_ == -1 ? channels_last_ : state_channels_last_ != 0;
+    ORT_RETURN_IF_NOT(batch > 0 && batch <= 65535 && length > 0 && length <= 65535 &&
+                          channels > 0 && kernel > 0 && input_channel_offset_ <= width &&
+                          channels <= width - input_channel_offset_ && weight_shape[0] == channels &&
+                          state_length >= pad,
+                      "invalid extended CausalConvWithState dimensions");
+    ORT_RETURN_IF_NOT(bias_tensor == nullptr || bias_tensor->Shape() == TensorShape({channels}),
+                      "bias must match selected input channels");
+    const TensorShape state_shape = state_last ? TensorShape({batch, state_length, channels})
+                                               : TensorShape({batch, channels, state_length});
+    ORT_RETURN_IF_NOT(past_state_tensor == nullptr || past_state_tensor->Shape() == state_shape,
+                      "past_state must match configured state layout and length");
+    Tensor* output = context->Output(0, TensorShape({batch, length, channels}));
+    Tensor* present = context->Output(1, state_shape);
+    Tensor* first = context->OutputCount() > 2 ? context->Output(2, state_shape) : nullptr;
+    using CudaT = typename OrtToCudaType<T>::type;
+    return LaunchCausalConvProjectionKernel<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,
+        past_state_tensor ? reinterpret_cast<const CudaT*>(past_state_tensor->Data<T>()) : nullptr,
+        reinterpret_cast<CudaT*>(output->MutableData<T>()),
+        reinterpret_cast<CudaT*>(present->MutableData<T>()),
+        first ? reinterpret_cast<CudaT*>(first->MutableData<T>()) : nullptr,
+        batch, channels, length, kernel, dilation_, width, input_channel_offset_, state_length,
+        state_last, activation_ != "none");
+  }
+
   // Validate input rank and weight rank
   if (channels_last_) {
     // (batch_size, sequence_length, d_1, ..., d_n): any number of trailing channel axes, so a

onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h

@@ -25,6 +25,10 @@ class CausalConvWithState final : public onnxruntime::cuda::CudaKernel {
   std::string activation_;
   // Leading (axis-0) extent of past_state / present_state; 0 means no window axis (single state).
   int state_window_;
+  int input_channels_;
+  int input_channel_offset_;
+  int state_length_;
+  int state_channels_last_;
 };
 
 }  // namespace cuda

onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu

@@ -27,6 +27,44 @@ namespace cuda {
 
 namespace {
 
+template <typename T>
+__global__ void CausalConvProjectionKernel(
+    const T* input, const T* weight, const T* bias, const T* past,
+    T* output, T* present, T* first_state, int channels, int length,
+    int kernel_size, int dilation, int input_width, int channel_offset,
+    int state_length, CausalConvLayout state_layout, bool apply_silu) {
+  const int b = blockIdx.y;
+  const int c = blockIdx.x * blockDim.x + threadIdx.x;
+  if (c >= channels) return;
+  const int pad = (kernel_size - 1) * dilation;
+  const int position = blockIdx.z;
+  {
+    float sum = bias == nullptr ? 0.0f : to_float(bias[c]);
+    for (int tap = 0; tap < kernel_size; ++tap) {
+      const int source = position - pad + tap * dilation;
+      const float value = source >= 0
+                              ? to_float(input[(static_cast<int64_t>(b) * length + source) * input_width + channel_offset + c])
+                              : (past == nullptr ? 0.0f : to_float(past[state_layout.Offset(b, state_length + source, c)]));
+      sum += value * to_float(weight[c * kernel_size + tap]);
+    }
+    if (apply_silu) sum = sum / (1.0f + expf(-sum));
+    output[(static_cast<int64_t>(b) * length + position) * channels + c] = from_float<T>(sum);
+  }
+  if (position != 0) return;
+  for (int position = 0; position < state_length; ++position) {
+    const int source = length - state_length + position;
+    present[state_layout.Offset(b, position, c)] =
+        source >= 0 ? input[(static_cast<int64_t>(b) * length + source) * input_width + channel_offset + c]
+                    : (past == nullptr ? from_float<T>(0.0f) : past[state_layout.Offset(b, state_length + source, c)]);
+    if (first_state != nullptr) {
+      const int first_source = 1 - state_length + position;
+      first_state[state_layout.Offset(b, position, c)] =
+          first_source >= 0 ? input[static_cast<int64_t>(b) * length * input_width + channel_offset + c]
+                            : (past == nullptr ? from_float<T>(0.0f) : past[state_layout.Offset(b, state_length + first_source, c)]);
+    }
+  }
+}
+
 __device__ __forceinline__ float silu_fn(float x) {
   return x / (1.0f + expf(-x));
 }
@@ -428,6 +466,29 @@ __global__ void CausalConvPrefillKernelChannelsLast(
 
 }  // anonymous namespace
 
+template <typename T>
+Status LaunchCausalConvProjectionKernel(
+    cudaStream_t stream, const T* input, const T* weight, const T* bias, const T* past,
+    T* output, T* present, T* first_state, int batch, int channels, int length,
+    int kernel_size, int dilation, int input_width, int channel_offset,
+    int state_length, bool state_channels_last, bool apply_silu) {
+  const dim3 grid((channels - 1) / 128 + 1, batch, length);
+  CausalConvProjectionKernel<T><<<grid, 128, 0, stream>>>(
+      input, weight, bias, past, output, present, first_state, channels, length,
+      kernel_size, dilation, input_width, channel_offset, state_length,
+      MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu);
+  return CUDA_CALL(cudaGetLastError());
+}
+
+#define INSTANTIATE_PROJECTION(T)                                       \
+  template Status LaunchCausalConvProjectionKernel<T>(                  \
+      cudaStream_t, const T*, const T*, const T*, const T*, T*, T*, T*, \
+      int, int, int, int, int, int, int, int, bool, bool);
+INSTANTIATE_PROJECTION(float)
+INSTANTIATE_PROJECTION(half)
+INSTANTIATE_PROJECTION(__nv_bfloat16)
+#undef INSTANTIATE_PROJECTION
+
 template <typename T>
 Status LaunchCausalConvWithStateKernel(
     cudaStream_t stream,

onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h

@@ -31,6 +31,13 @@ inline CausalConvLayout MakeCausalConvLayout(bool channels_last, int64_t channel
                        : CausalConvLayout{channels * length, 1, length};
 }
 
+template <typename T>
+Status LaunchCausalConvProjectionKernel(
+    cudaStream_t stream, const T* input, const T* weight, const T* bias, const T* past,
+    T* output, T* present, T* first_state, int batch, int channels, int length,
+    int kernel_size, int dilation, int input_width, int channel_offset,
+    int state_length, bool state_channels_last, bool apply_silu);
+
 // Fused causal depthwise conv1d + activation + state management.
 // One thread block per (batch, channel). For decode (L=1), this is a simple
 // dot product from shared memory. For prefill (L>1), each thread handles