From ea8791f7dcf9207948f811690912af5e76096b9d Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 19:14:59 +0000 Subject: [PATCH 01/17] Extend GDN and causal convolution fused state handling --- docs/ContribOperators.md | 1 - .../cpu/bert/causal_conv_with_state.cc | 6 + .../cuda/bert/causal_conv_with_state.cc | 49 +++++ .../cuda/bert/causal_conv_with_state.h | 4 + .../cuda/bert/causal_conv_with_state_impl.cu | 61 ++++++ .../cuda/bert/causal_conv_with_state_impl.h | 7 + .../contrib_ops/cuda/bert/gated_delta_net.cc | 98 ++++++--- .../contrib_ops/cuda/bert/gated_delta_net.h | 3 +- .../cuda/bert/gated_delta_net_impl.cu | 127 +++++++---- .../cuda/bert/gated_delta_net_impl.h | 15 +- .../cuda/bert/gated_delta_net_mma.cuh | 2 +- .../cuda/bert/gated_delta_net_plan.h | 2 +- .../cuda/bert/gated_delta_net_split_impl.cu | 4 +- .../webgpu/bert/causal_conv_with_state.cc | 6 + .../webgpu/bert/gated_delta_net.cc | 11 +- .../core/graph/contrib_ops/bert_defs.cc | 207 +++++++++++++----- .../causal_conv_with_state_op_test.cc | 121 ++++++++++ .../contrib_ops/gated_delta_net_op_test.cc | 165 +++++++++++++- 18 files changed, 751 insertions(+), 138 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 2703de7cc1fae..3510b52792635 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -8771,4 +8771,3 @@ No versioning maintained for experimental ops.
T : tensor(float)
Constrain input and output types to float32 tensors.
- 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 7585993088805..bf81d935636fe 100644 --- a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc @@ -41,6 +41,12 @@ REGISTER_KERNEL_TYPED(float) template CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : OpKernel(info) { + ORT_ENFORCE(info.GetAttrOrDefault("input_channels", 0) == 0 && + info.GetAttrOrDefault("input_channel_offset", 0) == 0 && + info.GetAttrOrDefault("state_length", 0) == 0 && + info.GetAttrOrDefault("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("ndim", 1); ORT_ENFORCE(ndim == 1, "CPU CausalConvWithState only supports ndim=1"); ndim_ = static_cast(ndim); 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 c79859f735c7f..f7a8d94c2b80a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc +++ b/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::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(info.GetAttrOrDefault("input_channels", 0)); + input_channel_offset_ = SafeInt(info.GetAttrOrDefault("input_channel_offset", 0)); + state_length_ = SafeInt(info.GetAttrOrDefault("state_length", 0)); + state_channels_last_ = SafeInt(info.GetAttrOrDefault("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 @@ -60,6 +68,47 @@ Status CausalConvWithState::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(input_shape[0]); + const int length = SafeInt(input_shape[1]); + const int width = SafeInt(input_shape[2]); + const int channels = input_channels_ == 0 ? width : input_channels_; + const int kernel = SafeInt(weight_shape[2]); + const int pad = SafeInt(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::type; + return LaunchCausalConvProjectionKernel( + Stream(context), reinterpret_cast(input_tensor->Data()), + reinterpret_cast(weight_tensor->Data()), + bias_tensor ? reinterpret_cast(bias_tensor->Data()) : nullptr, + past_state_tensor ? reinterpret_cast(past_state_tensor->Data()) : nullptr, + reinterpret_cast(output->MutableData()), + reinterpret_cast(present->MutableData()), + first ? reinterpret_cast(first->MutableData()) : 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 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 0a1f23f3abfd3..f0807d69c7da3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h +++ b/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 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 0c0e1da836521..eea53e8bbe272 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 @@ -27,6 +27,44 @@ namespace cuda { namespace { +template +__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(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(b) * length + position) * channels + c] = from_float(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(b) * length + source) * input_width + channel_offset + c] + : (past == nullptr ? from_float(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(b) * length * input_width + channel_offset + c] + : (past == nullptr ? from_float(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 +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<<>>( + 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( \ + 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 Status LaunchCausalConvWithStateKernel( cudaStream_t stream, 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 968cb9507339d..95d98dc96cdcd 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 @@ -31,6 +31,13 @@ inline CausalConvLayout MakeCausalConvLayout(bool channels_last, int64_t channel : CausalConvLayout{channels * length, 1, length}; } +template +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 diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index ff7f485dc1284..380e99746af31 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -22,19 +22,22 @@ using onnxruntime::cuda::CudaKernel; using onnxruntime::cuda::OrtToCudaType; namespace gdn = gated_delta_net; -#define REGISTER_KERNEL_TYPED(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - GatedDeltaNet, \ - kMSDomain, \ - 1, \ - T, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("TS", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ - .InputMemoryType(OrtMemTypeCPUInput, 10) \ - .MayInplace(6, 1), \ +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + GatedDeltaNet, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TS", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TG", {DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType()}) \ + .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ + .InputMemoryType(OrtMemTypeCPUInput, 10) \ + .MayInplace(6, 1), \ GatedDeltaNet); REGISTER_KERNEL_TYPED(float) @@ -42,7 +45,14 @@ REGISTER_KERNEL_TYPED(MLFloat16) REGISTER_KERNEL_TYPED(BFloat16) template -GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : CudaKernel(info) { +GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) + : CudaKernel(info), + mid_state_output_(info.GetAttrOrDefault("state_output", "state_update") == + "mid_state") { + ORT_ENFORCE(mid_state_output_ || + info.GetAttrOrDefault("state_output", "state_update") == + "state_update", + "state_output must be 'state_update' or 'mid_state'"); const std::string rule = info.GetAttrOrDefault("update_rule", "gated_delta"); if (rule == "linear") { update_rule_ = gdn::UpdateRule::kLinear; @@ -59,10 +69,10 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : CudaKernel(info) { const std::string gate = info.GetAttrOrDefault("gate_activation", "none"); if (gate == "none") { gate_activation_ = gdn::GateActivation::kNone; - } else if (gate == "qwen") { - gate_activation_ = gdn::GateActivation::kQwen; + } else if (gate == "softplus" || gate == "qwen") { + gate_activation_ = gdn::GateActivation::kSoftplusDecay; } else { - ORT_THROW("gate_activation must be one of: none, qwen"); + ORT_THROW("gate_activation must be one of: none, softplus"); } const std::string beta = info.GetAttrOrDefault("beta_activation", "none"); @@ -88,6 +98,9 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : CudaKernel(info) { forced_engine_ = gdn::EngineFromName( ParseEnvironmentVariableWithDefault("ORT_GDN_PLAN", "")); + if (mid_state_output_) { + forced_engine_ = gdn::Engine::kRecurrent; + } } template @@ -106,6 +119,15 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { const Tensor* capture_count = context->Input(9); const Tensor* state_update_active = context->Input(10); + ORT_RETURN_IF_NOT(query != nullptr, "query is required"); + for (const Tensor* gate : {decay, beta}) { + ORT_RETURN_IF_NOT(gate == nullptr || gate->IsDataType() || + gate->GetElementType() == query->GetElementType(), + "decay and beta must be float or match the query element type"); + } + ORT_RETURN_IF_NOT(!mid_state_output_ || state_update_capacity_ == 0, + "state_output='mid_state' cannot be combined with state_update_capacity"); + ORT_RETURN_IF_NOT((state_update_capacity_ > 0) == (capture_count != nullptr), "capture_count must be present exactly when state_update_capacity is positive"); @@ -120,7 +142,6 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(needs_beta == (beta != nullptr), "beta input presence must match update_rule"); - ORT_RETURN_IF_NOT(query != nullptr, "query is required"); const bool is_packed_qkv = key == nullptr && value == nullptr; ORT_RETURN_IF_NOT(is_packed_qkv || (key != nullptr && value != nullptr), "key and value must be both present or both absent"); @@ -191,6 +212,10 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(num_heads_q == num_heads_k, "num_heads_q (", num_heads_q, ") must equal num_heads_k (", num_heads_k, ")"); + if (mid_state_output_) { + ORT_RETURN_IF_NOT(num_heads_v >= num_heads_q, + "mid-state output requires value_heads >= query_heads"); + } ORT_RETURN_IF_NOT(num_heads_v > 0 && num_heads_q > 0 && num_heads_v % num_heads_q == 0, "num_heads_v (", num_heads_v, ") must be a positive multiple of num_heads_q (", num_heads_q, ")"); @@ -260,9 +285,9 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { b[token_dims] == num_heads_v, "beta must be [...tokens, num_heads_v]"); } - if (gate_activation_ == gdn::GateActivation::kQwen) { + if (gate_activation_ == gdn::GateActivation::kSoftplusDecay) { ORT_RETURN_IF_NOT(a_log != nullptr && dt_bias != nullptr, - "gate_activation=qwen requires a_log and dt_bias"); + "gate_activation=softplus requires a_log and dt_bias"); ORT_RETURN_IF_NOT(a_log->Shape().NumDimensions() == 1 && a_log->Shape()[0] == num_heads_v, "a_log must be [num_heads_v]"); ORT_RETURN_IF_NOT(dt_bias->Shape().NumDimensions() == 1 && @@ -270,7 +295,7 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { "dt_bias must be [num_heads_v]"); } else { ORT_RETURN_IF_NOT(a_log == nullptr && dt_bias == nullptr, - "a_log and dt_bias require gate_activation=qwen"); + "a_log and dt_bias require gate_activation=softplus"); } if (capture_count != nullptr) { @@ -293,10 +318,16 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { TensorShape state_shape{batch, num_heads_v, head_size_v, head_size_qk}; Tensor* final_state = context->Output(1, state_shape); - const int64_t state_update_width = static_cast(state_update_capacity_) * - (num_heads_v + num_heads_k * head_size_qk + - num_heads_v * head_size_v); - Tensor* state_update = context->Output(2, TensorShape{batch, state_update_width}); + Tensor* state_update = nullptr; + Tensor* mid_state = nullptr; + if (mid_state_output_) { + mid_state = context->Output(2, state_shape); + } else { + const int64_t state_update_width = static_cast(state_update_capacity_) * + (num_heads_v + num_heads_k * head_size_qk + + num_heads_v * head_size_v); + state_update = context->Output(2, TensorShape{batch, state_update_width}); + } gdn::Descriptor desc{}; desc.total_tokens = total_tokens; @@ -363,7 +394,9 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { const size_t query_row_bytes = static_cast(num_heads_q) * head_size_qk * sizeof(T); const size_t key_row_bytes = static_cast(num_heads_k) * head_size_qk * sizeof(T); const size_t value_row_bytes = static_cast(num_heads_v) * head_size_v * sizeof(T); - const size_t unpacked_qkv_bytes = is_packed_qkv + const bool direct_packed_decode = is_packed_qkv && plan.engine == gdn::Engine::kRecurrent && + plan.warp_specialized && total_tokens <= 2LL * batch; + const size_t unpacked_qkv_bytes = is_packed_qkv && !direct_packed_decode ? static_cast(total_tokens) * (query_row_bytes + key_row_bytes + value_row_bytes) : 0; @@ -377,7 +410,10 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { const CudaT* query_data = reinterpret_cast(query->Data()); const CudaT* key_data = key != nullptr ? reinterpret_cast(key->Data()) : nullptr; const CudaT* value_data = value != nullptr ? reinterpret_cast(value->Data()) : nullptr; - if (is_packed_qkv) { + if (direct_packed_decode) { + key_data = query_data + query_row_bytes / sizeof(T); + value_data = key_data + key_row_bytes / sizeof(T); + } else if (is_packed_qkv) { const size_t packed_row_bytes = query_row_bytes + key_row_bytes + value_row_bytes; auto* unpacked_query = workspace.get(); auto* unpacked_key = unpacked_query + static_cast(total_tokens) * query_row_bytes; @@ -402,15 +438,19 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { pack.query = query_data; pack.key = key_data; pack.value = value_data; + pack.packed_qkv = direct_packed_decode; pack.cu_seqlens = cu_seqlens != nullptr ? cu_seqlens->Data() : nullptr; pack.capture_count = desc.state_update_active ? capture_count->Data() : nullptr; - pack.decay = decay != nullptr ? decay->Data() : nullptr; - pack.beta = beta != nullptr ? beta->Data() : nullptr; + pack.decay = decay != nullptr ? decay->DataRaw() : nullptr; + pack.beta = beta != nullptr ? beta->DataRaw() : nullptr; + const Tensor* gate = decay != nullptr ? decay : beta; + pack.gates_in_input_type = gate != nullptr && !gate->IsDataType(); pack.initial_state = initial_state != nullptr ? initial_state->Data() : nullptr; pack.a_log = a_log != nullptr ? a_log->Data() : nullptr; pack.dt_bias = dt_bias != nullptr ? dt_bias->Data() : nullptr; pack.output = reinterpret_cast(output->MutableData()); pack.final_state = final_state != nullptr ? final_state->MutableData() : nullptr; + pack.mid_state = mid_state != nullptr ? mid_state->MutableData() : nullptr; pack.state_update = state_update != nullptr ? state_update->MutableData() : nullptr; // Only the inactive case is cleared. While capture is active the kernels write just the diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h index 26243a7e60ed6..1982ae2de7da4 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h @@ -14,7 +14,7 @@ namespace contrib { namespace cuda { template -class GatedDeltaNet final : public onnxruntime::cuda::CudaKernel { +class GatedDeltaNet : public onnxruntime::cuda::CudaKernel { public: GatedDeltaNet(const OpKernelInfo& info); Status ComputeInternal(OpKernelContext* context) const override; @@ -28,6 +28,7 @@ class GatedDeltaNet final : public onnxruntime::cuda::CudaKernel { int chunk_size_; int state_update_capacity_; bool qk_l2_norm_; + bool mid_state_output_; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu index 4729e31bd2e9d..048faa9c2f3f5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -172,7 +172,7 @@ __device__ __forceinline__ ChunkedSmem CarveChunked(char* raw) { // --------------------------------------------------------------------------- // Chunked engine // --------------------------------------------------------------------------- -template +template __global__ __launch_bounds__(kChunkedThreads) void GatedDeltaNetChunkedKernel( VariantPack pack, KernelParams p) { using L = Ld; @@ -237,10 +237,10 @@ __global__ __launch_bounds__(kChunkedThreads) void GatedDeltaNetChunkedKernel( const int64_t gi = tok * p.num_heads_v + hv; float g = 0.0f, be = 0.0f; if (tid < len) { - g = needs_decay ? EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv, + g = needs_decay ? EffectiveDecay(LoadGate(pack.decay, gi, InputTypeGates), pack.a_log, pack.dt_bias, hv, p.gate_activation) : 0.0f; - be = pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f; + be = pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, InputTypeGates), p.beta_activation) : 1.0f; } s.gc[tid] = g; s.beta[tid] = be; @@ -379,7 +379,7 @@ __device__ __forceinline__ float WarpReduceK(float v) { return v; } -template +template __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( VariantPack pack, KernelParams p, int d_v) { constexpr int kRegs = DK / 32; // state elements per lane @@ -401,9 +401,12 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( } const int capture_count = CaptureCount(pack, p, b, seq_len); - const int64_t q_stride = static_cast(p.num_heads_q) * DK; - const int64_t k_stride = static_cast(p.num_heads_k) * DK; + const int64_t packed_stride = (static_cast(p.num_heads_q) + p.num_heads_k) * DK + + static_cast(p.num_heads_v) * d_v; + const int64_t q_stride = PackedQkv ? packed_stride : static_cast(p.num_heads_q) * DK; + const int64_t k_stride = PackedQkv ? packed_stride : static_cast(p.num_heads_k) * DK; const int64_t v_stride = static_cast(p.num_heads_v) * d_v; + const int64_t v_input_stride = PackedQkv ? packed_stride : v_stride; const int64_t st_off = (static_cast(b) * p.num_heads_v + hv) * static_cast(d_v) * DK + static_cast(col) * DK; @@ -456,14 +459,14 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( float decay_scalar = 1.0f; if (needs_decay && !decay_per_key_dim) { - decay_scalar = __expf(EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv, + decay_scalar = __expf(EffectiveDecay(LoadGate(pack.decay, gi, InputTypeGates), pack.a_log, pack.dt_bias, hv, p.gate_activation)); } #pragma unroll for (int i = 0; i < kRegs; ++i) { if (needs_decay && decay_per_key_dim) { const int k = lane + i * 32; - const float raw = pack.decay[(tok * p.num_heads_v + hv) * DK + k]; + const float raw = LoadGate(pack.decay, (tok * p.num_heads_v + hv) * DK + k, InputTypeGates); s[i] *= __expf(EffectiveDecay(raw, pack.a_log, pack.dt_bias, hv, p.gate_activation)); } else { s[i] *= decay_scalar; @@ -478,8 +481,8 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( } const float beta_t = - pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f; - const float vv = static_cast(pack.value[tok * v_stride + hv * d_v + col]); + pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, InputTypeGates), p.beta_activation) : 1.0f; + const float vv = static_cast(pack.value[tok * v_input_stride + hv * d_v + col]); const float delta = beta_t * (vv - r); if (t < capture_count) { @@ -512,6 +515,12 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( pack.output[tok * v_stride + hv * d_v + col] = static_cast(p.scale * o); } + if (t == 0 && pack.mid_state != nullptr) { +#pragma unroll + for (int i = 0; i < kRegs; ++i) { + pack.mid_state[st_off + lane + i * 32] = s[i]; + } + } } if (pack.final_state != nullptr) { @@ -526,7 +535,7 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( // Generic recurrent engine: any head geometry the warp kernel cannot take. // One CTA per (sequence, v-head); the [DK x DV] state lives in shared memory. // --------------------------------------------------------------------------- -template +template __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(VariantPack pack, KernelParams p, int d_k, int d_v) { @@ -585,7 +594,7 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant } if (needs_decay && decay_per_key_dim) { for (int i = tid; i < d_k; i += kThreads) { - const float raw = pack.decay[(tok * p.num_heads_v + hv) * d_k + i]; + const float raw = LoadGate(pack.decay, (tok * p.num_heads_v + hv) * d_k + i, InputTypeGates); gbuf[i] = __expf(EffectiveDecay(raw, pack.a_log, pack.dt_bias, hv, p.gate_activation)); } } @@ -612,11 +621,11 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant const float g_scalar = (needs_decay && !decay_per_key_dim) - ? __expf(EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv, + ? __expf(EffectiveDecay(LoadGate(pack.decay, gi, InputTypeGates), pack.a_log, pack.dt_bias, hv, p.gate_activation)) : 1.0f; const float beta_t = - pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f; + pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, InputTypeGates), p.beta_activation) : 1.0f; // Decay, then retrieval against the decayed state. for (int idx = tid; idx < d_k * d_v; idx += kThreads) { @@ -665,6 +674,12 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant for (int r = 0; r < d_k; ++r) acc += S[r * d_v + c] * qbuf[r]; pack.output[tok * v_stride + hv * d_v + c] = (T)(p.scale * acc); } + if (t == 0 && pack.mid_state != nullptr) { + for (int idx = tid; idx < d_k * d_v; idx += kThreads) { + const int r = idx / d_v, c = idx % d_v; + pack.mid_state[st_off + static_cast(c) * d_k + r] = S[idx]; + } + } __syncthreads(); } @@ -680,27 +695,37 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant constexpr int kRecurrentThreads = 256; constexpr int kDecodeWarps = kRecurrentThreads / 32; -template +template +void LaunchDecodeWarp(const Descriptor& desc, const VariantPack& pack, + const KernelParams& p, cudaStream_t stream) { + const dim3 grid(desc.batch, desc.num_heads_v, + (desc.head_size_v + kDecodeWarps - 1) / kDecodeWarps); + const dim3 block(32, kDecodeWarps, 1); + switch (desc.head_size_qk) { + case 64: + GatedDeltaNetDecodeWarpKernel + <<>>(pack, p, desc.head_size_v); + break; + case 128: + GatedDeltaNetDecodeWarpKernel + <<>>(pack, p, desc.head_size_v); + break; + default: + GatedDeltaNetDecodeWarpKernel + <<>>(pack, p, desc.head_size_v); + break; + } +} + +template Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const VariantPack& pack, const KernelParams& p, size_t max_shared_memory_per_block, cudaStream_t stream) { if (warp_specialized) { - const dim3 grid(desc.batch, desc.num_heads_v, - (desc.head_size_v + kDecodeWarps - 1) / kDecodeWarps); - const dim3 block(32, kDecodeWarps, 1); - switch (desc.head_size_qk) { - case 64: - GatedDeltaNetDecodeWarpKernel - <<>>(pack, p, desc.head_size_v); - break; - case 128: - GatedDeltaNetDecodeWarpKernel - <<>>(pack, p, desc.head_size_v); - break; - default: - GatedDeltaNetDecodeWarpKernel - <<>>(pack, p, desc.head_size_v); - break; + if (pack.packed_qkv) { + LaunchDecodeWarp(desc, pack, p, stream); + } else { + LaunchDecodeWarp(desc, pack, p, stream); } return CUDA_CALL(cudaGetLastError()); } @@ -708,20 +733,20 @@ Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const Vari const size_t smem = RecurrentSmemBytes(desc.head_size_qk, desc.head_size_v); static DynamicSmemConfig recurrent_smem_config; ORT_RETURN_IF_ERROR(SetMaxDynamicSmem( - GatedDeltaNetRecurrentKernel, + GatedDeltaNetRecurrentKernel, max_shared_memory_per_block, recurrent_smem_config)); const dim3 grid(desc.batch, desc.num_heads_v, 1); - GatedDeltaNetRecurrentKernel + GatedDeltaNetRecurrentKernel <<>>(pack, p, desc.head_size_qk, desc.head_size_v); return CUDA_CALL(cudaGetLastError()); } } // namespace -template -Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const VariantPack& pack, - float scale, int max_threads_per_block, - size_t max_shared_memory_per_block, cudaStream_t stream) { +template +static Status LaunchGatedDeltaNetTypedGates(const Descriptor& desc, const Plan& plan, const VariantPack& pack, + float scale, int max_threads_per_block, + size_t max_shared_memory_per_block, cudaStream_t stream) { KernelParams p{}; p.total_tokens = desc.total_tokens; p.uniform_len = desc.batch > 0 ? desc.total_tokens / desc.batch : 0; @@ -754,15 +779,15 @@ Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const Varia // engine fit devices with a 99 KB opt-in limit (SM120) instead of SM90's 227 KB. if (plan.chunk_size == 32) { static DynamicSmemConfig chunk32_smem_config; - ORT_RETURN_IF_ERROR(SetMaxDynamicSmem(GatedDeltaNetChunkedKernel, + ORT_RETURN_IF_ERROR(SetMaxDynamicSmem(GatedDeltaNetChunkedKernel, plan.smem_bytes, chunk32_smem_config)); - GatedDeltaNetChunkedKernel + GatedDeltaNetChunkedKernel <<>>(pack, p); } else { static DynamicSmemConfig chunk64_smem_config; - ORT_RETURN_IF_ERROR(SetMaxDynamicSmem(GatedDeltaNetChunkedKernel, + ORT_RETURN_IF_ERROR(SetMaxDynamicSmem(GatedDeltaNetChunkedKernel, plan.smem_bytes, chunk64_smem_config)); - GatedDeltaNetChunkedKernel + GatedDeltaNetChunkedKernel <<>>(pack, p); } ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); @@ -774,14 +799,26 @@ Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const Varia // the two row sets are disjoint, so the merged result is deterministic. // The chunked engine only accepts head_size_qk == head_size_v == 128, so the warp // kernel is always the right shape here. - return LaunchRecurrent(desc, /*warp_specialized=*/true, pack, p, - max_shared_memory_per_block, stream); + return LaunchRecurrent(desc, /*warp_specialized=*/true, pack, p, + max_shared_memory_per_block, stream); } } (void)max_threads_per_block; - return LaunchRecurrent(desc, plan.warp_specialized, pack, p, max_shared_memory_per_block, - stream); + return LaunchRecurrent(desc, plan.warp_specialized, pack, p, max_shared_memory_per_block, + stream); +} + +template +Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const VariantPack& pack, + float scale, int max_threads_per_block, + size_t max_shared_memory_per_block, cudaStream_t stream) { + if (pack.gates_in_input_type) { + return LaunchGatedDeltaNetTypedGates( + desc, plan, pack, scale, max_threads_per_block, max_shared_memory_per_block, stream); + } + return LaunchGatedDeltaNetTypedGates( + desc, plan, pack, scale, max_threads_per_block, max_shared_memory_per_block, stream); } template Status LaunchGatedDeltaNet(const Descriptor&, const Plan&, const VariantPack&, diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h index feb75bf9b1fbf..3035ca02fa083 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h @@ -22,17 +22,28 @@ struct VariantPack { const T* value = nullptr; // [total_tokens, Hv, V] const int32_t* cu_seqlens = nullptr; // [B+1] or null for uniform packing const int32_t* capture_count = nullptr; // [B] or null when transition capture is disabled - const float* decay = nullptr; // [total_tokens, Hv] or [total_tokens, Hv, K] - const float* beta = nullptr; // [total_tokens, Hv] + const void* decay = nullptr; // [total_tokens, Hv] or [total_tokens, Hv, K] + const void* beta = nullptr; // [total_tokens, Hv] + bool gates_in_input_type = false; // gate buffers contain T instead of float + bool packed_qkv = false; // short decode reads packed Q/K/V without unpacking const float* initial_state = nullptr; // [B, Hv, V, K] V-major, may alias final_state const float* a_log = nullptr; // [Hv] const float* dt_bias = nullptr; // [Hv] T* output = nullptr; // [total_tokens, max(Hq,Hv), V] float* final_state = nullptr; // [B, Hv, V, K] + float* mid_state = nullptr; // [B, Hv, V, K], state after local row 0 float* state_update = nullptr; // [B, C * (Hv + Hk*K + Hv*V)] void* workspace = nullptr; }; +#ifdef __CUDACC__ +template +__device__ __forceinline__ float LoadGate(const void* data, int64_t index, bool gates_in_input_type) { + return gates_in_input_type ? static_cast(static_cast(data)[index]) + : static_cast(data)[index]; +} +#endif + template Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const VariantPack& pack, float scale, int max_threads_per_block, diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh index 33a841df9a4f0..582e3474bed1e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh @@ -52,7 +52,7 @@ __device__ __forceinline__ void SequenceRange(const int32_t* cu_seqlens, int b, // Effective per-token decay and beta, after the optional fused activations. __device__ __forceinline__ float EffectiveDecay(float raw, const float* a_log, const float* dt_bias, int h, GateActivation act) { - if (act == GateActivation::kQwen) { + if (act == GateActivation::kSoftplusDecay) { const float bias = dt_bias != nullptr ? dt_bias[h] : 0.0f; const float scale = a_log != nullptr ? -__expf(a_log[h]) : -1.0f; return scale * SoftPlus(raw + bias); diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h index 613663fccb3b0..898afaf4699fc 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h @@ -26,7 +26,7 @@ enum class UpdateRule : int { kLinear = 0, kDelta = 2, kGatedDelta = 3 }; enum class GateActivation : int { kNone = 0, - kQwen = 1 }; + kSoftplusDecay = 1 }; enum class BetaActivation : int { kNone = 0, kSigmoid = 1 }; diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu index 0837fbde8f3bd..fb7c5f2711c2d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu @@ -184,10 +184,10 @@ __global__ __launch_bounds__(kChunkedThreads) void GatedDeltaNetPrepareKernel( const int64_t gi = tok * p.num_heads_v + hv; float g = 0.0f, be = 0.0f; if (tid < len) { - g = needs_decay ? EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv, + g = needs_decay ? EffectiveDecay(LoadGate(pack.decay, gi, pack.gates_in_input_type), pack.a_log, pack.dt_bias, hv, p.gate_activation) : 0.0f; - be = pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f; + be = pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, pack.gates_in_input_type), p.beta_activation) : 1.0f; } s.gc[tid] = g; s.beta[tid] = be; 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 a47b313e482fd..ca7888b7c178b 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc @@ -38,6 +38,12 @@ ONNX_OPERATOR_KERNEL_EX( CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : WebGpuKernel(info) { + ORT_ENFORCE(info.GetAttrOrDefault("input_channels", 0) == 0 && + info.GetAttrOrDefault("input_channel_offset", 0) == 0 && + info.GetAttrOrDefault("state_length", 0) == 0 && + info.GetAttrOrDefault("state_channels_last", -1) == -1 && + info.node().OutputDefs().size() <= 2, + "WebGPU CausalConvWithState does not support extended channel/state options (CUDA only)"); std::string activation_str = info.GetAttrOrDefault("activation", "none"); activation_ = ParseCausalConvActivation(activation_str); ORT_ENFORCE(info.GetAttr("ndim", &ndim_).IsOK(), "Attribute 'ndim' is required"); diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc index 2294389a9728c..d8030e7412ecd 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc @@ -37,6 +37,7 @@ ONNX_OPERATOR_KERNEL_EX( (*KernelDefBuilder::Create()) .TypeConstraint("T", WebGpuSupportedFloatTypes()) .TypeConstraint("TS", DataTypeImpl::GetTensorType()) + .TypeConstraint("TG", DataTypeImpl::GetTensorType()) .TypeConstraint("TI", DataTypeImpl::GetTensorType()) .InputMemoryType(OrtMemTypeCPUInput, 10) .MayInplace(6, 1), @@ -47,9 +48,10 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : WebGpuKernel(info) { ORT_ENFORCE(update_rule_ != GatedDeltaNetUpdateRule::Invalid, "update_rule must be one of: linear, gated, delta, gated_delta"); const auto gate_activation = info.GetAttrOrDefault("gate_activation", "none"); - ORT_ENFORCE(gate_activation == "none" || gate_activation == "qwen", - "gate_activation must be one of: none, qwen"); - qwen_gate_ = gate_activation == "qwen"; + ORT_ENFORCE(gate_activation == "none" || gate_activation == "softplus" || + gate_activation == "qwen", + "gate_activation must be one of: none, softplus"); + qwen_gate_ = gate_activation == "softplus" || gate_activation == "qwen"; const auto beta_activation = info.GetAttrOrDefault("beta_activation", "none"); ORT_ENFORCE(beta_activation == "none" || beta_activation == "sigmoid", "beta_activation must be one of: none, sigmoid"); @@ -60,6 +62,9 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : WebGpuKernel(info) { const auto state_update_capacity = info.GetAttrOrDefault("state_update_capacity", 0); ORT_ENFORCE(state_update_capacity >= 0 && state_update_capacity <= 8, "state_update_capacity must be in [0, 8], got ", state_update_capacity); + ORT_ENFORCE(info.GetAttrOrDefault("state_output", "state_update") == + "state_update", + "WebGPU GatedDeltaNet supports only state_output='state_update'"); state_update_capacity_ = static_cast(state_update_capacity); } diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index fecac2f17f523..6ecab0397ab61 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -4650,6 +4650,20 @@ ONNX_MS_OPERATOR_SET_SCHEMA( CausalConvWithState, 1, OpSchema() .SetDoc(CausalConvWithState_ver1_doc) + .Attr("input_channels", + "CUDA channels-last input channel slice width; 0 selects all channels. " + "Nondefault extended attributes require rank-3 input and state_window=0.", + AttributeProto::INT, static_cast(0)) + .Attr("input_channel_offset", "Starting channel of the channels-last input slice (CUDA only).", + AttributeProto::INT, static_cast(0)) + .Attr("state_length", + "Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. " + "An explicit length must cover the receptive field; convolution reads its trailing history.", + AttributeProto::INT, static_cast(0)) + .Attr("state_channels_last", + "State layout override (CUDA only): -1 follows channels_last, " + "0 uses [B,C,state_length], 1 uses [B,state_length,C].", + AttributeProto::INT, static_cast(-1)) .Attr("activation", "Fused activation function. One of: 'silu', 'swish', 'none'. " "Default is 'none'.", @@ -4693,7 +4707,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "input", "Input tensor with shape (batch_size, channels, ...) in the default channels-first " "layout. Spatial dims: 1D: (L,); 2D: (H, W); 3D: (D, H, W). When channels_last = 1 " - "the shape is (batch_size, sequence_length, d_1, ..., d_n) instead.", + "the shape is (batch_size, sequence_length, d_1, ..., d_n) instead. Extended channel " + "selection uses rank-3 (batch_size, sequence_length, input_width).", "T") .Input(1, "weight", @@ -4711,12 +4726,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "or (W, batch_size, channels, state_length) when state_window = W > 0, in which case " "only slot W-1 is read, where state_length = (k_1 - 1) * dilation. When " "channels_last = 1 each slot is (batch_size, state_length, d_1, ..., d_n) instead. " - "If not provided, padding is zero.", + "The state_length and state_channels_last attributes may override the length and " + "layout for the CUDA extended path. If not provided, padding is zero.", "T", OpSchema::Optional) .Output(0, "output", - "Convolution output with same shape as input.", + "Convolution output with same shape as input, except the final dimension is " + "input_channels when extended channel selection is enabled.", "T") .Output(1, "present_state", @@ -4725,8 +4742,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "(batch_size, state_length, d_1, ..., d_n) per slot when channels_last = 1. Slot " "W-1 contains the last state_length values from the virtual input along the causal " "axis; slot j contains the same for the prefix ending at position " - "(seq_len - W + j).", + "(seq_len - W + j). Extended attributes override the state length and layout " + "as for past_state.", "T") + .Output(2, "first_state", + "Optional carry state after the first input position, with the same " + "layout and length as present_state. CUDA only; requires rank-3 channels-last input " + "and state_window=0.", + "T", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") @@ -4735,6 +4758,63 @@ ONNX_MS_OPERATOR_SET_SCHEMA( propagateElemTypeFromInputToOutput(ctx, 0, 1); const int64_t state_window = getAttribute(ctx, "state_window", 0); + const int64_t input_channels = getAttribute(ctx, "input_channels", 0); + const int64_t channel_offset = getAttribute(ctx, "input_channel_offset", 0); + const int64_t history = getAttribute(ctx, "state_length", 0); + const int64_t state_layout = getAttribute(ctx, "state_channels_last", -1); + if (input_channels < 0 || input_channels > std::numeric_limits::max() || + channel_offset < 0 || channel_offset > std::numeric_limits::max() || + history < 0 || history > std::numeric_limits::max() || + state_layout < -1 || state_layout > 1) { + fail_shape_inference("CausalConvWithState: invalid extended attributes"); + } + if (input_channels != 0 || channel_offset != 0 || history != 0 || state_layout != -1 || + ctx.getNumOutputs() > 2) { + if (getAttribute(ctx, "channels_last", 0) != 1 || getAttribute(ctx, "ndim", 1) != 1 || + state_window != 0 || getAttribute(ctx, "dilation", 1) < 1) { + fail_shape_inference("CausalConvWithState: extended attributes require channels_last=1, ndim=1, state_window=0"); + } + if (ctx.getNumOutputs() > 2) propagateElemTypeFromInputToOutput(ctx, 0, 2); + if (!hasInputShape(ctx, 0)) return; + const auto& input = getInputShape(ctx, 0); + if (input.dim_size() != 3) fail_shape_inference("CausalConvWithState: extended input must have rank 3"); + TensorShapeProto output = input; + auto* channels = output.mutable_dim(2); + if (input_channels > 0) channels->set_dim_value(input_channels); + if (input.dim(2).has_dim_value() && + (channel_offset > input.dim(2).dim_value() || + (channels->has_dim_value() && channels->dim_value() > input.dim(2).dim_value() - channel_offset))) { + fail_shape_inference("CausalConvWithState: channel slice exceeds input width"); + } + updateOutputShape(ctx, 0, output); + TensorShapeProto::Dimension length; + if (history > 0) length.set_dim_value(history); + if (hasInputShape(ctx, 1)) { + const auto& weight = getInputShape(ctx, 1); + if (weight.dim_size() != 3) fail_shape_inference("CausalConvWithState: weight must have rank 3"); + if (weight.dim(0).has_dim_value() && channels->has_dim_value() && + weight.dim(0).dim_value() != channels->dim_value()) { + fail_shape_inference("CausalConvWithState: weight channels mismatch"); + } + if (weight.dim(2).has_dim_value()) { + const int64_t kernel = weight.dim(2).dim_value(); + const int64_t dilation = getAttribute(ctx, "dilation", 1); + if (kernel <= 0 || kernel - 1 > std::numeric_limits::max() / dilation) { + fail_shape_inference("CausalConvWithState: invalid receptive field"); + } + const int64_t pad = (kernel - 1) * dilation; + if (history != 0 && history < pad) fail_shape_inference("CausalConvWithState: state_length is too short"); + if (history == 0) length.set_dim_value(pad); + } + } + TensorShapeProto state; + *state.add_dim() = input.dim(0); + *state.add_dim() = state_layout == 0 ? *channels : length; + *state.add_dim() = state_layout == 0 ? length : *channels; + updateOutputShape(ctx, 1, state); + if (ctx.getNumOutputs() > 2) updateOutputShape(ctx, 2, state); + return; + } if (state_window < 0 || state_window > kMaxStateWindow) { fail_shape_inference("CausalConvWithState: state_window must be in [0, ", kMaxStateWindow, "], got ", state_window); @@ -5351,7 +5431,7 @@ The delta family ('delta' and 'gated_delta') requires L2-normalized keys. Withou per-chunk system (I + M) is arbitrarily ill-conditioned and the recurrence diverges. Either normalize upstream or set `qk_l2_norm=1` to have the operator do it. -Fused activations. `gate_activation='qwen'` computes the effective decay in float32 from the +Fused activations. `gate_activation='softplus'` computes the effective decay in float32 from the raw projection carried by `decay`: g = -exp(a_log) * Softplus(decay + dt_bias) @@ -5359,6 +5439,9 @@ raw projection carried by `decay`: `beta_activation='sigmoid'` applies a sigmoid to `beta`, and `qk_l2_norm=1` L2-normalizes each query and key head vector. Folding these in avoids materializing the intermediates and keeps the gate arithmetic in float32 independent of the input type. +Decay and beta may be float or the query element type (both use the same type). CUDA widens +FP16/BF16 gate inputs to float before any gate arithmetic, avoiding standalone lossless casts. +State, a_log, and dt_bias remain float. WebGPU currently accepts float decay and beta only. )DOC"; @@ -5373,8 +5456,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Output scaling factor. When 0.0 (default) uses 1/sqrt(head_size_qk).", AttributeProto::FLOAT, 0.0f) .Attr("gate_activation", - "'none' (default) treats `decay` as the effective log-space decay. 'qwen' computes " - "-exp(a_log) * Softplus(decay + dt_bias) in float32.", + "'none' (default) treats `decay` as the effective log-space decay. 'softplus' computes " + "-exp(a_log) * Softplus(decay + dt_bias) in float32. 'qwen' is also accepted for " + "compatibility.", AttributeProto::STRING, std::string("none")) .Attr("beta_activation", "'none' (default) treats `beta` as the effective update rate. 'sigmoid' applies a " @@ -5394,6 +5478,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Capacity C for compact contiguous-prefix transition capture, in [0, 8]. " "0 (default) disables compact state-update outputs.", AttributeProto::INT, static_cast(0)) + .Attr("state_output", + "Third output is either 'state_update' (default) or the recurrent state after " + "the first token ('mid_state'). CUDA supports both modes; WebGPU supports " + "'state_update'.", + AttributeProto::STRING, std::string("state_update")) .Input(0, "query", "Query or packed QKV, shaped as described above.", "T") .Input(1, "key", "Key, shape (total_tokens, num_heads_k, head_size_qk)", "T", OpSchema::Optional) @@ -5405,20 +5494,21 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "TI", OpSchema::Optional) .Input(4, "decay", "Log-space decay, shape (total_tokens, num_heads_v) for a scalar per-head decay or " - "(total_tokens, num_heads_v, head_size_qk) for a per-key-dimension decay.", - "TS", OpSchema::Optional) - .Input(5, "beta", "Update rate, shape (total_tokens, num_heads_v)", "TS", + "(total_tokens, num_heads_v, head_size_qk) for a per-key-dimension decay. " + "Float or the query element type; converted to float before gate arithmetic.", + "TG", OpSchema::Optional) + .Input(5, "beta", "Update rate, shape (total_tokens, num_heads_v). Same element type as decay.", "TG", OpSchema::Optional) .Input(6, "initial_state", "Recurrent state, shape (batch_size, num_heads_v, head_size_v, head_size_qk), " "V-major. Required for packed QKV. May alias final_state.", "TS", OpSchema::Optional) .Input(7, "a_log", - "Per-head A_log, shape (num_heads_v). Requires gate_activation=qwen.", + "Per-head A_log, shape (num_heads_v). Requires gate_activation=softplus.", "TS", OpSchema::Optional) .Input(8, "dt_bias", "Per-head gate bias, shape (num_heads_v). " - "Requires gate_activation=qwen.", + "Requires gate_activation=softplus.", "TS", OpSchema::Optional) .Input(9, "capture_count", "Number of leading local token transitions to capture for each request, shape " @@ -5436,14 +5526,15 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "(batch_size, num_heads_v, head_size_v, head_size_qk)", "TS", OpSchema::Optional) .Output(2, "state_update", - "Struct-of-arrays compact transition factors, shape " - "(batch_size, state_update_capacity * (num_heads_v + num_heads_k * " - "head_size_qk + num_heads_v * head_size_v)).", + "Struct-of-arrays transition factors when state_output='state_update', or " + "the state after the first token when state_output='mid_state'.", "TS", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain query/key/value/output types.") .TypeConstraint("TS", {"tensor(float)"}, - "State, gate, beta and compact state-update tensors are always float.") + "State, gate parameters and compact state-update tensors are always float.") + .TypeConstraint("TG", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Decay and beta are float or match the query element type. Gate arithmetic is float.") .TypeConstraint("TI", {"tensor(int32)"}, "Constrain index and count tensors to int32.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { propagateElemTypeFromInputToOutput(ctx, 0, 0); @@ -5576,45 +5667,59 @@ ONNX_MS_OPERATOR_SET_SCHEMA( kMaxStateWindow, "], got ", state_update_capacity); } if (ctx.getNumOutputs() > 2) { - ONNX_NAMESPACE::TensorShapeProto capsule_shape; - add_batch_dim(capsule_shape); - auto* width = capsule_shape.add_dim(); - if (!packed_qkv && query_shape.dim(token_dims).has_dim_value() && - query_shape.dim(token_dims + 1).has_dim_value() && - value_shape->dim(token_dims).has_dim_value() && - value_shape->dim(token_dims + 1).has_dim_value()) { - // num_heads_k is constrained to equal num_heads_q, so query supplies it. - const int64_t num_heads_k = query_shape.dim(token_dims).dim_value(); - const int64_t head_size_qk = query_shape.dim(token_dims + 1).dim_value(); - const int64_t num_heads_v = value_shape->dim(token_dims).dim_value(); - const int64_t head_size_v = value_shape->dim(token_dims + 1).dim_value(); - if (num_heads_k <= 0 || head_size_qk <= 0 || num_heads_v <= 0 || head_size_v <= 0) { - fail_shape_inference( - "GatedDeltaNet: head counts and head sizes must be positive"); - } - - constexpr int64_t max_dimension = std::numeric_limits::max(); - if (num_heads_k > max_dimension / head_size_qk || - num_heads_v > max_dimension / head_size_v) { - fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); - } - const int64_t key_width = num_heads_k * head_size_qk; - const int64_t value_width = num_heads_v * head_size_v; - if (num_heads_v > max_dimension - key_width) { - fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); + const std::string state_output = + getAttribute(ctx, "state_output", std::string("state_update")); + if (state_output != "state_update" && state_output != "mid_state") { + fail_shape_inference("GatedDeltaNet: state_output must be 'state_update' or 'mid_state'"); + } + if (state_output == "mid_state") { + if (state_update_capacity != 0) { + fail_shape_inference("GatedDeltaNet: mid_state output requires state_update_capacity=0"); } - const int64_t width_without_capacity = num_heads_v + key_width; - if (value_width > max_dimension - width_without_capacity) { - fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); + if (state_shape != nullptr) { + updateOutputShape(ctx, 2, *state_shape); } - const int64_t per_token_width = width_without_capacity + value_width; - if (state_update_capacity > 0 && - per_token_width > max_dimension / state_update_capacity) { - fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); + } else { + ONNX_NAMESPACE::TensorShapeProto capsule_shape; + add_batch_dim(capsule_shape); + auto* width = capsule_shape.add_dim(); + if (!packed_qkv && query_shape.dim(token_dims).has_dim_value() && + query_shape.dim(token_dims + 1).has_dim_value() && + value_shape->dim(token_dims).has_dim_value() && + value_shape->dim(token_dims + 1).has_dim_value()) { + // num_heads_k is constrained to equal num_heads_q, so query supplies it. + const int64_t num_heads_k = query_shape.dim(token_dims).dim_value(); + const int64_t head_size_qk = query_shape.dim(token_dims + 1).dim_value(); + const int64_t num_heads_v = value_shape->dim(token_dims).dim_value(); + const int64_t head_size_v = value_shape->dim(token_dims + 1).dim_value(); + if (num_heads_k <= 0 || head_size_qk <= 0 || num_heads_v <= 0 || head_size_v <= 0) { + fail_shape_inference( + "GatedDeltaNet: head counts and head sizes must be positive"); + } + + constexpr int64_t max_dimension = std::numeric_limits::max(); + if (num_heads_k > max_dimension / head_size_qk || + num_heads_v > max_dimension / head_size_v) { + fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); + } + const int64_t key_width = num_heads_k * head_size_qk; + const int64_t value_width = num_heads_v * head_size_v; + if (num_heads_v > max_dimension - key_width) { + fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); + } + const int64_t width_without_capacity = num_heads_v + key_width; + if (value_width > max_dimension - width_without_capacity) { + fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); + } + const int64_t per_token_width = width_without_capacity + value_width; + if (state_update_capacity > 0 && + per_token_width > max_dimension / state_update_capacity) { + fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); + } + width->set_dim_value(state_update_capacity * per_token_width); } - width->set_dim_value(state_update_capacity * per_token_width); + updateOutputShape(ctx, 2, capsule_shape); } - updateOutputShape(ctx, 2, capsule_shape); } if (state_shape != nullptr) { diff --git a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc index 0adc85ac191b5..3ae0241955a71 100644 --- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc +++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc @@ -319,6 +319,127 @@ static void RunCausalConvWithStateTests( // Basic tests - simple cases // ============================================================================= +template +void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool use_bias, bool silu, int dilation, + bool first_output = true) { + constexpr int batch = 2, channels = 7, width = 13, offset = 3, kernel = 4; + const int history = (kernel - 1) * dilation + 1; + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) GTEST_SKIP() << "CUDA is required for extended convolution"; + const auto round = [](float value) { return static_cast(T(value)); }; + std::vector input(batch * length * width), weight(channels * kernel), bias(channels); + std::vector past(batch * channels * history), output(batch * length * channels); + std::vector present(past.size()), first(past.size()); + for (size_t i = 0; i < input.size(); ++i) input[i] = T(float(int(i % 17) - 8) / 16); + for (size_t i = 0; i < weight.size(); ++i) weight[i] = T(float(int(i % 11) - 5) / 32); + for (size_t i = 0; i < bias.size(); ++i) bias[i] = T(float(i) / 64); + for (size_t i = 0; i < past.size(); ++i) past[i] = T(float(int(i % 13) - 6) / 8); + const auto state_index = [&](int b, int t, int c) { + return state_last ? (b * history + t) * channels + c : (b * channels + c) * history + t; + }; + for (int b = 0; b < batch; ++b) { + for (int c = 0; c < channels; ++c) { + const auto value = [&](int t) { + return t >= 0 ? static_cast(input[(b * length + t) * width + offset + c]) + : (use_past ? static_cast(past[state_index(b, history + t, c)]) : 0.0f); + }; + for (int t = 0; t < length; ++t) { + float sum = use_bias ? static_cast(bias[c]) : 0.0f; + for (int j = 0; j < kernel; ++j) sum += value(t - (kernel - 1) * dilation + j * dilation) * static_cast(weight[c * kernel + j]); + if (silu) sum /= 1.0f + std::exp(-sum); + output[(b * length + t) * channels + c] = T(round(sum)); + } + for (int t = 0; t < history; ++t) { + present[state_index(b, t, c)] = T(value(length - history + t)); + first[state_index(b, t, c)] = T(value(1 - history + t)); + } + } + } + OpTester test("CausalConvWithState", 1, kMSDomain); + test.AddAttribute("channels_last", 1); + test.AddAttribute("input_channels", channels); + test.AddAttribute("input_channel_offset", offset); + test.AddAttribute("state_length", history); + test.AddAttribute("state_channels_last", state_last ? 1 : 0); + test.AddAttribute("dilation", dilation); + test.AddAttribute("activation", silu ? "silu" : "none"); + test.AddInput("input", {batch, length, width}, input); + test.AddInput("weight", {channels, 1, kernel}, weight); + if (use_bias) + test.AddInput("bias", {channels}, bias); + else + test.AddOptionalInputEdge(); + const std::vector state_shape = state_last ? std::vector{batch, history, channels} + : std::vector{batch, channels, history}; + if (use_past) + test.AddInput("past", state_shape, past); + else + test.AddOptionalInputEdge(); + test.AddOutput("output", {batch, length, channels}, output); + test.AddOutput("present", state_shape, present); + if (first_output) test.AddOutput("first", state_shape, first); + test.SetOutputTolerance(std::is_same_v ? 0.01f : 0.001f, 0.001f); + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(CausalConvWithStateTest, PackedChannelsExtendedHistoryAndFirstState) { + for (int length : {1, 2, 9}) { + for (bool state_last : {false, true}) { + for (bool past : {false, true}) { + for (int dilation : {1, 2}) { + TestPackedChannelHistory(length, state_last, past, true, true, dilation); + TestPackedChannelHistory(length, state_last, past, false, true, dilation); + TestPackedChannelHistory(length, state_last, past, true, false, dilation); + TestPackedChannelHistory(length, state_last, past, false, true, dilation); + } + } + } + } + TestPackedChannelHistory(2, false, true, true, true, 1, false); +} + +TEST(CausalConvWithStateTest, PackedChannelsInvalidConfigurations) { + const std::vector errors{ + "invalid extended attributes", "channel slice exceeds input width", "state_length is too short", + "extended attributes require", "past_state must match", "bias must match"}; + for (size_t i = 0; i < errors.size(); ++i) { + SCOPED_TRACE(errors[i]); + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) GTEST_SKIP() << "CUDA is required for extended convolution"; + OpTester test("CausalConvWithState", 1, kMSDomain); + test.AddAttribute("channels_last", i == 3 ? 0 : 1); + test.AddAttribute("input_channels", i == 1 ? 14 : 7); + test.AddAttribute("input_channel_offset", 3); + test.AddAttribute("state_length", i == 2 ? 2 : 4); + test.AddAttribute("state_channels_last", i == 0 ? 2 : 0); + test.AddInput("input", {2, 2, 13}, std::vector(52)); + test.AddInput("weight", {7, 1, 4}, std::vector(28)); + test.AddInput("bias", {i == 5 ? 6 : 7}, std::vector(i == 5 ? 6 : 7)); + test.AddInput("past", {2, 7, i == 4 ? 5 : 4}, std::vector(i == 4 ? 70 : 56)); + test.AddOutput("output", {2, 2, 7}, std::vector(28)); + test.AddOutput("present", {2, 7, 4}, std::vector(56)); + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectFailure, errors[i], {}, nullptr, &providers); + } +} + +TEST(CausalConvWithStateTest, PackedChannelsRejectsUnsupportedCpu) { + OpTester test("CausalConvWithState", 1, kMSDomain); + test.AddAttribute("channels_last", 1); + test.AddAttribute("input_channels", 2); + test.AddInput("input", {1, 1, 2}, {1.0f, 2.0f}); + test.AddInput("weight", {2, 1, 2}, {1.0f, 1.0f, 1.0f, 1.0f}); + test.AddOutput("output", {1, 1, 2}, {1.0f, 2.0f}); + test.AddOutput("present", {1, 1, 2}, {1.0f, 2.0f}); + std::vector> providers; + providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, "does not support extended channel/state options", + {}, nullptr, &providers); +} + TEST(CausalConvWithStateTest, BasicNoStateNoBias) { // B=1, D=2, L=4, K=3, activation=none int batch_size = 1, channels = 2, input_length = 4, kernel_size = 3; diff --git a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc index a4778180005f3..68a3b7c1a25fa 100644 --- a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -58,6 +59,7 @@ struct Options { int state_update_capacity = 0; int chunk_size = 0; float scale = 0.0f; + bool gates_in_input_type = false; }; struct Inputs { @@ -421,6 +423,10 @@ void RunTypedCase(const Geometry& g, const Options& o, const Inputs& in_raw, flo in.q = RoundToTensorType(in_raw.q); in.k = RoundToTensorType(in_raw.k); in.v = RoundToTensorType(in_raw.v); + if (o.gates_in_input_type) { + in.decay = RoundToTensorType(in_raw.decay); + in.beta = RoundToTensorType(in_raw.beta); + } std::vector ref_out, ref_state, ref_state_update; Reference(g, o, in, &ref_out, &ref_state, use_webgpu ? &ref_state_update : nullptr); @@ -490,12 +496,20 @@ void RunTypedCase(const Geometry& g, const Options& o, const Inputs& in_raw, flo const bool needs_beta = o.update_rule == "delta" || o.update_rule == "gated_delta"; if (needs_decay) { - test.AddInput("decay", shaped({g.hv}), in.decay); + if (o.gates_in_input_type) { + test.AddInput("decay", shaped({g.hv}), ToTensorType(in.decay)); + } else { + test.AddInput("decay", shaped({g.hv}), in.decay); + } } else { test.AddOptionalInputEdge(); } if (needs_beta) { - test.AddInput("beta", shaped({g.hv}), in.beta); + if (o.gates_in_input_type) { + test.AddInput("beta", shaped({g.hv}), ToTensorType(in.beta)); + } else { + test.AddInput("beta", shaped({g.hv}), in.beta); + } } else { test.AddOptionalInputEdge(); } @@ -605,12 +619,120 @@ void RunCase(const Geometry& g, const Options& o, const Inputs& in, float out_to RunTypedCase(g, o, in, out_tol, state_tol, rank4, fetches, use_webgpu); } +template +void TestInputTypeGates() { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA is required for input-type gates"; + } + for (const int tokens : {1, 2, 9, 65}) { + const Geometry geometry{tokens, 1, 2, 4, 128, 128}; + Inputs inputs = MakeInputs(geometry, 917 + tokens); + inputs.decay = RoundToTensorType(inputs.decay); + inputs.beta = RoundToTensorType(inputs.beta); + for (const std::string rule : {"gated_delta", "gated", "delta"}) { + Options options; + options.update_rule = rule; + options.gate_activation = rule == "delta" ? "none" : "qwen"; + options.beta_activation = "sigmoid"; + std::vector reference, narrowed; + RunTypedCase(geometry, options, inputs, 6e-3f, 2e-3f, true, &reference); + options.gates_in_input_type = true; + RunTypedCase(geometry, options, inputs, 6e-3f, 2e-3f, true, &narrowed); + ASSERT_EQ(reference.size(), narrowed.size()); + for (size_t i = 0; i < reference.size(); ++i) { + const Tensor& expected = reference[i].Get(); + const Tensor& actual = narrowed[i].Get(); + ASSERT_EQ(expected.Shape(), actual.Shape()); + ASSERT_EQ(expected.SizeInBytes(), actual.SizeInBytes()); + if (expected.SizeInBytes() != 0) { + EXPECT_EQ(std::memcmp(expected.DataRaw(), actual.DataRaw(), expected.SizeInBytes()), 0) + << "tokens=" << tokens << " rule=" << rule << " output=" << i; + } + } + } + } +} + +TEST(GatedDeltaNetTest, InputTypeGatesFloat16ExactParity) { + TestInputTypeGates(); +} + +TEST(GatedDeltaNetTest, InputTypeGatesBFloat16ExactParity) { + TestInputTypeGates(); +} + +TEST(GatedDeltaNetTest, RejectsGateTypeDifferentFromQuery) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA is required for input-type gates"; + } + OpTester test("GatedDeltaNet", 1, kMSDomain); + test.AddInput("query", {1, 1, 64}, std::vector(64)); + test.AddInput("key", {1, 1, 64}, std::vector(64)); + test.AddInput("value", {1, 1, 64}, std::vector(64)); + test.AddOptionalInputEdge(); + test.AddInput("decay", {1, 1}, {BFloat16(-0.1f)}); + test.AddInput("beta", {1, 1}, {BFloat16(0.5f)}); + test.AddOutput("output", {1, 1, 64}, std::vector(64)); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, + "decay and beta must be float or match the query element type", {}, nullptr, &providers); +} + constexpr int kDim = 128; // the chunked engine is specialised for head_size 128 bool NeedSkipGatedDeltaNetTest() { return !HasCudaEnvironment(800); } +TEST(GatedDeltaNetTest, RecurrentOutputsIncludeFirstRowState) { + if (NeedSkipGatedDeltaNetTest()) return; + + const Geometry g{2, 1, 1, 2, 128, 128}; + Inputs in = MakeInputs(g, 211); + Options options; + options.gate_activation = "softplus"; + options.beta_activation = "sigmoid"; + options.qk_l2_norm = 1; + + in.q = RoundToTensorType(in.q); + in.k = RoundToTensorType(in.k); + in.v = RoundToTensorType(in.v); + std::vector expected_output, expected_final_state; + Reference(g, options, in, &expected_output, &expected_final_state); + + Geometry first_row_geometry = g; + first_row_geometry.total_tokens = 1; + Inputs first_row = in; + first_row.q.resize(static_cast(g.hq) * g.dk); + first_row.k.resize(static_cast(g.hq) * g.dk); + first_row.v.resize(static_cast(g.hv) * g.dv); + first_row.decay.resize(g.hv); + first_row.beta.resize(g.hv); + std::vector unused_output, expected_mid_state; + Reference(first_row_geometry, options, first_row, &unused_output, &expected_mid_state); + + OpTester test("GatedDeltaNet", 1, onnxruntime::kMSDomain); + AddCommonAttrs(test, options); + test.AddAttribute("state_output", "mid_state"); + test.AddInput("query", {1, 2, g.hq, g.dk}, ToTensorType(in.q)); + test.AddInput("key", {1, 2, g.hq, g.dk}, ToTensorType(in.k)); + test.AddInput("value", {1, 2, g.hv, g.dv}, ToTensorType(in.v)); + test.AddOptionalInputEdge(); + test.AddInput("decay", {1, 2, g.hv}, in.decay); + test.AddInput("beta", {1, 2, g.hv}, in.beta); + test.AddInput("initial_state", {g.batch, g.hv, g.dv, g.dk}, in.state0); + test.AddInput("a_log", {g.hv}, in.a_log); + test.AddInput("dt_bias", {g.hv}, in.dt_bias); + test.AddOutput("output", {1, 2, g.hv, g.dv}, + ToTensorType(expected_output), false, 6e-3f, 6e-3f); + test.AddOutput("final_state", {g.batch, g.hv, g.dv, g.dk}, + expected_final_state, false, 2e-3f, 2e-3f); + test.AddOutput("mid_state", {g.batch, g.hv, g.dv, g.dk}, + expected_mid_state, false, 2e-3f, 2e-3f); + test.Run(); +} + bool NeedSkipGatedDeltaNetSplitTest() { if (NeedSkipGatedDeltaNetTest()) return true; @@ -1344,6 +1466,45 @@ TEST(GatedDeltaNetTest, Recurrent_PackedQkvDecode) { /*test_max_storage_buffer_binding_size=*/0, /*packed_qkv=*/true); } +template +void TestDirectPackedQkvParity() { + if (NeedSkipGatedDeltaNetTest()) return; + for (int width : {64, 128, 256}) { + for (int tokens : {1, 2, 9, 65}) { + const Geometry geometry{tokens * 2, 2, 2, 4, width, 128}; + Inputs inputs = MakeInputs(geometry, width + tokens); + for (bool input_type_gates : {false, true}) { + Options options; + options.qk_l2_norm = 1; + options.gate_activation = "qwen"; + options.beta_activation = "sigmoid"; + options.gates_in_input_type = input_type_gates; + std::vector reference, packed; + RunTypedCase(geometry, options, inputs, 2e-2f, 2e-2f, true, &reference); + RunTypedCase(geometry, options, inputs, 2e-2f, 2e-2f, true, &packed, + false, false, nullptr, 0, true); + ASSERT_EQ(reference.size(), packed.size()); + for (size_t index = 0; index < reference.size(); ++index) { + const auto& expected = reference[index].Get(); + const auto& actual = packed[index].Get(); + ASSERT_EQ(expected.Shape(), actual.Shape()); + ASSERT_EQ(expected.SizeInBytes(), actual.SizeInBytes()); + EXPECT_EQ(std::memcmp(expected.DataRaw(), actual.DataRaw(), expected.SizeInBytes()), 0) + << "tokens=" << tokens << " width=" << width << " input_type_gates=" << input_type_gates; + } + } + } + } +} + +TEST(GatedDeltaNetTest, DirectPackedQkvFloat16ExactParity) { + TestDirectPackedQkvParity(); +} + +TEST(GatedDeltaNetTest, DirectPackedQkvBFloat16ExactParity) { + TestDirectPackedQkvParity(); +} + TEST(GatedDeltaNetTest, Recurrent_VerifyBatch) { if (NeedSkipGatedDeltaNetTest()) return; Geometry g{4, 1, 1, 2, kDim, kDim}; From c1e138f9e7f5c05df679b618a3455114f623f779 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 20:18:30 +0000 Subject: [PATCH 02/17] Fix GDN empty-request state capture and WebGPU plugin compatibility Initialize mid_state for empty ragged requests in both recurrent kernels, correct the softplus reference alias, and add regression coverage. Use the adapter-supported output-count API and synchronize operator documentation with the extended schemas. --- docs/ContribOperators.md | 48 +++++++++++++------ docs/OperatorKernels.md | 6 +-- .../cuda/bert/gated_delta_net_impl.cu | 13 +++++ .../webgpu/bert/causal_conv_with_state.cc | 5 +- .../contrib_ops/gated_delta_net_op_test.cc | 45 ++++++++++++++++- 5 files changed, 95 insertions(+), 22 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 3510b52792635..ee504bfd1ed2f 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -1020,8 +1020,16 @@ This version of the operator has been available since version 1 of the 'com.micr
When 1, input, output, past_state and present_state use a sequence-major, channels-last layout: input and output are (batch_size, sequence_length, d_1, ..., d_n) and the state tensors are (batch_size, state_length, d_1, ..., d_n), where channels = d_1 * ... * d_n. weight and bias keep their channels-first shapes. Requires ndim = 1. Default is 0 (channels-first).
dilation : int
Spacing between kernel taps along the causal (last spatial) axis. The receptive field spans (k_1 - 1) * dilation positions before the current one, and past_state / present_state hold that many positions. Must be >= 1. Default is 1 (undilated).
+
input_channel_offset : int
+
Starting channel of the channels-last input slice (CUDA only).
+
input_channels : int
+
CUDA channels-last input channel slice width; 0 selects all channels. Nondefault extended attributes require rank-3 input and state_window=0.
ndim : int
Spatial dimensionality: 1, 2, or 3. Default is 1.
+
state_channels_last : int
+
State layout override (CUDA only): -1 follows channels_last, 0 uses [B,C,state_length], 1 uses [B,state_length,C].
+
state_length : int
+
Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. An explicit length must cover the receptive field; convolution reads its trailing history.
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, state_length) where state_length = (k_1 - 1) * dilation. 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, state_length) 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].
@@ -1030,22 +1038,24 @@ This version of the operator has been available since version 1 of the 'com.micr
input : T
-
Input tensor with shape (batch_size, channels, ...) in the default channels-first layout. Spatial dims: 1D: (L,); 2D: (H, W); 3D: (D, H, W). When channels_last = 1 the shape is (batch_size, sequence_length, d_1, ..., d_n) instead.
+
Input tensor with shape (batch_size, channels, ...) in the default channels-first layout. Spatial dims: 1D: (L,); 2D: (H, W); 3D: (D, H, W). When channels_last = 1 the shape is (batch_size, sequence_length, d_1, ..., d_n) instead. Extended channel selection uses rank-3 (batch_size, sequence_length, input_width).
weight : T
Depthwise convolution kernel with shape (channels, 1, k_1, ...). Spatial kernel sizes: (k_1, ..., k_ndim).
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, state_length), or (W, batch_size, channels, state_length) when state_window = W > 0, in which case only slot W-1 is read, where state_length = (k_1 - 1) * dilation. When channels_last = 1 each slot is (batch_size, state_length, d_1, ..., d_n) instead. If not provided, padding is zero.
+
Carry state from previous step. For ndim=1: (batch_size, channels, state_length), or (W, batch_size, channels, state_length) when state_window = W > 0, in which case only slot W-1 is read, where state_length = (k_1 - 1) * dilation. When channels_last = 1 each slot is (batch_size, state_length, d_1, ..., d_n) instead. The state_length and state_channels_last attributes may override the length and layout for the CUDA extended path. If not provided, padding is zero.
-#### Outputs +#### Outputs (2 - 3)
output : T
-
Convolution output with same shape as input.
+
Convolution output with same shape as input, except the final dimension is input_channels when extended channel selection is enabled.
present_state : T
-
Updated carry state. For ndim=1: (batch_size, channels, state_length), or (W, batch_size, channels, state_length) when state_window = W > 0, and (batch_size, state_length, d_1, ..., d_n) per slot when channels_last = 1. Slot W-1 contains the last state_length values from the virtual input along the causal axis; slot j contains the same for the prefix ending at position (seq_len - W + j).
+
Updated carry state. For ndim=1: (batch_size, channels, state_length), or (W, batch_size, channels, state_length) when state_window = W > 0, and (batch_size, state_length, d_1, ..., d_n) per slot when channels_last = 1. Slot W-1 contains the last state_length values from the virtual input along the causal axis; slot j contains the same for the prefix ending at position (seq_len - W + j). Extended attributes override the state length and layout as for past_state.
+
first_state (optional) : T
+
Optional carry state after the first input position, with the same layout and length as present_state. CUDA only; requires rank-3 channels-last input and state_window=0.
#### Type Constraints @@ -2425,7 +2435,7 @@ This version of the operator has been available since version 1 of the 'com.micr per-chunk system (I + M) is arbitrarily ill-conditioned and the recurrence diverges. Either normalize upstream or set `qk_l2_norm=1` to have the operator do it. - Fused activations. `gate_activation='qwen'` computes the effective decay in float32 from the + Fused activations. `gate_activation='softplus'` computes the effective decay in float32 from the raw projection carried by `decay`: g = -exp(a_log) * Softplus(decay + dt_bias) @@ -2433,6 +2443,10 @@ This version of the operator has been available since version 1 of the 'com.micr `beta_activation='sigmoid'` applies a sigmoid to `beta`, and `qk_l2_norm=1` L2-normalizes each query and key head vector. Folding these in avoids materializing the intermediates and keeps the gate arithmetic in float32 independent of the input type. + + Decay and beta may be float or the query element type (both use the same type). CUDA widens + FP16/BF16 gate inputs to float before any gate arithmetic, avoiding standalone lossless casts. + State, a_log, and dt_bias remain float. WebGPU currently accepts float decay and beta only. #### Version @@ -2447,13 +2461,15 @@ This version of the operator has been available since version 1 of the 'com.micr
chunk_size : int
Tuning hint for the chunk-parallel prefill algorithm. On CUDA, 32 pins the narrow chunk and any other value lets the implementation select the widest chunk the device can hold. On WebGPU, 16 opts into the chunkwise prefill route on supported Apple silicon; default 64 retains recurrent execution. Default 64.
gate_activation : string
-
'none' (default) treats `decay` as the effective log-space decay. 'qwen' computes -exp(a_log) * Softplus(decay + dt_bias) in float32.
+
'none' (default) treats `decay` as the effective log-space decay. 'softplus' computes -exp(a_log) * Softplus(decay + dt_bias) in float32. 'qwen' is also accepted for compatibility.
qk_l2_norm : int
When 1, L2-normalize each query and key head vector before the recurrence. Default 0.
scale : float
Output scaling factor. When 0.0 (default) uses 1/sqrt(head_size_qk).
state_update_capacity : int
Capacity C for compact contiguous-prefix transition capture, in [0, 8]. 0 (default) disables compact state-update outputs.
+
state_output : string
+
Third output is either 'state_update' (default) or the recurrent state after the first token ('mid_state'). CUDA supports both modes; WebGPU supports 'state_update'.
update_rule : string
One of: 'linear', 'gated', 'delta', 'gated_delta'. Default is 'gated_delta'.
@@ -2469,16 +2485,16 @@ This version of the operator has been available since version 1 of the 'com.micr
Value, shape (total_tokens, num_heads_v, head_size_v)
cu_seqlens (optional) : TI
Exclusive prefix sums of the per-request token counts, shape (batch_size + 1). Absent means uniform packing.
-
decay (optional) : TS
-
Log-space decay, shape (total_tokens, num_heads_v) for a scalar per-head decay or (total_tokens, num_heads_v, head_size_qk) for a per-key-dimension decay.
-
beta (optional) : TS
-
Update rate, shape (total_tokens, num_heads_v)
+
decay (optional) : TG
+
Log-space decay, shape (total_tokens, num_heads_v) for a scalar per-head decay or (total_tokens, num_heads_v, head_size_qk) for a per-key-dimension decay. Float or the query element type; converted to float before gate arithmetic.
+
beta (optional) : TG
+
Update rate, shape (total_tokens, num_heads_v). Same element type as decay.
initial_state (optional) : TS
Recurrent state, shape (batch_size, num_heads_v, head_size_v, head_size_qk), V-major. Required for packed QKV. May alias final_state.
a_log (optional) : TS
-
Per-head A_log, shape (num_heads_v). Requires gate_activation=qwen.
+
Per-head A_log, shape (num_heads_v). Requires gate_activation=softplus.
dt_bias (optional) : TS
-
Per-head gate bias, shape (num_heads_v). Requires gate_activation=qwen.
+
Per-head gate bias, shape (num_heads_v). Requires gate_activation=softplus.
capture_count (optional) : TI
Number of leading local token transitions to capture for each request, shape (batch_size). Clamped on device to [0, min(state_update_capacity, sequence_length)]. Required exactly when state_update_capacity is positive.
state_update_active (optional) : TI
@@ -2493,7 +2509,7 @@ This version of the operator has been available since version 1 of the 'com.micr
final_state (optional) : TS
State after the last token of each request, shape (batch_size, num_heads_v, head_size_v, head_size_qk)
state_update (optional) : TS
-
Struct-of-arrays compact transition factors, shape (batch_size, state_update_capacity * (num_heads_v + num_heads_k * head_size_qk + num_heads_v * head_size_v)).
+
Struct-of-arrays transition factors when state_output='state_update', or the state after the first token when state_output='mid_state'.
#### Type Constraints @@ -2502,7 +2518,9 @@ This version of the operator has been available since version 1 of the 'com.micr
T : tensor(float), tensor(float16), tensor(bfloat16)
Constrain query/key/value/output types.
TS : tensor(float)
-
State, gate, beta and compact state-update tensors are always float.
+
State, gate parameters and compact state-update tensors are always float.
+
TG : tensor(float), tensor(float16), tensor(bfloat16)
+
Decay and beta are float or match the query element type. Gate arithmetic is float.
TI : tensor(int32)
Constrain index and count tensors to int32.
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 669e3083f6b15..5d96963555321 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -577,7 +577,7 @@ The **OpSet Version** column uses the following notation: |BifurcationDetector|*in* src_tokens:**T**
*in* cur_tokens:**T**
*in* prev_suffix_match_idx:**T**
*in* pred_tokens:**T**
*out* tokens:**T**
*out* suffix_match_idx:**T**|1+|**T** = tensor(int64)| |BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |CDist|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(double), tensor(float)| -|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**|1+|**T** = tensor(float)| +|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* first_state:**T**|1+|**T** = tensor(float)| |ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |CropAndResize|*in* X:**T1**
*in* rois:**T1**
*in* batch_indices:**T2**
*in* crop_size:**T2**
*out* Y:**T1**|1+|**T1** = tensor(float)
**T2** = tensor(int32)| |DecoderMaskedMultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* mask_index:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* beam_width:**M**
*in* cache_indirection:**M**
*in* bias:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**T** = tensor(float)| @@ -1089,7 +1089,7 @@ The **OpSet Version** column uses the following notation: |BitmaskBiasDropout|*in* data:**T**
*in* bias:**T**
*in* residual:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| |BitmaskDropout|*in* data:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| |BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| -|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* first_state:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |ComplexMul|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ComplexMulConj|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)| @@ -1106,7 +1106,7 @@ The **OpSet Version** column uses the following notation: |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)| |GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| -|GatedDeltaNet|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* cu_seqlens:**TI**
*in* decay:**TS**
*in* beta:**TS**
*in* initial_state:**TS**
*in* a_log:**TS**
*in* dt_bias:**TS**
*in* capture_count:**TI**
*in* state_update_active:**TI**
*out* output:**T**
*out* final_state:**TS**
*out* state_update:**TS**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)
**TS** = tensor(float)| +|GatedDeltaNet|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* cu_seqlens:**TI**
*in* decay:**TG**
*in* beta:**TG**
*in* initial_state:**TS**
*in* a_log:**TS**
*in* dt_bias:**TS**
*in* capture_count:**TI**
*in* state_update_active:**TI**
*out* output:**T**
*out* final_state:**TS**
*out* state_update:**TS**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TG** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)
**TS** = tensor(float)| |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(float4e2m1), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int4), tensor(uint4), tensor(uint8)
**T2** = tensor(bfloat16), tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)| diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu index 048faa9c2f3f5..dbca80a231f47 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -428,6 +428,12 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( const int k = lane + i * 32; s[i] = pack.initial_state != nullptr ? pack.initial_state[st_off + k] : 0.0f; } + if (seq_len == 0 && pack.mid_state != nullptr) { +#pragma unroll + for (int i = 0; i < kRegs; ++i) { + pack.mid_state[st_off + lane + i * 32] = s[i]; + } + } for (int t = 0; t < seq_len; ++t) { const int64_t tok = tok_base + t; @@ -584,6 +590,13 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant } __syncthreads(); + if (seq_len == 0 && pack.mid_state != nullptr) { + for (int idx = tid; idx < d_k * d_v; idx += kThreads) { + const int r = idx / d_v, c = idx % d_v; + pack.mid_state[st_off + static_cast(c) * d_k + r] = S[idx]; + } + } + for (int t = 0; t < seq_len; ++t) { const int64_t tok = tok_base + t; const int64_t gi = tok * p.num_heads_v + hv; 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 ca7888b7c178b..88dca80758232 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc @@ -41,8 +41,7 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) ORT_ENFORCE(info.GetAttrOrDefault("input_channels", 0) == 0 && info.GetAttrOrDefault("input_channel_offset", 0) == 0 && info.GetAttrOrDefault("state_length", 0) == 0 && - info.GetAttrOrDefault("state_channels_last", -1) == -1 && - info.node().OutputDefs().size() <= 2, + info.GetAttrOrDefault("state_channels_last", -1) == -1, "WebGPU CausalConvWithState does not support extended channel/state options (CUDA only)"); std::string activation_str = info.GetAttrOrDefault("activation", "none"); activation_ = ParseCausalConvActivation(activation_str); @@ -127,6 +126,8 @@ Status CausalConvUpdateStateProgram::GenerateShaderCode(ShaderHelper& shader) co } Status CausalConvWithState::ComputeInternal(ComputeContext& context) const { + ORT_RETURN_IF(context.OutputCount() > 2, + "WebGPU CausalConvWithState does not support extended state outputs (CUDA only)"); const Tensor* input = context.Input(0); // (B, D, L) const Tensor* weight = context.Input(1); // (D, 1, K) const Tensor* bias = context.Input(2); // optional (D,) diff --git a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc index 68a3b7c1a25fa..6eaee59a4ab4c 100644 --- a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc @@ -302,7 +302,7 @@ void Reference(const Geometry& g, const Options& o, const Inputs& in, std::vecto double decay = 1.0; if (gated) { float raw = in.decay[gi]; - if (o.gate_activation == "qwen") { + if (o.gate_activation == "qwen" || o.gate_activation == "softplus") { raw = -std::exp(in.a_log[hv]) * SoftPlus(raw + in.dt_bias[hv]); } decay = std::exp(static_cast(raw)); @@ -518,7 +518,7 @@ void RunTypedCase(const Geometry& g, const Options& o, const Inputs& in_raw, flo } else { test.AddInput("initial_state", {g.batch, g.hv, g.dv, g.dk}, in.state0); } - if (o.gate_activation == "qwen") { + if (o.gate_activation == "qwen" || o.gate_activation == "softplus") { test.AddInput("a_log", {g.hv}, in.a_log); test.AddInput("dt_bias", {g.hv}, in.dt_bias); } else { @@ -733,6 +733,47 @@ TEST(GatedDeltaNetTest, RecurrentOutputsIncludeFirstRowState) { test.Run(); } +TEST(GatedDeltaNetTest, RecurrentMidStatePreservesEmptyRequests) { + if (NeedSkipGatedDeltaNetTest()) return; + for (int head_size : {128, 32}) { + for (bool with_state : {false, true}) { + SCOPED_TRACE(head_size); + SCOPED_TRACE(with_state); + const Geometry g{1, 2, 1, 1, head_size, head_size}; + Inputs in = MakeInputs(g, 917, with_state); + in.cu_seqlens = {0, 0, 1}; + in.q = RoundToTensorType(in.q); + in.k = RoundToTensorType(in.k); + in.v = RoundToTensorType(in.v); + Options options; + std::vector expected_output, expected_state; + Reference(g, options, in, &expected_output, &expected_state); + + OpTester test("GatedDeltaNet", 1, kMSDomain); + AddCommonAttrs(test, options); + test.AddAttribute("state_output", "mid_state"); + test.AddInput("query", {1, 1, head_size}, ToTensorType(in.q)); + test.AddInput("key", {1, 1, head_size}, ToTensorType(in.k)); + test.AddInput("value", {1, 1, head_size}, ToTensorType(in.v)); + test.AddInput("cu_seqlens", {3}, in.cu_seqlens); + test.AddInput("decay", {1, 1}, in.decay); + test.AddInput("beta", {1, 1}, in.beta); + if (with_state) { + test.AddInput("initial_state", {2, 1, head_size, head_size}, in.state0); + } else { + test.AddOptionalInputEdge(); + } + test.AddOutput("output", {1, 1, head_size}, + ToTensorType(expected_output), false, 6e-3f, 6e-3f); + test.AddOutput("final_state", {2, 1, head_size, head_size}, expected_state, false, 2e-3f, 2e-3f); + test.AddOutput("mid_state", {2, 1, head_size, head_size}, expected_state, false, 2e-3f, 2e-3f); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); + } + } +} + bool NeedSkipGatedDeltaNetSplitTest() { if (NeedSkipGatedDeltaNetTest()) return true; From 62bd923adb13df65127adccd4d7162bb539e41df Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 21 Oct 2026 22:13:53 +0000 Subject: [PATCH 03/17] Add opt-in two-token GDN and convolution CUDA specialization Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/contrib_ops/cuda/gdn.md | 19 ++++++ .../cuda/bert/causal_conv_with_state_impl.cu | 17 ++++- .../cuda/bert/gated_delta_net_impl.cu | 62 ++++++++++++++++--- .../contrib_ops/gated_delta_net_op_test.cc | 32 ++++++++++ 4 files changed, 120 insertions(+), 10 deletions(-) diff --git a/docs/contrib_ops/cuda/gdn.md b/docs/contrib_ops/cuda/gdn.md index fee553a7b26dc..b541e7d53d103 100644 --- a/docs/contrib_ops/cuda/gdn.md +++ b/docs/contrib_ops/cuda/gdn.md @@ -1,5 +1,24 @@ # GatedDeltaNet — Operator Documentation +## Experimental two-token CUDA dispatch + +Set process-level tuning options before the first applicable invocation: + +- `ORT_GDN_S2_SPECIALIZE=1` specializes packed QKV for a single request with + two tokens, 16 query/key heads, 48 value heads, and 128-wide heads. + It additionally requires Qwen softplus-decay gates, sigmoid beta, + gated-delta updates, Q/K L2 normalization, uniform non-packed sequence + metadata, and no per-key decay, state-update capture, or tail passes. +- `ORT_CAUSAL_CONV_S2_SPECIALIZE=1` specializes the companion + `CausalConvProjection` for length two, kernel/state length four, and dilation + one. Layout, optional bias/SiLU, and first-token state semantics are unchanged. +- `ORT_GDN_DECODE_WARPS` selects 2, 4, or 8 warps for the existing warp decode + family; its default is 8. + +Both specialization switches default to `0` and accept only `0` or `1`. +Unsupported geometry retains generic dispatch. These options do not alter +operator schemas, state layouts, or the default package configuration. + This document describes the `com.microsoft::GatedDeltaNet` contrib operator: its schema and state contract, the chunked and warp-specialised CUDA engines and how one is selected, the consumer-Blackwell shared-memory constraint, and measured comparisons against 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 eea53e8bbe272..1b5a232331723 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 @@ -19,6 +19,7 @@ #include #include #include "contrib_ops/cuda/bert/causal_conv_with_state_impl.h" +#include "core/platform/env_var_utils.h" #include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" namespace onnxruntime { @@ -27,12 +28,17 @@ namespace cuda { namespace { -template +template __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) { + if constexpr (S2) { + length = 2; + kernel_size = state_length = 4; + dilation = 1; + } const int b = blockIdx.y; const int c = blockIdx.x * blockDim.x + threadIdx.x; if (c >= channels) return; @@ -473,6 +479,15 @@ Status LaunchCausalConvProjectionKernel( 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); + static const int specialize = ParseEnvironmentVariableWithDefault("ORT_CAUSAL_CONV_S2_SPECIALIZE", 0); + ORT_ENFORCE(specialize == 0 || specialize == 1, "ORT_CAUSAL_CONV_S2_SPECIALIZE must be 0 or 1."); + if (specialize && length == 2 && kernel_size == 4 && dilation == 1 && state_length == 4) { + CausalConvProjectionKernel<<>>( + 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()); + } CausalConvProjectionKernel<<>>( input, weight, bias, past, output, present, first_state, channels, length, kernel_size, dilation, input_width, channel_offset, state_length, diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu index dbca80a231f47..27fdb760e41bb 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -44,6 +44,7 @@ #include #include #include +#include "core/platform/env_var_utils.h" #include "contrib_ops/cuda/bert/gated_delta_net_impl.h" #include "contrib_ops/cuda/bert/gated_delta_net_mma.cuh" @@ -379,9 +380,23 @@ __device__ __forceinline__ float WarpReduceK(float v) { return v; } -template +template __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( VariantPack pack, KernelParams p, int d_v) { + if constexpr (QwenS2) { + // Host dispatch checks every field before replacing it with a compile-time constant. + p.total_tokens = p.uniform_len = 2; + p.num_heads_q = p.num_heads_k = 16; + p.num_heads_v = 48; + p.gate_activation = GateActivation::kSoftplusDecay; + p.beta_activation = BetaActivation::kSigmoid; + p.update_rule = UpdateRule::kGatedDelta; + p.qk_l2_norm = true; + p.state_update_capacity = 0; + p.decay_per_key_dim_flag = 0; + p.state_update_tail_pass = p.short_row_tail_pass = false; + d_v = 128; + } constexpr int kRegs = DK / 32; // state elements per lane const int lane = threadIdx.x; const int warp = threadIdx.y; @@ -708,28 +723,57 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant constexpr int kRecurrentThreads = 256; constexpr int kDecodeWarps = kRecurrentThreads / 32; -template -void LaunchDecodeWarp(const Descriptor& desc, const VariantPack& pack, - const KernelParams& p, cudaStream_t stream) { +template +void LaunchDecodeWarpSized(const Descriptor& desc, const VariantPack& pack, + const KernelParams& p, cudaStream_t stream) { const dim3 grid(desc.batch, desc.num_heads_v, - (desc.head_size_v + kDecodeWarps - 1) / kDecodeWarps); - const dim3 block(32, kDecodeWarps, 1); + (desc.head_size_v + Warps - 1) / Warps); + const dim3 block(32, Warps, 1); switch (desc.head_size_qk) { case 64: - GatedDeltaNetDecodeWarpKernel + GatedDeltaNetDecodeWarpKernel <<>>(pack, p, desc.head_size_v); break; case 128: - GatedDeltaNetDecodeWarpKernel + if constexpr (PackedQkv) { + static const int specialize = ParseEnvironmentVariableWithDefault("ORT_GDN_S2_SPECIALIZE", 0); + ORT_ENFORCE(specialize == 0 || specialize == 1, "ORT_GDN_S2_SPECIALIZE must be 0 or 1."); + if (specialize && desc.batch == 1 && p.total_tokens == 2 && p.uniform_len == 2 && + p.num_heads_q == 16 && p.num_heads_k == 16 && p.num_heads_v == 48 && + desc.head_size_v == 128 && pack.cu_seqlens == nullptr && + p.gate_activation == GateActivation::kSoftplusDecay && + p.beta_activation == BetaActivation::kSigmoid && p.update_rule == UpdateRule::kGatedDelta && + p.qk_l2_norm && p.state_update_capacity == 0 && p.decay_per_key_dim_flag == 0 && + !p.state_update_tail_pass && !p.short_row_tail_pass) { + GatedDeltaNetDecodeWarpKernel + <<>>(pack, p, desc.head_size_v); + break; + } + } + GatedDeltaNetDecodeWarpKernel <<>>(pack, p, desc.head_size_v); break; default: - GatedDeltaNetDecodeWarpKernel + GatedDeltaNetDecodeWarpKernel <<>>(pack, p, desc.head_size_v); break; } } +template +void LaunchDecodeWarp(const Descriptor& desc, const VariantPack& pack, + const KernelParams& p, cudaStream_t stream) { + static const int warps = ParseEnvironmentVariableWithDefault("ORT_GDN_DECODE_WARPS", kDecodeWarps); + ORT_ENFORCE(warps == 2 || warps == 4 || warps == kDecodeWarps, "ORT_GDN_DECODE_WARPS must be 2, 4, or 8."); + if (warps == 2) { + LaunchDecodeWarpSized(desc, pack, p, stream); + } else if (warps == 4) { + LaunchDecodeWarpSized(desc, pack, p, stream); + } else { + LaunchDecodeWarpSized(desc, pack, p, stream); + } +} + template Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const VariantPack& pack, const KernelParams& p, size_t max_shared_memory_per_block, diff --git a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc index 6eaee59a4ab4c..0664feda322cf 100644 --- a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc @@ -1546,6 +1546,38 @@ TEST(GatedDeltaNetTest, DirectPackedQkvBFloat16ExactParity) { TestDirectPackedQkvParity(); } +template +void TestQwenS2PackedParity() { + if (NeedSkipGatedDeltaNetTest()) return; + Geometry geometry{2, 1, 16, 48, 128, 128}; + Options options; + options.qk_l2_norm = 1; + options.gate_activation = "qwen"; + options.beta_activation = "sigmoid"; + options.gates_in_input_type = true; + Inputs inputs = MakeInputs(geometry, 151); + std::vector reference, packed; + RunTypedCase(geometry, options, inputs, 2e-2f, 2e-2f, true, &reference); + RunTypedCase(geometry, options, inputs, 2e-2f, 2e-2f, true, &packed, + false, false, nullptr, 0, true); + ASSERT_EQ(reference.size(), packed.size()); + for (size_t index = 0; index < reference.size(); ++index) { + const auto& expected = reference[index].Get(); + const auto& actual = packed[index].Get(); + ASSERT_EQ(expected.Shape(), actual.Shape()); + ASSERT_EQ(expected.SizeInBytes(), actual.SizeInBytes()); + EXPECT_EQ(std::memcmp(expected.DataRaw(), actual.DataRaw(), expected.SizeInBytes()), 0); + } +} + +TEST(GatedDeltaNetTest, QwenS2PackedFloat16ExactParity) { + TestQwenS2PackedParity(); +} + +TEST(GatedDeltaNetTest, QwenS2PackedBFloat16ExactParity) { + TestQwenS2PackedParity(); +} + TEST(GatedDeltaNetTest, Recurrent_VerifyBatch) { if (NeedSkipGatedDeltaNetTest()) return; Geometry g{4, 1, 1, 2, kDim, kDim}; From 51cd77f6e0f7e40599dc9172da765d626238a0cb Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 23 Oct 2026 23:19:15 +0000 Subject: [PATCH 04/17] [CUDA] Add S=4 prefix-state outputs (present_mid2/3) to GatedDeltaNet and CausalConvWithState Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 12 +++++-- docs/OperatorKernels.md | 6 ++-- .../cuda/bert/causal_conv_with_state.cc | 4 +++ .../cuda/bert/causal_conv_with_state_impl.cu | 27 +++++++------- .../cuda/bert/causal_conv_with_state_impl.h | 2 +- .../contrib_ops/cuda/bert/gated_delta_net.cc | 8 +++++ .../cuda/bert/gated_delta_net_impl.cu | 36 ++++++++++++++++--- .../cuda/bert/gated_delta_net_impl.h | 2 ++ .../core/graph/contrib_ops/bert_defs.cc | 22 +++++++++--- 9 files changed, 92 insertions(+), 27 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index dd3b2164ae6e2..17cefd0b949a6 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -1048,7 +1048,7 @@ This version of the operator has been available since version 1 of the 'com.micr
Carry state from previous step. For ndim=1: (batch_size, channels, state_length), or (W, batch_size, channels, state_length) when state_window = W > 0, in which case only slot W-1 is read, where state_length = (k_1 - 1) * dilation. When channels_last = 1 each slot is (batch_size, state_length, d_1, ..., d_n) instead. The state_length and state_channels_last attributes may override the length and layout for the CUDA extended path. If not provided, padding is zero.
-#### Outputs (2 - 3) +#### Outputs (2 - 5)
output : T
@@ -1057,6 +1057,10 @@ This version of the operator has been available since version 1 of the 'com.micr
Updated carry state. For ndim=1: (batch_size, channels, state_length), or (W, batch_size, channels, state_length) when state_window = W > 0, and (batch_size, state_length, d_1, ..., d_n) per slot when channels_last = 1. Slot W-1 contains the last state_length values from the virtual input along the causal axis; slot j contains the same for the prefix ending at position (seq_len - W + j). Extended attributes override the state length and layout as for past_state.
first_state (optional) : T
Optional carry state after the first input position, with the same layout and length as present_state. CUDA only; requires rank-3 channels-last input and state_window=0.
+
second_state (optional) : T
+
Optional carry state after at most two input positions, with the same layout and length as present_state.
+
third_state (optional) : T
+
Optional carry state after at most three input positions, with the same layout and length as present_state.
#### Type Constraints @@ -2565,7 +2569,7 @@ This version of the operator has been available since version 1 of the 'com.micr
CPU int32 control with shape (1). Zero disables transition capture, ignores capture_count, and produces a zero-filled state_update. Omission is conservative.
-#### Outputs (1 - 3) +#### Outputs (1 - 5)
output : T
@@ -2574,6 +2578,10 @@ This version of the operator has been available since version 1 of the 'com.micr
State after the last token of each request, shape (batch_size, num_heads_v, head_size_v, head_size_qk)
state_update (optional) : TS
Struct-of-arrays transition factors when state_output='state_update', or the state after the first token when state_output='mid_state'.
+
mid_state2 (optional) : TS
+
CUDA state after at most two request tokens; requires state_output='mid_state'.
+
mid_state3 (optional) : TS
+
CUDA state after at most three request tokens; requires state_output='mid_state'.
#### Type Constraints diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 55e218741747a..f06631532841a 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -585,7 +585,7 @@ The **OpSet Version** column uses the following notation: |BifurcationDetector|*in* src_tokens:**T**
*in* cur_tokens:**T**
*in* prev_suffix_match_idx:**T**
*in* pred_tokens:**T**
*out* tokens:**T**
*out* suffix_match_idx:**T**|1+|**T** = tensor(int64)| |BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |CDist|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(double), tensor(float)| -|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* first_state:**T**|1+|**T** = tensor(float)| +|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* first_state:**T**
*out* second_state:**T**
*out* third_state:**T**|1+|**T** = tensor(float)| |ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |CropAndResize|*in* X:**T1**
*in* rois:**T1**
*in* batch_indices:**T2**
*in* crop_size:**T2**
*out* Y:**T1**|1+|**T1** = tensor(float)
**T2** = tensor(int32)| |DecoderMaskedMultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* mask_index:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* beam_width:**M**
*in* cache_indirection:**M**
*in* bias:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**T** = tensor(float)| @@ -1103,7 +1103,7 @@ The **OpSet Version** column uses the following notation: |BitmaskBiasDropout|*in* data:**T**
*in* bias:**T**
*in* residual:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| |BitmaskDropout|*in* data:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| |BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| -|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* first_state:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* first_state:**T**
*out* second_state:**T**
*out* third_state:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |ComplexMul|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ComplexMulConj|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)| @@ -1120,7 +1120,7 @@ The **OpSet Version** column uses the following notation: |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)| |GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| -|GatedDeltaNet|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* cu_seqlens:**TI**
*in* decay:**TG**
*in* beta:**TG**
*in* initial_state:**TS**
*in* a_log:**TS**
*in* dt_bias:**TS**
*in* capture_count:**TI**
*in* state_update_active:**TI**
*out* output:**T**
*out* final_state:**TS**
*out* state_update:**TS**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TG** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)
**TS** = tensor(float)| +|GatedDeltaNet|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* cu_seqlens:**TI**
*in* decay:**TG**
*in* beta:**TG**
*in* initial_state:**TS**
*in* a_log:**TS**
*in* dt_bias:**TS**
*in* capture_count:**TI**
*in* state_update_active:**TI**
*out* output:**T**
*out* final_state:**TS**
*out* state_update:**TS**
*out* mid_state2:**TS**
*out* mid_state3:**TS**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TG** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)
**TS** = tensor(float)| |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(float4e2m1), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int4), tensor(uint4), tensor(uint8)
**T2** = tensor(bfloat16), tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)| 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 f7a8d94c2b80a..fa12811f9caba 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc @@ -96,6 +96,8 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { 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; + Tensor* second = context->OutputCount() > 3 ? context->Output(3, state_shape) : nullptr; + Tensor* third = context->OutputCount() > 4 ? context->Output(4, state_shape) : nullptr; using CudaT = typename OrtToCudaType::type; return LaunchCausalConvProjectionKernel( Stream(context), reinterpret_cast(input_tensor->Data()), @@ -105,6 +107,8 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { reinterpret_cast(output->MutableData()), reinterpret_cast(present->MutableData()), first ? reinterpret_cast(first->MutableData()) : nullptr, + second ? reinterpret_cast(second->MutableData()) : nullptr, + third ? reinterpret_cast(third->MutableData()) : nullptr, batch, channels, length, kernel, dilation_, width, input_channel_offset_, state_length, state_last, activation_ != "none"); } 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 e18f13817fe37..bacfae35be6f6 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 @@ -31,7 +31,7 @@ namespace { template __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, + T* output, T* present, T* first_state, T* second_state, T* third_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) { if constexpr (S2) { @@ -62,11 +62,14 @@ __global__ void CausalConvProjectionKernel( present[state_layout.Offset(b, position, c)] = source >= 0 ? input[(static_cast(b) * length + source) * input_width + channel_offset + c] : (past == nullptr ? from_float(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(b) * length * input_width + channel_offset + c] - : (past == nullptr ? from_float(0.0f) : past[state_layout.Offset(b, state_length + first_source, c)]); + T* snapshots[3] = {first_state, second_state, third_state}; +#pragma unroll + for (int prefix = 0; prefix < 3; ++prefix) { + if (snapshots[prefix] == nullptr) continue; + const int prefix_source = min(length, prefix + 1) - state_length + position; + snapshots[prefix][state_layout.Offset(b, position, c)] = + prefix_source >= 0 ? input[(static_cast(b) * length + prefix_source) * input_width + channel_offset + c] + : (past == nullptr ? from_float(0.0f) : past[state_layout.Offset(b, state_length + prefix_source, c)]); } } } @@ -475,7 +478,7 @@ __global__ void CausalConvPrefillKernelChannelsLast( template 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, + T* output, T* present, T* first_state, T* second_state, T* third_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); @@ -484,21 +487,21 @@ Status LaunchCausalConvProjectionKernel( if (specialize && length >= 1 && length <= 8 && kernel_size == 4 && dilation == 1 && state_length == 4 && apply_silu) { CausalConvProjectionKernel<<>>( - input, weight, bias, past, output, present, first_state, channels, length, + input, weight, bias, past, output, present, first_state, second_state, third_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()); } CausalConvProjectionKernel<<>>( - input, weight, bias, past, output, present, first_state, channels, length, + input, weight, bias, past, output, present, first_state, second_state, third_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( \ - cudaStream_t, const T*, const T*, const T*, const T*, T*, T*, T*, \ +#define INSTANTIATE_PROJECTION(T) \ + template Status LaunchCausalConvProjectionKernel( \ + cudaStream_t, const T*, const T*, const T*, const T*, T*, T*, T*, T*, T*, \ int, int, int, int, int, int, int, int, bool, bool); INSTANTIATE_PROJECTION(float) INSTANTIATE_PROJECTION(half) 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 95d98dc96cdcd..d59982af947ba 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 @@ -34,7 +34,7 @@ inline CausalConvLayout MakeCausalConvLayout(bool channels_last, int64_t channel template 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, + T* output, T* present, T* first_state, T* second_state, T* third_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); diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index 7bd851bca9bc1..cdb907467b170 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -337,9 +337,15 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { Tensor* state_update = nullptr; Tensor* mid_state = nullptr; + Tensor* mid_state2 = nullptr; + Tensor* mid_state3 = nullptr; if (mid_state_output_) { mid_state = context->Output(2, state_shape); + if (context->OutputCount() > 3) mid_state2 = context->Output(3, state_shape); + if (context->OutputCount() > 4) mid_state3 = context->Output(4, state_shape); } else { + ORT_RETURN_IF_NOT(context->OutputCount() <= 3, + "additional prefix states require state_output='mid_state'"); const int64_t state_update_width = static_cast(state_update_capacity_) * (num_heads_v + num_heads_k * head_size_qk + num_heads_v * head_size_v); @@ -483,6 +489,8 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { pack.output = reinterpret_cast(output->MutableData()); pack.final_state = final_state != nullptr ? final_state->MutableData() : nullptr; pack.mid_state = mid_state != nullptr ? mid_state->MutableData() : nullptr; + pack.mid_state2 = mid_state2 != nullptr ? mid_state2->MutableData() : nullptr; + pack.mid_state3 = mid_state3 != nullptr ? mid_state3->MutableData() : nullptr; pack.state_update = state_update != nullptr ? state_update->MutableData() : nullptr; // Only the inactive case is cleared. While capture is active the kernels write just the diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu index de3c6168d2264..6b4edb89bbdaa 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -692,10 +692,14 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( pack.output[tok * v_stride + hv * d_v + col] = static_cast(p.scale * o); } - if (t == 0 && pack.mid_state != nullptr) { + float* snapshots[3] = {pack.mid_state, pack.mid_state2, pack.mid_state3}; +#pragma unroll + for (int prefix = 0; prefix < 3; ++prefix) { + float* prefix_state = snapshots[prefix]; + if (prefix_state == nullptr || t != min(prefix + 1, seq_len) - 1) continue; #pragma unroll for (int i = 0; i < kRegs; ++i) { - pack.mid_state[st_off + lane + i * 32] = s[i]; + prefix_state[st_off + lane + i * 32] = s[i]; } } } @@ -706,6 +710,15 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( pack.final_state[st_off + lane + i * 32] = s[i]; } } + if (seq_len == 0) { + float* snapshots[3] = {pack.mid_state, pack.mid_state2, pack.mid_state3}; +#pragma unroll + for (int prefix = 0; prefix < 3; ++prefix) { + if (snapshots[prefix] == nullptr) continue; +#pragma unroll + for (int i = 0; i < kRegs; ++i) snapshots[prefix][st_off + lane + i * 32] = s[i]; + } + } } // --------------------------------------------------------------------------- @@ -858,10 +871,14 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant for (int r = 0; r < d_k; ++r) acc += S[r * d_v + c] * qbuf[r]; pack.output[tok * v_stride + hv * d_v + c] = (T)(p.scale * acc); } - if (t == 0 && pack.mid_state != nullptr) { + float* snapshots[3] = {pack.mid_state, pack.mid_state2, pack.mid_state3}; +#pragma unroll + for (int prefix = 0; prefix < 3; ++prefix) { + float* prefix_state = snapshots[prefix]; + if (prefix_state == nullptr || t != min(prefix + 1, seq_len) - 1) continue; for (int idx = tid; idx < d_k * d_v; idx += kThreads) { const int r = idx / d_v, c = idx % d_v; - pack.mid_state[st_off + static_cast(c) * d_k + r] = S[idx]; + prefix_state[st_off + static_cast(c) * d_k + r] = S[idx]; } } @@ -874,6 +891,17 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant pack.final_state[st_off + static_cast(c) * d_k + r] = S[idx]; } } + if (seq_len == 0) { + float* snapshots[3] = {pack.mid_state, pack.mid_state2, pack.mid_state3}; +#pragma unroll + for (int prefix = 0; prefix < 3; ++prefix) { + if (snapshots[prefix] == nullptr) continue; + for (int idx = tid; idx < d_k * d_v; idx += kThreads) { + const int r = idx / d_v, c = idx % d_v; + snapshots[prefix][st_off + static_cast(c) * d_k + r] = S[idx]; + } + } + } } constexpr int kRecurrentThreads = 256; diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h index 2e3010c7138ba..eb82ba072f1fa 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h @@ -34,6 +34,8 @@ struct VariantPack { T* output = nullptr; // [total_tokens, max(Hq,Hv), V] float* final_state = nullptr; // [B, Hv, V, K] float* mid_state = nullptr; // [B, Hv, V, K], state after local row 0 + float* mid_state2 = nullptr; // state after local row 1 + float* mid_state3 = nullptr; // state after local row 2 float* state_update = nullptr; // [B, C * (Hv + Hk*K + Hv*V)] void* workspace = nullptr; }; diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 607237591f9cc..cfd4c6e91235d 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -4788,6 +4788,12 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "layout and length as present_state. CUDA only; requires rank-3 channels-last input " "and state_window=0.", "T", OpSchema::Optional) + .Output(3, "second_state", + "Optional carry state after at most two input positions, with the same layout and length as present_state.", + "T", OpSchema::Optional) + .Output(4, "third_state", + "Optional carry state after at most three input positions, with the same layout and length as present_state.", + "T", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") @@ -4812,7 +4818,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( state_window != 0 || getAttribute(ctx, "dilation", 1) < 1) { fail_shape_inference("CausalConvWithState: extended attributes require channels_last=1, ndim=1, state_window=0"); } - if (ctx.getNumOutputs() > 2) propagateElemTypeFromInputToOutput(ctx, 0, 2); + for (size_t i = 2; i < ctx.getNumOutputs(); ++i) propagateElemTypeFromInputToOutput(ctx, 0, i); if (!hasInputShape(ctx, 0)) return; const auto& input = getInputShape(ctx, 0); if (input.dim_size() != 3) fail_shape_inference("CausalConvWithState: extended input must have rank 3"); @@ -4850,7 +4856,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( *state.add_dim() = state_layout == 0 ? *channels : length; *state.add_dim() = state_layout == 0 ? length : *channels; updateOutputShape(ctx, 1, state); - if (ctx.getNumOutputs() > 2) updateOutputShape(ctx, 2, state); + for (size_t i = 2; i < ctx.getNumOutputs(); ++i) updateOutputShape(ctx, i, state); return; } if (state_window < 0 || state_window > kMaxStateWindow) { @@ -5624,6 +5630,10 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Struct-of-arrays transition factors when state_output='state_update', or " "the state after the first token when state_output='mid_state'.", "TS", OpSchema::Optional) + .Output(3, "mid_state2", "CUDA state after at most two request tokens; requires state_output='mid_state'.", + "TS", OpSchema::Optional) + .Output(4, "mid_state3", "CUDA state after at most three request tokens; requires state_output='mid_state'.", + "TS", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain query/key/value/output types.") .TypeConstraint("TS", {"tensor(float)"}, @@ -5636,8 +5646,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( if (ctx.getNumOutputs() > 1) { updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); } - if (ctx.getNumOutputs() > 2) { - updateOutputElemType(ctx, 2, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + for (size_t i = 2; i < ctx.getNumOutputs(); ++i) { + updateOutputElemType(ctx, i, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); } const bool has_key = ctx.getNumInputs() > 1 && ctx.getInputType(1) != nullptr; @@ -5772,9 +5782,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA( fail_shape_inference("GatedDeltaNet: mid_state output requires state_update_capacity=0"); } if (state_shape != nullptr) { - updateOutputShape(ctx, 2, *state_shape); + for (size_t i = 2; i < ctx.getNumOutputs(); ++i) updateOutputShape(ctx, i, *state_shape); } } else { + if (ctx.getNumOutputs() > 3) + fail_shape_inference("GatedDeltaNet: additional prefix states require state_output='mid_state'"); ONNX_NAMESPACE::TensorShapeProto capsule_shape; add_batch_dim(capsule_shape); auto* width = capsule_shape.add_dim(); From bbeb06e27da9f25bcc01b739a750f308542ea48b Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 00:32:03 +0000 Subject: [PATCH 05/17] [CUDA] Remove unused ORT_GDN_DECODE_WARPS instantiations and duplicate seq_len==0 blocks in GatedDeltaNet Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/contrib_ops/cuda/gdn.md | 2 -- .../cuda/bert/gated_delta_net_impl.cu | 31 ++----------------- .../cuda/bert/gated_delta_net_mma.cuh | 2 +- 3 files changed, 3 insertions(+), 32 deletions(-) diff --git a/docs/contrib_ops/cuda/gdn.md b/docs/contrib_ops/cuda/gdn.md index 6501a340525fb..8037400dfb5b9 100644 --- a/docs/contrib_ops/cuda/gdn.md +++ b/docs/contrib_ops/cuda/gdn.md @@ -13,8 +13,6 @@ Set process-level tuning options before the first applicable invocation: `CausalConvProjection` for lengths 1 to 8, kernel/state length four, dilation one, and SiLU activation. Layout, optional bias, and first-token state semantics are unchanged. -- `ORT_GDN_DECODE_WARPS` selects 2, 4, or 8 warps for the existing warp decode - family; its default is 8. Both specialization switches default to `1` and accept only `0` or `1`. Unsupported geometry retains generic dispatch. These options do not alter diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu index 6b4edb89bbdaa..4337c49c848af 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -599,12 +599,6 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( const int k = lane + i * 32; s[i] = pack.initial_state != nullptr ? pack.initial_state[st_off + k] : 0.0f; } - if (seq_len == 0 && pack.mid_state != nullptr) { -#pragma unroll - for (int i = 0; i < kRegs; ++i) { - pack.mid_state[st_off + lane + i * 32] = s[i]; - } - } for (int t = 0; t < seq_len; ++t) { const int64_t tok = tok_base + t; @@ -774,13 +768,6 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant } __syncthreads(); - if (seq_len == 0 && pack.mid_state != nullptr) { - for (int idx = tid; idx < d_k * d_v; idx += kThreads) { - const int r = idx / d_v, c = idx % d_v; - pack.mid_state[st_off + static_cast(c) * d_k + r] = S[idx]; - } - } - for (int t = 0; t < seq_len; ++t) { const int64_t tok = tok_base + t; const int64_t gi = tok * p.num_heads_v + hv; @@ -945,29 +932,15 @@ void LaunchDecodeWarpSized(const Descriptor& desc, const VariantPack& pack, } } -template -void LaunchDecodeWarp(const Descriptor& desc, const VariantPack& pack, - const KernelParams& p, cudaStream_t stream) { - static const int warps = ParseEnvironmentVariableWithDefault("ORT_GDN_DECODE_WARPS", kDecodeWarps); - ORT_ENFORCE(warps == 2 || warps == 4 || warps == kDecodeWarps, "ORT_GDN_DECODE_WARPS must be 2, 4, or 8."); - if (warps == 2) { - LaunchDecodeWarpSized(desc, pack, p, stream); - } else if (warps == 4) { - LaunchDecodeWarpSized(desc, pack, p, stream); - } else { - LaunchDecodeWarpSized(desc, pack, p, stream); - } -} - template Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const VariantPack& pack, const KernelParams& p, size_t max_shared_memory_per_block, cudaStream_t stream) { if (warp_specialized) { if (pack.packed_qkv) { - LaunchDecodeWarp(desc, pack, p, stream); + LaunchDecodeWarpSized(desc, pack, p, stream); } else { - LaunchDecodeWarp(desc, pack, p, stream); + LaunchDecodeWarpSized(desc, pack, p, stream); } return CUDA_CALL(cudaGetLastError()); } diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh index 7c3bf5b4b59d5..67c4e8be40e6e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh @@ -122,7 +122,7 @@ __device__ __forceinline__ uint32_t PackHalf2(__nv_bfloat16 lo, __nv_bfloat16 hi // mma.sync.m16n8k16 is Ampere and later, so ptxas rejects it for a pre-Ampere target. The only // caller is the chunked engine, and SelectPlan requires sm_major >= 8 before choosing it, so the // pre-Ampere body below exists purely to let those targets compile and is never executed. -template +template __device__ __forceinline__ void MmaM16N8K16(float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]); From 87cc338d5af5e23214f4635e99f292d442e61b83 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 04:09:37 +0000 Subject: [PATCH 06/17] fix(cuda): include env_var_utils.h after cuda_common.h and sync GatedDeltaNet docs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 6 +++--- onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu | 5 ++++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 01be02a5e874a..347b892d5ba82 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2514,7 +2514,6 @@ This version of the operator has been available since version 1 of the 'com.micr `beta_activation='sigmoid'` applies a sigmoid to `beta`, and `qk_l2_norm=1` L2-normalizes each query and key head vector. Folding these in avoids materializing the intermediates and keeps the gate arithmetic in float32 independent of the input type. - Decay and beta may be float or the query element type (both use the same type). CUDA widens FP16/BF16 gate inputs to float before any gate arithmetic, avoiding standalone lossless casts. State, a_log, and dt_bias remain float. WebGPU currently accepts float decay and beta only. @@ -2537,10 +2536,10 @@ This version of the operator has been available since version 1 of the 'com.micr
When 1, L2-normalize each query and key head vector before the recurrence. Default 0.
scale : float
Output scaling factor. When 0.0 (default) uses 1/sqrt(head_size_qk).
-
state_update_capacity : int
-
Capacity C for compact contiguous-prefix transition capture, in [0, 8]. 0 (default) disables compact state-update outputs.
state_output : string
Third output is either 'state_update' (default) or the recurrent state after the first token ('mid_state'). CUDA supports both modes; WebGPU supports 'state_update'.
+
state_update_capacity : int
+
Capacity C for compact contiguous-prefix transition capture, in [0, 8]. 0 (default) disables compact state-update outputs.
update_rule : string
One of: 'linear', 'gated', 'delta', 'gated_delta'. Default is 'gated_delta'.
@@ -8975,3 +8974,4 @@ No versioning maintained for experimental ops.
Constrain input and output types to float32 tensors.
+ diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu index 4337c49c848af..811a9079262e8 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -44,13 +44,16 @@ #include #include #include -#include "core/platform/env_var_utils.h" #include "contrib_ops/cuda/bert/gated_delta_net_impl.h" #include "contrib_ops/cuda/bert/gated_delta_net_mma.cuh" #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cuda_common.h" +// Include env_var_utils.h after cuda_common.h: the latter transitively pulls in provider_api.h, +// which defines SHARED_PROVIDER and suppresses env_var_utils.h's own logging.h include. +#include "core/platform/env_var_utils.h" + namespace onnxruntime { namespace contrib { namespace cuda { From 4439bb3075585fe369000c1fab75367538a6b286 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 15:17:49 +0000 Subject: [PATCH 07/17] feat(cuda): support state_update_capacity in CausalConvWithState and QwenS2 GatedDeltaNet decode - Add state_update_capacity (0..8) attribute, shape inference, and device emission to extended CausalConvWithState (output 2 of shape [batch, state_update_capacity, channels]). - Enable state_update_capacity > 0 in the QwenS2 warp-specialized GatedDeltaNetDecodeWarpKernel and gate state_update_active on cu_seqlens != nullptr || total_tokens <= (state_update_capacity + 1) * batch. - Preserve packed-QKV uniform prefill engine/precision selection in GatedDeltaNet and pass DefaultCudaExecutionProvider() in RecurrentOutputsIncludeFirstRowState for TensorRT CI compatibility. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 2 ++ .../cuda/bert/causal_conv_with_state.cc | 28 +++++++++++++++---- .../cuda/bert/causal_conv_with_state.h | 1 + .../cuda/bert/causal_conv_with_state_impl.cu | 26 +++++++++++------ .../cuda/bert/causal_conv_with_state_impl.h | 3 +- .../contrib_ops/cuda/bert/gated_delta_net.cc | 11 ++++++-- .../cuda/bert/gated_delta_net_impl.cu | 3 +- .../cuda/bert/gated_delta_net_mma.cuh | 2 +- .../core/graph/contrib_ops/bert_defs.cc | 26 +++++++++++++++-- .../causal_conv_with_state_op_test.cc | 26 +++++++++++++++++ .../contrib_ops/gated_delta_net_op_test.cc | 4 ++- 11 files changed, 106 insertions(+), 26 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 347b892d5ba82..e6bbbf8667d82 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -1034,6 +1034,8 @@ This version of the operator has been available since version 1 of the 'com.micr
State layout override (CUDA only): -1 follows channels_last, 0 uses [B,C,state_length], 1 uses [B,state_length,C].
state_length : int
Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. An explicit length must cover the receptive field; convolution reads its trailing history.
+
state_update_capacity : int
+
Compact speculative state-update capture capacity (CUDA extended path only, 0..8); when > 0, output 2 has shape (batch_size, state_update_capacity, channels) and records the first min(sequence_length, state_update_capacity) projected input rows.
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, state_length) where state_length = (k_1 - 1) * dilation. 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, state_length) 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].
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 fa12811f9caba..da07a543b1c0d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc @@ -50,8 +50,10 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : CudaKern input_channel_offset_ = SafeInt(info.GetAttrOrDefault("input_channel_offset", 0)); state_length_ = SafeInt(info.GetAttrOrDefault("state_length", 0)); state_channels_last_ = SafeInt(info.GetAttrOrDefault("state_channels_last", -1)); + state_update_capacity_ = SafeInt(info.GetAttrOrDefault("state_update_capacity", 0)); ORT_ENFORCE(input_channels_ >= 0 && input_channel_offset_ >= 0 && state_length_ >= 0 && - state_channels_last_ >= -1 && state_channels_last_ <= 1, + state_channels_last_ >= -1 && state_channels_last_ <= 1 && + state_update_capacity_ >= 0 && state_update_capacity_ <= 8, "invalid packed-channel or state attributes"); } @@ -69,11 +71,13 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { 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; + state_channels_last_ != -1 || state_update_capacity_ != 0 || 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"); + ORT_RETURN_IF_NOT(state_update_capacity_ == 0 || context->OutputCount() <= 3, + "extended CausalConvWithState with state_update_capacity > 0 supports at most 3 outputs"); const int batch = SafeInt(input_shape[0]); const int length = SafeInt(input_shape[1]); const int width = SafeInt(input_shape[2]); @@ -95,9 +99,19 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { "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; - Tensor* second = context->OutputCount() > 3 ? context->Output(3, state_shape) : nullptr; - Tensor* third = context->OutputCount() > 4 ? context->Output(4, state_shape) : nullptr; + Tensor* state_update = nullptr; + Tensor* first = nullptr; + Tensor* second = nullptr; + Tensor* third = nullptr; + if (state_update_capacity_ > 0) { + state_update = context->OutputCount() > 2 + ? context->Output(2, TensorShape({batch, state_update_capacity_, channels})) + : nullptr; + } else { + first = context->OutputCount() > 2 ? context->Output(2, state_shape) : nullptr; + second = context->OutputCount() > 3 ? context->Output(3, state_shape) : nullptr; + third = context->OutputCount() > 4 ? context->Output(4, state_shape) : nullptr; + } using CudaT = typename OrtToCudaType::type; return LaunchCausalConvProjectionKernel( Stream(context), reinterpret_cast(input_tensor->Data()), @@ -110,7 +124,9 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { second ? reinterpret_cast(second->MutableData()) : nullptr, third ? reinterpret_cast(third->MutableData()) : nullptr, batch, channels, length, kernel, dilation_, width, input_channel_offset_, state_length, - state_last, activation_ != "none"); + state_last, activation_ != "none", + state_update ? reinterpret_cast(state_update->MutableData()) : nullptr, + state_update_capacity_); } // Validate input rank and weight rank 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 f0807d69c7da3..a9fd12f83ad6a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h @@ -29,6 +29,7 @@ class CausalConvWithState final : public onnxruntime::cuda::CudaKernel { int input_channel_offset_; int state_length_; int state_channels_last_; + int state_update_capacity_; }; } // 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 bacfae35be6f6..719192aa4da26 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 @@ -33,11 +33,11 @@ __global__ void CausalConvProjectionKernel( const T* input, const T* weight, const T* bias, const T* past, T* output, T* present, T* first_state, T* second_state, T* third_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) { + int state_length, CausalConvLayout state_layout, bool apply_silu, + T* state_update = nullptr, int state_update_capacity = 0) { if constexpr (S2) { kernel_size = state_length = 4; dilation = 1; - apply_silu = true; } const int b = blockIdx.y; const int c = blockIdx.x * blockDim.x + threadIdx.x; @@ -45,11 +45,17 @@ __global__ void CausalConvProjectionKernel( const int pad = (kernel_size - 1) * dilation; const int position = blockIdx.z; { + const T cur_in = input[(static_cast(b) * length + position) * input_width + channel_offset + c]; + if (state_update != nullptr && position < state_update_capacity) { + state_update[(static_cast(b) * state_update_capacity + position) * channels + c] = cur_in; + } 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(b) * length + source) * input_width + channel_offset + c]) + ? (source == position + ? to_float(cur_in) + : to_float(input[(static_cast(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]); } @@ -480,29 +486,31 @@ Status LaunchCausalConvProjectionKernel( cudaStream_t stream, const T* input, const T* weight, const T* bias, const T* past, T* output, T* present, T* first_state, T* second_state, T* third_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) { + int state_length, bool state_channels_last, bool apply_silu, + T* state_update, int state_update_capacity) { const dim3 grid((channels - 1) / 128 + 1, batch, length); static const int specialize = ParseEnvironmentVariableWithDefault("ORT_CAUSAL_CONV_S2_SPECIALIZE", 1); ORT_ENFORCE(specialize == 0 || specialize == 1, "ORT_CAUSAL_CONV_S2_SPECIALIZE must be 0 or 1."); - if (specialize && length >= 1 && length <= 8 && kernel_size == 4 && dilation == 1 && - state_length == 4 && apply_silu) { + if (specialize && length >= 1 && length <= 8 && kernel_size == 4 && dilation == 1 && state_length == 4) { CausalConvProjectionKernel<<>>( input, weight, bias, past, output, present, first_state, second_state, third_state, channels, length, kernel_size, dilation, input_width, channel_offset, state_length, - MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu); + MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu, + state_update, state_update_capacity); return CUDA_CALL(cudaGetLastError()); } CausalConvProjectionKernel<<>>( input, weight, bias, past, output, present, first_state, second_state, third_state, channels, length, kernel_size, dilation, input_width, channel_offset, state_length, - MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu); + MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu, + state_update, state_update_capacity); return CUDA_CALL(cudaGetLastError()); } #define INSTANTIATE_PROJECTION(T) \ template Status LaunchCausalConvProjectionKernel( \ cudaStream_t, const T*, const T*, const T*, const T*, T*, T*, T*, T*, T*, \ - int, int, int, int, int, int, int, int, bool, bool); + int, int, int, int, int, int, int, int, bool, bool, T*, int); INSTANTIATE_PROJECTION(float) INSTANTIATE_PROJECTION(half) INSTANTIATE_PROJECTION(__nv_bfloat16) 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 d59982af947ba..436e00278efda 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 @@ -36,7 +36,8 @@ Status LaunchCausalConvProjectionKernel( cudaStream_t stream, const T* input, const T* weight, const T* bias, const T* past, T* output, T* present, T* first_state, T* second_state, T* third_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); + int state_length, bool state_channels_last, bool apply_silu, + T* state_update = nullptr, int state_update_capacity = 0); // Fused causal depthwise conv1d + activation + state management. // One thread block per (batch, channel). For decode (L=1), this is a simple diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index cdb907467b170..4f5418662b573 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -364,7 +364,9 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { desc.state_update_capacity = state_update_capacity_; desc.state_update_active = state_update_capacity_ > 0 && - (state_update_active == nullptr || state_update_active->Data()[0] != 0); + (state_update_active == nullptr || state_update_active->Data()[0] != 0) && + (cu_seqlens != nullptr || + total_tokens <= static_cast(state_update_capacity_ + 1) * batch); desc.update_rule = update_rule_; desc.gate_activation = gate_activation_; desc.beta_activation = beta_activation_; @@ -378,12 +380,15 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { desc.qk_l2_norm = qk_l2_norm_; desc.decay_per_key_dim = decay_per_key_dim; desc.ragged = cu_seqlens != nullptr; + const bool packed_uniform_qkv = is_packed_qkv && cu_seqlens == nullptr; const gdn::Engine effective_engine = - (mid_state_output_ && total_tokens <= 512LL * batch && forced_engine_ == gdn::Engine::kAuto) + ((mid_state_output_ || packed_uniform_qkv) && total_tokens <= 512LL * batch && + forced_engine_ == gdn::Engine::kAuto) ? gdn::Engine::kRecurrent : forced_engine_; desc.preferred_engine = effective_engine; - desc.bf16_precision = bf16_precision_; + desc.bf16_precision = + (packed_uniform_qkv && !has_session_bf16_precision_) ? gdn::Bf16Precision::kBf16 : bf16_precision_; if (!has_session_bf16_precision_) { const std::string env_val = ParseEnvironmentVariableWithDefault(gdn::kGatedDeltaNetBf16PrecisionEnv, ""); diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu index 811a9079262e8..74f728a362334 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -551,7 +551,6 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( p.beta_activation = BetaActivation::kSigmoid; p.update_rule = UpdateRule::kGatedDelta; p.qk_l2_norm = true; - p.state_update_capacity = 0; p.decay_per_key_dim_flag = 0; p.state_update_tail_pass = p.short_row_tail_pass = false; d_v = 128; @@ -918,7 +917,7 @@ void LaunchDecodeWarpSized(const Descriptor& desc, const VariantPack& pack, desc.head_size_v == 128 && pack.cu_seqlens == nullptr && p.gate_activation == GateActivation::kSoftplusDecay && p.beta_activation == BetaActivation::kSigmoid && p.update_rule == UpdateRule::kGatedDelta && - p.qk_l2_norm && p.state_update_capacity == 0 && p.decay_per_key_dim_flag == 0 && + p.qk_l2_norm && p.decay_per_key_dim_flag == 0 && !p.state_update_tail_pass && !p.short_row_tail_pass) { GatedDeltaNetDecodeWarpKernel <<>>(pack, p, desc.head_size_v); diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh index 67c4e8be40e6e..7c3bf5b4b59d5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh @@ -122,7 +122,7 @@ __device__ __forceinline__ uint32_t PackHalf2(__nv_bfloat16 lo, __nv_bfloat16 hi // mma.sync.m16n8k16 is Ampere and later, so ptxas rejects it for a pre-Ampere target. The only // caller is the chunked engine, and SelectPlan requires sm_major >= 8 before choosing it, so the // pre-Ampere body below exists purely to let those targets compile and is never executed. -template +template __device__ __forceinline__ void MmaM16N8K16(float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]); diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 70da10a4b1e51..4b422c6255848 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -4704,6 +4704,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "State layout override (CUDA only): -1 follows channels_last, " "0 uses [B,C,state_length], 1 uses [B,state_length,C].", AttributeProto::INT, static_cast(-1)) + .Attr("state_update_capacity", + "Compact speculative state-update capture capacity (CUDA extended path only, 0..8); " + "when > 0, output 2 has shape (batch_size, state_update_capacity, channels) and records " + "the first min(sequence_length, state_update_capacity) projected input rows.", + AttributeProto::INT, static_cast(0)) .Attr("activation", "Fused activation function. One of: 'silu', 'swish', 'none'. " "Default is 'none'.", @@ -4808,18 +4813,23 @@ ONNX_MS_OPERATOR_SET_SCHEMA( const int64_t channel_offset = getAttribute(ctx, "input_channel_offset", 0); const int64_t history = getAttribute(ctx, "state_length", 0); const int64_t state_layout = getAttribute(ctx, "state_channels_last", -1); + const int64_t state_update_capacity = getAttribute(ctx, "state_update_capacity", 0); if (input_channels < 0 || input_channels > std::numeric_limits::max() || channel_offset < 0 || channel_offset > std::numeric_limits::max() || history < 0 || history > std::numeric_limits::max() || - state_layout < -1 || state_layout > 1) { + state_layout < -1 || state_layout > 1 || + state_update_capacity < 0 || state_update_capacity > 8) { fail_shape_inference("CausalConvWithState: invalid extended attributes"); } if (input_channels != 0 || channel_offset != 0 || history != 0 || state_layout != -1 || - ctx.getNumOutputs() > 2) { + state_update_capacity != 0 || ctx.getNumOutputs() > 2) { if (getAttribute(ctx, "channels_last", 0) != 1 || getAttribute(ctx, "ndim", 1) != 1 || state_window != 0 || getAttribute(ctx, "dilation", 1) < 1) { fail_shape_inference("CausalConvWithState: extended attributes require channels_last=1, ndim=1, state_window=0"); } + if (state_update_capacity > 0 && ctx.getNumOutputs() > 3) { + fail_shape_inference("CausalConvWithState: state_update_capacity > 0 supports at most 3 outputs"); + } for (size_t i = 2; i < ctx.getNumOutputs(); ++i) propagateElemTypeFromInputToOutput(ctx, 0, i); if (!hasInputShape(ctx, 0)) return; const auto& input = getInputShape(ctx, 0); @@ -4858,7 +4868,17 @@ ONNX_MS_OPERATOR_SET_SCHEMA( *state.add_dim() = state_layout == 0 ? *channels : length; *state.add_dim() = state_layout == 0 ? length : *channels; updateOutputShape(ctx, 1, state); - for (size_t i = 2; i < ctx.getNumOutputs(); ++i) updateOutputShape(ctx, i, state); + if (state_update_capacity > 0) { + if (ctx.getNumOutputs() > 2 && ctx.hasOutput(2)) { + TensorShapeProto update_shape; + *update_shape.add_dim() = input.dim(0); + update_shape.add_dim()->set_dim_value(state_update_capacity); + *update_shape.add_dim() = *channels; + updateOutputShape(ctx, 2, update_shape); + } + } else { + for (size_t i = 2; i < ctx.getNumOutputs(); ++i) updateOutputShape(ctx, i, state); + } return; } if (state_window < 0 || state_window > kMaxStateWindow) { diff --git a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc index b8d3c144bc7af..f49710474888d 100644 --- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc +++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc @@ -400,6 +400,32 @@ TEST(CausalConvWithStateTest, PackedChannelsExtendedHistoryAndFirstState) { TestPackedChannelHistory(2, false, true, true, true, 1, false); } +TEST(CausalConvWithStateTest, PackedChannelsStateUpdateCapacity) { + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) GTEST_SKIP() << "CUDA is required for extended convolution"; + OpTester test("CausalConvWithState", 1, kMSDomain); + test.AddAttribute("channels_last", 1); + test.AddAttribute("input_channels", 2); + test.AddAttribute("input_channel_offset", 1); + test.AddAttribute("state_length", 3); + test.AddAttribute("state_channels_last", 0); + test.AddAttribute("state_update_capacity", 2); + test.AddAttribute("activation", "none"); + test.AddInput("input", {1, 3, 4}, + {10.f, 1.f, 2.f, 99.f, + 20.f, 3.f, 4.f, 99.f, + 30.f, 7.f, 8.f, 99.f}); + test.AddInput("weight", {2, 1, 2}, {1.f, 1.f, 1.f, 1.f}); + test.AddOptionalInputEdge(); + test.AddInput("past", {1, 2, 3}, {0.f, 0.f, 5.f, 0.f, 0.f, 6.f}); + test.AddOutput("output", {1, 3, 2}, {6.f, 8.f, 4.f, 6.f, 10.f, 12.f}); + test.AddOutput("present", {1, 2, 3}, {1.f, 3.f, 7.f, 2.f, 4.f, 8.f}); + test.AddOutput("state_update", {1, 2, 2}, {1.f, 2.f, 3.f, 4.f}, false, 1e-4f, 1e-4f); + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + TEST(CausalConvWithStateTest, PackedChannelsInvalidConfigurations) { const std::vector errors{ "invalid extended attributes", "channel slice exceeds input width", "state_length is too short", diff --git a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc index b592d321a8228..1f5520ba09c38 100644 --- a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc @@ -730,7 +730,9 @@ TEST(GatedDeltaNetTest, RecurrentOutputsIncludeFirstRowState) { expected_final_state, false, 2e-3f, 2e-3f); test.AddOutput("mid_state", {g.batch, g.hv, g.dv, g.dk}, expected_mid_state, false, 2e-3f, 2e-3f); - test.Run(); + std::vector> providers; + providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } TEST(GatedDeltaNetTest, RecurrentMidStatePreservesEmptyRequests) { From fe6c41c3709cfa71842a133f172aae92b93c74ff Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 25 Oct 2026 05:01:21 +0000 Subject: [PATCH 08/17] Remove legacy explicit mid-state outputs from CausalConvWithState and GatedDeltaNet - Remove unused first_state/second_state/third_state outputs from CausalConvWithState and name output 2 state_update (used when state_update_capacity > 0). - Remove unused state_output attribute and mid_state/mid_state2/mid_state3 outputs from GatedDeltaNet, keeping compact state_update capture. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 20 +-- docs/OperatorKernels.md | 6 +- .../cuda/bert/causal_conv_with_state.cc | 23 +-- .../cuda/bert/causal_conv_with_state_impl.cu | 23 +-- .../cuda/bert/causal_conv_with_state_impl.h | 2 +- .../contrib_ops/cuda/bert/gated_delta_net.cc | 43 +----- .../contrib_ops/cuda/bert/gated_delta_net.h | 1 - .../cuda/bert/gated_delta_net_impl.cu | 40 ----- .../cuda/bert/gated_delta_net_impl.h | 3 - .../webgpu/bert/gated_delta_net.cc | 3 - .../core/graph/contrib_ops/bert_defs.cc | 138 +++++++----------- .../causal_conv_with_state_op_test.cc | 10 +- .../contrib_ops/gated_delta_net_op_test.cc | 91 ------------ 13 files changed, 82 insertions(+), 321 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index e6bbbf8667d82..ce522a2415c8c 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -1053,19 +1053,15 @@ This version of the operator has been available since version 1 of the 'com.micr
Carry state from previous step. For ndim=1: (batch_size, channels, state_length), or (W, batch_size, channels, state_length) when state_window = W > 0, in which case only slot W-1 is read, where state_length = (k_1 - 1) * dilation. When channels_last = 1 each slot is (batch_size, state_length, d_1, ..., d_n) instead. The state_length and state_channels_last attributes may override the length and layout for the CUDA extended path. If not provided, padding is zero.
-#### Outputs (2 - 5) +#### Outputs (2 - 3)
output : T
Convolution output with same shape as input, except the final dimension is input_channels when extended channel selection is enabled.
present_state : T
Updated carry state. For ndim=1: (batch_size, channels, state_length), or (W, batch_size, channels, state_length) when state_window = W > 0, and (batch_size, state_length, d_1, ..., d_n) per slot when channels_last = 1. Slot W-1 contains the last state_length values from the virtual input along the causal axis; slot j contains the same for the prefix ending at position (seq_len - W + j). Extended attributes override the state length and layout as for past_state.
-
first_state (optional) : T
-
Optional carry state after the first input position, with the same layout and length as present_state. CUDA only; requires rank-3 channels-last input and state_window=0.
-
second_state (optional) : T
-
Optional carry state after at most two input positions, with the same layout and length as present_state.
-
third_state (optional) : T
-
Optional carry state after at most three input positions, with the same layout and length as present_state.
+
state_update (optional) : T
+
Optional compact state-update capture with shape (batch_size, state_update_capacity, channels) when state_update_capacity > 0. CUDA extended path only; requires rank-3 channels-last input and state_window=0.
#### Type Constraints @@ -2538,8 +2534,6 @@ This version of the operator has been available since version 1 of the 'com.micr
When 1, L2-normalize each query and key head vector before the recurrence. Default 0.
scale : float
Output scaling factor. When 0.0 (default) uses 1/sqrt(head_size_qk).
-
state_output : string
-
Third output is either 'state_update' (default) or the recurrent state after the first token ('mid_state'). CUDA supports both modes; WebGPU supports 'state_update'.
state_update_capacity : int
Capacity C for compact contiguous-prefix transition capture, in [0, 8]. 0 (default) disables compact state-update outputs.
update_rule : string
@@ -2573,7 +2567,7 @@ This version of the operator has been available since version 1 of the 'com.micr
CPU int32 control with shape (1). Zero disables transition capture, ignores capture_count, and produces a zero-filled state_update. Omission is conservative.
-#### Outputs (1 - 5) +#### Outputs (1 - 3)
output : T
@@ -2581,11 +2575,7 @@ This version of the operator has been available since version 1 of the 'com.micr
final_state (optional) : TS
State after the last token of each request, shape (batch_size, num_heads_v, head_size_v, head_size_qk)
state_update (optional) : TS
-
Struct-of-arrays transition factors when state_output='state_update', or the state after the first token when state_output='mid_state'.
-
mid_state2 (optional) : TS
-
CUDA state after at most two request tokens; requires state_output='mid_state'.
-
mid_state3 (optional) : TS
-
CUDA state after at most three request tokens; requires state_output='mid_state'.
+
Struct-of-arrays compact transition factors, shape (batch_size, state_update_capacity * (num_heads_v + num_heads_k * head_size_qk + num_heads_v * head_size_v)).
#### Type Constraints diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 6ab401fbaa648..3b0511bd3eb91 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -585,7 +585,7 @@ The **OpSet Version** column uses the following notation: |BifurcationDetector|*in* src_tokens:**T**
*in* cur_tokens:**T**
*in* prev_suffix_match_idx:**T**
*in* pred_tokens:**T**
*out* tokens:**T**
*out* suffix_match_idx:**T**|1+|**T** = tensor(int64)| |BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| |CDist|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(double), tensor(float)| -|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* first_state:**T**
*out* second_state:**T**
*out* third_state:**T**|1+|**T** = tensor(float)| +|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* state_update:**T**|1+|**T** = tensor(float)| |ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |CropAndResize|*in* X:**T1**
*in* rois:**T1**
*in* batch_indices:**T2**
*in* crop_size:**T2**
*out* Y:**T1**|1+|**T1** = tensor(float)
**T2** = tensor(int32)| |DecoderMaskedMultiHeadAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* mask_index:**M**
*in* attention_bias:**T**
*in* past_key:**T**
*in* past_value:**T**
*in* past_sequence_length:**M**
*in* beam_width:**M**
*in* cache_indirection:**M**
*in* bias:**T**
*out* output:**T**
*out* present_key:**T**
*out* present_value:**T**
*out* qk:**QK**|1+|**T** = tensor(float)| @@ -1103,7 +1103,7 @@ The **OpSet Version** column uses the following notation: |BitmaskBiasDropout|*in* data:**T**
*in* bias:**T**
*in* residual:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| |BitmaskDropout|*in* data:**T**
*in* ratio:**T1**
*in* training_mode:**T2**
*out* output:**T**
*out* mask:**T3**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T1** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)
**T2** = tensor(bool)
**T3** = tensor(uint32)| |BranchwiseRMSNorm|*in* X:**T**
*in* scale:**M**
*out* Y:**T**|1+|**M** = tensor(bfloat16), tensor(float), tensor(float16)
**T** = tensor(bfloat16), tensor(float), tensor(float16)| -|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* first_state:**T**
*out* second_state:**T**
*out* third_state:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| +|CausalConvWithState|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* past_state:**T**
*out* output:**T**
*out* present_state:**T**
*out* state_update:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| |ComplexMul|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ComplexMulConj|*in* A:**T**
*in* B:**T**
*out* C:**T**|1+|**T** = tensor(float), tensor(float16)| |ConvTransposeWithDynamicPads|*in* X:**T**
*in* W:**T**
*in* Pads:**tensor(int64)**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)| @@ -1120,7 +1120,7 @@ The **OpSet Version** column uses the following notation: |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)| |GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)| -|GatedDeltaNet|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* cu_seqlens:**TI**
*in* decay:**TG**
*in* beta:**TG**
*in* initial_state:**TS**
*in* a_log:**TS**
*in* dt_bias:**TS**
*in* capture_count:**TI**
*in* state_update_active:**TI**
*out* output:**T**
*out* final_state:**TS**
*out* state_update:**TS**
*out* mid_state2:**TS**
*out* mid_state3:**TS**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TG** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)
**TS** = tensor(float)| +|GatedDeltaNet|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* cu_seqlens:**TI**
*in* decay:**TG**
*in* beta:**TG**
*in* initial_state:**TS**
*in* a_log:**TS**
*in* dt_bias:**TS**
*in* capture_count:**TI**
*in* state_update_active:**TI**
*out* output:**T**
*out* final_state:**TS**
*out* state_update:**TS**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TG** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)
**TS** = tensor(float)| |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(float4e2m1), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int4), tensor(uint4), tensor(uint8)
**T2** = tensor(bfloat16), tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)| 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 da07a543b1c0d..f03fb42cc4819 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc @@ -76,8 +76,8 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { 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"); - ORT_RETURN_IF_NOT(state_update_capacity_ == 0 || context->OutputCount() <= 3, - "extended CausalConvWithState with state_update_capacity > 0 supports at most 3 outputs"); + ORT_RETURN_IF_NOT(context->OutputCount() <= 2 || state_update_capacity_ > 0, + "extended CausalConvWithState state_update output requires state_update_capacity > 0"); const int batch = SafeInt(input_shape[0]); const int length = SafeInt(input_shape[1]); const int width = SafeInt(input_shape[2]); @@ -99,19 +99,9 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { "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* state_update = nullptr; - Tensor* first = nullptr; - Tensor* second = nullptr; - Tensor* third = nullptr; - if (state_update_capacity_ > 0) { - state_update = context->OutputCount() > 2 - ? context->Output(2, TensorShape({batch, state_update_capacity_, channels})) - : nullptr; - } else { - first = context->OutputCount() > 2 ? context->Output(2, state_shape) : nullptr; - second = context->OutputCount() > 3 ? context->Output(3, state_shape) : nullptr; - third = context->OutputCount() > 4 ? context->Output(4, state_shape) : nullptr; - } + Tensor* state_update = state_update_capacity_ > 0 && context->OutputCount() > 2 + ? context->Output(2, TensorShape({batch, state_update_capacity_, channels})) + : nullptr; using CudaT = typename OrtToCudaType::type; return LaunchCausalConvProjectionKernel( Stream(context), reinterpret_cast(input_tensor->Data()), @@ -120,9 +110,6 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { past_state_tensor ? reinterpret_cast(past_state_tensor->Data()) : nullptr, reinterpret_cast(output->MutableData()), reinterpret_cast(present->MutableData()), - first ? reinterpret_cast(first->MutableData()) : nullptr, - second ? reinterpret_cast(second->MutableData()) : nullptr, - third ? reinterpret_cast(third->MutableData()) : nullptr, batch, channels, length, kernel, dilation_, width, input_channel_offset_, state_length, state_last, activation_ != "none", state_update ? reinterpret_cast(state_update->MutableData()) : nullptr, 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 719192aa4da26..c697d2fbdac0f 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 @@ -31,7 +31,7 @@ namespace { template __global__ void CausalConvProjectionKernel( const T* input, const T* weight, const T* bias, const T* past, - T* output, T* present, T* first_state, T* second_state, T* third_state, int channels, int length, + T* output, T* present, int channels, int length, int kernel_size, int dilation, int input_width, int channel_offset, int state_length, CausalConvLayout state_layout, bool apply_silu, T* state_update = nullptr, int state_update_capacity = 0) { @@ -68,15 +68,6 @@ __global__ void CausalConvProjectionKernel( present[state_layout.Offset(b, position, c)] = source >= 0 ? input[(static_cast(b) * length + source) * input_width + channel_offset + c] : (past == nullptr ? from_float(0.0f) : past[state_layout.Offset(b, state_length + source, c)]); - T* snapshots[3] = {first_state, second_state, third_state}; -#pragma unroll - for (int prefix = 0; prefix < 3; ++prefix) { - if (snapshots[prefix] == nullptr) continue; - const int prefix_source = min(length, prefix + 1) - state_length + position; - snapshots[prefix][state_layout.Offset(b, position, c)] = - prefix_source >= 0 ? input[(static_cast(b) * length + prefix_source) * input_width + channel_offset + c] - : (past == nullptr ? from_float(0.0f) : past[state_layout.Offset(b, state_length + prefix_source, c)]); - } } } @@ -484,7 +475,7 @@ __global__ void CausalConvPrefillKernelChannelsLast( template Status LaunchCausalConvProjectionKernel( cudaStream_t stream, const T* input, const T* weight, const T* bias, const T* past, - T* output, T* present, T* first_state, T* second_state, T* third_state, int batch, int channels, int length, + T* output, T* present, 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, T* state_update, int state_update_capacity) { @@ -493,23 +484,23 @@ Status LaunchCausalConvProjectionKernel( ORT_ENFORCE(specialize == 0 || specialize == 1, "ORT_CAUSAL_CONV_S2_SPECIALIZE must be 0 or 1."); if (specialize && length >= 1 && length <= 8 && kernel_size == 4 && dilation == 1 && state_length == 4) { CausalConvProjectionKernel<<>>( - input, weight, bias, past, output, present, first_state, second_state, third_state, channels, length, + input, weight, bias, past, output, present, channels, length, kernel_size, dilation, input_width, channel_offset, state_length, MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu, state_update, state_update_capacity); return CUDA_CALL(cudaGetLastError()); } CausalConvProjectionKernel<<>>( - input, weight, bias, past, output, present, first_state, second_state, third_state, channels, length, + input, weight, bias, past, output, present, channels, length, kernel_size, dilation, input_width, channel_offset, state_length, MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu, state_update, state_update_capacity); return CUDA_CALL(cudaGetLastError()); } -#define INSTANTIATE_PROJECTION(T) \ - template Status LaunchCausalConvProjectionKernel( \ - cudaStream_t, const T*, const T*, const T*, const T*, T*, T*, T*, T*, T*, \ +#define INSTANTIATE_PROJECTION(T) \ + template Status LaunchCausalConvProjectionKernel( \ + cudaStream_t, const T*, const T*, const T*, const T*, T*, T*, \ int, int, int, int, int, int, int, int, bool, bool, T*, int); INSTANTIATE_PROJECTION(float) INSTANTIATE_PROJECTION(half) 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 436e00278efda..8339eefb7efbe 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 @@ -34,7 +34,7 @@ inline CausalConvLayout MakeCausalConvLayout(bool channels_last, int64_t channel template Status LaunchCausalConvProjectionKernel( cudaStream_t stream, const T* input, const T* weight, const T* bias, const T* past, - T* output, T* present, T* first_state, T* second_state, T* third_state, int batch, int channels, int length, + T* output, T* present, 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, T* state_update = nullptr, int state_update_capacity = 0); diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index 4f5418662b573..cb04fe8037b71 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -45,14 +45,7 @@ REGISTER_KERNEL_TYPED(MLFloat16) REGISTER_KERNEL_TYPED(BFloat16) template -GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) - : CudaKernel(info), - mid_state_output_(info.GetAttrOrDefault("state_output", "state_update") == - "mid_state") { - ORT_ENFORCE(mid_state_output_ || - info.GetAttrOrDefault("state_output", "state_update") == - "state_update", - "state_output must be 'state_update' or 'mid_state'"); +GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : CudaKernel(info) { const std::string rule = info.GetAttrOrDefault("update_rule", "gated_delta"); if (rule == "linear") { update_rule_ = gdn::UpdateRule::kLinear; @@ -98,9 +91,6 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) forced_engine_ = gdn::EngineFromName( ParseEnvironmentVariableWithDefault("ORT_GDN_PLAN", "")); - if (mid_state_output_) { - bf16_precision_ = gdn::Bf16Precision::kBf16; - } const auto from_config = info.GetConfigOptions().GetConfigEntry(gdn::kGatedDeltaNetBf16PrecisionConfig); @@ -142,8 +132,6 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { gate->GetElementType() == query->GetElementType(), "decay and beta must be float or match the query element type"); } - ORT_RETURN_IF_NOT(!mid_state_output_ || state_update_capacity_ == 0, - "state_output='mid_state' cannot be combined with state_update_capacity"); ORT_RETURN_IF_NOT((state_update_capacity_ > 0) == (capture_count != nullptr), "capture_count must be present exactly when state_update_capacity is positive"); @@ -229,10 +217,6 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(num_heads_q == num_heads_k, "num_heads_q (", num_heads_q, ") must equal num_heads_k (", num_heads_k, ")"); - if (mid_state_output_) { - ORT_RETURN_IF_NOT(num_heads_v >= num_heads_q, - "mid-state output requires value_heads >= query_heads"); - } ORT_RETURN_IF_NOT(num_heads_v > 0 && num_heads_q > 0 && num_heads_v % num_heads_q == 0, "num_heads_v (", num_heads_v, ") must be a positive multiple of num_heads_q (", num_heads_q, ")"); @@ -335,22 +319,10 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { TensorShape state_shape{batch, num_heads_v, head_size_v, head_size_qk}; Tensor* final_state = context->Output(1, state_shape); - Tensor* state_update = nullptr; - Tensor* mid_state = nullptr; - Tensor* mid_state2 = nullptr; - Tensor* mid_state3 = nullptr; - if (mid_state_output_) { - mid_state = context->Output(2, state_shape); - if (context->OutputCount() > 3) mid_state2 = context->Output(3, state_shape); - if (context->OutputCount() > 4) mid_state3 = context->Output(4, state_shape); - } else { - ORT_RETURN_IF_NOT(context->OutputCount() <= 3, - "additional prefix states require state_output='mid_state'"); - const int64_t state_update_width = static_cast(state_update_capacity_) * - (num_heads_v + num_heads_k * head_size_qk + - num_heads_v * head_size_v); - state_update = context->Output(2, TensorShape{batch, state_update_width}); - } + const int64_t state_update_width = static_cast(state_update_capacity_) * + (num_heads_v + num_heads_k * head_size_qk + + num_heads_v * head_size_v); + Tensor* state_update = context->Output(2, TensorShape{batch, state_update_width}); gdn::Descriptor desc{}; desc.total_tokens = total_tokens; @@ -382,7 +354,7 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { desc.ragged = cu_seqlens != nullptr; const bool packed_uniform_qkv = is_packed_qkv && cu_seqlens == nullptr; const gdn::Engine effective_engine = - ((mid_state_output_ || packed_uniform_qkv) && total_tokens <= 512LL * batch && + (packed_uniform_qkv && total_tokens <= 512LL * batch && forced_engine_ == gdn::Engine::kAuto) ? gdn::Engine::kRecurrent : forced_engine_; @@ -493,9 +465,6 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { pack.dt_bias = dt_bias != nullptr ? dt_bias->Data() : nullptr; pack.output = reinterpret_cast(output->MutableData()); pack.final_state = final_state != nullptr ? final_state->MutableData() : nullptr; - pack.mid_state = mid_state != nullptr ? mid_state->MutableData() : nullptr; - pack.mid_state2 = mid_state2 != nullptr ? mid_state2->MutableData() : nullptr; - pack.mid_state3 = mid_state3 != nullptr ? mid_state3->MutableData() : nullptr; pack.state_update = state_update != nullptr ? state_update->MutableData() : nullptr; // Only the inactive case is cleared. While capture is active the kernels write just the diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h index 8cdbfb67d8287..7d321e932a9c6 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h @@ -30,7 +30,6 @@ class GatedDeltaNet : public onnxruntime::cuda::CudaKernel { int chunk_size_; int state_update_capacity_; bool qk_l2_norm_; - bool mid_state_output_; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu index 74f728a362334..1c6830ae49873 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -688,16 +688,6 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( pack.output[tok * v_stride + hv * d_v + col] = static_cast(p.scale * o); } - float* snapshots[3] = {pack.mid_state, pack.mid_state2, pack.mid_state3}; -#pragma unroll - for (int prefix = 0; prefix < 3; ++prefix) { - float* prefix_state = snapshots[prefix]; - if (prefix_state == nullptr || t != min(prefix + 1, seq_len) - 1) continue; -#pragma unroll - for (int i = 0; i < kRegs; ++i) { - prefix_state[st_off + lane + i * 32] = s[i]; - } - } } if (pack.final_state != nullptr) { @@ -706,15 +696,6 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( pack.final_state[st_off + lane + i * 32] = s[i]; } } - if (seq_len == 0) { - float* snapshots[3] = {pack.mid_state, pack.mid_state2, pack.mid_state3}; -#pragma unroll - for (int prefix = 0; prefix < 3; ++prefix) { - if (snapshots[prefix] == nullptr) continue; -#pragma unroll - for (int i = 0; i < kRegs; ++i) snapshots[prefix][st_off + lane + i * 32] = s[i]; - } - } } // --------------------------------------------------------------------------- @@ -860,16 +841,6 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant for (int r = 0; r < d_k; ++r) acc += S[r * d_v + c] * qbuf[r]; pack.output[tok * v_stride + hv * d_v + c] = (T)(p.scale * acc); } - float* snapshots[3] = {pack.mid_state, pack.mid_state2, pack.mid_state3}; -#pragma unroll - for (int prefix = 0; prefix < 3; ++prefix) { - float* prefix_state = snapshots[prefix]; - if (prefix_state == nullptr || t != min(prefix + 1, seq_len) - 1) continue; - for (int idx = tid; idx < d_k * d_v; idx += kThreads) { - const int r = idx / d_v, c = idx % d_v; - prefix_state[st_off + static_cast(c) * d_k + r] = S[idx]; - } - } __syncthreads(); } @@ -880,17 +851,6 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant pack.final_state[st_off + static_cast(c) * d_k + r] = S[idx]; } } - if (seq_len == 0) { - float* snapshots[3] = {pack.mid_state, pack.mid_state2, pack.mid_state3}; -#pragma unroll - for (int prefix = 0; prefix < 3; ++prefix) { - if (snapshots[prefix] == nullptr) continue; - for (int idx = tid; idx < d_k * d_v; idx += kThreads) { - const int r = idx / d_v, c = idx % d_v; - snapshots[prefix][st_off + static_cast(c) * d_k + r] = S[idx]; - } - } - } } constexpr int kRecurrentThreads = 256; diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h index eb82ba072f1fa..8f2a0eb3d7acc 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h @@ -33,9 +33,6 @@ struct VariantPack { const float* dt_bias = nullptr; // [Hv] T* output = nullptr; // [total_tokens, max(Hq,Hv), V] float* final_state = nullptr; // [B, Hv, V, K] - float* mid_state = nullptr; // [B, Hv, V, K], state after local row 0 - float* mid_state2 = nullptr; // state after local row 1 - float* mid_state3 = nullptr; // state after local row 2 float* state_update = nullptr; // [B, C * (Hv + Hk*K + Hv*V)] void* workspace = nullptr; }; diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc index d8030e7412ecd..4f83cd61c65ea 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc @@ -62,9 +62,6 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : WebGpuKernel(info) { const auto state_update_capacity = info.GetAttrOrDefault("state_update_capacity", 0); ORT_ENFORCE(state_update_capacity >= 0 && state_update_capacity <= 8, "state_update_capacity must be in [0, 8], got ", state_update_capacity); - ORT_ENFORCE(info.GetAttrOrDefault("state_output", "state_update") == - "state_update", - "WebGPU GatedDeltaNet supports only state_output='state_update'"); state_update_capacity_ = static_cast(state_update_capacity); } diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 4b422c6255848..1df9680608701 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -4790,17 +4790,11 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "(seq_len - W + j). Extended attributes override the state length and layout " "as for past_state.", "T") - .Output(2, "first_state", - "Optional carry state after the first input position, with the same " - "layout and length as present_state. CUDA only; requires rank-3 channels-last input " + .Output(2, "state_update", + "Optional compact state-update capture with shape (batch_size, state_update_capacity, channels) " + "when state_update_capacity > 0. CUDA extended path only; requires rank-3 channels-last input " "and state_window=0.", "T", OpSchema::Optional) - .Output(3, "second_state", - "Optional carry state after at most two input positions, with the same layout and length as present_state.", - "T", OpSchema::Optional) - .Output(4, "third_state", - "Optional carry state after at most three input positions, with the same layout and length as present_state.", - "T", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") @@ -4827,10 +4821,10 @@ ONNX_MS_OPERATOR_SET_SCHEMA( state_window != 0 || getAttribute(ctx, "dilation", 1) < 1) { fail_shape_inference("CausalConvWithState: extended attributes require channels_last=1, ndim=1, state_window=0"); } - if (state_update_capacity > 0 && ctx.getNumOutputs() > 3) { - fail_shape_inference("CausalConvWithState: state_update_capacity > 0 supports at most 3 outputs"); + if (ctx.getNumOutputs() > 2 && state_update_capacity <= 0) { + fail_shape_inference("CausalConvWithState: state_update output requires state_update_capacity > 0"); } - for (size_t i = 2; i < ctx.getNumOutputs(); ++i) propagateElemTypeFromInputToOutput(ctx, 0, i); + if (ctx.getNumOutputs() > 2) propagateElemTypeFromInputToOutput(ctx, 0, 2); if (!hasInputShape(ctx, 0)) return; const auto& input = getInputShape(ctx, 0); if (input.dim_size() != 3) fail_shape_inference("CausalConvWithState: extended input must have rank 3"); @@ -4868,16 +4862,12 @@ ONNX_MS_OPERATOR_SET_SCHEMA( *state.add_dim() = state_layout == 0 ? *channels : length; *state.add_dim() = state_layout == 0 ? length : *channels; updateOutputShape(ctx, 1, state); - if (state_update_capacity > 0) { - if (ctx.getNumOutputs() > 2 && ctx.hasOutput(2)) { - TensorShapeProto update_shape; - *update_shape.add_dim() = input.dim(0); - update_shape.add_dim()->set_dim_value(state_update_capacity); - *update_shape.add_dim() = *channels; - updateOutputShape(ctx, 2, update_shape); - } - } else { - for (size_t i = 2; i < ctx.getNumOutputs(); ++i) updateOutputShape(ctx, i, state); + if (state_update_capacity > 0 && ctx.getNumOutputs() > 2 && ctx.hasOutput(2)) { + TensorShapeProto update_shape; + *update_shape.add_dim() = input.dim(0); + update_shape.add_dim()->set_dim_value(state_update_capacity); + *update_shape.add_dim() = *channels; + updateOutputShape(ctx, 2, update_shape); } return; } @@ -5601,11 +5591,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "Capacity C for compact contiguous-prefix transition capture, in [0, 8]. " "0 (default) disables compact state-update outputs.", AttributeProto::INT, static_cast(0)) - .Attr("state_output", - "Third output is either 'state_update' (default) or the recurrent state after " - "the first token ('mid_state'). CUDA supports both modes; WebGPU supports " - "'state_update'.", - AttributeProto::STRING, std::string("state_update")) .Input(0, "query", "Query or packed QKV, shaped as described above.", "T") .Input(1, "key", "Key, shape (total_tokens, num_heads_k, head_size_qk)", "T", OpSchema::Optional) @@ -5649,12 +5634,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "(batch_size, num_heads_v, head_size_v, head_size_qk)", "TS", OpSchema::Optional) .Output(2, "state_update", - "Struct-of-arrays transition factors when state_output='state_update', or " - "the state after the first token when state_output='mid_state'.", - "TS", OpSchema::Optional) - .Output(3, "mid_state2", "CUDA state after at most two request tokens; requires state_output='mid_state'.", - "TS", OpSchema::Optional) - .Output(4, "mid_state3", "CUDA state after at most three request tokens; requires state_output='mid_state'.", + "Struct-of-arrays compact transition factors, shape " + "(batch_size, state_update_capacity * (num_heads_v + num_heads_k * " + "head_size_qk + num_heads_v * head_size_v)).", "TS", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain query/key/value/output types.") @@ -5668,8 +5650,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( if (ctx.getNumOutputs() > 1) { updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); } - for (size_t i = 2; i < ctx.getNumOutputs(); ++i) { - updateOutputElemType(ctx, i, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + if (ctx.getNumOutputs() > 2) { + updateOutputElemType(ctx, 2, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); } const bool has_key = ctx.getNumInputs() > 1 && ctx.getInputType(1) != nullptr; @@ -5794,61 +5776,45 @@ ONNX_MS_OPERATOR_SET_SCHEMA( kMaxStateWindow, "], got ", state_update_capacity); } if (ctx.getNumOutputs() > 2) { - const std::string state_output = - getAttribute(ctx, "state_output", std::string("state_update")); - if (state_output != "state_update" && state_output != "mid_state") { - fail_shape_inference("GatedDeltaNet: state_output must be 'state_update' or 'mid_state'"); - } - if (state_output == "mid_state") { - if (state_update_capacity != 0) { - fail_shape_inference("GatedDeltaNet: mid_state output requires state_update_capacity=0"); - } - if (state_shape != nullptr) { - for (size_t i = 2; i < ctx.getNumOutputs(); ++i) updateOutputShape(ctx, i, *state_shape); + ONNX_NAMESPACE::TensorShapeProto capsule_shape; + add_batch_dim(capsule_shape); + auto* width = capsule_shape.add_dim(); + if (!packed_qkv && query_shape.dim(token_dims).has_dim_value() && + query_shape.dim(token_dims + 1).has_dim_value() && + value_shape->dim(token_dims).has_dim_value() && + value_shape->dim(token_dims + 1).has_dim_value()) { + // num_heads_k is constrained to equal num_heads_q, so query supplies it. + const int64_t num_heads_k = query_shape.dim(token_dims).dim_value(); + const int64_t head_size_qk = query_shape.dim(token_dims + 1).dim_value(); + const int64_t num_heads_v = value_shape->dim(token_dims).dim_value(); + const int64_t head_size_v = value_shape->dim(token_dims + 1).dim_value(); + if (num_heads_k <= 0 || head_size_qk <= 0 || num_heads_v <= 0 || head_size_v <= 0) { + fail_shape_inference( + "GatedDeltaNet: head counts and head sizes must be positive"); } - } else { - if (ctx.getNumOutputs() > 3) - fail_shape_inference("GatedDeltaNet: additional prefix states require state_output='mid_state'"); - ONNX_NAMESPACE::TensorShapeProto capsule_shape; - add_batch_dim(capsule_shape); - auto* width = capsule_shape.add_dim(); - if (!packed_qkv && query_shape.dim(token_dims).has_dim_value() && - query_shape.dim(token_dims + 1).has_dim_value() && - value_shape->dim(token_dims).has_dim_value() && - value_shape->dim(token_dims + 1).has_dim_value()) { - // num_heads_k is constrained to equal num_heads_q, so query supplies it. - const int64_t num_heads_k = query_shape.dim(token_dims).dim_value(); - const int64_t head_size_qk = query_shape.dim(token_dims + 1).dim_value(); - const int64_t num_heads_v = value_shape->dim(token_dims).dim_value(); - const int64_t head_size_v = value_shape->dim(token_dims + 1).dim_value(); - if (num_heads_k <= 0 || head_size_qk <= 0 || num_heads_v <= 0 || head_size_v <= 0) { - fail_shape_inference( - "GatedDeltaNet: head counts and head sizes must be positive"); - } - constexpr int64_t max_dimension = std::numeric_limits::max(); - if (num_heads_k > max_dimension / head_size_qk || - num_heads_v > max_dimension / head_size_v) { - fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); - } - const int64_t key_width = num_heads_k * head_size_qk; - const int64_t value_width = num_heads_v * head_size_v; - if (num_heads_v > max_dimension - key_width) { - fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); - } - const int64_t width_without_capacity = num_heads_v + key_width; - if (value_width > max_dimension - width_without_capacity) { - fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); - } - const int64_t per_token_width = width_without_capacity + value_width; - if (state_update_capacity > 0 && - per_token_width > max_dimension / state_update_capacity) { - fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); - } - width->set_dim_value(state_update_capacity * per_token_width); + constexpr int64_t max_dimension = std::numeric_limits::max(); + if (num_heads_k > max_dimension / head_size_qk || + num_heads_v > max_dimension / head_size_v) { + fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); + } + const int64_t key_width = num_heads_k * head_size_qk; + const int64_t value_width = num_heads_v * head_size_v; + if (num_heads_v > max_dimension - key_width) { + fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); + } + const int64_t width_without_capacity = num_heads_v + key_width; + if (value_width > max_dimension - width_without_capacity) { + fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); + } + const int64_t per_token_width = width_without_capacity + value_width; + if (state_update_capacity > 0 && + per_token_width > max_dimension / state_update_capacity) { + fail_shape_inference("GatedDeltaNet: state_update width overflows int64"); } - updateOutputShape(ctx, 2, capsule_shape); + width->set_dim_value(state_update_capacity * per_token_width); } + updateOutputShape(ctx, 2, capsule_shape); } if (state_shape != nullptr) { diff --git a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc index f49710474888d..c1e85816cafa9 100644 --- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc +++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc @@ -320,8 +320,7 @@ static void RunCausalConvWithStateTests( // ============================================================================= template -void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool use_bias, bool silu, int dilation, - bool first_output = true) { +void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool use_bias, bool silu, int dilation) { constexpr int batch = 2, channels = 7, width = 13, offset = 3, kernel = 4; const int history = (kernel - 1) * dilation + 1; auto cuda = DefaultCudaExecutionProvider(); @@ -329,7 +328,7 @@ void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool u const auto round = [](float value) { return static_cast(T(value)); }; std::vector input(batch * length * width), weight(channels * kernel), bias(channels); std::vector past(batch * channels * history), output(batch * length * channels); - std::vector present(past.size()), first(past.size()); + std::vector present(past.size()); for (size_t i = 0; i < input.size(); ++i) input[i] = T(float(int(i % 17) - 8) / 16); for (size_t i = 0; i < weight.size(); ++i) weight[i] = T(float(int(i % 11) - 5) / 32); for (size_t i = 0; i < bias.size(); ++i) bias[i] = T(float(i) / 64); @@ -351,7 +350,6 @@ void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool u } for (int t = 0; t < history; ++t) { present[state_index(b, t, c)] = T(value(length - history + t)); - first[state_index(b, t, c)] = T(value(1 - history + t)); } } } @@ -377,14 +375,13 @@ void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool u test.AddOptionalInputEdge(); test.AddOutput("output", {batch, length, channels}, output); test.AddOutput("present", state_shape, present); - if (first_output) test.AddOutput("first", state_shape, first); test.SetOutputTolerance(std::is_same_v ? 0.01f : 0.001f, 0.001f); std::vector> providers; providers.push_back(std::move(cuda)); test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); } -TEST(CausalConvWithStateTest, PackedChannelsExtendedHistoryAndFirstState) { +TEST(CausalConvWithStateTest, PackedChannelsExtendedHistory) { for (int length : {1, 2, 8, 9}) { for (bool state_last : {false, true}) { for (bool past : {false, true}) { @@ -397,7 +394,6 @@ TEST(CausalConvWithStateTest, PackedChannelsExtendedHistoryAndFirstState) { } } } - TestPackedChannelHistory(2, false, true, true, true, 1, false); } TEST(CausalConvWithStateTest, PackedChannelsStateUpdateCapacity) { diff --git a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc index 1f5520ba09c38..72cf073a82abb 100644 --- a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc @@ -685,97 +685,6 @@ bool NeedSkipGatedDeltaNetTest() { return !HasCudaEnvironment(800); } -TEST(GatedDeltaNetTest, RecurrentOutputsIncludeFirstRowState) { - if (NeedSkipGatedDeltaNetTest()) return; - - const Geometry g{2, 1, 1, 2, 128, 128}; - Inputs in = MakeInputs(g, 211); - Options options; - options.gate_activation = "softplus"; - options.beta_activation = "sigmoid"; - options.qk_l2_norm = 1; - - in.q = RoundToTensorType(in.q); - in.k = RoundToTensorType(in.k); - in.v = RoundToTensorType(in.v); - std::vector expected_output, expected_final_state; - Reference(g, options, in, &expected_output, &expected_final_state); - - Geometry first_row_geometry = g; - first_row_geometry.total_tokens = 1; - Inputs first_row = in; - first_row.q.resize(static_cast(g.hq) * g.dk); - first_row.k.resize(static_cast(g.hq) * g.dk); - first_row.v.resize(static_cast(g.hv) * g.dv); - first_row.decay.resize(g.hv); - first_row.beta.resize(g.hv); - std::vector unused_output, expected_mid_state; - Reference(first_row_geometry, options, first_row, &unused_output, &expected_mid_state); - - OpTester test("GatedDeltaNet", 1, onnxruntime::kMSDomain); - AddCommonAttrs(test, options); - test.AddAttribute("state_output", "mid_state"); - test.AddInput("query", {1, 2, g.hq, g.dk}, ToTensorType(in.q)); - test.AddInput("key", {1, 2, g.hq, g.dk}, ToTensorType(in.k)); - test.AddInput("value", {1, 2, g.hv, g.dv}, ToTensorType(in.v)); - test.AddOptionalInputEdge(); - test.AddInput("decay", {1, 2, g.hv}, in.decay); - test.AddInput("beta", {1, 2, g.hv}, in.beta); - test.AddInput("initial_state", {g.batch, g.hv, g.dv, g.dk}, in.state0); - test.AddInput("a_log", {g.hv}, in.a_log); - test.AddInput("dt_bias", {g.hv}, in.dt_bias); - test.AddOutput("output", {1, 2, g.hv, g.dv}, - ToTensorType(expected_output), false, 6e-3f, 6e-3f); - test.AddOutput("final_state", {g.batch, g.hv, g.dv, g.dk}, - expected_final_state, false, 2e-3f, 2e-3f); - test.AddOutput("mid_state", {g.batch, g.hv, g.dv, g.dk}, - expected_mid_state, false, 2e-3f, 2e-3f); - std::vector> providers; - providers.push_back(DefaultCudaExecutionProvider()); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); -} - -TEST(GatedDeltaNetTest, RecurrentMidStatePreservesEmptyRequests) { - if (NeedSkipGatedDeltaNetTest()) return; - for (int head_size : {128, 32}) { - for (bool with_state : {false, true}) { - SCOPED_TRACE(head_size); - SCOPED_TRACE(with_state); - const Geometry g{1, 2, 1, 1, head_size, head_size}; - Inputs in = MakeInputs(g, 917, with_state); - in.cu_seqlens = {0, 0, 1}; - in.q = RoundToTensorType(in.q); - in.k = RoundToTensorType(in.k); - in.v = RoundToTensorType(in.v); - Options options; - std::vector expected_output, expected_state; - Reference(g, options, in, &expected_output, &expected_state); - - OpTester test("GatedDeltaNet", 1, kMSDomain); - AddCommonAttrs(test, options); - test.AddAttribute("state_output", "mid_state"); - test.AddInput("query", {1, 1, head_size}, ToTensorType(in.q)); - test.AddInput("key", {1, 1, head_size}, ToTensorType(in.k)); - test.AddInput("value", {1, 1, head_size}, ToTensorType(in.v)); - test.AddInput("cu_seqlens", {3}, in.cu_seqlens); - test.AddInput("decay", {1, 1}, in.decay); - test.AddInput("beta", {1, 1}, in.beta); - if (with_state) { - test.AddInput("initial_state", {2, 1, head_size, head_size}, in.state0); - } else { - test.AddOptionalInputEdge(); - } - test.AddOutput("output", {1, 1, head_size}, - ToTensorType(expected_output), false, 6e-3f, 6e-3f); - test.AddOutput("final_state", {2, 1, head_size, head_size}, expected_state, false, 2e-3f, 2e-3f); - test.AddOutput("mid_state", {2, 1, head_size, head_size}, expected_state, false, 2e-3f, 2e-3f); - std::vector> providers; - providers.push_back(DefaultCudaExecutionProvider()); - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); - } - } -} - bool NeedSkipGatedDeltaNetSplitTest() { if (NeedSkipGatedDeltaNetTest()) return true; From 08f2f5a8393a1d55a2ccbf947007debe4e29d5c6 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 25 Oct 2026 23:28:38 +0000 Subject: [PATCH 09/17] fix(cuda): restore GatedDeltaNet packed-QKV prefill plan parity and sync docs - Use standard SelectPlan engine and BF16 precision selection for unpacked packed-QKV prefill so packed and unpacked QKV take the same chunked path at tokens >= 32 (fixing GatedDeltaNetTest.DirectPackedQkv{Float16,BFloat16}ExactParity). - Sync docs/ContribOperators.md with gen_contrib_doc.py output for CausalConvWithState state_update. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 2 +- .../contrib_ops/cuda/bert/gated_delta_net.cc | 23 +++++++------------ 2 files changed, 9 insertions(+), 16 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index ce522a2415c8c..877cb70deb478 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -1061,7 +1061,7 @@ This version of the operator has been available since version 1 of the 'com.micr
present_state : T
Updated carry state. For ndim=1: (batch_size, channels, state_length), or (W, batch_size, channels, state_length) when state_window = W > 0, and (batch_size, state_length, d_1, ..., d_n) per slot when channels_last = 1. Slot W-1 contains the last state_length values from the virtual input along the causal axis; slot j contains the same for the prefix ending at position (seq_len - W + j). Extended attributes override the state length and layout as for past_state.
state_update (optional) : T
-
Optional compact state-update capture with shape (batch_size, state_update_capacity, channels) when state_update_capacity > 0. CUDA extended path only; requires rank-3 channels-last input and state_window=0.
+
Optional compact state-update capture with shape (batch_size, state_update_capacity, channels) when state_update_capacity > 0. CUDA extended path only; requires rank-3 channels-last input and state_window=0.
#### Type Constraints diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index cb04fe8037b71..19e4d8e6a0de2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -352,15 +352,8 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { desc.qk_l2_norm = qk_l2_norm_; desc.decay_per_key_dim = decay_per_key_dim; desc.ragged = cu_seqlens != nullptr; - const bool packed_uniform_qkv = is_packed_qkv && cu_seqlens == nullptr; - const gdn::Engine effective_engine = - (packed_uniform_qkv && total_tokens <= 512LL * batch && - forced_engine_ == gdn::Engine::kAuto) - ? gdn::Engine::kRecurrent - : forced_engine_; - desc.preferred_engine = effective_engine; - desc.bf16_precision = - (packed_uniform_qkv && !has_session_bf16_precision_) ? gdn::Bf16Precision::kBf16 : bf16_precision_; + desc.preferred_engine = forced_engine_; + desc.bf16_precision = bf16_precision_; if (!has_session_bf16_precision_) { const std::string env_val = ParseEnvironmentVariableWithDefault(gdn::kGatedDeltaNetBf16PrecisionEnv, ""); @@ -376,21 +369,21 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { desc.sm_major = prop.major; gdn::Plan plan = gdn::SelectPlan(desc, prop.sharedMemPerBlockOptin); - if (effective_engine == gdn::Engine::kChunked || - effective_engine == gdn::Engine::kChunkedSplit) { + if (forced_engine_ == gdn::Engine::kChunked || + forced_engine_ == gdn::Engine::kChunkedSplit) { ORT_RETURN_IF_NOT( - plan.engine == effective_engine, - "GatedDeltaNet: the ", gdn::EngineName(effective_engine), + plan.engine == forced_engine_, + "GatedDeltaNet: the ", gdn::EngineName(forced_engine_), " engine cannot serve this descriptor (", plan.reject_reason ? plan.reject_reason : "no reason recorded", ")"); - } else if (effective_engine == gdn::Engine::kRecurrent) { + } else if (forced_engine_ == gdn::Engine::kRecurrent) { plan.engine = gdn::Engine::kRecurrent; // Match the fallback path's kernel choice: the override must measure what the // heuristic would have run, not the generic kernel. plan.warp_specialized = gdn::RecurrentIsWarpSpecialized(head_size_qk); plan.state_update_tail_pass = false; plan.short_row_tail_pass = false; - } else if (effective_engine == gdn::Engine::kCudnn) { + } else if (forced_engine_ == gdn::Engine::kCudnn) { plan.engine = gdn::Engine::kCudnn; } ORT_RETURN_IF_NOT(plan.supported, "GatedDeltaNet: no supported plan (", From 0cdb489d004cb42e3dcb9d9c088358b7750fe791 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 05:34:02 +0000 Subject: [PATCH 10/17] perf(cuda): tile 8 positions per CTA in CausalConvProjection for K=4 prefill Add CausalConvProjectionTiledK4Kernel for length >= 16 when kernel_size == 4, dilation == 1, and state_length == 4, reusing weights, bias, and the 3-element input halo in registers across 8 consecutive token positions per CTA. Cuts 36-layer GDN CausalConvProjection prefill at L=2048 on H200 from 8.53 ms to 2.03 ms (4.10x faster). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/bert/causal_conv_with_state_impl.cu | 68 ++++++++++++++++++- 1 file changed, 67 insertions(+), 1 deletion(-) 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 c697d2fbdac0f..a38292ba164aa 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 @@ -28,6 +28,62 @@ namespace cuda { namespace { +template +__global__ void CausalConvProjectionTiledK4Kernel( + const T* __restrict__ input, const T* __restrict__ weight, const T* __restrict__ bias, + const T* __restrict__ past, T* __restrict__ output, T* __restrict__ present, + int channels, int length, int input_width, int channel_offset, + CausalConvLayout state_layout, bool apply_silu, + T* __restrict__ state_update = nullptr, int state_update_capacity = 0) { + const int b = blockIdx.y; + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= channels) return; + const int pos_base = static_cast(blockIdx.z) * kPosTile; + + const float b_val = bias == nullptr ? 0.0f : to_float(bias[c]); + const float w0 = to_float(weight[c * 4 + 0]); + const float w1 = to_float(weight[c * 4 + 1]); + const float w2 = to_float(weight[c * 4 + 2]); + const float w3 = to_float(weight[c * 4 + 3]); + + auto load_src = [&](int src) -> float { + if (src >= 0) { + return to_float(input[(static_cast(b) * length + src) * input_width + channel_offset + c]); + } + return past == nullptr ? 0.0f : to_float(past[state_layout.Offset(b, 4 + src, c)]); + }; + + float v0 = load_src(pos_base - 3); + float v1 = load_src(pos_base - 2); + float v2 = load_src(pos_base - 1); + +#pragma unroll + for (int p = 0; p < kPosTile; ++p) { + const int position = pos_base + p; + if (position >= length) break; + const T cur_in = input[(static_cast(b) * length + position) * input_width + channel_offset + c]; + if (state_update != nullptr && position < state_update_capacity) { + state_update[(static_cast(b) * state_update_capacity + position) * channels + c] = cur_in; + } + const float v3 = to_float(cur_in); + float sum = b_val + v0 * w0 + v1 * w1 + v2 * w2 + v3 * w3; + if (apply_silu) sum = sum / (1.0f + expf(-sum)); + output[(static_cast(b) * length + position) * channels + c] = from_float(sum); + v0 = v1; + v1 = v2; + v2 = v3; + } + + if (blockIdx.z != 0) return; +#pragma unroll + for (int position = 0; position < 4; ++position) { + const int source = length - 4 + position; + present[state_layout.Offset(b, position, c)] = + source >= 0 ? input[(static_cast(b) * length + source) * input_width + channel_offset + c] + : (past == nullptr ? from_float(0.0f) : past[state_layout.Offset(b, 4 + source, c)]); + } +} + template __global__ void CausalConvProjectionKernel( const T* input, const T* weight, const T* bias, const T* past, @@ -482,7 +538,17 @@ Status LaunchCausalConvProjectionKernel( const dim3 grid((channels - 1) / 128 + 1, batch, length); static const int specialize = ParseEnvironmentVariableWithDefault("ORT_CAUSAL_CONV_S2_SPECIALIZE", 1); ORT_ENFORCE(specialize == 0 || specialize == 1, "ORT_CAUSAL_CONV_S2_SPECIALIZE must be 0 or 1."); - if (specialize && length >= 1 && length <= 8 && kernel_size == 4 && dilation == 1 && state_length == 4) { + if (specialize && length >= 1 && kernel_size == 4 && dilation == 1 && state_length == 4) { + if (length >= 16) { + constexpr int kPosTile = 8; + const dim3 tiled_grid((channels - 1) / 128 + 1, batch, (length + kPosTile - 1) / kPosTile); + CausalConvProjectionTiledK4Kernel<<>>( + input, weight, bias, past, output, present, channels, length, + input_width, channel_offset, + MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu, + state_update, state_update_capacity); + return CUDA_CALL(cudaGetLastError()); + } CausalConvProjectionKernel<<>>( input, weight, bias, past, output, present, channels, length, kernel_size, dilation, input_width, channel_offset, state_length, From 2ba61c300b852ba725745999f06b587e37797842 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 05:46:54 +0000 Subject: [PATCH 11/17] fix(cuda): address review feedback on CausalConvWithState/GatedDeltaNet - Treat an empty trailing state_update output as omitted in shape inference and the CPU, CUDA and WebGPU CausalConvWithState kernels. - Remove the sequence-length veto that silently disabled uniform GatedDeltaNet state capture when S > state_update_capacity + 1. - Use 64-bit weight offsets in the extended CausalConv projection kernels. - Add regression tests for all three. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cpu/bert/causal_conv_with_state.cc | 2 +- .../cuda/bert/causal_conv_with_state.cc | 8 ++-- .../cuda/bert/causal_conv_with_state.h | 1 + .../cuda/bert/causal_conv_with_state_impl.cu | 10 ++--- .../contrib_ops/cuda/bert/gated_delta_net.cc | 4 +- .../webgpu/bert/causal_conv_with_state.cc | 5 +-- .../core/graph/contrib_ops/bert_defs.cc | 9 ++-- .../causal_conv_with_state_op_test.cc | 13 ++++++ .../contrib_ops/gated_delta_net_op_test.cc | 44 +++++++++++++++++++ 9 files changed, 77 insertions(+), 19 deletions(-) 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 bf81d935636fe..3ebf0eb65ce5c 100644 --- a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc @@ -45,7 +45,7 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : OpKernel info.GetAttrOrDefault("input_channel_offset", 0) == 0 && info.GetAttrOrDefault("state_length", 0) == 0 && info.GetAttrOrDefault("state_channels_last", -1) == -1 && - info.node().OutputDefs().size() <= 2, + (info.node().OutputDefs().size() <= 2 || !info.node().OutputDefs()[2]->Exists()), "CPU CausalConvWithState does not support extended channel/state options (CUDA only)"); int64_t ndim = info.GetAttrOrDefault("ndim", 1); ORT_ENFORCE(ndim == 1, "CPU CausalConvWithState only supports ndim=1"); 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 f03fb42cc4819..8cdd1eaf6addc 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc @@ -51,6 +51,8 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : CudaKern state_length_ = SafeInt(info.GetAttrOrDefault("state_length", 0)); state_channels_last_ = SafeInt(info.GetAttrOrDefault("state_channels_last", -1)); state_update_capacity_ = SafeInt(info.GetAttrOrDefault("state_update_capacity", 0)); + const auto& output_defs = info.node().OutputDefs(); + has_state_update_output_ = output_defs.size() > 2 && output_defs[2]->Exists(); ORT_ENFORCE(input_channels_ >= 0 && input_channel_offset_ >= 0 && state_length_ >= 0 && state_channels_last_ >= -1 && state_channels_last_ <= 1 && state_update_capacity_ >= 0 && state_update_capacity_ <= 8, @@ -71,12 +73,12 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { const auto& weight_shape = weight_tensor->Shape(); const bool extended = input_channels_ != 0 || input_channel_offset_ != 0 || state_length_ != 0 || - state_channels_last_ != -1 || state_update_capacity_ != 0 || context->OutputCount() > 2; + state_channels_last_ != -1 || state_update_capacity_ != 0 || has_state_update_output_; 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"); - ORT_RETURN_IF_NOT(context->OutputCount() <= 2 || state_update_capacity_ > 0, + ORT_RETURN_IF_NOT(!has_state_update_output_ || state_update_capacity_ > 0, "extended CausalConvWithState state_update output requires state_update_capacity > 0"); const int batch = SafeInt(input_shape[0]); const int length = SafeInt(input_shape[1]); @@ -99,7 +101,7 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { "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* state_update = state_update_capacity_ > 0 && context->OutputCount() > 2 + Tensor* state_update = state_update_capacity_ > 0 && has_state_update_output_ ? context->Output(2, TensorShape({batch, state_update_capacity_, channels})) : nullptr; using CudaT = typename OrtToCudaType::type; 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 a9fd12f83ad6a..aec21dfdb4928 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h @@ -30,6 +30,7 @@ class CausalConvWithState final : public onnxruntime::cuda::CudaKernel { int state_length_; int state_channels_last_; int state_update_capacity_; + bool has_state_update_output_; }; } // 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 a38292ba164aa..e9b508e191f20 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,10 +41,10 @@ __global__ void CausalConvProjectionTiledK4Kernel( const int pos_base = static_cast(blockIdx.z) * kPosTile; const float b_val = bias == nullptr ? 0.0f : to_float(bias[c]); - const float w0 = to_float(weight[c * 4 + 0]); - const float w1 = to_float(weight[c * 4 + 1]); - const float w2 = to_float(weight[c * 4 + 2]); - const float w3 = to_float(weight[c * 4 + 3]); + const float w0 = to_float(weight[static_cast(c) * 4 + 0]); + const float w1 = to_float(weight[static_cast(c) * 4 + 1]); + const float w2 = to_float(weight[static_cast(c) * 4 + 2]); + const float w3 = to_float(weight[static_cast(c) * 4 + 3]); auto load_src = [&](int src) -> float { if (src >= 0) { @@ -113,7 +113,7 @@ __global__ void CausalConvProjectionKernel( ? to_float(cur_in) : to_float(input[(static_cast(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]); + sum += value * to_float(weight[static_cast(c) * kernel_size + tap]); } if (apply_silu) sum = sum / (1.0f + expf(-sum)); output[(static_cast(b) * length + position) * channels + c] = from_float(sum); diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index 19e4d8e6a0de2..0fc61ca88911e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -336,9 +336,7 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { desc.state_update_capacity = state_update_capacity_; desc.state_update_active = state_update_capacity_ > 0 && - (state_update_active == nullptr || state_update_active->Data()[0] != 0) && - (cu_seqlens != nullptr || - total_tokens <= static_cast(state_update_capacity_ + 1) * batch); + (state_update_active == nullptr || state_update_active->Data()[0] != 0); desc.update_rule = update_rule_; desc.gate_activation = gate_activation_; desc.beta_activation = beta_activation_; 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 88dca80758232..733222f5e00f9 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc @@ -41,7 +41,8 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) ORT_ENFORCE(info.GetAttrOrDefault("input_channels", 0) == 0 && info.GetAttrOrDefault("input_channel_offset", 0) == 0 && info.GetAttrOrDefault("state_length", 0) == 0 && - info.GetAttrOrDefault("state_channels_last", -1) == -1, + info.GetAttrOrDefault("state_channels_last", -1) == -1 && + (info.node().OutputDefs().size() <= 2 || !info.node().OutputDefs()[2]->Exists()), "WebGPU CausalConvWithState does not support extended channel/state options (CUDA only)"); std::string activation_str = info.GetAttrOrDefault("activation", "none"); activation_ = ParseCausalConvActivation(activation_str); @@ -126,8 +127,6 @@ Status CausalConvUpdateStateProgram::GenerateShaderCode(ShaderHelper& shader) co } Status CausalConvWithState::ComputeInternal(ComputeContext& context) const { - ORT_RETURN_IF(context.OutputCount() > 2, - "WebGPU CausalConvWithState does not support extended state outputs (CUDA only)"); const Tensor* input = context.Input(0); // (B, D, L) const Tensor* weight = context.Input(1); // (D, 1, K) const Tensor* bias = context.Input(2); // optional (D,) diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 1df9680608701..64a5f697dba22 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -4808,6 +4808,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( const int64_t history = getAttribute(ctx, "state_length", 0); const int64_t state_layout = getAttribute(ctx, "state_channels_last", -1); const int64_t state_update_capacity = getAttribute(ctx, "state_update_capacity", 0); + const bool has_state_update_output = ctx.getNumOutputs() > 2 && ctx.hasOutput(2); if (input_channels < 0 || input_channels > std::numeric_limits::max() || channel_offset < 0 || channel_offset > std::numeric_limits::max() || history < 0 || history > std::numeric_limits::max() || @@ -4816,15 +4817,15 @@ ONNX_MS_OPERATOR_SET_SCHEMA( fail_shape_inference("CausalConvWithState: invalid extended attributes"); } if (input_channels != 0 || channel_offset != 0 || history != 0 || state_layout != -1 || - state_update_capacity != 0 || ctx.getNumOutputs() > 2) { + state_update_capacity != 0 || has_state_update_output) { if (getAttribute(ctx, "channels_last", 0) != 1 || getAttribute(ctx, "ndim", 1) != 1 || state_window != 0 || getAttribute(ctx, "dilation", 1) < 1) { fail_shape_inference("CausalConvWithState: extended attributes require channels_last=1, ndim=1, state_window=0"); } - if (ctx.getNumOutputs() > 2 && state_update_capacity <= 0) { + if (has_state_update_output && state_update_capacity <= 0) { fail_shape_inference("CausalConvWithState: state_update output requires state_update_capacity > 0"); } - if (ctx.getNumOutputs() > 2) propagateElemTypeFromInputToOutput(ctx, 0, 2); + if (has_state_update_output) propagateElemTypeFromInputToOutput(ctx, 0, 2); if (!hasInputShape(ctx, 0)) return; const auto& input = getInputShape(ctx, 0); if (input.dim_size() != 3) fail_shape_inference("CausalConvWithState: extended input must have rank 3"); @@ -4862,7 +4863,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( *state.add_dim() = state_layout == 0 ? *channels : length; *state.add_dim() = state_layout == 0 ? length : *channels; updateOutputShape(ctx, 1, state); - if (state_update_capacity > 0 && ctx.getNumOutputs() > 2 && ctx.hasOutput(2)) { + if (state_update_capacity > 0 && has_state_update_output) { TensorShapeProto update_shape; *update_shape.add_dim() = input.dim(0); update_shape.add_dim()->set_dim_value(state_update_capacity); diff --git a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc index c1e85816cafa9..4d7bc5423d341 100644 --- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc +++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc @@ -462,6 +462,19 @@ TEST(CausalConvWithStateTest, PackedChannelsRejectsUnsupportedCpu) { {}, nullptr, &providers); } +TEST(CausalConvWithStateTest, EmptyTrailingStateUpdateOutputIsOmitted) { + OpTester test("CausalConvWithState", 1, kMSDomain); + test.AddAttribute("activation", "none"); + test.AddInput("input", {1, 1, 2}, {1.0f, 2.0f}); + test.AddInput("weight", {1, 1, 2}, {1.0f, 1.0f}); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddOutput("output", {1, 1, 2}, {1.0f, 3.0f}); + test.AddOutput("present", {1, 1, 1}, {2.0f}); + test.AddOptionalOutputEdge(); + test.Run(); +} + TEST(CausalConvWithStateTest, BasicNoStateNoBias) { // B=1, D=2, L=4, K=3, activation=none int batch_size = 1, channels = 2, input_length = 4, kernel_size = 3; diff --git a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc index 72cf073a82abb..929794fd1c280 100644 --- a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc @@ -2597,6 +2597,50 @@ TEST(GatedDeltaNetTest, RequiresCaptureCountExactlyWhenCapacityIsPositive) { } } +// A uniform call longer than state_update_capacity + 1 must still capture the requested prefix. +TEST(GatedDeltaNetTest, UniformCompactCaptureLongerThanCapacity) { + if (NeedSkipGatedDeltaNetTest()) return; + constexpr int kCapacity = 4; + Options options; + options.gate_activation = "qwen"; + options.beta_activation = "sigmoid"; + options.qk_l2_norm = 1; + options.state_update_capacity = kCapacity; + + for (int tokens : {8, 64}) { + SCOPED_TRACE(tokens); + Geometry g{tokens, 1, 1, 2, kDim, kDim}; + Inputs in = MakeInputs(g, 83 + tokens); + in.capture_count = {kCapacity}; + + std::vector fetches; + RunCase(g, options, in, 3e-2f, 3e-2f, false, &fetches); + ASSERT_EQ(fetches.size(), 3u); + + std::vector recurrent_fetches; + { + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_GDN_PLAN", "recurrent"}}}; + RunCase(g, options, in, 3e-2f, 3e-2f, false, &recurrent_fetches); + } + ASSERT_EQ(recurrent_fetches.size(), 3u); + + const Tensor& state_update = fetches[2].Get(); + const Tensor& recurrent_state_update = recurrent_fetches[2].Get(); + ASSERT_EQ(state_update.Shape(), recurrent_state_update.Shape()); + const float* actual = state_update.Data(); + const float* expected = recurrent_state_update.Data(); + const int64_t decay_elements = static_cast(kCapacity) * g.hv; + const int64_t key_elements = static_cast(kCapacity) * g.hq * g.dk; + const int64_t delta_elements = static_cast(kCapacity) * g.hv * g.dv; + for (int64_t i = 0; i < decay_elements; ++i) { + EXPECT_GT(actual[i], 0.0f) << "decay element " << i; + } + for (int64_t i = 0; i < decay_elements + key_elements + delta_elements; ++i) { + EXPECT_NEAR(actual[i], expected[i], 3e-2f) << "state_update element " << i; + } + } +} + TEST(GatedDeltaNetTest, InactiveStateUpdateIgnoresCountsAndClearsOutput) { if (NeedSkipGatedDeltaNetTest()) return; Geometry g{4, 1, 1, 2, 64, 32}; From 7495ae574af6f4ad5f7533fc64eb54ccf61e3a70 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 16:52:30 +0000 Subject: [PATCH 12/17] fix(causal_conv): use plugin-compatible optional-output check The CUDA and WebGPU CausalConvWithState kernels called info.node().OutputDefs() to detect the optional state_update output. ep::adapter::Node (used by the CUDA plugin EP and WebGPU plugin/static plugin builds) has no OutputDefs(), so those builds failed to compile. Add causal_conv_with_state_helper::HasOutput, which uses the adapter's OutputCount()/OutputExists() when available and falls back to OutputDefs() for the in-tree Node, and use it in the CPU, CUDA and WebGPU kernels. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cpu/bert/causal_conv_with_state.cc | 2 +- .../cpu/bert/causal_conv_with_state_helper.h | 15 +++++++++++++++ .../cuda/bert/causal_conv_with_state.cc | 3 +-- .../webgpu/bert/causal_conv_with_state.cc | 2 +- 4 files changed, 18 insertions(+), 4 deletions(-) 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 3ebf0eb65ce5c..9e73c5a75f888 100644 --- a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc @@ -45,7 +45,7 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : OpKernel info.GetAttrOrDefault("input_channel_offset", 0) == 0 && info.GetAttrOrDefault("state_length", 0) == 0 && info.GetAttrOrDefault("state_channels_last", -1) == -1 && - (info.node().OutputDefs().size() <= 2 || !info.node().OutputDefs()[2]->Exists()), + !causal_conv_with_state_helper::HasOutput(info, 2), "CPU CausalConvWithState does not support extended channel/state options (CUDA only)"); int64_t ndim = info.GetAttrOrDefault("ndim", 1); ORT_ENFORCE(ndim == 1, "CPU CausalConvWithState only supports ndim=1"); 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 index 8f16d8ee71f49..0bc24d8cf9502 100644 --- a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state_helper.h @@ -16,6 +16,21 @@ namespace onnxruntime { namespace contrib { namespace causal_conv_with_state_helper { +// Returns true if output `output_index` is present and not an omitted optional output. +// +// The in-tree onnxruntime::Node exposes NodeArgs through OutputDefs(), while the plugin EP adapter +// node (ep::adapter::Node) only exposes OutputCount() and OutputExists(). +template +bool HasOutput(const TKernelInfo& info, size_t output_index) { + const auto& node = info.node(); + if constexpr (requires { node.OutputCount(); node.OutputExists(output_index); }) { + return node.OutputCount() > output_index && node.OutputExists(output_index); + } else { + const auto& output_defs = node.OutputDefs(); + return output_defs.size() > output_index && output_defs[output_index]->Exists(); + } +} + // Reads and validates the optional `state_window` attribute. // // 0 (the default, i.e. attribute absent) selects the legacy unwindowed state layout. Every model 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 8cdd1eaf6addc..cc2a3552de4dc 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc @@ -51,8 +51,7 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : CudaKern state_length_ = SafeInt(info.GetAttrOrDefault("state_length", 0)); state_channels_last_ = SafeInt(info.GetAttrOrDefault("state_channels_last", -1)); state_update_capacity_ = SafeInt(info.GetAttrOrDefault("state_update_capacity", 0)); - const auto& output_defs = info.node().OutputDefs(); - has_state_update_output_ = output_defs.size() > 2 && output_defs[2]->Exists(); + has_state_update_output_ = causal_conv_with_state_helper::HasOutput(info, 2); ORT_ENFORCE(input_channels_ >= 0 && input_channel_offset_ >= 0 && state_length_ >= 0 && state_channels_last_ >= -1 && state_channels_last_ <= 1 && state_update_capacity_ >= 0 && state_update_capacity_ <= 8, 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 733222f5e00f9..a4042ecd9340c 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc @@ -42,7 +42,7 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) info.GetAttrOrDefault("input_channel_offset", 0) == 0 && info.GetAttrOrDefault("state_length", 0) == 0 && info.GetAttrOrDefault("state_channels_last", -1) == -1 && - (info.node().OutputDefs().size() <= 2 || !info.node().OutputDefs()[2]->Exists()), + !causal_conv_with_state_helper::HasOutput(info, 2), "WebGPU CausalConvWithState does not support extended channel/state options (CUDA only)"); std::string activation_str = info.GetAttrOrDefault("activation", "none"); activation_ = ParseCausalConvActivation(activation_str); From e62ab5a3aac9d1d1eb4c181239f009b1ca9d5506 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:48:25 +0000 Subject: [PATCH 13/17] Remove gate_activation 'qwen' alias from GatedDeltaNet; use 'softplus' only The alias was introduced alongside 'softplus' in this PR. Use a single spelling across the schema, CUDA/WebGPU kernels, docs and tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 2 +- docs/contrib_ops/cuda/gdn.md | 4 +- .../contrib_ops/cuda/bert/gated_delta_net.cc | 2 +- .../webgpu/bert/gated_delta_net.cc | 9 ++-- .../core/graph/contrib_ops/bert_defs.cc | 3 +- .../contrib_ops/gated_delta_net_op_test.cc | 46 +++++++++---------- 6 files changed, 32 insertions(+), 34 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 877cb70deb478..83b67b03a532d 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2529,7 +2529,7 @@ This version of the operator has been available since version 1 of the 'com.micr
chunk_size : int
Tuning hint for the chunk-parallel prefill algorithm. On CUDA, 32 pins the narrow chunk and any other value lets the implementation select the widest chunk the device can hold. On WebGPU, 16 opts into the chunkwise prefill route on supported Apple silicon; default 64 retains recurrent execution. Default 64.
gate_activation : string
-
'none' (default) treats `decay` as the effective log-space decay. 'softplus' computes -exp(a_log) * Softplus(decay + dt_bias) in float32. 'qwen' is also accepted for compatibility.
+
'none' (default) treats `decay` as the effective log-space decay. 'softplus' computes -exp(a_log) * Softplus(decay + dt_bias) in float32.
qk_l2_norm : int
When 1, L2-normalize each query and key head vector before the recurrence. Default 0.
scale : float
diff --git a/docs/contrib_ops/cuda/gdn.md b/docs/contrib_ops/cuda/gdn.md index 8037400dfb5b9..53230ef74c5fb 100644 --- a/docs/contrib_ops/cuda/gdn.md +++ b/docs/contrib_ops/cuda/gdn.md @@ -74,7 +74,7 @@ Source: |---|---|---|---| | `update_rule` | string | `gated_delta` | One of `linear`, `gated`, `delta`, `gated_delta`. Selects which of the decay and delta-retrieval terms are present. | | `scale` | float | `0.0` | Output scale. `0.0` means `1/sqrt(head_size_qk)`. | -| `gate_activation` | string | `none` | `none` treats `decay` as the effective log-space decay. `qwen` computes `-exp(a_log) * Softplus(decay + dt_bias)` in float32. | +| `gate_activation` | string | `none` | `none` treats `decay` as the effective log-space decay. `softplus` computes `-exp(a_log) * Softplus(decay + dt_bias)` in float32. | | `beta_activation` | string | `none` | `none` treats `beta` as the effective update rate; `sigmoid` applies a sigmoid. | | `qk_l2_norm` | int | `0` | When `1`, L2-normalizes each query and key head vector in-kernel. | | `chunk_size` | int | `64` | Chunk length for the chunked engine. `32` pins the reduced shared-memory configuration (see §6). | @@ -525,7 +525,7 @@ support three intermediate precision modes controlled by session option Measured on NVIDIA H200 (`sm_90a`, CUDA 13.0) at the Qwen3.8-Flash-Next geometry (`batch = 1`, `num_heads_q = num_heads_k = 16`, `num_heads_v = 32`, `head_size_qk = head_size_v = 128`, -`update_rule = "gated_delta"`, `gate_activation = "qwen"`, `beta_activation = "sigmoid"`, +`update_rule = "gated_delta"`, `gate_activation = "softplus"`, `beta_activation = "sigmoid"`, `qk_l2_norm = 1`, 36 GDN layers) against the sequential `float32` `recurrent` engine reference: | Sequence Length | Mode | 1-Op Latency (us) | 36-Layer Time (ms) | Output RMSE vs. `recurrent` | State RMSE vs. `recurrent` | Exact BF16 Output Match | diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index 0fc61ca88911e..bf60fcb3d936b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -62,7 +62,7 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : CudaKernel(info) { const std::string gate = info.GetAttrOrDefault("gate_activation", "none"); if (gate == "none") { gate_activation_ = gdn::GateActivation::kNone; - } else if (gate == "softplus" || gate == "qwen") { + } else if (gate == "softplus") { gate_activation_ = gdn::GateActivation::kSoftplusDecay; } else { ORT_THROW("gate_activation must be one of: none, softplus"); diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc index 4f83cd61c65ea..968dfc3a10426 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc @@ -48,10 +48,9 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : WebGpuKernel(info) { ORT_ENFORCE(update_rule_ != GatedDeltaNetUpdateRule::Invalid, "update_rule must be one of: linear, gated, delta, gated_delta"); const auto gate_activation = info.GetAttrOrDefault("gate_activation", "none"); - ORT_ENFORCE(gate_activation == "none" || gate_activation == "softplus" || - gate_activation == "qwen", + ORT_ENFORCE(gate_activation == "none" || gate_activation == "softplus", "gate_activation must be one of: none, softplus"); - qwen_gate_ = gate_activation == "softplus" || gate_activation == "qwen"; + qwen_gate_ = gate_activation == "softplus"; const auto beta_activation = info.GetAttrOrDefault("beta_activation", "none"); ORT_ENFORCE(beta_activation == "none" || beta_activation == "sigmoid", "beta_activation must be one of: none, sigmoid"); @@ -453,9 +452,9 @@ Status GatedDeltaNet::ComputeInternal(onnxruntime::webgpu::ComputeContext& conte ORT_RETURN_IF_NOT(a_log != nullptr && dt_bias != nullptr && a_log->Shape().NumDimensions() == 1 && a_log->Shape()[0] == hv && dt_bias->Shape().NumDimensions() == 1 && dt_bias->Shape()[0] == hv, - "gate_activation=qwen requires a_log and dt_bias with shape [num_heads_v]"); + "gate_activation=softplus requires a_log and dt_bias with shape [num_heads_v]"); } else { - ORT_RETURN_IF_NOT(a_log == nullptr && dt_bias == nullptr, "a_log and dt_bias require gate_activation=qwen"); + ORT_RETURN_IF_NOT(a_log == nullptr && dt_bias == nullptr, "a_log and dt_bias require gate_activation=softplus"); } TensorShapeVector output_dims(q_shape.GetDims().begin(), q_shape.GetDims().begin() + token_dims); diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 64a5f697dba22..9ab1adfdba3c4 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -5571,8 +5571,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( AttributeProto::FLOAT, 0.0f) .Attr("gate_activation", "'none' (default) treats `decay` as the effective log-space decay. 'softplus' computes " - "-exp(a_log) * Softplus(decay + dt_bias) in float32. 'qwen' is also accepted for " - "compatibility.", + "-exp(a_log) * Softplus(decay + dt_bias) in float32.", AttributeProto::STRING, std::string("none")) .Attr("beta_activation", "'none' (default) treats `beta` as the effective update rate. 'sigmoid' applies a " diff --git a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc index 929794fd1c280..8c115196fda85 100644 --- a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc @@ -302,7 +302,7 @@ void Reference(const Geometry& g, const Options& o, const Inputs& in, std::vecto double decay = 1.0; if (gated) { float raw = in.decay[gi]; - if (o.gate_activation == "qwen" || o.gate_activation == "softplus") { + if (o.gate_activation == "softplus") { raw = -std::exp(in.a_log[hv]) * SoftPlus(raw + in.dt_bias[hv]); } decay = std::exp(static_cast(raw)); @@ -518,7 +518,7 @@ void RunTypedCase(const Geometry& g, const Options& o, const Inputs& in_raw, flo } else { test.AddInput("initial_state", {g.batch, g.hv, g.dv, g.dk}, in.state0); } - if (o.gate_activation == "qwen" || o.gate_activation == "softplus") { + if (o.gate_activation == "softplus") { test.AddInput("a_log", {g.hv}, in.a_log); test.AddInput("dt_bias", {g.hv}, in.dt_bias); } else { @@ -632,7 +632,7 @@ void TestInputTypeGates() { for (const std::string rule : {"gated_delta", "gated", "delta"}) { Options options; options.update_rule = rule; - options.gate_activation = rule == "delta" ? "none" : "qwen"; + options.gate_activation = rule == "delta" ? "none" : "softplus"; options.beta_activation = "sigmoid"; std::vector reference, narrowed; RunTypedCase(geometry, options, inputs, 6e-3f, 2e-3f, true, &reference); @@ -730,7 +730,7 @@ TEST(GatedDeltaNetWebGpuTest, Rank4AllRulesAndInverseGqa) { SCOPED_TRACE(rule); Options options; options.update_rule = rule; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; RunCase(g, options, MakeInputs(g, 201), 3e-3f, 3e-4f, /*rank4=*/true, @@ -756,7 +756,7 @@ TEST(GatedDeltaNetWebGpuTest, RecurrentVectorizedValueIoAndSharedGate) { inputs.cu_seqlens = {0, 3, 7}; Options options; options.update_rule = "gated"; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; RunTypedCase(g, options, inputs, 3e-3f, 3e-4f, /*rank4=*/false, /*fetches=*/nullptr, /*use_webgpu=*/true, /*omit_final_state=*/true); @@ -812,7 +812,7 @@ TEST(GatedDeltaNetWebGpuTest, PackedQkvAndParamsWithSegmentedBacking) { inputs.cu_seqlens = {0, 4, 9}; Options options; options.update_rule = "gated_delta"; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; const auto binding_count = [](uint64_t bytes) { return (bytes + max_binding_size - 1) / max_binding_size; @@ -919,7 +919,7 @@ TEST(GatedDeltaNetWebGpuTest, RaggedWithoutInitialState) { Inputs inputs = MakeInputs(g, 223, /*with_state=*/false); inputs.cu_seqlens = {0, 1, 5, 7}; Options options; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; options.scale = 0.37f; @@ -935,7 +935,7 @@ TEST(GatedDeltaNetWebGpuTest, RaggedQwenWithInitialStateAndNonDivisibleDv) { Inputs inputs = MakeInputs(g, 225); inputs.cu_seqlens = {0, 1, 5, 7}; Options options; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; RunTypedCase(g, options, inputs, 3e-4f, 3e-4f, @@ -948,7 +948,7 @@ TEST(GatedDeltaNetWebGpuTest, ChunkwiseQwenPrefillPartialChunk) { } Options options; options.chunk_size = 16; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; for (const Geometry geometry : {Geometry{65, 1, 2, 4, 128, 128}, Geometry{65, 1, 16, 48, 128, 128}}) { @@ -967,7 +967,7 @@ TEST(GatedDeltaNetWebGpuTest, ChunkwiseQwenPrefillSaturatedGates) { Geometry geometry{32, 1, 1, 2, 16, 7}; Options options; options.chunk_size = 16; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; for (bool isolated_resets : {false, true}) { @@ -996,7 +996,7 @@ TEST(GatedDeltaNetWebGpuTest, ChunkwiseQwenPrefillRaggedAndEmptySequence) { Geometry geometry{130, 3, 1, 3, 16, 7}; Options options; options.chunk_size = 16; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; options.scale = 0.37f; @@ -1023,7 +1023,7 @@ TEST(GatedDeltaNetWebGpuTest, ChunkwiseQwenPrefillLargeRaggedBatch) { Geometry geometry{2080, 64, 1, 1, 128, 128}; Options options; options.chunk_size = 16; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; Inputs inputs = MakeInputs(geometry, 234); @@ -1040,7 +1040,7 @@ TEST(GatedDeltaNetWebGpuTest, ChunkwiseQwenPrefillPackedAndUniform) { Geometry geometry{130, 2, 2, 4, 32, 32}; Options options; options.chunk_size = 16; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; const Inputs inputs = MakeInputs(geometry, 233); @@ -1062,7 +1062,7 @@ TEST(GatedDeltaNetWebGpuTest, CompactStateUpdates) { inputs.cu_seqlens = {0, 1, 5, 12}; inputs.capture_count = {-1, 2, 8}; Options options; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; options.state_update_capacity = 7; @@ -1079,7 +1079,7 @@ TEST(GatedDeltaNetWebGpuTest, CompactStateUpdatesFp16QwenGeometry) { inputs.cu_seqlens = {0, 1, 5, 12}; inputs.capture_count = {-1, 2, 8}; Options options; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; options.state_update_capacity = 7; @@ -1100,7 +1100,7 @@ TEST(GatedDeltaNetWebGpuTest, CapturedStateReplayMatchesIndependentPrefixes) { } inputs.capture_count = {geometry.total_tokens}; Options options; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; // Capture capacity intentionally exceeds both the full run length and every partial prefix. @@ -1288,7 +1288,7 @@ TEST(GatedDeltaNetTest, Chunked_FusedActivations) { if (NeedSkipGatedDeltaNetTest()) return; Geometry g{128, 1, 1, 2, kDim, kDim}; Options o; - o.gate_activation = "qwen"; + o.gate_activation = "softplus"; o.beta_activation = "sigmoid"; o.qk_l2_norm = 1; RunCase(g, o, MakeInputs(g, 7), 3e-2f, 3e-2f); @@ -1394,7 +1394,7 @@ TEST(GatedDeltaNetTest, ChunkedSplit_FusedActivationsMatchReference) { Geometry g{128, 1, 2, 6, kDim, kDim}; Options options; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; RunCase(g, options, MakeInputs(g, 137), 3e-2f, 3e-2f); @@ -1428,7 +1428,7 @@ void TestDirectPackedQkvParity() { for (bool input_type_gates : {false, true}) { Options options; options.qk_l2_norm = 1; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.gates_in_input_type = input_type_gates; std::vector reference, packed; @@ -1464,7 +1464,7 @@ void TestQwenS2PackedParity() { Geometry geometry{tokens, 1, 16, 48, 128, 128}; Options options; options.qk_l2_norm = 1; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.gates_in_input_type = true; Inputs inputs = MakeInputs(geometry, 151 + tokens); @@ -2002,7 +2002,7 @@ TEST(GatedDeltaNetTest, Ragged_CompactReplayMatchesSequentialPrefixes) { constexpr int kPartial = 4; Geometry g{kPrefillTokens + 5 + kCapacity, 4, 1, 2, kDim, kDim}; Options options; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; options.state_update_capacity = kCapacity; @@ -2456,7 +2456,7 @@ TEST(GatedDeltaNetTest, RejectsPerKeyDtBias) { Inputs in = MakeInputs(g, 53); OpTester test("GatedDeltaNet", 1, onnxruntime::kMSDomain); Options o; - o.gate_activation = "qwen"; + o.gate_activation = "softplus"; AddCommonAttrs(test, o); test.AddInput("query", {g.total_tokens, g.hq, g.dk}, ToFloat16(in.q)); test.AddInput("key", {g.total_tokens, g.hq, g.dk}, ToFloat16(in.k)); @@ -2602,7 +2602,7 @@ TEST(GatedDeltaNetTest, UniformCompactCaptureLongerThanCapacity) { if (NeedSkipGatedDeltaNetTest()) return; constexpr int kCapacity = 4; Options options; - options.gate_activation = "qwen"; + options.gate_activation = "softplus"; options.beta_activation = "sigmoid"; options.qk_l2_norm = 1; options.state_update_capacity = kCapacity; From 63ad89b0d90c131a91149e8d57a5ef3855c193fc Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:48:58 +0000 Subject: [PATCH 14/17] CausalConvWithState: zero unused state_update rows; fix S2 env var docs Rows of the fixed-capacity state_update output beyond min(length, capacity) were never written. Zero them with a graph-capturable strided memset in the launcher. Document that ORT_CAUSAL_CONV_S2_SPECIALIZE also gates the tiled length>=16 kernel and that the activation attribute is honored as configured. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 2 +- docs/contrib_ops/cuda/gdn.md | 13 +++++++++---- .../cuda/bert/causal_conv_with_state_impl.cu | 9 +++++++++ onnxruntime/core/graph/contrib_ops/bert_defs.cc | 3 ++- 4 files changed, 21 insertions(+), 6 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 83b67b03a532d..60083dd653d85 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -1035,7 +1035,7 @@ This version of the operator has been available since version 1 of the 'com.micr
state_length : int
Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. An explicit length must cover the receptive field; convolution reads its trailing history.
state_update_capacity : int
-
Compact speculative state-update capture capacity (CUDA extended path only, 0..8); when > 0, output 2 has shape (batch_size, state_update_capacity, channels) and records the first min(sequence_length, state_update_capacity) projected input rows.
+
Compact speculative state-update capture capacity (CUDA extended path only, 0..8); when > 0, output 2 has shape (batch_size, state_update_capacity, channels) and records the first min(sequence_length, state_update_capacity) projected input rows and any remaining rows are zero.
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, state_length) where state_length = (k_1 - 1) * dilation. 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, state_length) 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].
diff --git a/docs/contrib_ops/cuda/gdn.md b/docs/contrib_ops/cuda/gdn.md index 53230ef74c5fb..77fd8d40fbe77 100644 --- a/docs/contrib_ops/cuda/gdn.md +++ b/docs/contrib_ops/cuda/gdn.md @@ -8,11 +8,16 @@ Set process-level tuning options before the first applicable invocation: 1 to 8 tokens, 16 query/key heads, 48 value heads, and 128-wide heads. It additionally requires Qwen softplus-decay gates, sigmoid beta, gated-delta updates, Q/K L2 normalization, uniform non-packed sequence - metadata, and no per-key decay, state-update capture, or tail passes. + metadata, and no per-key decay, state-update tail pass, or short-row tail + pass. - `ORT_CAUSAL_CONV_S2_SPECIALIZE=1` specializes the companion - `CausalConvProjection` for lengths 1 to 8, kernel/state length four, dilation - one, and SiLU activation. Layout, optional bias, and first-token state - semantics are unchanged. + `CausalConvWithState` extended (packed-channel) projection path for kernel size + four, state length four and dilation one. Lengths 1 to 8 use the per-position + specialized kernel and lengths of 16 or more use the tiled kernel (eight + positions per CTA); other lengths keep the generic kernel. The `activation` + attribute (`none`, `silu` or `swish`; the last two are identical) is honored + as configured and is not part of the dispatch condition. Layout, optional + bias, and first-token state semantics are unchanged. Both specialization switches default to `1` and accept only `0` or `1`. Unsupported geometry retains generic dispatch. These options do not alter 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 e9b508e191f20..1747dcc2de1b0 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 @@ -536,6 +536,15 @@ Status LaunchCausalConvProjectionKernel( int state_length, bool state_channels_last, bool apply_silu, T* state_update, int state_update_capacity) { const dim3 grid((channels - 1) / 128 + 1, batch, length); + if (state_update != nullptr && length < state_update_capacity) { + // The kernels write only the first min(length, capacity) rows per batch; zero the rest so the + // output does not expose stale allocator contents. A strided memset keeps this graph-capturable. + const size_t row_bytes = static_cast(channels) * sizeof(T); + CUDA_RETURN_IF_ERROR(cudaMemset2DAsync( + state_update + static_cast(length) * channels, + static_cast(state_update_capacity) * row_bytes, 0, + static_cast(state_update_capacity - length) * row_bytes, static_cast(batch), stream)); + } static const int specialize = ParseEnvironmentVariableWithDefault("ORT_CAUSAL_CONV_S2_SPECIALIZE", 1); ORT_ENFORCE(specialize == 0 || specialize == 1, "ORT_CAUSAL_CONV_S2_SPECIALIZE must be 0 or 1."); if (specialize && length >= 1 && kernel_size == 4 && dilation == 1 && state_length == 4) { diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 9ab1adfdba3c4..8d73ee8a00644 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -4707,7 +4707,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Attr("state_update_capacity", "Compact speculative state-update capture capacity (CUDA extended path only, 0..8); " "when > 0, output 2 has shape (batch_size, state_update_capacity, channels) and records " - "the first min(sequence_length, state_update_capacity) projected input rows.", + "the first min(sequence_length, state_update_capacity) projected input rows and any " + "remaining rows are zero.", AttributeProto::INT, static_cast(0)) .Attr("activation", "Fused activation function. One of: 'silu', 'swish', 'none'. " From 7fa1b491494afc03f016dd3410c36e0a835104e6 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:50:06 +0000 Subject: [PATCH 15/17] Test CausalConvWithState tiled K4 projection and state_update tail zeroing Cover lengths 16/17/24 (CausalConvProjectionTiledK4Kernel) across both state layouts, present/absent history, full and partial channel blocks, and FP32/FP16/BF16, with positive state_update capacity in short-K4 and tiled cases. Add a regression where capacity exceeds the sequence length. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../causal_conv_with_state_op_test.cc | 71 ++++++++++++++++++- 1 file changed, 69 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc index 4d7bc5423d341..bb9bcf03de1ad 100644 --- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc +++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include #include #include @@ -320,8 +321,10 @@ static void RunCausalConvWithStateTests( // ============================================================================= template -void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool use_bias, bool silu, int dilation) { - constexpr int batch = 2, channels = 7, width = 13, offset = 3, kernel = 4; +void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool use_bias, bool silu, int dilation, + int capacity = 0, int channels = 7) { + constexpr int batch = 2, offset = 3, kernel = 4; + const int width = channels + 6; const int history = (kernel - 1) * dilation + 1; auto cuda = DefaultCudaExecutionProvider(); if (!cuda) GTEST_SKIP() << "CUDA is required for extended convolution"; @@ -353,8 +356,18 @@ void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool u } } } + // Captured rows are the raw projected input; rows beyond min(length, capacity) are defined as zero. + std::vector state_update(static_cast(batch) * capacity * channels, T(0.0f)); + for (int b = 0; b < batch; ++b) { + for (int t = 0; t < std::min(length, capacity); ++t) { + for (int c = 0; c < channels; ++c) { + state_update[(static_cast(b) * capacity + t) * channels + c] = input[(b * length + t) * width + offset + c]; + } + } + } OpTester test("CausalConvWithState", 1, kMSDomain); test.AddAttribute("channels_last", 1); + if (capacity > 0) test.AddAttribute("state_update_capacity", capacity); test.AddAttribute("input_channels", channels); test.AddAttribute("input_channel_offset", offset); test.AddAttribute("state_length", history); @@ -375,6 +388,7 @@ void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool u test.AddOptionalInputEdge(); test.AddOutput("output", {batch, length, channels}, output); test.AddOutput("present", state_shape, present); + if (capacity > 0) test.AddOutput("state_update", {batch, capacity, channels}, state_update); test.SetOutputTolerance(std::is_same_v ? 0.01f : 0.001f, 0.001f); std::vector> providers; providers.push_back(std::move(cuda)); @@ -396,6 +410,59 @@ TEST(CausalConvWithStateTest, PackedChannelsExtendedHistory) { } } +// Lengths 16, 17 and 24 select CausalConvProjectionTiledK4Kernel (kernel 4, dilation 1, history 4): +// two full eight-position tiles, a partial trailing tile, and three full tiles. 130 channels spans a full and a +// partial 128-channel block. Positive capacities cover the captured rows in the tiled path and in +// the per-position short-K4 path (lengths 1..8), including capacity above the sequence length. +TEST(CausalConvWithStateTest, PackedChannelsTiledK4AndCapture) { + for (int length : {16, 17, 24}) { + for (bool state_last : {false, true}) { + for (bool past : {false, true}) { + for (int channels : {7, 130}) { + TestPackedChannelHistory(length, state_last, past, true, true, 1, 0, channels); + TestPackedChannelHistory(length, state_last, past, false, true, 1, 7, channels); + TestPackedChannelHistory(length, state_last, past, true, false, 1, 8, channels); + } + } + } + } + for (int length : {1, 3, 8}) { + for (bool state_last : {false, true}) { + for (bool past : {false, true}) { + TestPackedChannelHistory(length, state_last, past, true, true, 1, 2); + TestPackedChannelHistory(length, state_last, past, false, true, 1, 8, 130); + TestPackedChannelHistory(length, state_last, past, true, false, 1, 7); + } + } + } +} + +// Capacity above the sequence length leaves rows the kernels never write; they are defined as zero. +TEST(CausalConvWithStateTest, PackedChannelsStateUpdateCapacityExceedsLength) { + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) GTEST_SKIP() << "CUDA is required for extended convolution"; + OpTester test("CausalConvWithState", 1, kMSDomain); + test.AddAttribute("channels_last", 1); + test.AddAttribute("input_channels", 2); + test.AddAttribute("state_length", 1); + test.AddAttribute("state_channels_last", 0); + test.AddAttribute("state_update_capacity", 8); + test.AddAttribute("activation", "none"); + test.AddInput("input", {1, 1, 2}, {3.f, 5.f}); + test.AddInput("weight", {2, 1, 2}, {1.f, 1.f, 1.f, 1.f}); + test.AddOptionalInputEdge(); + test.AddOptionalInputEdge(); + test.AddOutput("output", {1, 1, 2}, {3.f, 5.f}); + test.AddOutput("present", {1, 2, 1}, {3.f, 5.f}); + std::vector expected_update(16, 0.f); + expected_update[0] = 3.f; + expected_update[1] = 5.f; + test.AddOutput("state_update", {1, 8, 2}, expected_update, false, 1e-4f, 1e-4f); + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + TEST(CausalConvWithStateTest, PackedChannelsStateUpdateCapacity) { auto cuda = DefaultCudaExecutionProvider(); if (!cuda) GTEST_SKIP() << "CUDA is required for extended convolution"; From 5bdeab4267d910ee4daa51c1118b995eb0f67f67 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:59:03 +0000 Subject: [PATCH 16/17] Restore gate_activation 'qwen' as a legacy alias of 'softplus' Revert the alias removal: CUDA and WebGPU accept 'qwen' again with behavior identical to 'softplus'. Schema and docs now name 'softplus' canonical and 'qwen' a legacy alias; error messages use the canonical spelling. Tests use 'softplus' and add alias coverage (CUDA bit-exact vs softplus, WebGPU ragged). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/ContribOperators.md | 2 +- docs/contrib_ops/cuda/gdn.md | 2 +- .../contrib_ops/cuda/bert/gated_delta_net.cc | 4 +- .../webgpu/bert/gated_delta_net.cc | 7 +-- .../core/graph/contrib_ops/bert_defs.cc | 3 +- .../contrib_ops/gated_delta_net_op_test.cc | 45 ++++++++++++++++++- 6 files changed, 53 insertions(+), 10 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 60083dd653d85..49008e1768e6b 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2529,7 +2529,7 @@ This version of the operator has been available since version 1 of the 'com.micr
chunk_size : int
Tuning hint for the chunk-parallel prefill algorithm. On CUDA, 32 pins the narrow chunk and any other value lets the implementation select the widest chunk the device can hold. On WebGPU, 16 opts into the chunkwise prefill route on supported Apple silicon; default 64 retains recurrent execution. Default 64.
gate_activation : string
-
'none' (default) treats `decay` as the effective log-space decay. 'softplus' computes -exp(a_log) * Softplus(decay + dt_bias) in float32.
+
'none' (default) treats `decay` as the effective log-space decay. 'softplus' computes -exp(a_log) * Softplus(decay + dt_bias) in float32. 'qwen' is accepted as a legacy alias of 'softplus' with identical behavior.
qk_l2_norm : int
When 1, L2-normalize each query and key head vector before the recurrence. Default 0.
scale : float
diff --git a/docs/contrib_ops/cuda/gdn.md b/docs/contrib_ops/cuda/gdn.md index 77fd8d40fbe77..68dec15293772 100644 --- a/docs/contrib_ops/cuda/gdn.md +++ b/docs/contrib_ops/cuda/gdn.md @@ -79,7 +79,7 @@ Source: |---|---|---|---| | `update_rule` | string | `gated_delta` | One of `linear`, `gated`, `delta`, `gated_delta`. Selects which of the decay and delta-retrieval terms are present. | | `scale` | float | `0.0` | Output scale. `0.0` means `1/sqrt(head_size_qk)`. | -| `gate_activation` | string | `none` | `none` treats `decay` as the effective log-space decay. `softplus` computes `-exp(a_log) * Softplus(decay + dt_bias)` in float32. | +| `gate_activation` | string | `none` | `none` treats `decay` as the effective log-space decay. `softplus` computes `-exp(a_log) * Softplus(decay + dt_bias)` in float32. `qwen` is accepted as a legacy alias of `softplus` with identical behavior. | | `beta_activation` | string | `none` | `none` treats `beta` as the effective update rate; `sigmoid` applies a sigmoid. | | `qk_l2_norm` | int | `0` | When `1`, L2-normalizes each query and key head vector in-kernel. | | `chunk_size` | int | `64` | Chunk length for the chunked engine. `32` pins the reduced shared-memory configuration (see §6). | diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index bf60fcb3d936b..d326e3b83be25 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -62,10 +62,10 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : CudaKernel(info) { const std::string gate = info.GetAttrOrDefault("gate_activation", "none"); if (gate == "none") { gate_activation_ = gdn::GateActivation::kNone; - } else if (gate == "softplus") { + } else if (gate == "softplus" || gate == "qwen") { gate_activation_ = gdn::GateActivation::kSoftplusDecay; } else { - ORT_THROW("gate_activation must be one of: none, softplus"); + ORT_THROW("gate_activation must be one of: none, softplus (legacy alias: qwen)"); } const std::string beta = info.GetAttrOrDefault("beta_activation", "none"); diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc index 968dfc3a10426..e1c9560342a6e 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc @@ -48,9 +48,10 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : WebGpuKernel(info) { ORT_ENFORCE(update_rule_ != GatedDeltaNetUpdateRule::Invalid, "update_rule must be one of: linear, gated, delta, gated_delta"); const auto gate_activation = info.GetAttrOrDefault("gate_activation", "none"); - ORT_ENFORCE(gate_activation == "none" || gate_activation == "softplus", - "gate_activation must be one of: none, softplus"); - qwen_gate_ = gate_activation == "softplus"; + ORT_ENFORCE(gate_activation == "none" || gate_activation == "softplus" || + gate_activation == "qwen", + "gate_activation must be one of: none, softplus (legacy alias: qwen)"); + qwen_gate_ = gate_activation == "softplus" || gate_activation == "qwen"; const auto beta_activation = info.GetAttrOrDefault("beta_activation", "none"); ORT_ENFORCE(beta_activation == "none" || beta_activation == "sigmoid", "beta_activation must be one of: none, sigmoid"); diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 8d73ee8a00644..4008e97ea820a 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -5572,7 +5572,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( AttributeProto::FLOAT, 0.0f) .Attr("gate_activation", "'none' (default) treats `decay` as the effective log-space decay. 'softplus' computes " - "-exp(a_log) * Softplus(decay + dt_bias) in float32.", + "-exp(a_log) * Softplus(decay + dt_bias) in float32. 'qwen' is accepted as a legacy " + "alias of 'softplus' with identical behavior.", AttributeProto::STRING, std::string("none")) .Attr("beta_activation", "'none' (default) treats `beta` as the effective update rate. 'sigmoid' applies a " diff --git a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc index 8c115196fda85..caf02bfd067da 100644 --- a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc +++ b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc @@ -302,7 +302,7 @@ void Reference(const Geometry& g, const Options& o, const Inputs& in, std::vecto double decay = 1.0; if (gated) { float raw = in.decay[gi]; - if (o.gate_activation == "softplus") { + if ((o.gate_activation == "softplus" || o.gate_activation == "qwen")) { raw = -std::exp(in.a_log[hv]) * SoftPlus(raw + in.dt_bias[hv]); } decay = std::exp(static_cast(raw)); @@ -518,7 +518,7 @@ void RunTypedCase(const Geometry& g, const Options& o, const Inputs& in_raw, flo } else { test.AddInput("initial_state", {g.batch, g.hv, g.dv, g.dk}, in.state0); } - if (o.gate_activation == "softplus") { + if ((o.gate_activation == "softplus" || o.gate_activation == "qwen")) { test.AddInput("a_log", {g.hv}, in.a_log); test.AddInput("dt_bias", {g.hv}, in.dt_bias); } else { @@ -1294,6 +1294,47 @@ TEST(GatedDeltaNetTest, Chunked_FusedActivations) { RunCase(g, o, MakeInputs(g, 7), 3e-2f, 3e-2f); } +// "qwen" is a legacy alias of "softplus" with identical behavior; the reference applies the fused +// softplus decay for either spelling, and the two spellings must also match each other bit-for-bit. +TEST(GatedDeltaNetTest, QwenGateActivationIsAliasOfSoftplus) { + if (NeedSkipGatedDeltaNetTest()) return; + for (int tokens : {1, 128}) { + Geometry g{tokens, 1, 1, 2, kDim, kDim}; + Inputs inputs = MakeInputs(g, 8); + Options o; + o.beta_activation = "sigmoid"; + o.qk_l2_norm = 1; + std::vector softplus_fetches, qwen_fetches; + o.gate_activation = "softplus"; + RunCase(g, o, inputs, 3e-2f, 3e-2f, false, &softplus_fetches); + o.gate_activation = "qwen"; + RunCase(g, o, inputs, 3e-2f, 3e-2f, false, &qwen_fetches); + ASSERT_EQ(softplus_fetches.size(), qwen_fetches.size()); + for (size_t i = 0; i < softplus_fetches.size(); ++i) { + const Tensor& expected = softplus_fetches[i].Get(); + const Tensor& actual = qwen_fetches[i].Get(); + ASSERT_EQ(expected.Shape(), actual.Shape()); + ASSERT_EQ(0, memcmp(expected.DataRaw(), actual.DataRaw(), expected.SizeInBytes())) + << "output " << i << " tokens=" << tokens; + } + } +} + +TEST(GatedDeltaNetWebGpuTest, QwenGateActivationAlias) { + if (NeedSkipGatedDeltaNetWebGpuTest()) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + Geometry g{7, 3, 1, 3, 8, 5}; + Inputs inputs = MakeInputs(g, 226); + inputs.cu_seqlens = {0, 1, 5, 7}; + Options options; + options.gate_activation = "qwen"; + options.beta_activation = "sigmoid"; + options.qk_l2_norm = 1; + RunTypedCase(g, options, inputs, 3e-4f, 3e-4f, + /*rank4=*/false, /*fetches=*/nullptr, /*use_webgpu=*/true); +} + TEST(GatedDeltaNetTest, Chunked_NoInitialState) { if (NeedSkipGatedDeltaNetTest()) return; Geometry g{128, 1, 1, 2, kDim, kDim}; From c4babd00d5f664d53fba374089974464bb154da5 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 27 Oct 2026 00:15:31 +0000 Subject: [PATCH 17/17] Document QwenS2 / CausalConvWithState S2 op-level and end-to-end benchmarks Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- docs/contrib_ops/cuda/gdn.md | 67 ++++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/docs/contrib_ops/cuda/gdn.md b/docs/contrib_ops/cuda/gdn.md index 68dec15293772..c283f20dbf5f4 100644 --- a/docs/contrib_ops/cuda/gdn.md +++ b/docs/contrib_ops/cuda/gdn.md @@ -23,6 +23,73 @@ Both specialization switches default to `1` and accept only `0` or `1`. Unsupported geometry retains generic dispatch. These options do not alter operator schemas, state layouts, or the default package configuration. +### Benchmark + +Op-level GPU kernel time per call on NVIDIA H200 (SM90, driver 580.159.04, CUDA 13.0, ORT 1.32.0 integration build +of this branch), measured 2026-10-27 at the Qwen3.8-Flash-Next rc3 production geometry +(taken from the rc3 DFlash2 graph): `GatedDeltaNet` with packed BF16 QKV `(B, S, 10240)`, 16 q/k heads, 48 v heads, +head size 128, FP32 state `(B, 48, 128, 128)`, BF16 decay/beta, `gated_delta`/`softplus`/`sigmoid`, +`qk_l2_norm=1`, `state_update_capacity=7`; `CausalConvWithState` with BF16 `(B, S, 10240)` channels-last input, +K=4, `state_length=4`, `dilation=1`, `silu`, `state_update_capacity=7`. Median kernel duration in microseconds +(lower is better), specialized (`=1`, default) versus generic (`=0`) in a fresh process per setting. Outputs are +bit-identical between the two settings for every row (max abs diff 0 on every output, including the state and +`state_update`). + +`GatedDeltaNet` (`ORT_GDN_S2_SPECIALIZE`; S is tokens per request, S=8 is the production verification step): + +| Batch | S | QwenS2 (us) | Generic (us) | Speedup | +|---:|---:|---:|---:|---:| +| 1 | 1 | 4.83 | 6.05 | 1.25x | +| 1 | 2 | 6.82 | 8.16 | 1.20x | +| 1 | 4 | 10.94 | 12.74 | 1.16x | +| 1 | 8 | 18.82 | 21.60 | 1.15x | +| 2 | 8 | 40.74 | 40.74 | 1.00x (generic kernel; not eligible) | +| 4 | 8 | 74.88 | 74.88 | 1.00x (generic kernel; not eligible) | +| 8 | 8 | 142.24 | 142.24 | 1.00x (generic kernel; not eligible) | + +`CausalConvWithState` (`ORT_CAUSAL_CONV_S2_SPECIALIZE`): + +| Batch | S | Specialized (us) | Generic (us) | Speedup | +|---:|---:|---:|---:|---:| +| 1 | 1 | 1.76 | 2.21 | 1.25x | +| 1 | 2 | 2.02 | 2.46 | 1.22x | +| 1 | 4 | 2.30 | 2.75 | 1.19x | +| 1 | 8 | 2.46 | 2.98 | 1.21x | +| 2 | 8 | 2.94 | 3.49 | 1.18x | +| 4 | 8 | 4.13 | 5.60 | 1.36x | +| 8 | 8 | 6.78 | 9.28 | 1.37x | +| 1 | 16 (tiled) | 3.90 | 3.36 | 0.86x | +| 1 | 2048 (tiled) | 55.13 | 223.97 | 4.06x | + +QwenS2 trims 1.2-2.8 us (13-20%) from each single-request decode call; at S=8 that is 2.8 us x 36 GDN layers, or +roughly 0.1 ms per verification cycle; the S<=8 convolution kernel saves about 0.5 us per call at batch 1 and +more at larger batch. The tiled convolution kernel is the main win for prefill (4x at 2048 tokens), but at the +smallest tiled length (S=16) it is slightly slower than the generic kernel (0.86x, +0.5 us), an absolute cost that is +negligible next to the S>=2048 gain. + +End to end on the rc3 model (INT4 `all_int4`, decode ms per DFlash2 S=8 cycle, H200, median of 5 repetitions): + +| Context | Default | `ORT_GDN_S2_SPECIALIZE=0` | +|---:|---:|---:| +| 2K | 14.93 | 15.03 | +| 8K | 15.07 | 15.17 | +| 32K | 15.30 | 15.45 | + +That is +0.10 to +0.15 ms per cycle without QwenS2, i.e. QwenS2 decodes about 0.7-1.0% faster. +`ORT_CAUSAL_CONV_S2_SPECIALIZE=0` leaves decode unchanged but raises time to first token by 2.9-3.1% (2K / 8K / 32K +prompts: 0.1987 -> 0.2044 s, 0.8506 -> 0.8791 s, 3.5843 -> 3.6942 s). + +Methodology: each op is a single-node ONNX model with random BF16/FP32 inputs (fixed seeds), `state_update` +capture enabled (`capture_count = min(S, 7)`), 48 copies of the node (distinct state tensors, like the 48 layers) in +one CUDA-graph-captured session with IOBinding and device-resident inputs; 30 warm-up replays and 300 timed replays +(14,400 kernel instances per row). Kernel durations come from Nsight Systems (`--cuda-graph-trace=node`), so host +launch and inter-kernel gaps are excluded; the table gives the median of the three per-process medians (spread across +processes <= 0.15 us, every row run on the same GPU as its counterpart). Wall-clock back-to-back graph replay +(slope between 16- and 48-node graphs) agrees for the large kernels (conv S=2048: 54 vs 224 us; GDN S=8: 20.5 vs +23.6 us, which includes about 1.8 us of per-node launch gap) but is noisy for kernels of a few microseconds, so it is +not used here. Run `ORT_GDN_S2_SPECIALIZE=0/1` and `ORT_CAUSAL_CONV_S2_SPECIALIZE=0/1` in separate processes because +the switches are cached on first use. + This document describes the `com.microsoft::GatedDeltaNet` contrib operator: its schema and state contract, the chunked and warp-specialised CUDA engines and how one is selected, the consumer-Blackwell shared-memory constraint, and measured comparisons against