diff --git a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc index 7585993088805..bf81d935636fe 100644 --- a/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc @@ -41,6 +41,12 @@ REGISTER_KERNEL_TYPED(float) template CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : OpKernel(info) { + ORT_ENFORCE(info.GetAttrOrDefault("input_channels", 0) == 0 && + info.GetAttrOrDefault("input_channel_offset", 0) == 0 && + info.GetAttrOrDefault("state_length", 0) == 0 && + info.GetAttrOrDefault("state_channels_last", -1) == -1 && + info.node().OutputDefs().size() <= 2, + "CPU CausalConvWithState does not support extended channel/state options (CUDA only)"); int64_t ndim = info.GetAttrOrDefault("ndim", 1); ORT_ENFORCE(ndim == 1, "CPU CausalConvWithState only supports ndim=1"); ndim_ = static_cast(ndim); diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc index c79859f735c7f..f7a8d94c2b80a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc @@ -6,6 +6,7 @@ #include "contrib_ops/cpu/bert/causal_conv_with_state_helper.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cuda_type_conversion.h" +#include "core/common/safeint.h" namespace onnxruntime { namespace contrib { @@ -45,6 +46,13 @@ CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : CudaKern // See LinearAttention: only the trailing per-position states are ever consumed, so a window // caps the allocation and the write traffic for long prompts. 0 keeps the plain single state. ORT_THROW_IF_ERROR(causal_conv_with_state_helper::ParseStateWindow(info, state_window_)); + input_channels_ = SafeInt(info.GetAttrOrDefault("input_channels", 0)); + input_channel_offset_ = SafeInt(info.GetAttrOrDefault("input_channel_offset", 0)); + state_length_ = SafeInt(info.GetAttrOrDefault("state_length", 0)); + state_channels_last_ = SafeInt(info.GetAttrOrDefault("state_channels_last", -1)); + ORT_ENFORCE(input_channels_ >= 0 && input_channel_offset_ >= 0 && state_length_ >= 0 && + state_channels_last_ >= -1 && state_channels_last_ <= 1, + "invalid packed-channel or state attributes"); } template @@ -60,6 +68,47 @@ Status CausalConvWithState::ComputeInternal(OpKernelContext* context) const { const auto& input_shape = input_tensor->Shape(); const auto& weight_shape = weight_tensor->Shape(); + const bool extended = input_channels_ != 0 || input_channel_offset_ != 0 || state_length_ != 0 || + state_channels_last_ != -1 || context->OutputCount() > 2; + if (extended) { + ORT_RETURN_IF_NOT(channels_last_ && state_window_ == 0 && input_shape.NumDimensions() == 3 && + weight_shape.NumDimensions() == 3 && weight_shape[1] == 1, + "extended CausalConvWithState requires rank-3 channels-last input, rank-3 weight and state_window=0"); + const int batch = SafeInt(input_shape[0]); + const int length = SafeInt(input_shape[1]); + const int width = SafeInt(input_shape[2]); + const int channels = input_channels_ == 0 ? width : input_channels_; + const int kernel = SafeInt(weight_shape[2]); + const int pad = SafeInt(kernel - 1) * dilation_; + const int state_length = state_length_ == 0 ? pad : state_length_; + const bool state_last = state_channels_last_ == -1 ? channels_last_ : state_channels_last_ != 0; + ORT_RETURN_IF_NOT(batch > 0 && batch <= 65535 && length > 0 && length <= 65535 && + channels > 0 && kernel > 0 && input_channel_offset_ <= width && + channels <= width - input_channel_offset_ && weight_shape[0] == channels && + state_length >= pad, + "invalid extended CausalConvWithState dimensions"); + ORT_RETURN_IF_NOT(bias_tensor == nullptr || bias_tensor->Shape() == TensorShape({channels}), + "bias must match selected input channels"); + const TensorShape state_shape = state_last ? TensorShape({batch, state_length, channels}) + : TensorShape({batch, channels, state_length}); + ORT_RETURN_IF_NOT(past_state_tensor == nullptr || past_state_tensor->Shape() == state_shape, + "past_state must match configured state layout and length"); + Tensor* output = context->Output(0, TensorShape({batch, length, channels})); + Tensor* present = context->Output(1, state_shape); + Tensor* first = context->OutputCount() > 2 ? context->Output(2, state_shape) : nullptr; + using CudaT = typename OrtToCudaType::type; + return LaunchCausalConvProjectionKernel( + Stream(context), reinterpret_cast(input_tensor->Data()), + reinterpret_cast(weight_tensor->Data()), + bias_tensor ? reinterpret_cast(bias_tensor->Data()) : nullptr, + past_state_tensor ? reinterpret_cast(past_state_tensor->Data()) : nullptr, + reinterpret_cast(output->MutableData()), + reinterpret_cast(present->MutableData()), + first ? reinterpret_cast(first->MutableData()) : nullptr, + batch, channels, length, kernel, dilation_, width, input_channel_offset_, state_length, + state_last, activation_ != "none"); + } + // Validate input rank and weight rank if (channels_last_) { // (batch_size, sequence_length, d_1, ..., d_n): any number of trailing channel axes, so a diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h index 0a1f23f3abfd3..f0807d69c7da3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h @@ -25,6 +25,10 @@ class CausalConvWithState final : public onnxruntime::cuda::CudaKernel { std::string activation_; // Leading (axis-0) extent of past_state / present_state; 0 means no window axis (single state). int state_window_; + int input_channels_; + int input_channel_offset_; + int state_length_; + int state_channels_last_; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu index 0c0e1da836521..eea53e8bbe272 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu @@ -27,6 +27,44 @@ namespace cuda { namespace { +template +__global__ void CausalConvProjectionKernel( + const T* input, const T* weight, const T* bias, const T* past, + T* output, T* present, T* first_state, int channels, int length, + int kernel_size, int dilation, int input_width, int channel_offset, + int state_length, CausalConvLayout state_layout, bool apply_silu) { + const int b = blockIdx.y; + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= channels) return; + const int pad = (kernel_size - 1) * dilation; + const int position = blockIdx.z; + { + float sum = bias == nullptr ? 0.0f : to_float(bias[c]); + for (int tap = 0; tap < kernel_size; ++tap) { + const int source = position - pad + tap * dilation; + const float value = source >= 0 + ? to_float(input[(static_cast(b) * length + source) * input_width + channel_offset + c]) + : (past == nullptr ? 0.0f : to_float(past[state_layout.Offset(b, state_length + source, c)])); + sum += value * to_float(weight[c * kernel_size + tap]); + } + if (apply_silu) sum = sum / (1.0f + expf(-sum)); + output[(static_cast(b) * length + position) * channels + c] = from_float(sum); + } + if (position != 0) return; + for (int position = 0; position < state_length; ++position) { + const int source = length - state_length + position; + present[state_layout.Offset(b, position, c)] = + source >= 0 ? input[(static_cast(b) * length + source) * input_width + channel_offset + c] + : (past == nullptr ? from_float(0.0f) : past[state_layout.Offset(b, state_length + source, c)]); + if (first_state != nullptr) { + const int first_source = 1 - state_length + position; + first_state[state_layout.Offset(b, position, c)] = + first_source >= 0 ? input[static_cast(b) * length * input_width + channel_offset + c] + : (past == nullptr ? from_float(0.0f) : past[state_layout.Offset(b, state_length + first_source, c)]); + } + } +} + __device__ __forceinline__ float silu_fn(float x) { return x / (1.0f + expf(-x)); } @@ -428,6 +466,29 @@ __global__ void CausalConvPrefillKernelChannelsLast( } // anonymous namespace +template +Status LaunchCausalConvProjectionKernel( + cudaStream_t stream, const T* input, const T* weight, const T* bias, const T* past, + T* output, T* present, T* first_state, int batch, int channels, int length, + int kernel_size, int dilation, int input_width, int channel_offset, + int state_length, bool state_channels_last, bool apply_silu) { + const dim3 grid((channels - 1) / 128 + 1, batch, length); + CausalConvProjectionKernel<<>>( + input, weight, bias, past, output, present, first_state, channels, length, + kernel_size, dilation, input_width, channel_offset, state_length, + MakeCausalConvLayout(state_channels_last, channels, state_length), apply_silu); + return CUDA_CALL(cudaGetLastError()); +} + +#define INSTANTIATE_PROJECTION(T) \ + template Status LaunchCausalConvProjectionKernel( \ + cudaStream_t, const T*, const T*, const T*, const T*, T*, T*, T*, \ + int, int, int, int, int, int, int, int, bool, bool); +INSTANTIATE_PROJECTION(float) +INSTANTIATE_PROJECTION(half) +INSTANTIATE_PROJECTION(__nv_bfloat16) +#undef INSTANTIATE_PROJECTION + template Status LaunchCausalConvWithStateKernel( cudaStream_t stream, diff --git a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h index 968cb9507339d..95d98dc96cdcd 100644 --- a/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h @@ -31,6 +31,13 @@ inline CausalConvLayout MakeCausalConvLayout(bool channels_last, int64_t channel : CausalConvLayout{channels * length, 1, length}; } +template +Status LaunchCausalConvProjectionKernel( + cudaStream_t stream, const T* input, const T* weight, const T* bias, const T* past, + T* output, T* present, T* first_state, int batch, int channels, int length, + int kernel_size, int dilation, int input_width, int channel_offset, + int state_length, bool state_channels_last, bool apply_silu); + // Fused causal depthwise conv1d + activation + state management. // One thread block per (batch, channel). For decode (L=1), this is a simple // dot product from shared memory. For prefill (L>1), each thread handles diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc index ff7f485dc1284..380e99746af31 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.cc @@ -22,19 +22,22 @@ using onnxruntime::cuda::CudaKernel; using onnxruntime::cuda::OrtToCudaType; namespace gdn = gated_delta_net; -#define REGISTER_KERNEL_TYPED(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX( \ - GatedDeltaNet, \ - kMSDomain, \ - 1, \ - T, \ - kCudaExecutionProvider, \ - (*KernelDefBuilder::Create()) \ - .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("TS", DataTypeImpl::GetTensorType()) \ - .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ - .InputMemoryType(OrtMemTypeCPUInput, 10) \ - .MayInplace(6, 1), \ +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + GatedDeltaNet, \ + kMSDomain, \ + 1, \ + T, \ + kCudaExecutionProvider, \ + (*KernelDefBuilder::Create()) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TS", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("TG", {DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType(), \ + DataTypeImpl::GetTensorType()}) \ + .TypeConstraint("TI", DataTypeImpl::GetTensorType()) \ + .InputMemoryType(OrtMemTypeCPUInput, 10) \ + .MayInplace(6, 1), \ GatedDeltaNet); REGISTER_KERNEL_TYPED(float) @@ -42,7 +45,14 @@ REGISTER_KERNEL_TYPED(MLFloat16) REGISTER_KERNEL_TYPED(BFloat16) template -GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : CudaKernel(info) { +GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) + : CudaKernel(info), + mid_state_output_(info.GetAttrOrDefault("state_output", "state_update") == + "mid_state") { + ORT_ENFORCE(mid_state_output_ || + info.GetAttrOrDefault("state_output", "state_update") == + "state_update", + "state_output must be 'state_update' or 'mid_state'"); const std::string rule = info.GetAttrOrDefault("update_rule", "gated_delta"); if (rule == "linear") { update_rule_ = gdn::UpdateRule::kLinear; @@ -59,10 +69,10 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : CudaKernel(info) { const std::string gate = info.GetAttrOrDefault("gate_activation", "none"); if (gate == "none") { gate_activation_ = gdn::GateActivation::kNone; - } else if (gate == "qwen") { - gate_activation_ = gdn::GateActivation::kQwen; + } else if (gate == "softplus" || gate == "qwen") { + gate_activation_ = gdn::GateActivation::kSoftplusDecay; } else { - ORT_THROW("gate_activation must be one of: none, qwen"); + ORT_THROW("gate_activation must be one of: none, softplus"); } const std::string beta = info.GetAttrOrDefault("beta_activation", "none"); @@ -88,6 +98,9 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : CudaKernel(info) { forced_engine_ = gdn::EngineFromName( ParseEnvironmentVariableWithDefault("ORT_GDN_PLAN", "")); + if (mid_state_output_) { + forced_engine_ = gdn::Engine::kRecurrent; + } } template @@ -106,6 +119,15 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { const Tensor* capture_count = context->Input(9); const Tensor* state_update_active = context->Input(10); + ORT_RETURN_IF_NOT(query != nullptr, "query is required"); + for (const Tensor* gate : {decay, beta}) { + ORT_RETURN_IF_NOT(gate == nullptr || gate->IsDataType() || + gate->GetElementType() == query->GetElementType(), + "decay and beta must be float or match the query element type"); + } + ORT_RETURN_IF_NOT(!mid_state_output_ || state_update_capacity_ == 0, + "state_output='mid_state' cannot be combined with state_update_capacity"); + ORT_RETURN_IF_NOT((state_update_capacity_ > 0) == (capture_count != nullptr), "capture_count must be present exactly when state_update_capacity is positive"); @@ -120,7 +142,6 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(needs_beta == (beta != nullptr), "beta input presence must match update_rule"); - ORT_RETURN_IF_NOT(query != nullptr, "query is required"); const bool is_packed_qkv = key == nullptr && value == nullptr; ORT_RETURN_IF_NOT(is_packed_qkv || (key != nullptr && value != nullptr), "key and value must be both present or both absent"); @@ -191,6 +212,10 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { ORT_RETURN_IF_NOT(num_heads_q == num_heads_k, "num_heads_q (", num_heads_q, ") must equal num_heads_k (", num_heads_k, ")"); + if (mid_state_output_) { + ORT_RETURN_IF_NOT(num_heads_v >= num_heads_q, + "mid-state output requires value_heads >= query_heads"); + } ORT_RETURN_IF_NOT(num_heads_v > 0 && num_heads_q > 0 && num_heads_v % num_heads_q == 0, "num_heads_v (", num_heads_v, ") must be a positive multiple of num_heads_q (", num_heads_q, ")"); @@ -260,9 +285,9 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { b[token_dims] == num_heads_v, "beta must be [...tokens, num_heads_v]"); } - if (gate_activation_ == gdn::GateActivation::kQwen) { + if (gate_activation_ == gdn::GateActivation::kSoftplusDecay) { ORT_RETURN_IF_NOT(a_log != nullptr && dt_bias != nullptr, - "gate_activation=qwen requires a_log and dt_bias"); + "gate_activation=softplus requires a_log and dt_bias"); ORT_RETURN_IF_NOT(a_log->Shape().NumDimensions() == 1 && a_log->Shape()[0] == num_heads_v, "a_log must be [num_heads_v]"); ORT_RETURN_IF_NOT(dt_bias->Shape().NumDimensions() == 1 && @@ -270,7 +295,7 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { "dt_bias must be [num_heads_v]"); } else { ORT_RETURN_IF_NOT(a_log == nullptr && dt_bias == nullptr, - "a_log and dt_bias require gate_activation=qwen"); + "a_log and dt_bias require gate_activation=softplus"); } if (capture_count != nullptr) { @@ -293,10 +318,16 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { TensorShape state_shape{batch, num_heads_v, head_size_v, head_size_qk}; Tensor* final_state = context->Output(1, state_shape); - const int64_t state_update_width = static_cast(state_update_capacity_) * - (num_heads_v + num_heads_k * head_size_qk + - num_heads_v * head_size_v); - Tensor* state_update = context->Output(2, TensorShape{batch, state_update_width}); + Tensor* state_update = nullptr; + Tensor* mid_state = nullptr; + if (mid_state_output_) { + mid_state = context->Output(2, state_shape); + } else { + const int64_t state_update_width = static_cast(state_update_capacity_) * + (num_heads_v + num_heads_k * head_size_qk + + num_heads_v * head_size_v); + state_update = context->Output(2, TensorShape{batch, state_update_width}); + } gdn::Descriptor desc{}; desc.total_tokens = total_tokens; @@ -363,7 +394,9 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { const size_t query_row_bytes = static_cast(num_heads_q) * head_size_qk * sizeof(T); const size_t key_row_bytes = static_cast(num_heads_k) * head_size_qk * sizeof(T); const size_t value_row_bytes = static_cast(num_heads_v) * head_size_v * sizeof(T); - const size_t unpacked_qkv_bytes = is_packed_qkv + const bool direct_packed_decode = is_packed_qkv && plan.engine == gdn::Engine::kRecurrent && + plan.warp_specialized && total_tokens <= 2LL * batch; + const size_t unpacked_qkv_bytes = is_packed_qkv && !direct_packed_decode ? static_cast(total_tokens) * (query_row_bytes + key_row_bytes + value_row_bytes) : 0; @@ -377,7 +410,10 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { const CudaT* query_data = reinterpret_cast(query->Data()); const CudaT* key_data = key != nullptr ? reinterpret_cast(key->Data()) : nullptr; const CudaT* value_data = value != nullptr ? reinterpret_cast(value->Data()) : nullptr; - if (is_packed_qkv) { + if (direct_packed_decode) { + key_data = query_data + query_row_bytes / sizeof(T); + value_data = key_data + key_row_bytes / sizeof(T); + } else if (is_packed_qkv) { const size_t packed_row_bytes = query_row_bytes + key_row_bytes + value_row_bytes; auto* unpacked_query = workspace.get(); auto* unpacked_key = unpacked_query + static_cast(total_tokens) * query_row_bytes; @@ -402,15 +438,19 @@ Status GatedDeltaNet::ComputeInternal(OpKernelContext* context) const { pack.query = query_data; pack.key = key_data; pack.value = value_data; + pack.packed_qkv = direct_packed_decode; pack.cu_seqlens = cu_seqlens != nullptr ? cu_seqlens->Data() : nullptr; pack.capture_count = desc.state_update_active ? capture_count->Data() : nullptr; - pack.decay = decay != nullptr ? decay->Data() : nullptr; - pack.beta = beta != nullptr ? beta->Data() : nullptr; + pack.decay = decay != nullptr ? decay->DataRaw() : nullptr; + pack.beta = beta != nullptr ? beta->DataRaw() : nullptr; + const Tensor* gate = decay != nullptr ? decay : beta; + pack.gates_in_input_type = gate != nullptr && !gate->IsDataType(); pack.initial_state = initial_state != nullptr ? initial_state->Data() : nullptr; pack.a_log = a_log != nullptr ? a_log->Data() : nullptr; pack.dt_bias = dt_bias != nullptr ? dt_bias->Data() : nullptr; pack.output = reinterpret_cast(output->MutableData()); pack.final_state = final_state != nullptr ? final_state->MutableData() : nullptr; + pack.mid_state = mid_state != nullptr ? mid_state->MutableData() : nullptr; pack.state_update = state_update != nullptr ? state_update->MutableData() : nullptr; // Only the inactive case is cleared. While capture is active the kernels write just the diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h index 26243a7e60ed6..1982ae2de7da4 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net.h @@ -14,7 +14,7 @@ namespace contrib { namespace cuda { template -class GatedDeltaNet final : public onnxruntime::cuda::CudaKernel { +class GatedDeltaNet : public onnxruntime::cuda::CudaKernel { public: GatedDeltaNet(const OpKernelInfo& info); Status ComputeInternal(OpKernelContext* context) const override; @@ -28,6 +28,7 @@ class GatedDeltaNet final : public onnxruntime::cuda::CudaKernel { int chunk_size_; int state_update_capacity_; bool qk_l2_norm_; + bool mid_state_output_; }; } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu index 4729e31bd2e9d..048faa9c2f3f5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.cu @@ -172,7 +172,7 @@ __device__ __forceinline__ ChunkedSmem CarveChunked(char* raw) { // --------------------------------------------------------------------------- // Chunked engine // --------------------------------------------------------------------------- -template +template __global__ __launch_bounds__(kChunkedThreads) void GatedDeltaNetChunkedKernel( VariantPack pack, KernelParams p) { using L = Ld; @@ -237,10 +237,10 @@ __global__ __launch_bounds__(kChunkedThreads) void GatedDeltaNetChunkedKernel( const int64_t gi = tok * p.num_heads_v + hv; float g = 0.0f, be = 0.0f; if (tid < len) { - g = needs_decay ? EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv, + g = needs_decay ? EffectiveDecay(LoadGate(pack.decay, gi, InputTypeGates), pack.a_log, pack.dt_bias, hv, p.gate_activation) : 0.0f; - be = pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f; + be = pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, InputTypeGates), p.beta_activation) : 1.0f; } s.gc[tid] = g; s.beta[tid] = be; @@ -379,7 +379,7 @@ __device__ __forceinline__ float WarpReduceK(float v) { return v; } -template +template __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( VariantPack pack, KernelParams p, int d_v) { constexpr int kRegs = DK / 32; // state elements per lane @@ -401,9 +401,12 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( } const int capture_count = CaptureCount(pack, p, b, seq_len); - const int64_t q_stride = static_cast(p.num_heads_q) * DK; - const int64_t k_stride = static_cast(p.num_heads_k) * DK; + const int64_t packed_stride = (static_cast(p.num_heads_q) + p.num_heads_k) * DK + + static_cast(p.num_heads_v) * d_v; + const int64_t q_stride = PackedQkv ? packed_stride : static_cast(p.num_heads_q) * DK; + const int64_t k_stride = PackedQkv ? packed_stride : static_cast(p.num_heads_k) * DK; const int64_t v_stride = static_cast(p.num_heads_v) * d_v; + const int64_t v_input_stride = PackedQkv ? packed_stride : v_stride; const int64_t st_off = (static_cast(b) * p.num_heads_v + hv) * static_cast(d_v) * DK + static_cast(col) * DK; @@ -456,14 +459,14 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( float decay_scalar = 1.0f; if (needs_decay && !decay_per_key_dim) { - decay_scalar = __expf(EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv, + decay_scalar = __expf(EffectiveDecay(LoadGate(pack.decay, gi, InputTypeGates), pack.a_log, pack.dt_bias, hv, p.gate_activation)); } #pragma unroll for (int i = 0; i < kRegs; ++i) { if (needs_decay && decay_per_key_dim) { const int k = lane + i * 32; - const float raw = pack.decay[(tok * p.num_heads_v + hv) * DK + k]; + const float raw = LoadGate(pack.decay, (tok * p.num_heads_v + hv) * DK + k, InputTypeGates); s[i] *= __expf(EffectiveDecay(raw, pack.a_log, pack.dt_bias, hv, p.gate_activation)); } else { s[i] *= decay_scalar; @@ -478,8 +481,8 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( } const float beta_t = - pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f; - const float vv = static_cast(pack.value[tok * v_stride + hv * d_v + col]); + pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, InputTypeGates), p.beta_activation) : 1.0f; + const float vv = static_cast(pack.value[tok * v_input_stride + hv * d_v + col]); const float delta = beta_t * (vv - r); if (t < capture_count) { @@ -512,6 +515,12 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( pack.output[tok * v_stride + hv * d_v + col] = static_cast(p.scale * o); } + if (t == 0 && pack.mid_state != nullptr) { +#pragma unroll + for (int i = 0; i < kRegs; ++i) { + pack.mid_state[st_off + lane + i * 32] = s[i]; + } + } } if (pack.final_state != nullptr) { @@ -526,7 +535,7 @@ __global__ __launch_bounds__(32 * kWarps) void GatedDeltaNetDecodeWarpKernel( // Generic recurrent engine: any head geometry the warp kernel cannot take. // One CTA per (sequence, v-head); the [DK x DV] state lives in shared memory. // --------------------------------------------------------------------------- -template +template __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(VariantPack pack, KernelParams p, int d_k, int d_v) { @@ -585,7 +594,7 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant } if (needs_decay && decay_per_key_dim) { for (int i = tid; i < d_k; i += kThreads) { - const float raw = pack.decay[(tok * p.num_heads_v + hv) * d_k + i]; + const float raw = LoadGate(pack.decay, (tok * p.num_heads_v + hv) * d_k + i, InputTypeGates); gbuf[i] = __expf(EffectiveDecay(raw, pack.a_log, pack.dt_bias, hv, p.gate_activation)); } } @@ -612,11 +621,11 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant const float g_scalar = (needs_decay && !decay_per_key_dim) - ? __expf(EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv, + ? __expf(EffectiveDecay(LoadGate(pack.decay, gi, InputTypeGates), pack.a_log, pack.dt_bias, hv, p.gate_activation)) : 1.0f; const float beta_t = - pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f; + pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, InputTypeGates), p.beta_activation) : 1.0f; // Decay, then retrieval against the decayed state. for (int idx = tid; idx < d_k * d_v; idx += kThreads) { @@ -665,6 +674,12 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant for (int r = 0; r < d_k; ++r) acc += S[r * d_v + c] * qbuf[r]; pack.output[tok * v_stride + hv * d_v + c] = (T)(p.scale * acc); } + if (t == 0 && pack.mid_state != nullptr) { + for (int idx = tid; idx < d_k * d_v; idx += kThreads) { + const int r = idx / d_v, c = idx % d_v; + pack.mid_state[st_off + static_cast(c) * d_k + r] = S[idx]; + } + } __syncthreads(); } @@ -680,27 +695,37 @@ __global__ __launch_bounds__(kThreads) void GatedDeltaNetRecurrentKernel(Variant constexpr int kRecurrentThreads = 256; constexpr int kDecodeWarps = kRecurrentThreads / 32; -template +template +void LaunchDecodeWarp(const Descriptor& desc, const VariantPack& pack, + const KernelParams& p, cudaStream_t stream) { + const dim3 grid(desc.batch, desc.num_heads_v, + (desc.head_size_v + kDecodeWarps - 1) / kDecodeWarps); + const dim3 block(32, kDecodeWarps, 1); + switch (desc.head_size_qk) { + case 64: + GatedDeltaNetDecodeWarpKernel + <<>>(pack, p, desc.head_size_v); + break; + case 128: + GatedDeltaNetDecodeWarpKernel + <<>>(pack, p, desc.head_size_v); + break; + default: + GatedDeltaNetDecodeWarpKernel + <<>>(pack, p, desc.head_size_v); + break; + } +} + +template Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const VariantPack& pack, const KernelParams& p, size_t max_shared_memory_per_block, cudaStream_t stream) { if (warp_specialized) { - const dim3 grid(desc.batch, desc.num_heads_v, - (desc.head_size_v + kDecodeWarps - 1) / kDecodeWarps); - const dim3 block(32, kDecodeWarps, 1); - switch (desc.head_size_qk) { - case 64: - GatedDeltaNetDecodeWarpKernel - <<>>(pack, p, desc.head_size_v); - break; - case 128: - GatedDeltaNetDecodeWarpKernel - <<>>(pack, p, desc.head_size_v); - break; - default: - GatedDeltaNetDecodeWarpKernel - <<>>(pack, p, desc.head_size_v); - break; + if (pack.packed_qkv) { + LaunchDecodeWarp(desc, pack, p, stream); + } else { + LaunchDecodeWarp(desc, pack, p, stream); } return CUDA_CALL(cudaGetLastError()); } @@ -708,20 +733,20 @@ Status LaunchRecurrent(const Descriptor& desc, bool warp_specialized, const Vari const size_t smem = RecurrentSmemBytes(desc.head_size_qk, desc.head_size_v); static DynamicSmemConfig recurrent_smem_config; ORT_RETURN_IF_ERROR(SetMaxDynamicSmem( - GatedDeltaNetRecurrentKernel, + GatedDeltaNetRecurrentKernel, max_shared_memory_per_block, recurrent_smem_config)); const dim3 grid(desc.batch, desc.num_heads_v, 1); - GatedDeltaNetRecurrentKernel + GatedDeltaNetRecurrentKernel <<>>(pack, p, desc.head_size_qk, desc.head_size_v); return CUDA_CALL(cudaGetLastError()); } } // namespace -template -Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const VariantPack& pack, - float scale, int max_threads_per_block, - size_t max_shared_memory_per_block, cudaStream_t stream) { +template +static Status LaunchGatedDeltaNetTypedGates(const Descriptor& desc, const Plan& plan, const VariantPack& pack, + float scale, int max_threads_per_block, + size_t max_shared_memory_per_block, cudaStream_t stream) { KernelParams p{}; p.total_tokens = desc.total_tokens; p.uniform_len = desc.batch > 0 ? desc.total_tokens / desc.batch : 0; @@ -754,15 +779,15 @@ Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const Varia // engine fit devices with a 99 KB opt-in limit (SM120) instead of SM90's 227 KB. if (plan.chunk_size == 32) { static DynamicSmemConfig chunk32_smem_config; - ORT_RETURN_IF_ERROR(SetMaxDynamicSmem(GatedDeltaNetChunkedKernel, + ORT_RETURN_IF_ERROR(SetMaxDynamicSmem(GatedDeltaNetChunkedKernel, plan.smem_bytes, chunk32_smem_config)); - GatedDeltaNetChunkedKernel + GatedDeltaNetChunkedKernel <<>>(pack, p); } else { static DynamicSmemConfig chunk64_smem_config; - ORT_RETURN_IF_ERROR(SetMaxDynamicSmem(GatedDeltaNetChunkedKernel, + ORT_RETURN_IF_ERROR(SetMaxDynamicSmem(GatedDeltaNetChunkedKernel, plan.smem_bytes, chunk64_smem_config)); - GatedDeltaNetChunkedKernel + GatedDeltaNetChunkedKernel <<>>(pack, p); } ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); @@ -774,14 +799,26 @@ Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const Varia // the two row sets are disjoint, so the merged result is deterministic. // The chunked engine only accepts head_size_qk == head_size_v == 128, so the warp // kernel is always the right shape here. - return LaunchRecurrent(desc, /*warp_specialized=*/true, pack, p, - max_shared_memory_per_block, stream); + return LaunchRecurrent(desc, /*warp_specialized=*/true, pack, p, + max_shared_memory_per_block, stream); } } (void)max_threads_per_block; - return LaunchRecurrent(desc, plan.warp_specialized, pack, p, max_shared_memory_per_block, - stream); + return LaunchRecurrent(desc, plan.warp_specialized, pack, p, max_shared_memory_per_block, + stream); +} + +template +Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const VariantPack& pack, + float scale, int max_threads_per_block, + size_t max_shared_memory_per_block, cudaStream_t stream) { + if (pack.gates_in_input_type) { + return LaunchGatedDeltaNetTypedGates( + desc, plan, pack, scale, max_threads_per_block, max_shared_memory_per_block, stream); + } + return LaunchGatedDeltaNetTypedGates( + desc, plan, pack, scale, max_threads_per_block, max_shared_memory_per_block, stream); } template Status LaunchGatedDeltaNet(const Descriptor&, const Plan&, const VariantPack&, diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h index feb75bf9b1fbf..3035ca02fa083 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_impl.h @@ -22,17 +22,28 @@ struct VariantPack { const T* value = nullptr; // [total_tokens, Hv, V] const int32_t* cu_seqlens = nullptr; // [B+1] or null for uniform packing const int32_t* capture_count = nullptr; // [B] or null when transition capture is disabled - const float* decay = nullptr; // [total_tokens, Hv] or [total_tokens, Hv, K] - const float* beta = nullptr; // [total_tokens, Hv] + const void* decay = nullptr; // [total_tokens, Hv] or [total_tokens, Hv, K] + const void* beta = nullptr; // [total_tokens, Hv] + bool gates_in_input_type = false; // gate buffers contain T instead of float + bool packed_qkv = false; // short decode reads packed Q/K/V without unpacking const float* initial_state = nullptr; // [B, Hv, V, K] V-major, may alias final_state const float* a_log = nullptr; // [Hv] const float* dt_bias = nullptr; // [Hv] T* output = nullptr; // [total_tokens, max(Hq,Hv), V] float* final_state = nullptr; // [B, Hv, V, K] + float* mid_state = nullptr; // [B, Hv, V, K], state after local row 0 float* state_update = nullptr; // [B, C * (Hv + Hk*K + Hv*V)] void* workspace = nullptr; }; +#ifdef __CUDACC__ +template +__device__ __forceinline__ float LoadGate(const void* data, int64_t index, bool gates_in_input_type) { + return gates_in_input_type ? static_cast(static_cast(data)[index]) + : static_cast(data)[index]; +} +#endif + template Status LaunchGatedDeltaNet(const Descriptor& desc, const Plan& plan, const VariantPack& pack, float scale, int max_threads_per_block, diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh index 33a841df9a4f0..582e3474bed1e 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_mma.cuh @@ -52,7 +52,7 @@ __device__ __forceinline__ void SequenceRange(const int32_t* cu_seqlens, int b, // Effective per-token decay and beta, after the optional fused activations. __device__ __forceinline__ float EffectiveDecay(float raw, const float* a_log, const float* dt_bias, int h, GateActivation act) { - if (act == GateActivation::kQwen) { + if (act == GateActivation::kSoftplusDecay) { const float bias = dt_bias != nullptr ? dt_bias[h] : 0.0f; const float scale = a_log != nullptr ? -__expf(a_log[h]) : -1.0f; return scale * SoftPlus(raw + bias); diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h index 613663fccb3b0..898afaf4699fc 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_plan.h @@ -26,7 +26,7 @@ enum class UpdateRule : int { kLinear = 0, kDelta = 2, kGatedDelta = 3 }; enum class GateActivation : int { kNone = 0, - kQwen = 1 }; + kSoftplusDecay = 1 }; enum class BetaActivation : int { kNone = 0, kSigmoid = 1 }; diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu index 0837fbde8f3bd..fb7c5f2711c2d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/gated_delta_net_split_impl.cu @@ -184,10 +184,10 @@ __global__ __launch_bounds__(kChunkedThreads) void GatedDeltaNetPrepareKernel( const int64_t gi = tok * p.num_heads_v + hv; float g = 0.0f, be = 0.0f; if (tid < len) { - g = needs_decay ? EffectiveDecay(pack.decay[gi], pack.a_log, pack.dt_bias, hv, + g = needs_decay ? EffectiveDecay(LoadGate(pack.decay, gi, pack.gates_in_input_type), pack.a_log, pack.dt_bias, hv, p.gate_activation) : 0.0f; - be = pack.beta != nullptr ? EffectiveBeta(pack.beta[gi], p.beta_activation) : 1.0f; + be = pack.beta != nullptr ? EffectiveBeta(LoadGate(pack.beta, gi, pack.gates_in_input_type), p.beta_activation) : 1.0f; } s.gc[tid] = g; s.beta[tid] = be; diff --git a/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc b/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc index a47b313e482fd..ca7888b7c178b 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc @@ -38,6 +38,12 @@ ONNX_OPERATOR_KERNEL_EX( CausalConvWithState::CausalConvWithState(const OpKernelInfo& info) : WebGpuKernel(info) { + ORT_ENFORCE(info.GetAttrOrDefault("input_channels", 0) == 0 && + info.GetAttrOrDefault("input_channel_offset", 0) == 0 && + info.GetAttrOrDefault("state_length", 0) == 0 && + info.GetAttrOrDefault("state_channels_last", -1) == -1 && + info.node().OutputDefs().size() <= 2, + "WebGPU CausalConvWithState does not support extended channel/state options (CUDA only)"); std::string activation_str = info.GetAttrOrDefault("activation", "none"); activation_ = ParseCausalConvActivation(activation_str); ORT_ENFORCE(info.GetAttr("ndim", &ndim_).IsOK(), "Attribute 'ndim' is required"); diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc index 2294389a9728c..d8030e7412ecd 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/gated_delta_net.cc @@ -37,6 +37,7 @@ ONNX_OPERATOR_KERNEL_EX( (*KernelDefBuilder::Create()) .TypeConstraint("T", WebGpuSupportedFloatTypes()) .TypeConstraint("TS", DataTypeImpl::GetTensorType()) + .TypeConstraint("TG", DataTypeImpl::GetTensorType()) .TypeConstraint("TI", DataTypeImpl::GetTensorType()) .InputMemoryType(OrtMemTypeCPUInput, 10) .MayInplace(6, 1), @@ -47,9 +48,10 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : WebGpuKernel(info) { ORT_ENFORCE(update_rule_ != GatedDeltaNetUpdateRule::Invalid, "update_rule must be one of: linear, gated, delta, gated_delta"); const auto gate_activation = info.GetAttrOrDefault("gate_activation", "none"); - ORT_ENFORCE(gate_activation == "none" || gate_activation == "qwen", - "gate_activation must be one of: none, qwen"); - qwen_gate_ = gate_activation == "qwen"; + ORT_ENFORCE(gate_activation == "none" || gate_activation == "softplus" || + gate_activation == "qwen", + "gate_activation must be one of: none, softplus"); + qwen_gate_ = gate_activation == "softplus" || gate_activation == "qwen"; const auto beta_activation = info.GetAttrOrDefault("beta_activation", "none"); ORT_ENFORCE(beta_activation == "none" || beta_activation == "sigmoid", "beta_activation must be one of: none, sigmoid"); @@ -60,6 +62,9 @@ GatedDeltaNet::GatedDeltaNet(const OpKernelInfo& info) : WebGpuKernel(info) { const auto state_update_capacity = info.GetAttrOrDefault("state_update_capacity", 0); ORT_ENFORCE(state_update_capacity >= 0 && state_update_capacity <= 8, "state_update_capacity must be in [0, 8], got ", state_update_capacity); + ORT_ENFORCE(info.GetAttrOrDefault("state_output", "state_update") == + "state_update", + "WebGPU GatedDeltaNet supports only state_output='state_update'"); state_update_capacity_ = static_cast(state_update_capacity); } diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 0751df5d5de64..5b88eae6b079a 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -4660,6 +4660,20 @@ ONNX_MS_OPERATOR_SET_SCHEMA( CausalConvWithState, 1, OpSchema() .SetDoc(CausalConvWithState_ver1_doc) + .Attr("input_channels", + "CUDA channels-last input channel slice width; 0 selects all channels. " + "Nondefault extended attributes require rank-3 input and state_window=0.", + AttributeProto::INT, static_cast(0)) + .Attr("input_channel_offset", "Starting channel of the channels-last input slice (CUDA only).", + AttributeProto::INT, static_cast(0)) + .Attr("state_length", + "Carry history length (CUDA only); 0 uses (kernel_size-1)*dilation. " + "An explicit length must cover the receptive field; convolution reads its trailing history.", + AttributeProto::INT, static_cast(0)) + .Attr("state_channels_last", + "State layout override (CUDA only): -1 follows channels_last, " + "0 uses [B,C,state_length], 1 uses [B,state_length,C].", + AttributeProto::INT, static_cast(-1)) .Attr("activation", "Fused activation function. One of: 'silu', 'swish', 'none'. " "Default is 'none'.", @@ -4703,7 +4717,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", @@ -4721,12 +4736,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", @@ -4735,8 +4752,14 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "(batch_size, state_length, d_1, ..., d_n) per slot when channels_last = 1. Slot " "W-1 contains the last state_length values from the virtual input along the causal " "axis; slot j contains the same for the prefix ending at position " - "(seq_len - W + j).", + "(seq_len - W + j). Extended attributes override the state length and layout " + "as for past_state.", "T") + .Output(2, "first_state", + "Optional carry state after the first input position, with the same " + "layout and length as present_state. CUDA only; requires rank-3 channels-last input " + "and state_window=0.", + "T", OpSchema::Optional) .TypeConstraint("T", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain input and output types to float tensors.") @@ -4745,6 +4768,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); @@ -5361,7 +5441,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) @@ -5369,6 +5449,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"; @@ -5383,8 +5466,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 " @@ -5404,6 +5488,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) @@ -5415,20 +5504,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 " @@ -5446,14 +5536,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); @@ -5586,45 +5677,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 f857d895570bb..86755fc0c2869 100644 --- a/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc +++ b/onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc @@ -319,6 +319,127 @@ static void RunCausalConvWithStateTests( // Basic tests - simple cases // ============================================================================= +template +void TestPackedChannelHistory(int length, bool state_last, bool use_past, bool use_bias, bool silu, int dilation, + bool first_output = true) { + constexpr int batch = 2, channels = 7, width = 13, offset = 3, kernel = 4; + const int history = (kernel - 1) * dilation + 1; + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) GTEST_SKIP() << "CUDA is required for extended convolution"; + const auto round = [](float value) { return static_cast(T(value)); }; + std::vector input(batch * length * width), weight(channels * kernel), bias(channels); + std::vector past(batch * channels * history), output(batch * length * channels); + std::vector present(past.size()), first(past.size()); + for (size_t i = 0; i < input.size(); ++i) input[i] = T(float(int(i % 17) - 8) / 16); + for (size_t i = 0; i < weight.size(); ++i) weight[i] = T(float(int(i % 11) - 5) / 32); + for (size_t i = 0; i < bias.size(); ++i) bias[i] = T(float(i) / 64); + for (size_t i = 0; i < past.size(); ++i) past[i] = T(float(int(i % 13) - 6) / 8); + const auto state_index = [&](int b, int t, int c) { + return state_last ? (b * history + t) * channels + c : (b * channels + c) * history + t; + }; + for (int b = 0; b < batch; ++b) { + for (int c = 0; c < channels; ++c) { + const auto value = [&](int t) { + return t >= 0 ? static_cast(input[(b * length + t) * width + offset + c]) + : (use_past ? static_cast(past[state_index(b, history + t, c)]) : 0.0f); + }; + for (int t = 0; t < length; ++t) { + float sum = use_bias ? static_cast(bias[c]) : 0.0f; + for (int j = 0; j < kernel; ++j) sum += value(t - (kernel - 1) * dilation + j * dilation) * static_cast(weight[c * kernel + j]); + if (silu) sum /= 1.0f + std::exp(-sum); + output[(b * length + t) * channels + c] = T(round(sum)); + } + for (int t = 0; t < history; ++t) { + present[state_index(b, t, c)] = T(value(length - history + t)); + first[state_index(b, t, c)] = T(value(1 - history + t)); + } + } + } + OpTester test("CausalConvWithState", 1, kMSDomain); + test.AddAttribute("channels_last", 1); + test.AddAttribute("input_channels", channels); + test.AddAttribute("input_channel_offset", offset); + test.AddAttribute("state_length", history); + test.AddAttribute("state_channels_last", state_last ? 1 : 0); + test.AddAttribute("dilation", dilation); + test.AddAttribute("activation", silu ? "silu" : "none"); + test.AddInput("input", {batch, length, width}, input); + test.AddInput("weight", {channels, 1, kernel}, weight); + if (use_bias) + test.AddInput("bias", {channels}, bias); + else + test.AddOptionalInputEdge(); + const std::vector state_shape = state_last ? std::vector{batch, history, channels} + : std::vector{batch, channels, history}; + if (use_past) + test.AddInput("past", state_shape, past); + else + test.AddOptionalInputEdge(); + test.AddOutput("output", {batch, length, channels}, output); + test.AddOutput("present", state_shape, present); + if (first_output) test.AddOutput("first", state_shape, first); + test.SetOutputTolerance(std::is_same_v ? 0.01f : 0.001f, 0.001f); + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers); +} + +TEST(CausalConvWithStateTest, PackedChannelsExtendedHistoryAndFirstState) { + for (int length : {1, 2, 9}) { + for (bool state_last : {false, true}) { + for (bool past : {false, true}) { + for (int dilation : {1, 2}) { + TestPackedChannelHistory(length, state_last, past, true, true, dilation); + TestPackedChannelHistory(length, state_last, past, false, true, dilation); + TestPackedChannelHistory(length, state_last, past, true, false, dilation); + TestPackedChannelHistory(length, state_last, past, false, true, dilation); + } + } + } + } + TestPackedChannelHistory(2, false, true, true, true, 1, false); +} + +TEST(CausalConvWithStateTest, PackedChannelsInvalidConfigurations) { + const std::vector errors{ + "invalid extended attributes", "channel slice exceeds input width", "state_length is too short", + "extended attributes require", "past_state must match", "bias must match"}; + for (size_t i = 0; i < errors.size(); ++i) { + SCOPED_TRACE(errors[i]); + auto cuda = DefaultCudaExecutionProvider(); + if (!cuda) GTEST_SKIP() << "CUDA is required for extended convolution"; + OpTester test("CausalConvWithState", 1, kMSDomain); + test.AddAttribute("channels_last", i == 3 ? 0 : 1); + test.AddAttribute("input_channels", i == 1 ? 14 : 7); + test.AddAttribute("input_channel_offset", 3); + test.AddAttribute("state_length", i == 2 ? 2 : 4); + test.AddAttribute("state_channels_last", i == 0 ? 2 : 0); + test.AddInput("input", {2, 2, 13}, std::vector(52)); + test.AddInput("weight", {7, 1, 4}, std::vector(28)); + test.AddInput("bias", {i == 5 ? 6 : 7}, std::vector(i == 5 ? 6 : 7)); + test.AddInput("past", {2, 7, i == 4 ? 5 : 4}, std::vector(i == 4 ? 70 : 56)); + test.AddOutput("output", {2, 2, 7}, std::vector(28)); + test.AddOutput("present", {2, 7, 4}, std::vector(56)); + std::vector> providers; + providers.push_back(std::move(cuda)); + test.Run(OpTester::ExpectResult::kExpectFailure, errors[i], {}, nullptr, &providers); + } +} + +TEST(CausalConvWithStateTest, PackedChannelsRejectsUnsupportedCpu) { + OpTester test("CausalConvWithState", 1, kMSDomain); + test.AddAttribute("channels_last", 1); + test.AddAttribute("input_channels", 2); + test.AddInput("input", {1, 1, 2}, {1.0f, 2.0f}); + test.AddInput("weight", {2, 1, 2}, {1.0f, 1.0f, 1.0f, 1.0f}); + test.AddOutput("output", {1, 1, 2}, {1.0f, 2.0f}); + test.AddOutput("present", {1, 1, 2}, {1.0f, 2.0f}); + std::vector> providers; + providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectFailure, "does not support extended channel/state options", + {}, nullptr, &providers); +} + TEST(CausalConvWithStateTest, BasicNoStateNoBias) { // B=1, D=2, L=4, K=3, activation=none int batch_size = 1, channels = 2, input_length = 4, kernel_size = 3; diff --git a/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc b/onnxruntime/test/contrib_ops/gated_delta_net_op_test.cc index 541a4ecf6502e..275348fd69132 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; @@ -1263,6 +1385,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};