From 62dbc36ff46e2fc9d025f653e7f124b1ede6bbbd Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 17:56:51 +0000 Subject: [PATCH 1/2] Extend GDN and causal convolution fused state handling Part of the grouped generic CUDA operator stack. Preserve the qualified source snapshot. --- .../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 | 72 +++++-- .../contrib_ops/cuda/bert/gated_delta_net.h | 3 +- .../cuda/bert/gated_delta_net_impl.cu | 119 +++++++---- .../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 | 202 +++++++++++++----- .../causal_conv_with_state_op_test.cc | 117 ++++++++++ .../contrib_ops/gated_delta_net_op_test.cc | 165 +++++++++++++- 17 files changed, 725 insertions(+), 120 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 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..c5284935a5e04 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..f346bd3f6027d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -32,6 +32,9 @@ namespace gdn = gated_delta_net; (*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), \ @@ -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..f308442329c5d 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,18 +733,18 @@ 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, +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{}; @@ -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,16 +799,28 @@ 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, + 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, + 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&, float, int, size_t, cudaStream_t); 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..2bd82be0b53bf 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 f3e2fa310427e..e3eb187d2966b 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -4620,6 +4620,17 @@ 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'.", @@ -4663,7 +4674,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", @@ -4681,12 +4693,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", @@ -4695,8 +4709,12 @@ 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.") @@ -4705,6 +4723,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); @@ -5270,7 +5345,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) @@ -5278,6 +5353,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"; @@ -5292,8 +5370,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 " @@ -5313,6 +5392,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) @@ -5324,20 +5408,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 " @@ -5355,14 +5440,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); @@ -5454,45 +5540,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 eeb48a0579b58..8e015e80d5b32 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,123 @@ 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 4db622e45ad2b..72ad942d6eaa0 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 { @@ -340,6 +342,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); @@ -409,12 +415,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(); } @@ -524,12 +538,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; @@ -1129,6 +1251,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 c64cce886b05360f5023911b0a207924b9b4d381 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 19 Oct 2026 18:10:29 +0000 Subject: [PATCH 2/2] Fix lint --- .../cuda/bert/branchwise_rms_norm_impl.cu | 4 +- .../cuda/bert/causal_conv_with_state_impl.cu | 4 +- .../contrib_ops/cuda/bert/fused_tensor_ops.cc | 46 ++-- .../contrib_ops/cuda/bert/fused_tensor_ops.cu | 18 +- .../contrib_ops/cuda/bert/gated_delta_net.cc | 28 +- .../cuda/bert/gated_delta_net_impl.cu | 8 +- .../cuda/bert/gated_delta_net_impl.h | 2 +- .../cuda/bert/rotary_embedding_impl.cu | 9 +- .../bert/packed_sparse_attention_indexer.cc | 6 +- .../webgpu/bert/sparse_attention_indexer.h | 8 +- .../core/graph/contrib_ops/bert_defs.cc | 242 +++++++++--------- .../cuda/reduction/reduction_functions.cu | 2 +- .../cuda/reduction/reduction_functions.h | 2 +- .../causal_conv_with_state_op_test.cc | 14 +- .../test/contrib_ops/fused_tensor_ops_test.cc | 15 +- .../contrib_ops/hyper_connection_op_test.cc | 10 +- .../contrib_ops/rotary_embedding_op_test.cc | 10 +- .../sparse_attention_indexer_op_test.cc | 5 +- ...untime_test_python_symbolic_shape_infer.py | 3 - 19 files changed, 228 insertions(+), 208 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu index c07197fa3bf14..240b057b827e0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu @@ -168,9 +168,9 @@ INSTANTIATE(__nv_bfloat16) #undef INSTANTIATE -#define INSTANTIATE_STANDARD(T) \ +#define INSTANTIATE_STANDARD(T) \ template Status LaunchStandardBranchwiseRMSNorm(cudaStream_t, const T*, const void*, int, T*, \ - int, int, int, bool, float, float*, float*, float*, void*, size_t, int); + int, int, int, bool, float, float*, float*, float*, void*, size_t, int); INSTANTIATE_STANDARD(float) INSTANTIATE_STANDARD(half) INSTANTIATE_STANDARD(__nv_bfloat16) 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 c5284935a5e04..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 @@ -480,8 +480,8 @@ Status LaunchCausalConvProjectionKernel( return CUDA_CALL(cudaGetLastError()); } -#define INSTANTIATE_PROJECTION(T) \ - template Status LaunchCausalConvProjectionKernel( \ +#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) diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc index 24d2a153bf450..555c98f1fb16d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cc @@ -13,12 +13,12 @@ namespace onnxruntime::contrib::cuda { -#define REGISTER_ARGMAX_LAST_AXIS(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - ArgMaxLastAxis, kMSDomain, 1, T, kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("I", DataTypeImpl::GetTensorType()), \ +#define REGISTER_ARGMAX_LAST_AXIS(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + ArgMaxLastAxis, kMSDomain, 1, T, kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("I", DataTypeImpl::GetTensorType()), \ ArgMaxLastAxis); REGISTER_ARGMAX_LAST_AXIS(float) @@ -30,38 +30,38 @@ REGISTER_ARGMAX_LAST_AXIS(BFloat16) ONNX_OPERATOR_KERNEL_EX( EmbeddingRMSNorm, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("I", DataTypeImpl::GetTensorType()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()) - .TypeConstraint("F", DataTypeImpl::GetTensorType()), + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()) + .TypeConstraint("F", DataTypeImpl::GetTensorType()), EmbeddingRMSNorm); ONNX_OPERATOR_KERNEL_EX( RowSelect, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("I", DataTypeImpl::GetTensorType()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), + .TypeConstraint("I", DataTypeImpl::GetTensorType()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), RowSelect); ONNX_OPERATOR_KERNEL_EX( SigmoidGateMul, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), + .TypeConstraint("T", DataTypeImpl::GetTensorType()), SigmoidGateMul); #if !defined(DISABLE_FLOAT8_TYPES) ONNX_OPERATOR_TYPED_KERNEL_EX( DequantizeLinear, kMSDomain, 1, Float8E4M3FN_BFloat16, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T1", DataTypeImpl::GetTensorType()) - .TypeConstraint("T2", DataTypeImpl::GetTensorType()) - .TypeConstraint("TS", DataTypeImpl::GetTensorType()), + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) + .TypeConstraint("TS", DataTypeImpl::GetTensorType()), DequantizeLinearFp8Bf16); #endif ONNX_OPERATOR_KERNEL_EX( UntileHyperConnectionWeights, kMSDomain, 1, kCudaExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", DataTypeImpl::GetTensorType()), + .TypeConstraint("T", DataTypeImpl::GetTensorType()), UntileHyperConnectionWeights); template @@ -153,12 +153,10 @@ SigmoidGateMul::SigmoidGateMul(const OpKernelInfo& info) const int64_t silu = info.GetAttrOrDefault("silu_gate", 0); ORT_ENFORCE(silu == 0 || silu == 1, "SigmoidGateMul silu_gate must be zero or one"); ORT_ENFORCE(!(broadcast_gate_ && silu_gate_), "SigmoidGateMul broadcast_gate and silu_gate are exclusive"); - ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 && - q_stride_ >= 2 * head_dim_ && - tok_stride_ > 0 && q_offset_ >= 0), + ORT_ENFORCE(broadcast_gate_ || silu_gate_ || (head_dim_ > 0 && heads_ > 0 && q_stride_ >= 2 * head_dim_ && tok_stride_ > 0 && q_offset_ >= 0), "SigmoidGateMul received invalid layout attributes"); ORT_ENFORCE(!(broadcast_gate_ || silu_gate_) || (head_dim_ == 0 && heads_ == 0 && q_stride_ == 0 && - tok_stride_ == 0 && q_offset_ == 0), + tok_stride_ == 0 && q_offset_ == 0), "SigmoidGateMul broadcast_gate/silu_gate cannot be combined with packed layout attributes"); } @@ -184,12 +182,12 @@ Status SigmoidGateMul::ComputeInternal(OpKernelContext* context) const { } else { const int64_t row_width = static_cast(heads_) * head_dim_; ORT_RETURN_IF_NOT(count % row_width == 0, - "SigmoidGateMul attention size must be divisible by heads * head_dim"); + "SigmoidGateMul attention size must be divisible by heads * head_dim"); const int64_t tokens = count / row_width; ORT_RETURN_IF_NOT(projection->Shape().Size() >= - (tokens - 1) * tok_stride_ + q_offset_ + - static_cast(heads_) * q_stride_, - "SigmoidGateMul projection is too small for the configured layout"); + (tokens - 1) * tok_stride_ + q_offset_ + + static_cast(heads_) * q_stride_, + "SigmoidGateMul projection is too small for the configured layout"); } Tensor* output = context->Output(0, attention->Shape()); return LaunchSigmoidGateMul( diff --git a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu index 6db5544942e1f..789edc4a26615 100644 --- a/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu +++ b/onnxruntime/contrib_ops/cuda/bert/fused_tensor_ops.cu @@ -169,8 +169,8 @@ __global__ void SigmoidGateMulKernel(const __nv_bfloat16* attention, const BFloat16 a(gate); // Match CUDA Sigmoid's input-type intermediates, including its negative-input subtraction. const BFloat16 value = a > BFloat16(0.0f) - ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) - : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a)); + ? BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(-a)) + : BFloat16(1.0f) - BFloat16(1.0f) / (BFloat16(1.0f) + _Exp(a)); sigmoid_bf16 = __float2bfloat16_rn(static_cast(value)); } if constexpr (SiluGate) { @@ -254,13 +254,13 @@ Status LaunchArgMaxLastAxis(cudaStream_t stream, const void* input, int input_ty switch (input_type) { case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); + partial_values, partial_indices); case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); + partial_values, partial_indices); case ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16: return LaunchTypedArgmax(stream, static_cast(input), rows, width, output, - partial_values, partial_indices); + partial_values, partial_indices); default: return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "ArgMaxLastAxis input type is not supported"); } @@ -299,11 +299,11 @@ Status LaunchSigmoidGateMul(cudaStream_t stream, const uint16_t* attention, const int64_t block_count = (count - 1) / kThreads + 1; ORT_RETURN_IF_NOT(block_count <= std::numeric_limits::max(), "SigmoidGateMul requires too many CUDA blocks"); -#define LAUNCH_GATE(BROADCAST, RESIDUAL, SILU) \ +#define LAUNCH_GATE(BROADCAST, RESIDUAL, SILU) \ SigmoidGateMulKernel<<(block_count), kThreads, 0, stream>>>( \ - reinterpret_cast(attention), \ - reinterpret_cast(projection), \ - reinterpret_cast<__nv_bfloat16*>(output), count, head_dim, heads, q_stride, \ + reinterpret_cast(attention), \ + reinterpret_cast(projection), \ + reinterpret_cast<__nv_bfloat16*>(output), count, head_dim, heads, q_stride, \ token_stride, q_offset, gate_row_width, reinterpret_cast(residual), sigmoid_in_fp32) if (silu_gate) { if (residual != nullptr) { diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index f346bd3f6027d..380e99746af31 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -22,22 +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("TG", {DataTypeImpl::GetTensorType(), \ +#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), \ + .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ + .InputMemoryType(OrtMemTypeCPUInput, 10) \ + .MayInplace(6, 1), \ GatedDeltaNet); REGISTER_KERNEL_TYPED(float) 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 f308442329c5d..048faa9c2f3f5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -745,8 +745,8 @@ Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const Vari 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) { + 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; @@ -800,13 +800,13 @@ static Status LaunchGatedDeltaNetTypedGates(const Descriptor& desc, const Plan& // 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); + 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); + stream); } template 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 2bd82be0b53bf..3035ca02fa083 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h @@ -40,7 +40,7 @@ struct VariantPack { 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]; + : static_cast(data)[index]; } #endif diff --git a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu index 9f162783b3a30..55648cc944249 100644 --- a/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu @@ -189,8 +189,7 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu const int position_ids_format, const bool interleaved, const int max_threads_per_block, int4 in_strides, int4 out_strides, // strides in bnsh coord - const bool use_fp32_arithmetic -) { + const bool use_fp32_arithmetic) { // Note: Current implementation assumes head_size <= max_threads_per_block // because head_size is currently large for LLaMA-2. For smaller head_size // and num_heads values, we can create a block as `block(num_heads, head_size, 1)` @@ -221,8 +220,10 @@ Status LaunchRotaryEmbeddingKernel(cudaStream_t stream, T* output, const T* inpu interleaved, in_strides, out_strides); } }; - if (use_fp32_arithmetic) launch(std::true_type{}); - else launch(std::false_type{}); + if (use_fp32_arithmetic) + launch(std::true_type{}); + else + launch(std::false_type{}); return CUDA_CALL(cudaGetLastError()); } diff --git a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc index 93850dffc6d74..e25c593a149df 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/packed_sparse_attention_indexer.cc @@ -1158,9 +1158,9 @@ Status PackedSparseAttentionIndexer::ComputeQsa(onnxruntime::webgpu::ComputeCont ORT_RETURN_IF_ERROR(CheckStorageBindingLimit( context, "QSA selection", {&query_metadata, present_key_state, &rotary_cache, &sequence_metadata, present_state_lengths, &overflow_flags, - selected_indices, selected_counts, position_ids})); - PackedSparseAttentionIndexerQsaSelectProgram select{rotary.batched, position_ids != nullptr}; - select.CacheHint(rotary.batched, position_ids != nullptr) + selected_indices, selected_counts, position_ids})); + PackedSparseAttentionIndexerQsaSelectProgram select{rotary.batched, position_ids != nullptr}; + select.CacheHint(rotary.batched, position_ids != nullptr) .SetWorkgroupSize(kWorkgroupSize) .AddInputs({{&query_metadata, ProgramTensorMetadataDependency::Type}, {present_key_state, ProgramTensorMetadataDependency::Type}, diff --git a/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h b/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h index f34579543a4cb..0fbfd764eb276 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h +++ b/onnxruntime/contrib_ops/webgpu/bert/sparse_attention_indexer.h @@ -32,9 +32,9 @@ class SparseAttentionIndexerQsaConcatProgram final {"sequence_length", ProgramUniformVariableDataType::Uint32}, {"past_sequence_length", ProgramUniformVariableDataType::Uint32}, {"total_sequence_length", ProgramUniformVariableDataType::Uint32}, - {"head_size", ProgramUniformVariableDataType::Uint32}, - {"key_row_stride", ProgramUniformVariableDataType::Uint32}, - {"key_offset", ProgramUniformVariableDataType::Uint32}); + {"head_size", ProgramUniformVariableDataType::Uint32}, + {"key_row_stride", ProgramUniformVariableDataType::Uint32}, + {"key_offset", ProgramUniformVariableDataType::Uint32}); private: bool has_past_; @@ -52,7 +52,7 @@ class SparseAttentionIndexerQsaSelectProgram final {"sequence_length", ProgramUniformVariableDataType::Uint32}, {"num_heads", ProgramUniformVariableDataType::Uint32}, {"head_size", ProgramUniformVariableDataType::Uint32}, - {"query_row_stride", ProgramUniformVariableDataType::Uint32}, + {"query_row_stride", ProgramUniformVariableDataType::Uint32}, {"rotary_width", ProgramUniformVariableDataType::Uint32}, {"max_rotary_length", ProgramUniformVariableDataType::Uint32}, {"rotary_cache_batch_stride", ProgramUniformVariableDataType::Uint32}, diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index e3eb187d2966b..8cb045352889b 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -3408,55 +3408,55 @@ that are multiplied to query and key before the inner product of query and key i )DOC"; static OpSchema RotaryEmbeddingSchema() { return OpSchema() - .SetDoc(RotaryEmbedding_ver1_doc) - .Attr("scale", - "Custom scale will be used if specified. Default value is 1.0", - AttributeProto::FLOAT, - OPTIONAL_VALUE) - .Attr("interleaved", - "Indicates whether the input has real and imaginary parts interleaved. " - "Default value is 0 (False), meaning the first half of the input consists of real values " - "and the second half consists of imaginary values.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("rotary_embedding_dim", - "Rotary embedding dimension. Default value is 0.", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("num_heads", - "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", - AttributeProto::INT, - OPTIONAL_VALUE) - .Attr("is_packed_batching", - "ragged batch inputs or not. Default value is 0", - AttributeProto::INT, - OPTIONAL_VALUE) - .Input(0, - "input", - "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", - "T") - .Input(1, - "position_ids", - "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", - "M") - .Input(2, - "cos_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Input(3, - "sin_cache", - "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", - "T") - .Output(0, - "output", - "tensor with same shape as input.", - "T") - .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") - .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - propagateShapeFromInputToOutput(ctx, 0, 0); - }); + .SetDoc(RotaryEmbedding_ver1_doc) + .Attr("scale", + "Custom scale will be used if specified. Default value is 1.0", + AttributeProto::FLOAT, + OPTIONAL_VALUE) + .Attr("interleaved", + "Indicates whether the input has real and imaginary parts interleaved. " + "Default value is 0 (False), meaning the first half of the input consists of real values " + "and the second half consists of imaginary values.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("rotary_embedding_dim", + "Rotary embedding dimension. Default value is 0.", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("num_heads", + "Number of attention heads. Default value is 0. Must use with rotary_embedding_dim", + AttributeProto::INT, + OPTIONAL_VALUE) + .Attr("is_packed_batching", + "ragged batch inputs or not. Default value is 0", + AttributeProto::INT, + OPTIONAL_VALUE) + .Input(0, + "input", + "3D tensor with shape (batch_size, sequence_length, hidden_size) or 4D with shape (batch_size, num_heads, sequence_length, head_size)", + "T") + .Input(1, + "position_ids", + "1D tensor with shape (1) or 2D tensor with shape (batch_size, sequence_length)", + "M") + .Input(2, + "cos_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Input(3, + "sin_cache", + "2D tensor with shape (max_sequence_length, head_size / 2) or (max_sequence_length, rotary_embedding_dim / 2)", + "T") + .Output(0, + "output", + "tensor with same shape as input.", + "T") + .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") + .TypeConstraint("M", {"tensor(int64)"}, "Constrain input and output types to integer tensors") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + propagateShapeFromInputToOutput(ctx, 0, 0); + }); } ONNX_MS_OPERATOR_SET_SCHEMA(RotaryEmbedding, 1, RotaryEmbeddingSchema()); @@ -4620,15 +4620,18 @@ 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. " + .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. " + .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, " + .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", @@ -4712,9 +4715,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 " + .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) + "and state_window=0.", + "T", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") @@ -5694,20 +5699,26 @@ ONNX_MS_OPERATOR_SET_SCHEMA( .Attr("q_stride", "Per-head stride in the packed projection.", AttributeProto::INT, static_cast(0)) .Attr("tok_stride", "Per-token stride in the packed projection.", AttributeProto::INT, static_cast(0)) .Attr("q_offset", "Offset of the query projection within each token.", AttributeProto::INT, static_cast(0)) - .Attr("broadcast_gate", "When one, projection has the attention shape with final dimension one. " - "Packed layout attributes must be zero. CUDA only.", AttributeProto::INT, static_cast(0)) - .Attr("sigmoid_in_fp32", "One computes sigmoid in FP32 before BF16 conversion. Zero preserves " + .Attr("broadcast_gate", + "When one, projection has the attention shape with final dimension one. " + "Packed layout attributes must be zero. CUDA only.", + AttributeProto::INT, static_cast(0)) + .Attr("sigmoid_in_fp32", + "One computes sigmoid in FP32 before BF16 conversion. Zero preserves " "the input-type intermediate arithmetic of CUDA Sigmoid, including subtraction on negative inputs.", AttributeProto::INT, static_cast(1)) - .Attr("silu_gate", "When one, projection and attention have identical shapes. Multiply the " + .Attr("silu_gate", + "When one, projection and attention have identical shapes. Multiply the " "BF16-rounded sigmoid by projection and round to BF16 before multiplying attention. " "Cannot be combined with broadcast_gate or packed layout attributes.", AttributeProto::INT, static_cast(0)) .Input(0, "attention", "BF16 attention values.", "T") .Input(1, "projection", "BF16 packed projection containing gate values.", "T") .Input(2, "residual", "Optional BF16 residual with the attention shape.", "T", OpSchema::Optional) - .Output(0, "output", "Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 " - "before optionally adding residual and rounding again.", "T") + .Output(0, "output", + "Attention multiplied by BF16-rounded sigmoid gates, rounded to BF16 " + "before optionally adding residual and rounding again.", + "T") .TypeConstraint("T", {"tensor(bfloat16)"}, "Inputs and output use BF16.") .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { propagateElemTypeFromInputToOutput(ctx, 0, 0); @@ -5878,64 +5889,66 @@ WebGPU currently supports only the default mode without Y_fp32. static OpSchema BranchwiseRMSNormSchema() { return OpSchema() - .SetDoc(BranchwiseRMSNorm_ver1_doc) - .Attr("epsilon", "Epsilon added before reciprocal square root.", - AttributeProto::FLOAT, 1e-5f) - .Attr("num_branches", - "Number of branches for flattened input. Omit or set to zero for grouped input.", - AttributeProto::INT, static_cast(0)) - .Attr("use_onnx_arithmetic", - "Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. " - "Supported values are 0 (default) and 1.", - AttributeProto::INT, static_cast(0)) - .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") - .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", - "M", OpSchema::Optional) - .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") - .Output(1, "Y_fp32", "Optional float32 widening of the already T-rounded Y. " - "Requires use_onnx_arithmetic=1.", "F", OpSchema::Optional) - .TypeConstraint("T", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain input and output to floating-point tensors.") - .TypeConstraint("M", - {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, - "Constrain scale to floating-point tensors.") - .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.") - .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { - propagateElemTypeFromInputToOutput(ctx, 0, 0); - const auto start = getAttribute(ctx, "input_start", 0); - const auto width = getAttribute(ctx, "input_width", 0); - if (start < 0 || width < 0 || (width == 0 && start != 0)) { - fail_shape_inference("input_start and input_width must describe a nonnegative slice"); + .SetDoc(BranchwiseRMSNorm_ver1_doc) + .Attr("epsilon", "Epsilon added before reciprocal square root.", + AttributeProto::FLOAT, 1e-5f) + .Attr("num_branches", + "Number of branches for flattened input. Omit or set to zero for grouped input.", + AttributeProto::INT, static_cast(0)) + .Attr("use_onnx_arithmetic", + "Use explicit FP32 mean/sqrt/divide/multiply arithmetic instead of reciprocal RMS. " + "Supported values are 0 (default) and 1.", + AttributeProto::INT, static_cast(0)) + .Input(0, "X", "Grouped (..., C, H) or flattened (..., C * H) input.", "T") + .Input(1, "scale", "Optional scale with shape (C * H), (C, H), or (H).", + "M", OpSchema::Optional) + .Output(0, "Y", "RMS-normalized output with the same shape as X.", "T") + .Output(1, "Y_fp32", + "Optional float32 widening of the already T-rounded Y. " + "Requires use_onnx_arithmetic=1.", + "F", OpSchema::Optional) + .TypeConstraint("T", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain input and output to floating-point tensors.") + .TypeConstraint("M", + {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, + "Constrain scale to floating-point tensors.") + .TypeConstraint("F", {"tensor(float)"}, "Widened output is float32.") + .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) { + propagateElemTypeFromInputToOutput(ctx, 0, 0); + const auto start = getAttribute(ctx, "input_start", 0); + const auto width = getAttribute(ctx, "input_width", 0); + if (start < 0 || width < 0 || (width == 0 && start != 0)) { + fail_shape_inference("input_start and input_width must describe a nonnegative slice"); + } + if (width != 0 && (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1 || + getAttribute(ctx, "num_branches", 0) != 0)) { + fail_shape_inference("input_width requires use_onnx_arithmetic=1 and grouped input"); + } + if (hasInputShape(ctx, 0)) { + auto shape = getInputShape(ctx, 0); + if (width != 0) { + if (shape.dim_size() < 2) fail_shape_inference("grouped input requires rank at least two"); + auto* last = shape.mutable_dim(shape.dim_size() - 1); + if (last->has_dim_value() && (start > last->dim_value() || width > last->dim_value() - start)) { + fail_shape_inference("input slice exceeds the last dimension"); + } + last->set_dim_value(width); } - if (width != 0 && (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1 || - getAttribute(ctx, "num_branches", 0) != 0)) { - fail_shape_inference("input_width requires use_onnx_arithmetic=1 and grouped input"); + updateOutputShape(ctx, 0, shape); + } + if (ctx.getNumOutputs() > 1) { + if (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1) { + fail_shape_inference("Y_fp32 requires use_onnx_arithmetic=1"); } + updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); if (hasInputShape(ctx, 0)) { auto shape = getInputShape(ctx, 0); - if (width != 0) { - if (shape.dim_size() < 2) fail_shape_inference("grouped input requires rank at least two"); - auto* last = shape.mutable_dim(shape.dim_size() - 1); - if (last->has_dim_value() && (start > last->dim_value() || width > last->dim_value() - start)) { - fail_shape_inference("input slice exceeds the last dimension"); - } - last->set_dim_value(width); - } - updateOutputShape(ctx, 0, shape); + if (width != 0) shape.mutable_dim(shape.dim_size() - 1)->set_dim_value(width); + updateOutputShape(ctx, 1, shape); } - if (ctx.getNumOutputs() > 1) { - if (getAttribute(ctx, "use_onnx_arithmetic", 0) != 1) { - fail_shape_inference("Y_fp32 requires use_onnx_arithmetic=1"); - } - updateOutputElemType(ctx, 1, ONNX_NAMESPACE::TensorProto_DataType_FLOAT); - if (hasInputShape(ctx, 0)) { - auto shape = getInputShape(ctx, 0); - if (width != 0) shape.mutable_dim(shape.dim_size() - 1)->set_dim_value(width); - updateOutputShape(ctx, 1, shape); - } - } - }); + } + }); } ONNX_MS_OPERATOR_SET_SCHEMA(BranchwiseRMSNorm, 1, BranchwiseRMSNormSchema()); @@ -5944,7 +5957,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA( BranchwiseRMSNormSchema() .Attr("input_start", "CUDA-only starting offset in the last input dimension.", AttributeProto::INT, static_cast(0)) - .Attr("input_width", "CUDA-only width selected from the last input dimension before normalization. " + .Attr("input_width", + "CUDA-only width selected from the last input dimension before normalization. " "Zero retains the full input. Requires grouped input and use_onnx_arithmetic=1. " "Outputs retain the leading dimensions and use this last dimension.", AttributeProto::INT, static_cast(0))); diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index bf960242503d4..75024a420d3a9 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -539,7 +539,7 @@ struct RoundedSquare { template Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, - int m, int n, void* buffer, size_t buffer_size, int input_row_stride) { + int m, int n, void* buffer, size_t buffer_size, int input_row_stride) { return detail::call_reduce_matrix_columns( stream, input, output, m, n, buffer, buffer_size, input_row_stride); } diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h index e5bb1b03820f0..9c877c0743b31 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h @@ -109,7 +109,7 @@ Status reduce_matrix_columns(cudaStream_t stream, const TIn* input, TOut* output /** Sum FP32-rounded squares across columns, preserving reduce_matrix_columns' summation order. */ template Status reduce_square_matrix_columns(cudaStream_t stream, const TIn* input, float* output, - int m, int n, void* buffer, size_t buffer_size, int input_row_stride = 0); + int m, int n, void* buffer, size_t buffer_size, int input_row_stride = 0); /** * Computes ReduceSum for an arbitrary set of axes. This is the general fallback 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 8e015e80d5b32..7cac5d1d53c16 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 @@ -365,12 +365,16 @@ void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool u 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(); + 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(); + : 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); diff --git a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc index 441a7bf62ccce..2173434746ce9 100644 --- a/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc +++ b/onnxruntime/test/contrib_ops/fused_tensor_ops_test.cc @@ -146,8 +146,8 @@ void RunSigmoidGateMulBroadcastResidualRounding(bool sigmoid_in_fp32, bool silu_ const BFloat16 denominator(1.0f + BFloat16(exp_value).ToFloat()); const BFloat16 reciprocal(1.0f / denominator.ToFloat()); const BFloat16 rounded_sigmoid = sigmoid_in_fp32 - ? BFloat16(sigmoid) - : (gate > 0 ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat())); + ? BFloat16(sigmoid) + : (gate > 0 ? reciprocal : BFloat16(1.0f - reciprocal.ToFloat())); const BFloat16 activation = silu_gate ? BFloat16(gate * rounded_sigmoid.ToFloat()) : rounded_sigmoid; const BFloat16 product(values[i].ToFloat() * activation.ToFloat()); expected[i] = residual ? BFloat16(product.ToFloat() + residuals[i].ToFloat()) : product; @@ -186,7 +186,7 @@ TEST(FusedTensorOpsTest, SigmoidGateMulRejectsBroadcastAndResidualShapes) { test.AddAttribute("broadcast_gate", 1); test.AddInput("attention", {2, 3}, ToBFloat16(std::vector(6, 1.0f))); test.AddInput("projection", bad_residual ? std::vector{2, 1} : std::vector{1, 2}, - ToBFloat16({1.0f, 1.0f})); + ToBFloat16({1.0f, 1.0f})); if (bad_residual) test.AddInput("residual", {1}, ToBFloat16({1.0f})); test.AddOutput("output", {2, 3}, ToBFloat16(std::vector(6, 0.0f))); test.Run(OpTester::ExpectResult::kExpectFailure, @@ -208,8 +208,8 @@ TEST(FusedTensorOpsTest, SigmoidGateMulBroadcastEmptyAndInvalidAttributes) { test.AddInput("residual", {0, 3}, {}); test.AddOutput("output", {0, 3}, {}); const std::vector errors = {"", "broadcast_gate must be zero or one", - "sigmoid_in_fp32 must be zero or one", - "cannot be combined with packed layout attributes"}; + "sigmoid_in_fp32 must be zero or one", + "cannot be combined with packed layout attributes"}; test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, errors[mode], {}, nullptr, &providers); } @@ -227,14 +227,13 @@ TEST(FusedTensorOpsTest, SigmoidGateMulSiluEmptyAndInvalidAttributes) { test.AddInput("projection", {0, mode == 4 ? 1 : 3}, {}); test.AddOutput("output", {0, 3}, {}); const std::vector errors = {"", "silu_gate must be zero or one", "are exclusive", - "cannot be combined with packed layout attributes", - "silu gate must have the same shape"}; + "cannot be combined with packed layout attributes", + "silu gate must have the same shape"}; test.Run(mode == 0 ? OpTester::ExpectResult::kExpectSuccess : OpTester::ExpectResult::kExpectFailure, errors[mode], {}, nullptr, &providers); } } - TEST(FusedTensorOpsTest, UntileHyperConnectionWeights) { constexpr int64_t kInputElements = 160ll * 20480; constexpr int64_t kVectorCount = kInputElements / 16; diff --git a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc index 7942e8fc22425..d7e1bc7acdd72 100644 --- a/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc +++ b/onnxruntime/test/contrib_ops/hyper_connection_op_test.cc @@ -113,7 +113,10 @@ class StandardRMSChainTester : public OpTester { public: StandardRMSChainTester(int output_type, float epsilon, int64_t start = 0, int64_t width = 0) : OpTester("StandardRMSChain", 17, kOnnxDomain, false), - output_type_(output_type), epsilon_(epsilon), start_(start), width_(width) {} + output_type_(output_type), + epsilon_(epsilon), + start_(start), + width_(width) {} void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, std::vector>&) override { @@ -132,7 +135,8 @@ class StandardRMSChainTester : public OpTester { source = value("selected"); graph.AddNode("slice", "Slice", "", {inputs[0], constant("starts", start_), constant("ends", start_ + width_), - constant("axes", -1)}, {source}); + constant("axes", -1)}, + {source}); } auto* x = value("float_x"); graph.AddNode("cast", "Cast", "", {source}, {x}) @@ -200,7 +204,7 @@ void TestStandardRMSArithmetic(int type, bool packed = false) { std::vector expected; for (OpTester* tester : {static_cast(&reference), &fused}) { const std::vector input_shape = rows == 48 ? std::vector{1, 2, 24, stride} - : std::vector{1, rows, stride}; + : std::vector{1, rows, stride}; auto output_shape = input_shape; output_shape.back() = width; tester->AddInput("X", input_shape, input); diff --git a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc index 9ebbc10c97bf2..a6957e18bc194 100644 --- a/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc +++ b/onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc @@ -27,7 +27,10 @@ class FloatRotaryChainTester : public OpTester { public: FloatRotaryChainTester(int type, int heads, int rotary_dim, bool interleaved) : OpTester("RotaryEmbedding", 1, kMSDomain, false), - type_(type), heads_(heads), rotary_dim_(rotary_dim), interleaved_(interleaved) {} + type_(type), + heads_(heads), + rotary_dim_(rotary_dim), + interleaved_(interleaved) {} void AddNodes(Graph& graph, std::vector& inputs, std::vector& outputs, std::vector>&) override { @@ -74,7 +77,7 @@ void TestFloatRotaryArithmetic(int type) { const int rotary_dim = 64; const int cache_rows = batch * sequence + 3; std::vector shape = transposed ? std::vector{batch, heads, sequence, width} - : std::vector{batch, sequence, heads * width}; + : std::vector{batch, sequence, heads * width}; std::vector input(batch * sequence * heads * width), cosine(cache_rows * rotary_dim / 2), sine(cosine.size()); const bool scalar_position = !transposed; @@ -94,8 +97,7 @@ void TestFloatRotaryArithmetic(int type) { std::vector expected; for (OpTester* tester : {static_cast(&reference), &fused}) { tester->AddInput("input", shape, input); - tester->AddInput("positions", scalar_position ? std::vector{1} - : std::vector{batch, sequence}, positions); + tester->AddInput("positions", scalar_position ? std::vector{1} : std::vector{batch, sequence}, positions); tester->AddInput("cosine", {cache_rows, rotary_dim / 2}, cosine); tester->AddInput("sine", {cache_rows, rotary_dim / 2}, sine); tester->AddOutput("output", shape, std::vector(input.size())); diff --git a/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc b/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc index 3cf1bb9e2ae93..0561c46571879 100644 --- a/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc +++ b/onnxruntime/test/contrib_ops/sparse_attention_indexer_op_test.cc @@ -386,7 +386,9 @@ void QsaReference(const QsaProblem& problem, std::vector& selected, std } else { pooled[static_cast(d)] += problem.key[(static_cast(b) * problem.sequence_length + - position - problem.past_sequence_length) * head_size + d]; + position - problem.past_sequence_length) * + head_size + + d]; } } } @@ -492,7 +494,6 @@ void QsaReference(const QsaProblem& problem, std::vector& selected, std } } } - } struct CsaProblem { diff --git a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py index 96bb35672a6cf..6ac581d72cab8 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py +++ b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py @@ -882,9 +882,6 @@ def infer(query_width, packed): self.assertEqual(separate_graph.output[0].type.tensor_type.shape.dim[2].dim_value, 33) self.assertFalse(separate_graph.output[1].type.tensor_type.HasField("shape")) - - - def test_unsqueeze_opset_11(self): graph = helper.make_graph( [