From 007142f6194893227aa1b54fc574806f43baab8e Mon Sep 17 00:00:00 2001 From: Shiyi Zou Date: Tue, 15 Sep 2026 16:37:01 +0800 Subject: [PATCH 1/3] bert shape inference --- .../core/graph/contrib_ops/bert_defs.cc | 55 ++++++++++++++++--- 1 file changed, 47 insertions(+), 8 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index b9229035170e9..952a3872c0f9f 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include +#include #include #include @@ -303,6 +304,9 @@ void BaseGroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceConte const auto* total_sequence_length_data = ctx.getInputData(6); if (total_sequence_length_data != nullptr) { const auto& data = ParseData(total_sequence_length_data); + if (data.empty()) { + fail_shape_inference("total_sequence_length input must contain a single element"); + } total_sequence_length_value = static_cast(data[0]); } @@ -2895,13 +2899,26 @@ ONNX_MS_OPERATOR_SET_SCHEMA( if (hasInputShape(ctx, 0) && hasInputShape(ctx, 1)) { auto& input_shape = getInputShape(ctx, 0); auto& weight_shape = getInputShape(ctx, 1); - if (input_shape.dim_size() < 2) { - fail_shape_inference("CausalConvWithState: input must have rank >= 2"); + int64_t ndim = getAttribute(ctx, "ndim", 1); + if (ndim < 1 || ndim > 3) { + fail_shape_inference("CausalConvWithState: ndim must be 1, 2, or 3, got ", ndim); } - if (weight_shape.dim_size() < 2) { - fail_shape_inference("CausalConvWithState: weight must have rank >= 2"); + // weight is always channels-first: (channels, 1, k_1, ..., k_ndim), rank == ndim + 2. + if (weight_shape.dim_size() != ndim + 2) { + fail_shape_inference("CausalConvWithState: weight must have rank ndim + 2 (", + ndim + 2, "), got rank ", weight_shape.dim_size()); + } + if (channels_last == 1) { + // channels_last input is (batch_size, sequence_length, d_1, ..., d_n), where + // d_1 * ... * d_n = channels. n is independent of ndim (always 1 here), so only a + // lower bound applies. + if (input_shape.dim_size() < 3) { + fail_shape_inference("CausalConvWithState: channels_last input must have rank >= 3"); + } + } else if (input_shape.dim_size() != ndim + 2) { + fail_shape_inference("CausalConvWithState: input must have rank ndim + 2 (", + ndim + 2, "), got rank ", input_shape.dim_size()); } - int64_t ndim = getAttribute(ctx, "ndim", 1); // (kernel_size - 1) * dilation, or an unset dim when kernel_size is symbolic. const int last_kernel_dim = weight_shape.dim_size() - 1; TensorShapeProto::Dimension state_length; @@ -3555,9 +3572,31 @@ ONNX_MS_OPERATOR_SET_SCHEMA( 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(); - width->set_dim_value(state_update_capacity * - (num_heads_v + num_heads_k * head_size_qk + - num_heads_v * head_size_v)); + 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); } updateOutputShape(ctx, 2, capsule_shape); } From 645a7f2e848cf0d5028fe04e9ecf8366a18248d7 Mon Sep 17 00:00:00 2001 From: Shiyi Zou Date: Tue, 15 Sep 2026 17:02:36 +0800 Subject: [PATCH 2/3] address copilot --- onnxruntime/core/graph/contrib_ops/bert_defs.cc | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 952a3872c0f9f..e487aa3ac81a6 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -2885,7 +2885,12 @@ ONNX_MS_OPERATOR_SET_SCHEMA( fail_shape_inference("CausalConvWithState: channels_last must be 0 or 1, got ", channels_last); } - if (channels_last == 1 && getAttribute(ctx, "ndim", 1) != 1) { + + const int64_t ndim = getAttribute(ctx, "ndim", 1); + if (ndim < 1 || ndim > 3) { + fail_shape_inference("CausalConvWithState: ndim must be 1, 2, or 3, got ", ndim); + } + if (channels_last == 1 && ndim != 1) { fail_shape_inference("CausalConvWithState: channels_last requires ndim = 1"); } @@ -2899,10 +2904,6 @@ ONNX_MS_OPERATOR_SET_SCHEMA( if (hasInputShape(ctx, 0) && hasInputShape(ctx, 1)) { auto& input_shape = getInputShape(ctx, 0); auto& weight_shape = getInputShape(ctx, 1); - int64_t ndim = getAttribute(ctx, "ndim", 1); - if (ndim < 1 || ndim > 3) { - fail_shape_inference("CausalConvWithState: ndim must be 1, 2, or 3, got ", ndim); - } // weight is always channels-first: (channels, 1, k_1, ..., k_ndim), rank == ndim + 2. if (weight_shape.dim_size() != ndim + 2) { fail_shape_inference("CausalConvWithState: weight must have rank ndim + 2 (", From d9764b6c8df4af01e70ea3b48dbec9001adaec6c Mon Sep 17 00:00:00 2001 From: Shiyi Zou Date: Tue, 15 Sep 2026 17:14:58 +0800 Subject: [PATCH 3/3] fix comment --- onnxruntime/core/graph/contrib_ops/bert_defs.cc | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index e487aa3ac81a6..a4b247895c16f 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -2910,9 +2910,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA( ndim + 2, "), got rank ", weight_shape.dim_size()); } if (channels_last == 1) { - // channels_last input is (batch_size, sequence_length, d_1, ..., d_n), where - // d_1 * ... * d_n = channels. n is independent of ndim (always 1 here), so only a - // lower bound applies. + // (batch_size, sequence_length, d_1, ..., d_n). Check the lower bound. if (input_shape.dim_size() < 3) { fail_shape_inference("CausalConvWithState: channels_last input must have rank >= 3"); }