From 55366fa9849b2c6ebbf82f812ba318307a464621 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Tue, 23 Jun 2026 23:49:09 +0000 Subject: [PATCH 01/36] Require all present outputs before populating attention shape inference DecoderAttention and MultiHeadAttention shape-inference functions guarded population of present_key (output 1) and present_value (output 2) with getNumOutputs() > 1, but write output index 2. present_key and present_value are produced as a both-or-neither pair, so require all three outputs (> 2) before populating them, matching BaseGroupQueryAttention (>= 3) and the EmbedLayerNorm guard. Also add a bounds check in InferenceContextImpl::getOutputType so an out-of-range output index fails inference cleanly instead of indexing past the end, mirroring DataPropagationContextImpl and getInputType. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/graph/contrib_ops/bert_defs.cc | 6 +++--- onnxruntime/core/graph/graph.cc | 4 ++++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 896774fb5c8d8..7ce714a954290 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -29,7 +29,7 @@ namespace contrib { void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { // Type inference ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); - if (ctx.getNumOutputs() > 1) { + if (ctx.getNumOutputs() > 2) { // present_key and present_value outputs ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 1); ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 2); } @@ -38,7 +38,7 @@ void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx auto& query_shape = getInputShape(ctx, 0); updateOutputShape(ctx, 0, query_shape); } - if (ctx.getNumOutputs() > 1) { + if (ctx.getNumOutputs() > 2) { // present_key and present_value outputs if (hasInputShape(ctx, 6) && hasInputShape(ctx, 7)) { auto& cache_shape = getInputShape(ctx, 6); auto& cache_dims = cache_shape.dim(); @@ -199,7 +199,7 @@ void MultiHeadAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& c } } - if (ctx.getNumOutputs() > 1) { // has present output + if (ctx.getNumOutputs() > 2) { // has present_key and present_value outputs if (hasInputShape(ctx, past_key_index)) { auto& past_shape = getInputShape(ctx, past_key_index); auto& past_dims = past_shape.dim(); diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index fe2df6a87d124..ff60f6cb5f741 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -2739,6 +2739,10 @@ class InferenceContextImpl : public ONNX_NAMESPACE::InferenceContext { } TypeProto* getOutputType(size_t index) override { + if (index >= node_output_types_.size()) { + fail_type_inference("output index ", index, " is out of range; node has ", + node_output_types_.size(), " outputs"); + } return &node_output_types_[index]; } From 2ad360cfbd9894be65fb4f7edb5481fbf0f28839 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Wed, 24 Jun 2026 00:09:20 +0000 Subject: [PATCH 02/36] Add regression tests for attention contrib ops with optional present outputs omitted Cover DecoderAttention, MultiHeadAttention and DecoderMaskedMultiHeadAttention nodes declared with exactly two outputs (present_key kept, present_value omitted). Each test builds the node and asserts Graph::Resolve() shape inference completes cleanly. Tests are execution-provider independent and throw-free, so they run on the default CPU build and in no-exception (ORT_NO_EXCEPTIONS) builds. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...n_optional_outputs_shape_inference_test.cc | 101 ++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100644 onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc diff --git a/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc new file mode 100644 index 0000000000000..ac2df0ab6a6e0 --- /dev/null +++ b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc @@ -0,0 +1,101 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Regression coverage for attention contrib ops declared with their optional "present" outputs +// partially omitted. DecoderAttention, MultiHeadAttention and DecoderMaskedMultiHeadAttention each +// expose an optional present_key (output 1) and present_value (output 2). A node that declares +// exactly two outputs (present_key kept, present_value omitted) is valid per the op schemas, and +// graph resolution / shape inference must complete cleanly without referencing the absent third +// output. +// +// These tests build each op with exactly two outputs and assert that Graph::Resolve() succeeds. +// They exercise only graph-load shape inference, which is execution-provider independent, so they +// run on the default CPU build with no provider-specific handling. The resolve path for these +// models is throw-free, so the tests are valid in builds compiled without exceptions +// (ORT_NO_EXCEPTIONS) and need no exception-specific guarding. + +#include "gtest/gtest.h" + +#include "core/graph/constants.h" +#include "core/graph/model.h" +#include "test/test_environment.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" + +namespace onnxruntime { +namespace test { + +namespace { + +constexpr int kOnnxOpsetVersion = 17; + +// Builds a single-node model via the supplied callback and verifies that graph resolution, which +// runs the node's type/shape inference, completes without error. +void ResolveSingleNodeModel(const std::function& add_node) { + std::unordered_map domain_to_version; + domain_to_version[kOnnxDomain] = kOnnxOpsetVersion; + domain_to_version[kMSDomain] = 1; + + Model model("attention_optional_outputs", /*is_onnx_domain_only=*/false, ModelMetaData(), + PathString(), IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, {}, + DefaultLoggingManager().DefaultLogger()); + + ModelTestBuilder builder(model.MainGraph()); + add_node(builder); + builder.SetGraphOutputs(); + + ASSERT_STATUS_OK(model.MainGraph().Resolve()); +} + +} // namespace + +// MultiHeadAttention with present_key kept and present_value omitted (exactly two outputs). +TEST(AttentionOptionalOutputsShapeInferenceTest, MultiHeadAttentionPresentValueOmitted) { + ResolveSingleNodeModel([](ModelTestBuilder& builder) { + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* output = builder.MakeOutput(std::nullopt); + NodeArg* present_key = builder.MakeOutput(std::nullopt); + Node& node = builder.AddNode("MultiHeadAttention", {query}, {output, present_key}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }); +} + +// DecoderMaskedMultiHeadAttention with present_key kept and present_value omitted. +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderMaskedMultiHeadAttentionPresentValueOmitted) { + ResolveSingleNodeModel([](ModelTestBuilder& builder) { + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* output = builder.MakeOutput(std::nullopt); + NodeArg* present_key = builder.MakeOutput(std::nullopt); + Node& node = builder.AddNode("DecoderMaskedMultiHeadAttention", {query}, {output, present_key}, + kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }); +} + +// DecoderAttention with new_key_cache kept and new_value_cache omitted (exactly two outputs). +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderAttentionNewValueCacheOmitted) { + ResolveSingleNodeModel([](ModelTestBuilder& builder) { + // DecoderAttention requires inputs 0-4 and 8-11; inputs 5-7 are optional and left empty here. + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* key = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* q_weight = builder.MakeInput(std::vector{4, 4}); + NodeArg* kv_weight = builder.MakeInput(std::vector{4, 8}); + NodeArg* bias = builder.MakeInput(std::vector{12}); + NodeArg* static_kv = builder.MakeInput(std::vector{1}); + NodeArg* use_past = builder.MakeInput(std::vector{1}); + NodeArg* has_layer_state = builder.MakeInput(std::vector{1}); + NodeArg* has_key_padding_mask = builder.MakeInput(std::vector{1}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + NodeArg* new_key_cache = builder.MakeOutput(std::nullopt); + + std::vector inputs = {query, key, q_weight, kv_weight, bias, &empty, &empty, &empty, + static_kv, use_past, has_layer_state, has_key_padding_mask}; + Node& node = builder.AddNode("DecoderAttention", inputs, {output, new_key_cache}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }); +} + +} // namespace test +} // namespace onnxruntime From d5ef725de0a3a62ef2dc7ed39c565cc7a68bc202 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Wed, 24 Jun 2026 00:14:51 +0000 Subject: [PATCH 03/36] Add 3-output positive cases for attention contrib op shape inference Extend the optional-present-output regression suite with cases that declare all three outputs for DecoderAttention, MultiHeadAttention and DecoderMaskedMultiHeadAttention and assert the present_key/present_value branch still runs and infers their element types. Together with the two-output cases this pins the output-count guard to exactly three. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...n_optional_outputs_shape_inference_test.cc | 142 ++++++++++++++++-- 1 file changed, 127 insertions(+), 15 deletions(-) diff --git a/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc index ac2df0ab6a6e0..52792529d2418 100644 --- a/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc +++ b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc @@ -1,16 +1,20 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -// Regression coverage for attention contrib ops declared with their optional "present" outputs -// partially omitted. DecoderAttention, MultiHeadAttention and DecoderMaskedMultiHeadAttention each -// expose an optional present_key (output 1) and present_value (output 2). A node that declares -// exactly two outputs (present_key kept, present_value omitted) is valid per the op schemas, and -// graph resolution / shape inference must complete cleanly without referencing the absent third -// output. +// Shape-inference correctness coverage for attention contrib ops with optional "present" outputs. +// DecoderAttention, MultiHeadAttention and DecoderMaskedMultiHeadAttention each expose an optional +// present_key (output 1) and present_value (output 2). These outputs are produced as a pair, so a +// node may declare either one output, or all three; declaring exactly two (present_key kept, +// present_value omitted) is also valid per the schemas. // -// These tests build each op with exactly two outputs and assert that Graph::Resolve() succeeds. -// They exercise only graph-load shape inference, which is execution-provider independent, so they -// run on the default CPU build with no provider-specific handling. The resolve path for these +// The "...Omitted" tests build each op with exactly two outputs and assert that Graph::Resolve() +// completes cleanly without referencing the absent third output. The "...AllPresentOutputs" tests +// build each op with all three outputs and assert that the present_key / present_value branch still +// runs and propagates their element types. Together they pin the guard to exactly "> 2": fewer +// outputs must not touch the missing one, and three outputs must still be inferred. +// +// These tests exercise only graph-load shape inference, which is execution-provider independent, so +// they run on the default CPU build with no provider-specific handling. The resolve path for these // models is throw-free, so the tests are valid in builds compiled without exceptions // (ORT_NO_EXCEPTIONS) and need no exception-specific guarding. @@ -29,9 +33,10 @@ namespace { constexpr int kOnnxOpsetVersion = 17; -// Builds a single-node model via the supplied callback and verifies that graph resolution, which -// runs the node's type/shape inference, completes without error. -void ResolveSingleNodeModel(const std::function& add_node) { +// Builds a single-node model via add_node, resolves it (running the node's type/shape inference), +// asserts success, and runs the optional verifier against the resolved graph. +void BuildResolveAndVerify(const std::function& add_node, + const std::function& verify = nullptr) { std::unordered_map domain_to_version; domain_to_version[kOnnxDomain] = kOnnxOpsetVersion; domain_to_version[kMSDomain] = 1; @@ -45,13 +50,25 @@ void ResolveSingleNodeModel(const std::function builder.SetGraphOutputs(); ASSERT_STATUS_OK(model.MainGraph().Resolve()); + if (verify) { + verify(model.MainGraph()); + } +} + +// Asserts that the given output NodeArg received a tensor element type from shape inference. +void ExpectInferredElemType(const NodeArg* output) { + ASSERT_NE(output, nullptr); + const ONNX_NAMESPACE::TypeProto* type = output->TypeAsProto(); + ASSERT_NE(type, nullptr); + EXPECT_TRUE(type->has_tensor_type()); + EXPECT_TRUE(type->tensor_type().has_elem_type()); } } // namespace // MultiHeadAttention with present_key kept and present_value omitted (exactly two outputs). TEST(AttentionOptionalOutputsShapeInferenceTest, MultiHeadAttentionPresentValueOmitted) { - ResolveSingleNodeModel([](ModelTestBuilder& builder) { + BuildResolveAndVerify([](ModelTestBuilder& builder) { NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); NodeArg* output = builder.MakeOutput(std::nullopt); NodeArg* present_key = builder.MakeOutput(std::nullopt); @@ -62,7 +79,7 @@ TEST(AttentionOptionalOutputsShapeInferenceTest, MultiHeadAttentionPresentValueO // DecoderMaskedMultiHeadAttention with present_key kept and present_value omitted. TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderMaskedMultiHeadAttentionPresentValueOmitted) { - ResolveSingleNodeModel([](ModelTestBuilder& builder) { + BuildResolveAndVerify([](ModelTestBuilder& builder) { NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); NodeArg* output = builder.MakeOutput(std::nullopt); NodeArg* present_key = builder.MakeOutput(std::nullopt); @@ -74,7 +91,7 @@ TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderMaskedMultiHeadAttention // DecoderAttention with new_key_cache kept and new_value_cache omitted (exactly two outputs). TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderAttentionNewValueCacheOmitted) { - ResolveSingleNodeModel([](ModelTestBuilder& builder) { + BuildResolveAndVerify([](ModelTestBuilder& builder) { // DecoderAttention requires inputs 0-4 and 8-11; inputs 5-7 are optional and left empty here. NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); @@ -97,5 +114,100 @@ TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderAttentionNewValueCacheOm }); } +// MultiHeadAttention with all three outputs: the present_key / present_value branch must still run. +TEST(AttentionOptionalOutputsShapeInferenceTest, MultiHeadAttentionAllPresentOutputs) { + NodeArg* present_key = nullptr; + NodeArg* present_value = nullptr; + BuildResolveAndVerify( + [&](ModelTestBuilder& builder) { + // present_key/present_value types are propagated from past_key (input 6) and past_value + // (input 7) when past buffer sharing is active, which the op detects from shaped past_key + // (input 6) and past_sequence_length (input 8). Leave inputs 1-5 empty to reach them. + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_sequence_length = builder.MakeInput(std::vector{1}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + present_key = builder.MakeOutput(); + present_value = builder.MakeOutput(); + + std::vector inputs = {query, &empty, &empty, &empty, &empty, &empty, + past_key, past_value, past_sequence_length}; + Node& node = builder.AddNode("MultiHeadAttention", inputs, + {output, present_key, present_value}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }, + [&](const Graph&) { + ExpectInferredElemType(present_key); + ExpectInferredElemType(present_value); + }); +} + +// DecoderMaskedMultiHeadAttention with all three outputs: shape inference must still populate them. +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderMaskedMultiHeadAttentionAllPresentOutputs) { + NodeArg* present_key = nullptr; + NodeArg* present_value = nullptr; + BuildResolveAndVerify( + [&](ModelTestBuilder& builder) { + // For this op past_key/past_value are inputs 5 and 6; leave inputs 1-4 empty to reach them. + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + present_key = builder.MakeOutput(); + present_value = builder.MakeOutput(); + + std::vector inputs = {query, &empty, &empty, &empty, &empty, past_key, past_value}; + Node& node = builder.AddNode("DecoderMaskedMultiHeadAttention", inputs, + {output, present_key, present_value}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + node.AddAttribute("past_present_share_buffer", static_cast(1)); + }, + [&](const Graph&) { + ExpectInferredElemType(present_key); + ExpectInferredElemType(present_value); + }); +} + +// DecoderAttention with all three outputs: new_key_cache / new_value_cache types must be inferred. +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderAttentionAllCacheOutputs) { + NodeArg* new_key_cache = nullptr; + NodeArg* new_value_cache = nullptr; + BuildResolveAndVerify( + [&](ModelTestBuilder& builder) { + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* key = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* q_weight = builder.MakeInput(std::vector{4, 4}); + NodeArg* kv_weight = builder.MakeInput(std::vector{4, 8}); + NodeArg* bias = builder.MakeInput(std::vector{12}); + NodeArg* key_cache = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* value_cache = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* static_kv = builder.MakeInput(std::vector{1}); + NodeArg* use_past = builder.MakeInput(std::vector{1}); + NodeArg* has_layer_state = builder.MakeInput(std::vector{1}); + NodeArg* has_key_padding_mask = builder.MakeInput(std::vector{1}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + new_key_cache = builder.MakeOutput(); + new_value_cache = builder.MakeOutput(); + + std::vector inputs = {query, key, q_weight, kv_weight, bias, &empty, key_cache, + value_cache, static_kv, use_past, has_layer_state, + has_key_padding_mask}; + Node& node = builder.AddNode("DecoderAttention", inputs, + {output, new_key_cache, new_value_cache}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }, + [&](const Graph&) { + ExpectInferredElemType(new_key_cache); + ExpectInferredElemType(new_value_cache); + }); +} + } // namespace test } // namespace onnxruntime From 7eb0a6ab87bf918ce8bef8a537ea764c0740fd3b Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Wed, 24 Jun 2026 00:17:41 +0000 Subject: [PATCH 04/36] Clarify optional-output guard comments in attention shape inference Use each op's actual output names (DecoderAttention: new_key_cache / new_value_cache; MultiHeadAttention: present_key / present_value), align the three guards to a consistent '// has outputs' phrasing, and note that the two optional cache outputs are produced as a pair, so they are present only when the node declares more than two outputs. Comment-only; no logic change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/graph/contrib_ops/bert_defs.cc | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 7ce714a954290..9fdcebf57081d 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -29,7 +29,7 @@ namespace contrib { void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { // Type inference ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); - if (ctx.getNumOutputs() > 2) { // present_key and present_value outputs + if (ctx.getNumOutputs() > 2) { // has new_key_cache and new_value_cache outputs; a pair, so present only when > 2 ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 1); ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 2); } @@ -38,7 +38,7 @@ void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx auto& query_shape = getInputShape(ctx, 0); updateOutputShape(ctx, 0, query_shape); } - if (ctx.getNumOutputs() > 2) { // present_key and present_value outputs + if (ctx.getNumOutputs() > 2) { // has new_key_cache and new_value_cache outputs; a pair, so present only when > 2 if (hasInputShape(ctx, 6) && hasInputShape(ctx, 7)) { auto& cache_shape = getInputShape(ctx, 6); auto& cache_dims = cache_shape.dim(); @@ -199,7 +199,7 @@ void MultiHeadAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& c } } - if (ctx.getNumOutputs() > 2) { // has present_key and present_value outputs + if (ctx.getNumOutputs() > 2) { // has present_key and present_value outputs; a pair, so present only when > 2 if (hasInputShape(ctx, past_key_index)) { auto& past_shape = getInputShape(ctx, past_key_index); auto& past_dims = past_shape.dim(); From 22411d1c84cbd1c467f165ba1c3631b45c96419e Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Wed, 24 Jun 2026 00:36:36 +0000 Subject: [PATCH 05/36] Strengthen MHA/DMMHA two-output shape-inference regression tests The MultiHeadAttention and DecoderMaskedMultiHeadAttention two-output cases only passed a query input, so the present-output branch (which references output index 2) was never entered and the tests could not detect a regression there. Supply shaped past_key / past_value (and past_sequence_length for MHA, buffer sharing for DMMHA) so the branch is exercised while only two outputs are declared, matching the DecoderAttention case which already reached that path. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...n_optional_outputs_shape_inference_test.cc | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc index 52792529d2418..5d6a76ba3f9b9 100644 --- a/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc +++ b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc @@ -67,25 +67,42 @@ void ExpectInferredElemType(const NodeArg* output) { } // namespace // MultiHeadAttention with present_key kept and present_value omitted (exactly two outputs). +// past_key (input 6), past_value (input 7) and past_sequence_length (input 8) are supplied with +// shapes so the present-output branch is active; with only two outputs declared, inference must not +// reference the absent present_value (output 2). Inputs 1-5 are optional and left empty. TEST(AttentionOptionalOutputsShapeInferenceTest, MultiHeadAttentionPresentValueOmitted) { BuildResolveAndVerify([](ModelTestBuilder& builder) { + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_sequence_length = builder.MakeInput(std::vector{1}); NodeArg* output = builder.MakeOutput(std::nullopt); NodeArg* present_key = builder.MakeOutput(std::nullopt); - Node& node = builder.AddNode("MultiHeadAttention", {query}, {output, present_key}, kMSDomain); + std::vector inputs = {query, &empty, &empty, &empty, &empty, &empty, + past_key, past_value, past_sequence_length}; + Node& node = builder.AddNode("MultiHeadAttention", inputs, {output, present_key}, kMSDomain); node.AddAttribute("num_heads", static_cast(2)); }); } // DecoderMaskedMultiHeadAttention with present_key kept and present_value omitted. +// past_key (input 5) and past_value (input 6) are supplied with shapes and past buffer sharing is +// enabled so the present-output branch is active; with only two outputs declared, inference must not +// reference the absent present_value (output 2). Inputs 1-4 are optional and left empty. TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderMaskedMultiHeadAttentionPresentValueOmitted) { BuildResolveAndVerify([](ModelTestBuilder& builder) { + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); NodeArg* output = builder.MakeOutput(std::nullopt); NodeArg* present_key = builder.MakeOutput(std::nullopt); - Node& node = builder.AddNode("DecoderMaskedMultiHeadAttention", {query}, {output, present_key}, + std::vector inputs = {query, &empty, &empty, &empty, &empty, past_key, past_value}; + Node& node = builder.AddNode("DecoderMaskedMultiHeadAttention", inputs, {output, present_key}, kMSDomain); node.AddAttribute("num_heads", static_cast(2)); + node.AddAttribute("past_present_share_buffer", static_cast(1)); }); } From bc3b5c4fe5cab647385f5b342a452987e3d7b26a Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Wed, 24 Jun 2026 18:18:13 +0000 Subject: [PATCH 06/36] Apply clang-format to attention optional-output shape inference tests Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../attention_optional_outputs_shape_inference_test.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc index 5d6a76ba3f9b9..d2d3a421a6560 100644 --- a/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc +++ b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc @@ -79,7 +79,7 @@ TEST(AttentionOptionalOutputsShapeInferenceTest, MultiHeadAttentionPresentValueO NodeArg* past_sequence_length = builder.MakeInput(std::vector{1}); NodeArg* output = builder.MakeOutput(std::nullopt); NodeArg* present_key = builder.MakeOutput(std::nullopt); - std::vector inputs = {query, &empty, &empty, &empty, &empty, &empty, + std::vector inputs = {query, &empty, &empty, &empty, &empty, &empty, past_key, past_value, past_sequence_length}; Node& node = builder.AddNode("MultiHeadAttention", inputs, {output, present_key}, kMSDomain); node.AddAttribute("num_heads", static_cast(2)); @@ -150,7 +150,7 @@ TEST(AttentionOptionalOutputsShapeInferenceTest, MultiHeadAttentionAllPresentOut present_key = builder.MakeOutput(); present_value = builder.MakeOutput(); - std::vector inputs = {query, &empty, &empty, &empty, &empty, &empty, + std::vector inputs = {query, &empty, &empty, &empty, &empty, &empty, past_key, past_value, past_sequence_length}; Node& node = builder.AddNode("MultiHeadAttention", inputs, {output, present_key, present_value}, kMSDomain); From fca589010271e7fa07872635e01c48a3394bfca0 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Wed, 24 Jun 2026 20:42:01 +0000 Subject: [PATCH 07/36] docs: add contrib-op shape-inference output-index safety skill + coding-convention note Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../SKILL.md | 237 ++++++++++++++++++ docs/Coding_Conventions_and_Standards.md | 1 + 2 files changed, 238 insertions(+) create mode 100644 .agents/skills/contrib-op-shape-inference-memory-safety/SKILL.md diff --git a/.agents/skills/contrib-op-shape-inference-memory-safety/SKILL.md b/.agents/skills/contrib-op-shape-inference-memory-safety/SKILL.md new file mode 100644 index 0000000000000..d39d800fb887c --- /dev/null +++ b/.agents/skills/contrib-op-shape-inference-memory-safety/SKILL.md @@ -0,0 +1,237 @@ +--- +name: contrib-op-shape-inference-memory-safety +description: "Audit and fix out-of-range output writes in ONNX Runtime operator shape-inference functions. Use when reviewing or fixing a contrib (or standard) op TypeAndShapeInference where a getNumOutputs() guard precedes a write to a higher output index - optional trailing outputs make a smaller output count schema-valid, so getOutputType(index) can run one past the declared outputs at Graph::Resolve." +--- + +# Contrib-Op Shape-Inference Output-Index Safety + +Reusable method for finding and fixing the bug class where an operator's +`TypeAndShapeInference` function guards an output write with `getNumOutputs() > N` but then +writes an output index **greater than** `N`. For a node that declares fewer outputs, the +written index is past the end of the inference context's output vector. + +> **Scope**: schema-level shape inference in `onnxruntime/core/graph/contrib_ops/*.cc` and +> `shape_inference_functions.cc`. This runs once during `Graph::Resolve` (model-load time), +> **EP-agnostic** - there is no per-EP (CPU/CUDA/ROCm) kernel duplicate of this code to +> chase. Op *kernels* allocate outputs via the bounds-safe `OpKernelContext::Output(index)` +> and are a separate concern. + +## 1. The pattern + +```cpp +// onnxruntime/core/graph/contrib_ops/bert_defs.cc (before) +propagateElemTypeFromInputToOutput(ctx, 0, 0); +if (ctx.getNumOutputs() > 1) { // guard says "> 1" + propagateElemTypeFromInputToOutput(ctx, 0, 1); + propagateElemTypeFromInputToOutput(ctx, 0, 2); // but writes index 2 +} +``` + +The guard `getNumOutputs() > 1` admits a node with **exactly 2 outputs** (indices 0, 1), yet +the body writes index **2**. The implication "`> 1` ⇒ index 2 exists" is false: `> 1` only +guarantees indices 0 and 1. + +### Why a smaller output count is valid + +Trailing outputs declared `OpSchema::Optional` **lower `min_output`**. ONNX derives +`min_output` = number of required outputs, `max_output` = total declared. The model checker +(`checker::check_node`) only enforces `min_output <= N <= max_output`. + +| Op | Output decls | min / max | A 2-output node? | +|---|---|---|---| +| `DecoderAttention` | out (req), new_key_cache (Opt), new_value_cache (Opt) | 1 / 3 | passes checker | +| `MultiHeadAttention` | out (req), present_key (Opt), present_value (Opt), qk (Opt) | 1 / 4 | passes checker | +| `DecoderMaskedMultiHeadAttention` | out (req) + 3 Optional | 1 / 4 | passes checker | + +So a node with `output=['out','present_key']` is schema-valid, passes the checker, and then +reaches the index-2 write. **A passing checker is not a guarantee the index is in range.** + +## 2. The sink (why the write is not caught) + +```cpp +// onnxruntime/core/graph/graph.cc - InferenceContextImpl +const TypeProto* getInputType(size_t index) const override { + return node_.InputDefs().at(index)->TypeAsProto(); // .at() -> bounds-checked +} +TypeProto* getOutputType(size_t index) override { + return &node_output_types_[index]; // operator[] -> NOT bounds-checked +} +``` + +- `node_output_types_` is sized to `node.OutputDefs().size()` in the `InferenceContextImpl` + ctor, so for a 2-output node it has 2 elements; `getOutputType(2)` returns one past the end. +- `getInputType` uses `.at()` (would throw on a bad index); `getOutputType` uses raw + `operator[]` (no check) - the asymmetry is the root cause. +- The call runs at `Graph::Resolve` → `InferAndVerifyTypeMatch` → `RunInferencing`. The + surrounding `ORT_TRY/ORT_CATCH(const std::exception&)` only catches *thrown* + `fail_shape_inference`; a raw out-of-range `operator[]` does not throw, so the catch does + not help. +- Because this is schema-level inference, it is **EP-independent** - no CUDA/ROCm copy. + +## 3. Audit technique — always sweep siblings + +Do not stop at the reported function. Grep **every** shape-inference guard and compare its +threshold against the **highest output index written before the next guard**. + +```bash +git grep -n 'getNumOutputs' -- \ + onnxruntime/core/graph/contrib_ops/*.cc \ + onnxruntime/core/graph/contrib_ops/shape_inference_functions.cc +``` + +For each `if (ctx.getNumOutputs() > N)` block, find the largest `index` passed to +`propagateElemTypeFromInputToOutput(ctx, _, index)` / `updateOutputShape(ctx, index, _)` / +`getOutputType(index)` inside it. **Rule: the guard must require strictly more outputs than +the highest index written** (write index `k` ⇒ guard must ensure `getNumOutputs() > k`). + +Correct exemplars already in the tree to copy: + +| Exemplar | Pattern | Why it is safe | +|---|---|---| +| `BaseGroupQueryAttention...` | `if (getNumOutputs() >= 3)` then writes idx 2 | guard covers highest index | +| `PagedAttention...` | nested `> 1` + inner `if (getNumOutputs() != 3) fail_shape_inference` | fails before any write | +| `EmbedLayerNormalizationShapeInference` | `> 2` then writes idx 2 | fixed by PR #28176 (precedent) | +| `SkipLayerNormalizationShapeInference` | each idx `k` guarded by `> k` | per-index guard | + +> **Gotcha — conditional writes can hide a vacuous audit.** A write may sit behind an inner +> condition (e.g. `hasInputShape(past_key_index)` before writing index 2). The site is still +> a bug, but you can only *observe* it when that inner condition is also satisfied. Keep this +> in mind both for the audit and for tests (§5). + +## 4. Fix patterns + +**Point fix (required): raise the guard to cover the highest index written.** + +```cpp +// before +if (ctx.getNumOutputs() > 1) { ... writes idx 2 ... } +// after +if (ctx.getNumOutputs() > 2) { // both present_key (idx 1) AND present_value (idx 2) + ... +} +``` + +Justify the threshold with the op's output semantics. For these attention ops the two trailing +outputs - `present_key` (idx 1) and `present_value` (idx 2) for `MultiHeadAttention`, +`new_key_cache` / `new_value_cache` for `DecoderAttention` (see the §1 table for each op's +exact output names) - are a **both-or-neither pair**: there is no valid configuration that +emits one without the other, so requiring all three outputs before populating indices 1 and 2 +is behavior-preserving. (`PagedAttention` encodes the same invariant via its nested `!= 3` +check.) + +**Defense-in-depth (recommended): bound the sink** so a future author cannot reintroduce the +class. + +```cpp +// onnxruntime/core/graph/graph.cc - InferenceContextImpl::getOutputType +TypeProto* getOutputType(size_t index) override { + if (index >= node_output_types_.size()) { + fail_type_inference("output index ", index, " is out of range; node has ", + node_output_types_.size(), " outputs"); + } + return &node_output_types_[index]; +} +``` + +This mirrors `getInputType`'s `.at()` and the existing bounds checks in the sibling +`DataPropagationContextImpl`. Placing it at the base layer transitively protects the NHWC and +quantization wrapper contexts. After the point fix this branch is unreachable through a normal +model (the guard already prevents the out-of-range index), so it is pure defense-in-depth. Its +failure mode is build-dependent: with exceptions enabled, `fail_type_inference` raises +`InferenceError` (a `std::exception`), caught by the existing `ORT_CATCH(const std::exception&)` +around `RunInferencing` and surfaced as a clean load-time error; under `ORT_NO_EXCEPTIONS` it is +**not** compiled out - ONNX's no-exceptions path prints the message to `std::cerr` and calls +`abort()`, a deterministic fail-fast (consistent with `getInputType`'s `.at()`, which likewise +terminates under no-exceptions). Either way the result is a controlled failure rather than an +out-of-range write. + +## 5. Test recipe + +Tests live in `onnxruntime/test/contrib_ops/*.cc` and are **auto-globbed** into the +`onnxruntime_provider_test` target by `cmake/onnxruntime_unittests.cmake` +(`test/contrib_ops/*.cc` pattern) - **no cmake edit needed** for a new file. See the +`ort-test` skill for the executable taxonomy (`onnxruntime_provider_test` vs +`onnxruntime_test_all`). + +Rules that make the regression test actually guard the fix: + +1. **Drive through `Model` + `Graph::Resolve`**, not ONNX's standalone `TestShapeInference`. + Only the full resolve path constructs the real `InferenceContextImpl` and hits the + `getOutputType` sink described in §2. A standalone ONNX shape-inference helper uses a + different context and **bypasses** the sink, so it cannot reproduce the bug. +2. **Negative tests must be NON-VACUOUS** - they must actually enter the write branch on + pre-fix source. If a write is gated by an inner condition (§3 gotcha), satisfy it: e.g. for + `MultiHeadAttention`/`DecoderMaskedMultiHeadAttention`, supply a **shaped `past_key`** + (and `past_sequence_length` / `past_present_share_buffer` as the op requires) so the + index-2 block runs. A negative test that only supplies `query` skips the block and passes + even on pre-fix source - regression-proof in name only. +3. **Add positive (all-outputs) cases**: a node with every output present must still infer the + trailing output types - proves the tightened guard did not over-restrict. +4. **Keep tests throw-free post-fix** so they are valid under `ORT_NO_EXCEPTIONS`. Any case + that is *expected* to `fail_shape_inference` (throws) must be excluded with + `#ifndef ORT_NO_EXCEPTIONS`. The "2 outputs must not go out of range" case is throw-free + after the point fix and is safe in all builds. + +**Verify the negative test is non-vacuous (sanitizer A/B)** - the most reliable way to prove a +negative test enters the previously-out-of-range branch: build the test at the **pre-fix** +commit with **AddressSanitizer** and confirm it flags the out-of-range output access; then +confirm it is clean after the fix. + +```bash +# Functional run (any Debug build): +cmake --build build/Linux/Debug --target onnxruntime_provider_test -j"$(nproc)" +./build/Linux/Debug/onnxruntime_provider_test \ + --gtest_filter='AttentionOptionalOutputsShapeInferenceTest.*' + +# A/B proof (isolated worktree at the pre-fix commit, CPU-only Debug + sanitizer): +git worktree add --detach ../ort-prefix-check ~1 +# copy the new test file in, then: +python3 tools/ci_build/build.py --build_dir build/asan --config Debug --parallel \ + --skip_tests --enable_address_sanitizer --skip_submodule_sync \ + --cmake_generator Ninja --target onnxruntime_provider_test +# Pre-fix: the negative tests fail (the sanitizer flags the out-of-range output access). +# Post-fix (cherry-pick the guard fix): all tests pass, no sanitizer report. +``` + +## 6. Process / wording conventions + +- Run **`lintrunner -a`** before pushing so the `CLANGFORMAT` / Python-format gate passes. See + the `ort-lint` skill. +- Use **correctness/robustness framing** in code, comments, commit messages, and the PR body + - describe the change as fixing an optional-output guard, not as a security fix. This + matches repo convention (compare `python-kwargs-setattr-security`) and keeps the PR neutral. + +## 7. Audit checklist (per-operator review) + +When reviewing or hardening any operator implementation or its shape inference: + +- [ ] Read the op's spec - ONNX standard op page, or for a contrib op its `OpSchema` + registration (`.Input/.Output/.Attr`, and `Optional`/`Variadic` markers). A local ONNX + checkout has the standard-op spec pages; contrib ops are defined only in ORT. +- [ ] Enumerate **all** inputs, attributes, and outputs, noting which are optional and the + resulting `min/max` input and output counts. +- [ ] Validate every input/attribute before indexing into it, to avoid out-of-range reads + (which can cascade into worse failures). Match each output-index write to a guard that + guarantees the index is in range (§3 rule). +- [ ] Prefer `ORT_RETURN_IF` / `ORT_RETURN_IF_NOT` for validation; use `ORT_ENFORCE` in + constructors. In shape inference use `fail_shape_inference` / `fail_type_inference`. +- [ ] Use `SafeInt<>` / `narrow<>()` for index and size arithmetic and casts to avoid overflow + or truncation that yields a wrong index. See `core/common/safeint.h` and + `docs/Coding_Conventions_and_Standards.md`. +- [ ] Ensure tests build and pass under **no-exceptions** builds; `#ifndef ORT_NO_EXCEPTIONS` + around any case expected to throw. +- [ ] Exclude EPs known not to support the op, with a comment explaining why. +- [ ] Check whether **other EPs (notably CUDA/ROCm)** implement the same op and whether the + same issue exists there. (For *shape inference* specifically, the logic is EP-agnostic + and single-source - confirm there is no kernel-side analogue.) + +## References + +- **PR #28176** - "Fix ... in EmbedLayerNormalizationShapeInference": the precedent that fixed + the identical `> 1` → `> 2` primitive in one site; the sibling attention sites were missed, + motivating the sweep in §3. +- **PR #29243** - this fix: guards corrected in `DecoderAttention` / `MultiHeadAttention` / + `DecoderMaskedMultiHeadAttention` shape inference, plus the `getOutputType` bounds check and + non-vacuous regression tests. +- Sibling skill: **`ort-test`** (test executables, `--gtest_filter`, contrib-op test layout); + **`ort-lint`** (`lintrunner -a`); **`ort-build`** (build flags, ASan). diff --git a/docs/Coding_Conventions_and_Standards.md b/docs/Coding_Conventions_and_Standards.md index 02af7ddaa49be..6dcf7e6dd43bd 100644 --- a/docs/Coding_Conventions_and_Standards.md +++ b/docs/Coding_Conventions_and_Standards.md @@ -109,6 +109,7 @@ void foo(gsl::span names) { * Use [SafeInt](https://github.com/dcleblanc/SafeInt) when calculating the size of memory to allocate to protect against overflow errors * `#include "core/common/safeint.h"` * search for `SafeInt` in the code for examples +* In operator shape inference, validate every output index against `getNumOutputs()` before writing it. Optional trailing outputs lower an op's `min_output`, so a node may legally declare fewer outputs than the schema's maximum; guard each optional output by the exact index it populates so a `getNumOutputs() > N` guard never writes an index greater than `N`. * The following C++ warnings should never be disabled in onnxruntime VC++ projects(Required by [Binskim](https://github.com/microsoft/binskim/blob/d9afb65c89a621411efded74c27999281d87867e/src/BinSkim.Rules/PERules/BA2007.EnableCriticalCompilerWarnings.cs)). 1. [4018](https://docs.microsoft.com/en-us/cpp/error-messages/compiler-warnings/compiler-warning-level-3-c4018) 'token' : signed/unsigned mismatch 2. [4146](https://docs.microsoft.com/en-us/cpp/error-messages/compiler-warnings/compiler-warning-level-2-c4146?view=msvc-160) unary minus operator applied to unsigned type, result still unsigned From 0d807d77e35203dcbc3be7752b0191d9365e8bdd Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Wed, 24 Jun 2026 22:34:05 +0000 Subject: [PATCH 08/36] contrib: validate MaxpoolWithMask kernel rank matches input spatial rank Residual hardening for MaxpoolWithMask::Compute: the 1D/2D/3D dispatch is selected by kernel_shape.size() and reads x_shape[2..4] / output_dims[2..4] accordingly, but the only rank guard was x_shape.NumDimensions() >= 3. A kernel_shape rank that exceeds the input spatial rank (e.g. rank-3 X with a 2D kernel_shape) caused out-of-bounds reads of x_shape[3]/x_shape[4] and output_dims[3]/output_dims[4]. ONNX shape inference catches this for shaped ONNX models, but the kernel must validate at Compute time for defense in depth (ORT-format / shapeless models bypass shape inference). Add ORT_RETURN_IF_NOT(kernel rank == input spatial rank) alongside the existing mask/input guards, and a kExpectFailure regression test that bypasses ONNX shape inference via AddShapeToTensorData(false) to exercise the Compute-time guard directly. Agent-signed-off: Developer (0b207188) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cpu/maxpool_with_mask.h | 5 +++ .../test/contrib_ops/maxpool_mask_test.cc | 33 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h index 7dfb2770a6979..7b7eda12d6ab2 100644 --- a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h +++ b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h @@ -216,6 +216,11 @@ class MaxpoolWithMask : public OpKernel, public PoolBase { "Mask and input spatial dimensions mismatch at dimension ", i, ": mask=", m_shape[i], " input=", x_shape[i]); } + // The pooling kernel rank drives the 1D/2D/3D dispatch below, which reads x_shape[2..4] and + // output_dims[2..4]. Require it to match the input spatial rank so those reads stay in bounds. + ORT_RETURN_IF_NOT(pool_attrs_.kernel_shape.size() == x_shape.NumDimensions() - 2, + "Pooling kernel rank must equal input spatial rank. Got kernel rank: ", + pool_attrs_.kernel_shape.size(), " input spatial rank: ", x_shape.NumDimensions() - 2); TensorShapeVector pads = pool_attrs_.pads; TensorShapeVector kernel_shape = pool_attrs_.kernel_shape; diff --git a/onnxruntime/test/contrib_ops/maxpool_mask_test.cc b/onnxruntime/test/contrib_ops/maxpool_mask_test.cc index ed65700ccd336..7da97302d0f19 100644 --- a/onnxruntime/test/contrib_ops/maxpool_mask_test.cc +++ b/onnxruntime/test/contrib_ops/maxpool_mask_test.cc @@ -161,5 +161,38 @@ TEST(ContribOpTest, MaxPoolWithMask_MaskEmptyBatchDim) { "Mask N and C dimensions must be greater than 0"); } +TEST(ContribOpTest, MaxPoolWithMask_KernelRankMismatch) { + OpTester test("MaxpoolWithMask", 1, onnxruntime::kMSDomain); + + // AddShapeToTensorData(false) omits input shape from the graph so ONNX shape inference is bypassed + // (convPoolShapeInference returns early when hasInputShape is false). This lets the model pass + // Graph::Resolve() and reach Compute() where the kernel-rank guard fires. + test.AddShapeToTensorData(false); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{1, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0}); + // 2D kernel_shape, but X has only one spatial dimension (rank 3). + test.AddAttribute("kernel_shape", std::vector{8, 8}); + + // Input X has shape {1, 1, 8} (rank 3 => one spatial dim) + std::vector x_dims = {1, 1, 8}; + std::vector x_vals(8, 1.0f); + + // Mask M matches X shape so the earlier spatial/rank guards pass and we reach the kernel-rank guard. + std::vector m_dims = {1, 1, 8}; + std::vector m_vals(8, 1); + + // Placeholder output shape and values (not validated since we expect failure) + std::vector expected_dims = {1, 1, 1, 1}; + std::vector expected_vals = {1.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddInput("M", m_dims, m_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(BaseTester::ExpectResult::kExpectFailure, + "Pooling kernel rank must equal input spatial rank"); +} + } // namespace test } // namespace onnxruntime From d1dccb9fbcc93631a2a1d4c29dd45e9e4512f44b Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Wed, 24 Jun 2026 22:40:33 +0000 Subject: [PATCH 09/36] contrib: reject unsupported MaxpoolWithMask pooling rank and DRY rank check Follow-up polish on the kernel-rank guard. The 1D/2D/3D dispatch switch has a default case that already returns INVALID_ARGUMENT, so an oversized kernel rank was never a silent/uninitialized-output bug, but the default produced a vague 'Unsupported pooling size :' message late in Compute after allocating the output. Add an explicit early guard requiring the input spatial rank (== pooling kernel rank) to be in {1, 2, 3}, with a clear message, alongside the existing guards and before output allocation. Extract the repeated x_shape.NumDimensions() - 2 into a single input_spatial_rank local (defined after the rank >= 3 guard so it cannot underflow) and reuse it. Add a MaxPoolWithMask_KernelRankTooLarge regression test (rank-6 X + 4D kernel_shape) that passes the equality guard but trips the new supported-rank guard, asserting the specific message substring. Agent-signed-off: Developer (0b207188) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cpu/maxpool_with_mask.h | 9 ++++-- .../test/contrib_ops/maxpool_mask_test.cc | 32 +++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h index 7b7eda12d6ab2..4c7f2268310ff 100644 --- a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h +++ b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h @@ -216,11 +216,16 @@ class MaxpoolWithMask : public OpKernel, public PoolBase { "Mask and input spatial dimensions mismatch at dimension ", i, ": mask=", m_shape[i], " input=", x_shape[i]); } + // x_shape.NumDimensions() >= 3 is guaranteed above, so this subtraction cannot underflow. + const size_t input_spatial_rank = x_shape.NumDimensions() - 2; // The pooling kernel rank drives the 1D/2D/3D dispatch below, which reads x_shape[2..4] and // output_dims[2..4]. Require it to match the input spatial rank so those reads stay in bounds. - ORT_RETURN_IF_NOT(pool_attrs_.kernel_shape.size() == x_shape.NumDimensions() - 2, + ORT_RETURN_IF_NOT(pool_attrs_.kernel_shape.size() == input_spatial_rank, "Pooling kernel rank must equal input spatial rank. Got kernel rank: ", - pool_attrs_.kernel_shape.size(), " input spatial rank: ", x_shape.NumDimensions() - 2); + pool_attrs_.kernel_shape.size(), " input spatial rank: ", input_spatial_rank); + // Only 1D/2D/3D pooling is implemented by the dispatch below; a larger rank would match no case. + ORT_RETURN_IF_NOT(input_spatial_rank >= 1 && input_spatial_rank <= 3, + "Only 1D, 2D, and 3D pooling are supported. Got input spatial rank: ", input_spatial_rank); TensorShapeVector pads = pool_attrs_.pads; TensorShapeVector kernel_shape = pool_attrs_.kernel_shape; diff --git a/onnxruntime/test/contrib_ops/maxpool_mask_test.cc b/onnxruntime/test/contrib_ops/maxpool_mask_test.cc index 7da97302d0f19..af095a141dfcd 100644 --- a/onnxruntime/test/contrib_ops/maxpool_mask_test.cc +++ b/onnxruntime/test/contrib_ops/maxpool_mask_test.cc @@ -194,5 +194,37 @@ TEST(ContribOpTest, MaxPoolWithMask_KernelRankMismatch) { "Pooling kernel rank must equal input spatial rank"); } +TEST(ContribOpTest, MaxPoolWithMask_KernelRankTooLarge) { + OpTester test("MaxpoolWithMask", 1, onnxruntime::kMSDomain); + + // Bypass ONNX shape inference (see MaxPoolWithMask_KernelRankMismatch) so the model reaches Compute(). + test.AddShapeToTensorData(false); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{1, 1, 1, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0, 0, 0, 0, 0}); + // 4D kernel_shape with a matching rank-6 X (4 spatial dims): passes the kernel-rank equality guard + // but exceeds the supported 1D/2D/3D pooling ranks. + test.AddAttribute("kernel_shape", std::vector{2, 2, 2, 2}); + + // Input X has shape {1, 1, 2, 2, 2, 2} (rank 6 => four spatial dims) + std::vector x_dims = {1, 1, 2, 2, 2, 2}; + std::vector x_vals(16, 1.0f); + + // Mask M matches X so the earlier guards and the kernel-rank equality guard all pass. + std::vector m_dims = {1, 1, 2, 2, 2, 2}; + std::vector m_vals(16, 1); + + // Placeholder output shape and values (not validated since we expect failure) + std::vector expected_dims = {1, 1, 1, 1}; + std::vector expected_vals = {1.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddInput("M", m_dims, m_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(BaseTester::ExpectResult::kExpectFailure, + "Only 1D, 2D, and 3D pooling are supported"); +} + } // namespace test } // namespace onnxruntime From 0a3b34529d45240e23ba3dea8e53ff122cf29a13 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 00:55:11 +0000 Subject: [PATCH 10/36] contrib: validate DynamicQuantizeLSTM recurrence quantization parameter shapes The recurrence zero-point (R_zero_point) and scale (R_scale) inputs were not run through the kernel's shape validation: R_zp_shape was bound to w_zp->Shape() instead of r_zp->Shape(), and the R_scale WeightCheck was passed W_scale_shape instead of R_scale_shape. As a result a malformed r_zp/r_scale (e.g. a shape inconsistent with R) was never rejected and downstream code iterated using the input parameter's element count over the recurrence parameter's buffer. Bind R_zp_shape to r_zp->Shape() and validate R_scale against R_scale_shape so the recurrence parameters are checked symmetrically with the input (W) ones. Add two expect-failure tests that supply R_zero_point / R_scale shapes whose first dimension does not match num_directions and assert the specific INVALID_ARGUMENT diagnostics. Status-based validation, so the tests need no ORT_NO_EXCEPTIONS guard. CPU-only op; no CUDA implementation exists. Agent-signed-off: Developer (0b207188) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cpu/quantization/dynamic_quantize_lstm.cc | 4 +- .../test/contrib_ops/quantize_lstm_op_test.cc | 79 +++++++++++++++++++ 2 files changed, 81 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc index 2094af78f40b7..ba3c8a2ada30b 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc @@ -178,14 +178,14 @@ Status DynamicQuantizeLSTM::Compute(OpKernelContext* context) const { const Tensor* r_zp = context->Input(11); const TensorShape& W_zp_shape = w_zp->Shape(); - const TensorShape& R_zp_shape = w_zp->Shape(); + const TensorShape& R_zp_shape = r_zp->Shape(); const TensorShape& W_scale_shape = w_scale->Shape(); const TensorShape& R_scale_shape = r_scale->Shape(); WeightCheck(W_zp_shape, W_zero_point); WeightCheck(R_zp_shape, R_zero_point); WeightCheck(W_scale_shape, W_scale); - WeightCheck(W_scale_shape, R_scale); + WeightCheck(R_scale_shape, R_scale); const bool is_W_signed = (W != nullptr) ? W->IsDataType() : is_W_signed_; const bool is_R_signed = (R != nullptr) ? R->IsDataType() : is_R_signed_; diff --git a/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc b/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc index 93a63ce19eb03..a0eceb8528b0a 100644 --- a/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc +++ b/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc @@ -525,5 +525,84 @@ TEST(DynamicQuantLSTMTest, SharedPrepackedWeights) { } #endif +// Builds a minimal per-tensor DynamicQuantizeLSTM and runs it expecting the kernel's input +// validation to reject the recurrence quantization parameters. The caller supplies the R_scale and +// R_zero_point shapes so a shape that is inconsistent with R (e.g. first dim != num_directions) can +// be exercised. The recurrence parameters must be validated symmetrically with the input (W) ones. +static void RunQuantLSTMExpectInvalidRecurrenceQuantParam(const std::vector& r_scale_dims, + const std::vector& r_zp_dims, + const std::string& expected_error) { + OpTester test("DynamicQuantizeLSTM", 1 /*opset_version*/, onnxruntime::kMSDomain /*domain*/); + + constexpr int64_t num_directions = 1; + constexpr int64_t input_size = 2; + constexpr int64_t hidden_size = 2; + constexpr int64_t batch_size = 1; + constexpr int64_t seq_len = 1; + + auto num_elements = [](const std::vector& dims) { + int64_t count = 1; + for (int64_t dim : dims) { + count *= dim; + } + return static_cast(count); + }; + + test.AddAttribute>("activations", {"sigmoid", "tanh", "tanh"}); + test.AddAttribute("direction", "forward"); + test.AddAttribute("hidden_size", hidden_size); + test.AddAttribute("input_forget", static_cast(0)); + + // X: [seq_length, batch_size, input_size] + test.AddInput("X", {seq_len, batch_size, input_size}, + std::vector(num_elements({seq_len, batch_size, input_size}), 0.0f)); + + // W / R quantized weight values are irrelevant: validation fails before any dequantization. + test.AddInput("W", {num_directions, input_size, 4 * hidden_size}, + std::vector(num_elements({num_directions, input_size, 4 * hidden_size}), 0)); + test.AddInput("R", {num_directions, hidden_size, 4 * hidden_size}, + std::vector(num_elements({num_directions, hidden_size, 4 * hidden_size}), 0)); + + test.AddOptionalInputEdge(); // B + test.AddOptionalInputEdge(); // sequence_lens + test.AddInput("initial_h", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + test.AddInput("initial_c", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + test.AddOptionalInputEdge(); // P + + // Valid per-tensor quantization parameters for the input weights. + test.AddInput("W_scale", {num_directions}, std::vector(num_directions, 1.0f)); + test.AddInput("W_zero_point", {num_directions}, std::vector(num_directions, 0)); + + // Recurrence parameters with caller-supplied (possibly inconsistent) shapes. + test.AddInput("R_scale", r_scale_dims, std::vector(num_elements(r_scale_dims), 1.0f)); + test.AddInput("R_zero_point", r_zp_dims, std::vector(num_elements(r_zp_dims), 0)); + + // Placeholder outputs (not validated: the run fails during input validation). + test.AddOutput("Y", {seq_len, num_directions, batch_size, hidden_size}, + std::vector(num_elements({seq_len, num_directions, batch_size, hidden_size}), 0.0f)); + test.AddOutput("Y_h", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + test.AddOutput("Y_c", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, expected_error); +} + +TEST(DynamicQuantLSTMTest, RejectsInconsistentRecurrenceZeroPointShape) { + // R_zero_point's first dim must equal num_directions (1); {2} is inconsistent and must be rejected + // rather than silently validated against the input zero point's shape. + RunQuantLSTMExpectInvalidRecurrenceQuantParam(/*r_scale_dims=*/{1}, /*r_zp_dims=*/{2}, + "Input R_zero_point must have shape"); +} + +TEST(DynamicQuantLSTMTest, RejectsInconsistentRecurrenceScaleShape) { + // R_scale's first dim must equal num_directions (1); {2} is inconsistent and must be rejected + // rather than silently validated against the input scale's shape. + RunQuantLSTMExpectInvalidRecurrenceQuantParam(/*r_scale_dims=*/{2}, /*r_zp_dims=*/{1}, + "Input R_scale must have shape"); +} + } // namespace test } // namespace onnxruntime From eed49565bf4b15b74fd95849ef7022be54fb93d9 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:01:58 +0000 Subject: [PATCH 11/36] contrib: validate GatherBlockQuantized CUDA inputs Port the host-side shape validations from the CPU implementation into the CUDA ComputeInternal: data/scales rank and per-dimension equality, scales/zero_points rank and per-dimension equality, and block_size > 0 (the constructor permits the unset default of 0, which is an invalid divisor for block mapping). Reuse a single packing-components value across the output-shape and quantization-parameter logic. Validate the gathered index value on the device: the kernel sets a device flag when an index falls outside [-gather_axis_dim, gather_axis_dim) and normalizes negative indices to match the CPU kernel; the host reads the flag back after the launch and returns a clean status. Index arithmetic stays in int64. Add CUDA-scoped expect-failure tests for mismatched scales shape, out-of-range index, and block_size == 0. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../quantization/gather_block_quantized.cc | 68 ++++++++++++++++--- .../quantization/gather_block_quantized.cu | 17 ++++- .../quantization/gather_block_quantized.cuh | 4 ++ .../gather_block_quantized_op_test.cc | 26 +++++++ 4 files changed, 103 insertions(+), 12 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc index 9b91215eba91d..de293506ff47e 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc @@ -71,6 +71,46 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) ORT_ENFORCE(quantize_axis_ == static_cast(data_rank) - 1); + // block_size is used as a divisor when mapping an element to its quantization block, + // so it must be strictly positive (the constructor permits the unset default of 0). + ORT_RETURN_IF_NOT(block_size_ > 0, "'block_size' must be greater than 0."); + + // When sub-byte data is packed into uint8_t, the last (quantize) dimension stores + // multiple values per stored element. components is the packing factor. + uint32_t components = 1; + if constexpr (std::is_same_v) { + components = 8 / static_cast(bits_); + } + + // Validate scales shape against data shape, matching the CPU kernel: every dimension + // must be equal except the quantize axis, which is divided into block_size blocks. + auto scales_shape = scales->Shape().GetDims(); + ORT_RETURN_IF_NOT(data_rank == static_cast(scales->Shape().NumDimensions()), + "data and scales must have the same rank."); + for (size_t i = 0; i < data_shape.size(); ++i) { + ORT_RETURN_IF_NOT(static_cast(i) == quantize_axis_ + ? (data_shape[i] * components + block_size_ - 1) / block_size_ == scales_shape[i] + : data_shape[i] == scales_shape[i], + "data and scales do not match shapes."); + } + + // Validate zero_points shape against scales shape when provided. + if (zero_points != nullptr) { + auto zero_points_shape = zero_points->Shape().GetDims(); + ORT_RETURN_IF_NOT(scales->Shape().NumDimensions() == zero_points->Shape().NumDimensions(), + "scales and zero_points must have the same rank."); + for (size_t i = 0; i < scales_shape.size(); ++i) { + if (components > 1 && static_cast(i) == quantize_axis_) { + // For uint8_t with bits < 8, zero_points are packed components per stored element. + ORT_RETURN_IF_NOT((scales_shape[i] + components - 1) / components == zero_points_shape[i], + "scales and zero_points shape does not match."); + } else { + ORT_RETURN_IF_NOT(scales_shape[i] == zero_points_shape[i], + "scales and zero_points must have the same shape."); + } + } + } + TensorShapeVector output_shape; output_shape.reserve(data_rank - 1 + indices_rank); @@ -98,11 +138,8 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) } // Special int4‐in‐uint8 packing tweak: expand the last dim by components - if constexpr (std::is_same_v) { - uint32_t components = 8 / static_cast(bits_); - if (components > 1) { - output_shape.back() *= components; - } + if (components > 1) { + output_shape.back() *= components; } Tensor* output = ctx->Output(0, TensorShape(output_shape)); @@ -123,11 +160,8 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) // after_gather_dim has to be adjusted to match // the unpacked output dims for correct kernel indexing int64_t after_gather_dim_unpacked = after_gather_dim; - if constexpr (std::is_same_v) { - uint32_t components = 8 / static_cast(bits_); - if (components > 1) { - after_gather_dim_unpacked *= components; - } + if (components > 1) { + after_gather_dim_unpacked *= components; } GatherBlockQuantizedParam param; @@ -140,6 +174,12 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) param.gather_axis = gather_axis_; param.N = N; + // Device flag the kernel sets when a gathered index is outside the valid range. + // It is initialized to 0 and read back after the launch to surface a clean status. + auto index_out_of_bounds = GetScratchBuffer(1, ctx->GetComputeStream()); + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(index_out_of_bounds.get(), 0, sizeof(int), param.stream)); + param.index_out_of_bounds = index_out_of_bounds.get(); + const auto dequantized_type = scales->GetElementType(); if (dequantized_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { const auto* scales_ptr = static_cast(scales->DataRaw()); @@ -155,6 +195,14 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) LaunchGatherBlockQuantizedKernel(data_ptr, indices_ptr, scales_ptr, zero_points_ptr, output_ptr, param); } + auto host_index_out_of_bounds = AllocateBufferOnCPUPinned(1); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(host_index_out_of_bounds.get(), index_out_of_bounds.get(), sizeof(int), + cudaMemcpyDeviceToHost, param.stream)); + CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(param.stream)); + ORT_RETURN_IF(*host_index_out_of_bounds.get() != 0, + "indices element out of data bounds. Each index must be within the inclusive range [", + -param.gather_axis_dim, ", ", param.gather_axis_dim - 1, "]."); + return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu index 39286c63e9a08..a40138f0d3bb3 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu @@ -44,7 +44,8 @@ __global__ void GatherBlockQuantizedKernel( int64_t block_size, int64_t gather_axis, int64_t N, - bool sign) { + bool sign, + int* index_out_of_bounds) { int64_t out_idx = blockDim.x * blockIdx.x + threadIdx.x; if (out_idx >= N) return; @@ -53,6 +54,18 @@ __global__ void GatherBlockQuantizedKernel( int64_t idx_after = out_idx % after_gather_dim; int64_t idx = (out_idx % (after_gather_dim * ind_dim)) / after_gather_dim; int64_t idx_at_g = indices[idx]; + + // Validate the gathered index and normalize negatives, mirroring the CPU kernel. + // A negative index counts from the end of the gather axis; anything outside + // [-gather_axis_dim, gather_axis_dim) is reported via the device error flag. + if (idx_at_g < -gather_axis_dim || idx_at_g >= gather_axis_dim) { + *index_out_of_bounds = 1; + return; + } + if (idx_at_g < 0) { + idx_at_g += gather_axis_dim; + } + int64_t in_idx = idx_before * gather_axis_dim * after_gather_dim + idx_at_g * after_gather_dim + idx_after; int64_t block_id = in_idx / block_size; @@ -82,7 +95,7 @@ void LaunchGatherBlockQuantizedKernel(const T1* data, bool sign = std::is_same::value; GatherBlockQuantizedKernel<<>>(data, indices, scales, zero_points, output, - param.after_gather_dim, param.gather_axis_dim, param.ind_dim, param.bits, param.block_size, param.gather_axis, param.N, sign); + param.after_gather_dim, param.gather_axis_dim, param.ind_dim, param.bits, param.block_size, param.gather_axis, param.N, sign, param.index_out_of_bounds); } template void LaunchGatherBlockQuantizedKernel(const uint8_t*, const int32_t*, const float*, const uint8_t*, float*, GatherBlockQuantizedParam); diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh index f5dea3b1f2d9d..36266121ac210 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh @@ -19,6 +19,10 @@ struct GatherBlockQuantizedParam { int64_t block_size; int64_t gather_axis; int64_t N; + // Device buffer of a single int. The kernel sets it to 1 if any gathered index + // is outside the valid range [-gather_axis_dim, gather_axis_dim). The host reads + // it back after the launch and returns an error status when it is set. + int* index_out_of_bounds; }; template diff --git a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc index 7fbe71de5251a..f7cce522c65eb 100644 --- a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc +++ b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc @@ -507,6 +507,32 @@ TEST(GatherBlockQuantizedOpTest, InvalidIndices) { } #endif +#ifdef USE_CUDA +// The CUDA EP validates quantization-parameter shapes on the host and gathered-index +// bounds on the device. These cases mirror the CPU coverage above but run on the CUDA +// EP, which previously skipped these checks. They fall back to the CPU EP (which rejects +// the same inputs) when no CUDA device is present. +TEST(GatherBlockQuantizedOpTest, CudaShapeMismatch) { + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); +} + +TEST(GatherBlockQuantizedOpTest, CudaInvalidIndices) { + Test_InvalidIndices_WithZeroPoints(); + Test_InvalidIndices_WithZeroPoints(); + Test_InvalidIndices_WithZeroPoints(); +} + +// block_size == 0 is accepted by the constructor but is an invalid divisor for block +// mapping, so ComputeInternal must reject it with a clean status. +TEST(GatherBlockQuantizedOpTest, CudaInvalidBlockSizeZero) { + Test_Fail_WithoutZeroPoints(0, 2, 0); + Test_Fail_WithoutZeroPoints(0, 2, 0); + Test_Fail_WithoutZeroPoints(0, 2, 0); +} +#endif + template void Test_GatherAxis0_WithZeroPoints(int bits = 4) { std::vector data = {-8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8, From af78a8e01029007c1e579e5297de64a7d9c55350 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:04:57 +0000 Subject: [PATCH 12/36] contrib: validate EmbedLayerNormalization CUDA embedding indices Harden EmbedLayerNormalization input handling so malformed inputs are rejected with a clear INVALID_ARGUMENT instead of reading past the embedding tables. - Shared CheckInputs helper: when position_ids are not supplied, require position_embedding to have at least sequence_length rows. - CUDA kernel: validate word/position/segment ids against their embedding table sizes, surfacing an out-of-range id via a device error flag instead of silently clamping; widen the embedding offset arithmetic to int64 to avoid overflow. - CUDA op: read back the device error flag and return INVALID_ARGUMENT. - Tests: enable EmbedLayerNormNegativePositionIds on CUDA and add EmbedLayerNormWordIdOutOfRange and EmbedLayerNormPositionEmbeddingTooFewRows expect-failure cases covering CPU and CUDA. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cpu/bert/embed_layer_norm_helper.h | 16 +++ .../contrib_ops/cuda/bert/embed_layer_norm.cc | 27 ++++- .../cuda/bert/embed_layer_norm_impl.cu | 42 ++++++-- .../cuda/bert/embed_layer_norm_impl.h | 6 +- .../contrib_ops/embed_layer_norm_op_test.cc | 98 ++++++++++++++++++- 5 files changed, 175 insertions(+), 14 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h index b0cf1fc976de0..64b5439b87c7b 100644 --- a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h @@ -24,8 +24,13 @@ inline Status CheckInputs(const TOpKernelContext* context, bool quantizedVersion const Tensor* beta = context->template Input(6); const Tensor* mask = context->template Input(7); // optional. nullptr if not provided + // Track whether explicit position ids are supplied. When they are not, positions default to + // [0, sequence_length) and must be covered by the position_embedding rows (checked below). + bool has_position_ids = false; + if (!quantizedVersion) { const Tensor* position_ids = context->template Input(8); // optional. nullptr if not provided + has_position_ids = (nullptr != position_ids); if (nullptr != position_ids) { if (input_ids->Shape()[1] != position_ids->Shape()[1]) { @@ -69,6 +74,17 @@ inline Status CheckInputs(const TOpKernelContext* context, bool quantizedVersion "position_embedding is expected to have 2 dimensions, got ", position_embedding_dims.size()); } + // When position_ids are not supplied, positions run over [0, sequence_length); the + // position_embedding table must have at least that many rows to be addressable. + if (!has_position_ids) { + int64_t sequence_length = input_ids->Shape()[1]; + if (position_embedding_dims[0] < sequence_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "position_embedding must have at least sequence_length (", sequence_length, + ") rows when position_ids is not provided, got ", position_embedding_dims[0]); + } + } + if (nullptr != segment_embedding) { const auto& segment_embedding_dims = segment_embedding->Shape().GetDims(); if (segment_embedding_dims.size() != 2) { diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc index 864e2d1623923..faf4cb805ed5a 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc @@ -61,9 +61,18 @@ Status EmbedLayerNorm::ComputeInternal(OpKernelContext* context) const { int sequence_length = static_cast(input_dims[1]); size_t element_size = sizeof(T); + int word_embedding_length = static_cast(word_embedding->Shape()[0]); + int position_embedding_length = static_cast(position_embedding->Shape()[0]); + int segment_embedding_length = + (nullptr == segment_embedding) ? 0 : static_cast(segment_embedding->Shape()[0]); + const bool broadcast_position_ids = (nullptr != position_ids && position_ids->Shape()[0] == 1); - return LaunchEmbedLayerNormKernel( + // Device flag raised by the kernel when an input id is outside its embedding table. + auto error_flag = GetScratchBuffer(1, context->GetComputeStream()); + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(error_flag.get(), 0, sizeof(int), Stream(context))); + + ORT_RETURN_IF_ERROR(LaunchEmbedLayerNormKernel( Stream(context), output->MutableData(), nullptr == mask_index ? nullptr : mask_index->MutableData(), @@ -82,7 +91,21 @@ Status EmbedLayerNorm::ComputeInternal(OpKernelContext* context) const { element_size, embedding_sum == nullptr ? nullptr : embedding_sum->MutableData(), position_ids == nullptr ? nullptr : position_ids->Data(), - broadcast_position_ids); + broadcast_position_ids, + word_embedding_length, + position_embedding_length, + segment_embedding_length, + error_flag.get())); + + // Surface any out-of-range id as a clean error instead of leaving an invalid output. + auto host_error_flag = AllocateBufferOnCPUPinned(1); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(host_error_flag.get(), error_flag.get(), sizeof(int), + cudaMemcpyDeviceToHost, Stream(context))); + CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(Stream(context))); + ORT_RETURN_IF(*host_error_flag.get() != 0, + "input id is out of range of the corresponding embedding table."); + + return Status::OK(); } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu index a6b97286d1733..ea06a0da309f2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu @@ -123,7 +123,8 @@ template __global__ void EmbedLayerNormKernel( int hidden_size, const int* input_ids, const int* segment_ids, const T* beta, const T* gamma, const T* word_embedding, const T* position_embedding, const T* segment_embedding, - float epsilon, T* output, T* embedding_sum, const int* position_ids, const bool broadcast_position_ids) { + float epsilon, T* output, T* embedding_sum, const int* position_ids, const bool broadcast_position_ids, + int word_embedding_length, int position_embedding_length, int segment_embedding_length, int* error_flag) { KeyValuePairSum pair_sum; // 1. lookup word and segment of the block // blockIdx.x = position in the sequence @@ -133,6 +134,7 @@ __global__ void EmbedLayerNormKernel( __shared__ int word_id; __shared__ int segment_id; __shared__ int position_id; + __shared__ bool is_valid_block; const float rld = 1.f / hidden_size; const int sequence_position = blockIdx.y * gridDim.x + blockIdx.x; @@ -150,15 +152,28 @@ __global__ void EmbedLayerNormKernel( } else { position_id = position_ids[sequence_position]; } + + // Reject ids that fall outside their embedding tables instead of reading past the buffers. + is_valid_block = (word_id >= 0 && word_id < word_embedding_length) && + (position_id >= 0 && position_id < position_embedding_length) && + (nullptr == segment_embedding || + (segment_id >= 0 && segment_id < segment_embedding_length)); + if (!is_valid_block) { + *error_flag = 1; + } } __syncthreads(); + if (!is_valid_block) { + return; + } + // 2. load pos/segment/word embeddings and add them together // offset into embeddings is given by word_id * hidden_size - const int position_offset = position_id * hidden_size; + const int64_t position_offset = static_cast(position_id) * hidden_size; - const int word_offset = word_id * hidden_size; - const int segment_offset = segment_id * hidden_size; + const int64_t word_offset = static_cast(word_id) * hidden_size; + const int64_t segment_offset = static_cast(segment_id) * hidden_size; // the output offset is given by b * (sequence_length * hidden_size) + s * hidden_size const int output_offset = sequence_position * hidden_size; @@ -192,7 +207,8 @@ Status EmbedSkipLayerNorm( const int* input_ids, const int* segment_ids, const T* beta, const T* gamma, const T* word_embedding, const T* position_embedding, const T* segment_embedding, float epsilon, T* output, T* embedding_sum, const int* position_ids, - const bool broadcast_position_ids) { + const bool broadcast_position_ids, int word_embedding_length, int position_embedding_length, + int segment_embedding_length, int* error_flag) { constexpr int tpb = 256; const dim3 grid(sequence_length, batch_size, 1); const dim3 block(tpb, 1, 1); @@ -200,7 +216,9 @@ Status EmbedSkipLayerNorm( EmbedLayerNormKernel <<>>(hidden_size, input_ids, segment_ids, beta, gamma, word_embedding, position_embedding, segment_embedding, - epsilon, output, embedding_sum, position_ids, broadcast_position_ids); + epsilon, output, embedding_sum, position_ids, broadcast_position_ids, + word_embedding_length, position_embedding_length, + segment_embedding_length, error_flag); return CUDA_CALL(cudaGetLastError()); } @@ -224,7 +242,11 @@ Status LaunchEmbedLayerNormKernel( const size_t element_size, void* embedding_sum, const int* position_ids, - const bool broadcast_position_ids) { + const bool broadcast_position_ids, + int word_embedding_length, + int position_embedding_length, + int segment_embedding_length, + int* error_flag) { if (mask_index != nullptr) { if (nullptr == input_mask) { CUDA_RETURN_IF_ERROR(cudaMemsetAsync(mask_index, 0, sizeof(int) * batch_size, stream)); @@ -241,7 +263,8 @@ Status LaunchEmbedLayerNormKernel( reinterpret_cast(word_embedding), reinterpret_cast(position_embedding), reinterpret_cast(segment_embedding), epsilon, reinterpret_cast(output), reinterpret_cast(embedding_sum), position_ids, - broadcast_position_ids); + broadcast_position_ids, word_embedding_length, position_embedding_length, + segment_embedding_length, error_flag); } else { return EmbedSkipLayerNorm( stream, hidden_size, batch_size, sequence_length, input_ids, segment_ids, @@ -249,7 +272,8 @@ Status LaunchEmbedLayerNormKernel( reinterpret_cast(word_embedding), reinterpret_cast(position_embedding), reinterpret_cast(segment_embedding), epsilon, reinterpret_cast(output), reinterpret_cast(embedding_sum), position_ids, - broadcast_position_ids); + broadcast_position_ids, word_embedding_length, position_embedding_length, + segment_embedding_length, error_flag); } } diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h index aaf1d891dab3a..b9d3bbc1d9304 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h @@ -26,7 +26,11 @@ Status LaunchEmbedLayerNormKernel(cudaStream_t stream, const size_t element_size, // size of output element: 2 for half, 4 for float. void* embedding_sum, // Optional output of sum of embeddings const int* position_ids, // Optional input of position ids - const bool broadcast_position_ids); // Whether to broadcast position ids + const bool broadcast_position_ids, // Whether to broadcast position ids + int word_embedding_length, // number of rows in word_embedding + int position_embedding_length, // number of rows in position_embedding + int segment_embedding_length, // number of rows in segment_embedding (0 if none) + int* error_flag); // device flag set when an id is out of range } // namespace cuda } // namespace contrib diff --git a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc b/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc index d32c7fc224fd8..5ce54eda4157c 100644 --- a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc +++ b/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc @@ -282,8 +282,102 @@ TEST(EmbedLayerNormTest, EmbedLayerNormNegativePositionIds) { tester.AddOutput("output", output_dims, std::vector(batch_size * sequence_size * hidden_size, 0.0f)); tester.AddOutput("mask_index", mask_index_dims, {0}); - // Run CPU only - expect failure due to negative position_ids - tester.Run(OpTester::ExpectResult::kExpectFailure, "", {kCudaExecutionProvider, kCudaNHWCExecutionProvider, kDmlExecutionProvider, kOpenVINOExecutionProvider}); + // Both CPU and CUDA reject the out-of-range position id via input validation; other EPs are + // not exercised here. + tester.Run(OpTester::ExpectResult::kExpectFailure, "", {kDmlExecutionProvider, kOpenVINOExecutionProvider}); +} + +// An input id that points outside the word_embedding table must be rejected rather than read. +// CPU validates per-index already; CUDA validates device-side and surfaces the same failure, so +// this case runs on both EPs. +TEST(EmbedLayerNormTest, EmbedLayerNormWordIdOutOfRange) { + int batch_size = 1; + int sequence_size = 2; + int hidden_size = 4; + + std::vector input_ids_dims = {batch_size, sequence_size}; + std::vector word_embedding_dims = {6, hidden_size}; + std::vector position_embedding_dims = {3, hidden_size}; + std::vector segment_embedding_dims = {2, hidden_size}; + std::vector gamma_dims = {hidden_size}; + std::vector beta_dims = {hidden_size}; + std::vector output_dims = {batch_size, sequence_size, hidden_size}; + std::vector mask_index_dims = {batch_size}; + + OpTester tester("EmbedLayerNormalization", 1, onnxruntime::kMSDomain); + // Second id (99) is outside the 6-row word_embedding table. + tester.AddInput("input_ids", input_ids_dims, {1, 99}); + tester.AddInput("segment_ids", input_ids_dims, {0, 1}); + tester.AddInput("word_embedding", word_embedding_dims, + {0.2f, 0.1f, 0.4f, -0.6f, + 0.3f, 0.2f, 0.5f, 0.6f, + 0.6f, 0.7f, 0.0f, -0.1f, + 0.8f, 0.6f, 0.9f, 1.2f, + 0.1f, 0.3f, 0.5f, 0.9f, + 1.0f, -2.0f, 1.1f, 0.8f}, + /*is_initializer=*/true); + tester.AddInput("position_embedding", position_embedding_dims, + {0.1f, 0.1f, 0.4f, 0.6f, + 0.6f, 0.0f, 0.8f, 0.6f, + 0.3f, 0.9f, -2.0f, 0.8f}, + /*is_initializer=*/true); + tester.AddInput("segment_embedding", segment_embedding_dims, + {0.3f, 0.4f, 0.9f, 0.1f, + 0.7f, 0.3f, 0.5f, 0.2f}, + /*is_initializer=*/true); + tester.AddInput("gamma", gamma_dims, {0.25f, 0.15f, 0.45f, -0.66f}, /*is_initializer=*/true); + tester.AddInput("beta", beta_dims, {0.6f, 0.2f, 0.5f, -0.6f}, /*is_initializer=*/true); + tester.AddAttribute("epsilon", embedlayernorm::kEpsilon); + + tester.AddOutput("output", output_dims, std::vector(batch_size * sequence_size * hidden_size, 0.0f)); + tester.AddOutput("mask_index", mask_index_dims, {0}); + + tester.Run(OpTester::ExpectResult::kExpectFailure, "", {kDmlExecutionProvider, kOpenVINOExecutionProvider}); +} + +// Without position_ids, positions default to [0, sequence_length); a position_embedding table with +// fewer rows than sequence_length is rejected by the shared input validation on both CPU and CUDA. +TEST(EmbedLayerNormTest, EmbedLayerNormPositionEmbeddingTooFewRows) { + int batch_size = 1; + int sequence_size = 2; + int hidden_size = 4; + + std::vector input_ids_dims = {batch_size, sequence_size}; + std::vector word_embedding_dims = {6, hidden_size}; + // Only 1 row, but sequence_size is 2 and no position_ids are supplied. + std::vector position_embedding_dims = {1, hidden_size}; + std::vector segment_embedding_dims = {2, hidden_size}; + std::vector gamma_dims = {hidden_size}; + std::vector beta_dims = {hidden_size}; + std::vector output_dims = {batch_size, sequence_size, hidden_size}; + std::vector mask_index_dims = {batch_size}; + + OpTester tester("EmbedLayerNormalization", 1, onnxruntime::kMSDomain); + tester.AddInput("input_ids", input_ids_dims, {1, 3}); + tester.AddInput("segment_ids", input_ids_dims, {0, 1}); + tester.AddInput("word_embedding", word_embedding_dims, + {0.2f, 0.1f, 0.4f, -0.6f, + 0.3f, 0.2f, 0.5f, 0.6f, + 0.6f, 0.7f, 0.0f, -0.1f, + 0.8f, 0.6f, 0.9f, 1.2f, + 0.1f, 0.3f, 0.5f, 0.9f, + 1.0f, -2.0f, 1.1f, 0.8f}, + /*is_initializer=*/true); + tester.AddInput("position_embedding", position_embedding_dims, + {0.1f, 0.1f, 0.4f, 0.6f}, + /*is_initializer=*/true); + tester.AddInput("segment_embedding", segment_embedding_dims, + {0.3f, 0.4f, 0.9f, 0.1f, + 0.7f, 0.3f, 0.5f, 0.2f}, + /*is_initializer=*/true); + tester.AddInput("gamma", gamma_dims, {0.25f, 0.15f, 0.45f, -0.66f}, /*is_initializer=*/true); + tester.AddInput("beta", beta_dims, {0.6f, 0.2f, 0.5f, -0.6f}, /*is_initializer=*/true); + tester.AddAttribute("epsilon", embedlayernorm::kEpsilon); + + tester.AddOutput("output", output_dims, std::vector(batch_size * sequence_size * hidden_size, 0.0f)); + tester.AddOutput("mask_index", mask_index_dims, {0}); + + tester.Run(OpTester::ExpectResult::kExpectFailure, "", {kDmlExecutionProvider, kOpenVINOExecutionProvider}); } } // namespace test From 459a12284ff62552e34c301e99684b7c2fc407ed Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:05:39 +0000 Subject: [PATCH 13/36] Validate GridSample input spatial dimensions are non-empty GridSample sizes its output from the grid, so a zero-size input spatial dimension skips the empty-output early return and reaches the clamp/reflect index computation with an invalid extent. Add an ORT_RETURN_IF_NOT check that every input spatial dimension is greater than zero in both the CPU Compute and the CUDA ComputeInternal (covering 2D and 3D-spatial inputs, NCHW and NHWC layouts), with a consistent diagnostic message. Add an expect-failure OpTester case that fans across the registered CPU and CUDA execution providers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Agent-signed-off: Developer (cc69a935) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cpu/tensor/grid_sample_test_custom.cc | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc b/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc index dd23b76d8275d..9abb362baca22 100644 --- a/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc +++ b/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc @@ -63,6 +63,17 @@ void RunCpuOnly(T& test) { RunTests(test, GetCpuExecutionProviders()); } +template +void RunTestsExpectFailure(T& test, const std::string& error_msg, + std::vector>&& execution_providers) { + for (size_t idx = 0; idx < execution_providers.size(); ++idx) { + test.Config(BaseTester::ExpectResult::kExpectFailure, error_msg) + .ConfigEp(std::move(execution_providers[idx])) + .RunWithConfig(); + } + execution_providers.clear(); +} + } // namespace template @@ -578,5 +589,30 @@ TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_reflection_mixed RunCpuOnly(test); } +TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_border_empty_spatial_rejected) { + // An input with a zero-size spatial dimension is rejected: the output is sized by the grid, + // so a non-empty grid skips the empty-output early return, and spatial dimensions must be + // non-empty for sampling to produce valid sample indices. + OpTester test("GridSample", 20); + test.AddAttribute("mode", std::string("nearest")); + test.AddAttribute("padding_mode", std::string("border")); + test.AddAttribute("align_corners", int64_t{0}); + + std::initializer_list X_shape{1, 1, 0, 5}; + std::initializer_list X_data{}; + std::initializer_list Grid_shape{1, 2, 2, 2}; + std::initializer_list Grid_data{ + TypeParam(0.0f), TypeParam(0.0f), TypeParam(-1.0f), TypeParam(-1.0f), + TypeParam(1.0f), TypeParam(1.0f), TypeParam(0.5f), TypeParam(-0.5f)}; + std::initializer_list Y_shape{1, 1, 2, 2}; + + test.AddInput("X", X_shape, X_data); + test.AddInput("Grid", Grid_shape, Grid_data); + test.AddOutput("Y", Y_shape, + {TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f)}); + RunTestsExpectFailure(test, "Input spatial dimensions must be non-empty for sampling", + GetExecutionProviders()); +} + } // namespace test } // namespace onnxruntime From 607dfb6d2023fbb62c5319c61fc9cafd01e5274d Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:06:12 +0000 Subject: [PATCH 14/36] Validate GridSample input spatial dimensions are non-empty (CPU and CUDA kernels) Add ORT_RETURN_IF_NOT guards that every input spatial dimension is greater than zero in GridSample::Compute (CPU) and GridSample::ComputeInternal (CUDA), covering 2D and 3D-spatial inputs in both NCHW and NHWC layouts. The output is sized by the grid, so a zero-size input spatial dimension would otherwise skip the empty-output early return and reach interpolation with an invalid extent. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/providers/cpu/tensor/grid_sample.cc | 8 ++++++++ onnxruntime/core/providers/cuda/tensor/grid_sample.cc | 11 +++++++++++ 2 files changed, 19 insertions(+) diff --git a/onnxruntime/core/providers/cpu/tensor/grid_sample.cc b/onnxruntime/core/providers/cpu/tensor/grid_sample.cc index 34135ca62e4a8..0299aca0057a3 100644 --- a/onnxruntime/core/providers/cpu/tensor/grid_sample.cc +++ b/onnxruntime/core/providers/cpu/tensor/grid_sample.cc @@ -373,6 +373,14 @@ Status GridSample::Compute(OpKernelContext* context) const { ORT_ENFORCE(input_dims.NumDimensions() == 4 || input_dims.NumDimensions() == 5, "Only 4-D or 5-D tensor is supported"); + // Spatial dimensions must be non-empty for sampling: the output is sized by the grid, so a zero-size + // input spatial dimension would otherwise lead to invalid index computations during interpolation. + for (size_t i = 2; i < input_dims.NumDimensions(); ++i) { + ORT_RETURN_IF_NOT(input_dims[i] > 0, + "Input spatial dimensions must be non-empty for sampling. Dimension ", i, + " has size ", input_dims[i]); + } + auto N = input_dims[0]; auto C = input_dims[1]; ORT_ENFORCE(grid_dims[0] == N, "Grid batch size ", grid_dims[0], " does not match input batch size ", N); diff --git a/onnxruntime/core/providers/cuda/tensor/grid_sample.cc b/onnxruntime/core/providers/cuda/tensor/grid_sample.cc index d97d5fcbb0b5b..e23b5bad573f7 100755 --- a/onnxruntime/core/providers/cuda/tensor/grid_sample.cc +++ b/onnxruntime/core/providers/cuda/tensor/grid_sample.cc @@ -124,6 +124,17 @@ Status GridSample::ComputeInternal(OpKernelContext* context) const { return Status(common::ONNXRUNTIME, common::NOT_IMPLEMENTED, "Cubic mode is only supported in 4-D cases."); } + // Spatial dimensions must be non-empty for sampling: the output is sized by the grid, so a zero-size + // input spatial dimension would otherwise lead to invalid index computations during interpolation. + // Spatial dims are [1, rank-1) for NHWC (N,...,C) and [2, rank) for NCHW (N,C,...). + const size_t spatial_begin = IsNHWC ? 1 : 2; + const size_t spatial_end = IsNHWC ? dims_input.size() - 1 : dims_input.size(); + for (size_t i = spatial_begin; i < spatial_end; ++i) { + ORT_RETURN_IF_NOT(dims_input[i] > 0, + "Input spatial dimensions must be non-empty for sampling. Dimension ", i, + " has size ", dims_input[i]); + } + using Ch = Channels; TensorShapeVector dims_output(dims_input.size()); From 3dc3c25f889c8c3834e05377b2f7c17cf0e746f7 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:07:30 +0000 Subject: [PATCH 15/36] contrib: reject non-positive block_size in GatherBlockQuantized shape inference The GatherBlockQuantized type/shape inference uses block_size as a divisor when relating the data and scales shapes, but only rejected negative values. A block_size of 0 reached the division and raised an integer divide-by-zero during Graph::Resolve. Reject block_size <= 0 with a clear status instead. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/graph/contrib_ops/contrib_defs.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index f3f2f521ecab2..b2a6b70dbb455 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -4072,8 +4072,10 @@ GatherBlockQuantized is a Gather with data quantized. It is similar to Gather (h if (quantize_axis < -r || quantize_axis >= r) { fail_shape_inference("quantize_axis must be in [-r, r-1]"); } - if (block_size < 0) { - fail_shape_inference("block_size must be non-negative"); + if (block_size <= 0) { + // block_size is used as a divisor below when relating data and scales shapes, + // so it must be strictly positive. + fail_shape_inference("block_size must be greater than 0"); } gather_axis = (gather_axis + r) % r; From 283f84ef818c25f5bd63a69785dedf6555cc5f28 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:13:36 +0000 Subject: [PATCH 16/36] contrib: guard bits and tidy GatherBlockQuantized CUDA kernel Validate the bits attribute in the CUDA constructor before it is used as the divisor in 8 / bits to derive the uint8 packing factor: uint8 data accepts bits == 4 or 8 and int4/uint4 accepts bits == 4. Wrap the kernel launch line to stay within the column limit and document the positional parameters of the Test_Fail_WithoutZeroPoints helper. Add CUDA-scoped expect-failure coverage for an unsupported bits value. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/quantization/gather_block_quantized.cc | 9 +++++++++ .../cuda/quantization/gather_block_quantized.cu | 6 ++++-- .../contrib_ops/gather_block_quantized_op_test.cc | 15 +++++++++++++++ 3 files changed, 28 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc index de293506ff47e..b2e5cffa5d54e 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc @@ -47,6 +47,15 @@ template GatherBlockQuantized::GatherBlockQuantized(const OpKernelInfo& info) : CudaKernel(info) { ORT_ENFORCE(info.GetAttr("bits", &bits_).IsOK()); + // bits_ is used as the divisor in 8 / bits_ to derive the uint8 packing factor, so it must + // be one of the supported widths. uint8 data packs 4 or 8 bit values; int4/uint4 are 4 bit. + if constexpr (std::is_same_v) { + ORT_ENFORCE(bits_ == 4 || bits_ == 8, + "GatherBlockQuantized with uint8 data only supports bits == 4 or 8."); + } else { + ORT_ENFORCE(bits_ == 4, "GatherBlockQuantized with int4/uint4 data only supports bits == 4."); + } + block_size_ = info.GetAttrOrDefault("block_size", 0); gather_axis_ = info.GetAttrOrDefault("gather_axis", 0); quantize_axis_ = info.GetAttrOrDefault("quantize_axis", 0); diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu index a40138f0d3bb3..8539b4e5d2869 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu @@ -94,8 +94,10 @@ void LaunchGatherBlockQuantizedKernel(const T1* data, int blocksPerGrid = (int)(ceil(static_cast(param.N) / GridDim::maxThreadsPerBlock)); bool sign = std::is_same::value; - GatherBlockQuantizedKernel<<>>(data, indices, scales, zero_points, output, - param.after_gather_dim, param.gather_axis_dim, param.ind_dim, param.bits, param.block_size, param.gather_axis, param.N, sign, param.index_out_of_bounds); + GatherBlockQuantizedKernel<<>>( + data, indices, scales, zero_points, output, + param.after_gather_dim, param.gather_axis_dim, param.ind_dim, param.bits, param.block_size, + param.gather_axis, param.N, sign, param.index_out_of_bounds); } template void LaunchGatherBlockQuantizedKernel(const uint8_t*, const int32_t*, const float*, const uint8_t*, float*, GatherBlockQuantizedParam); diff --git a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc index f7cce522c65eb..8e010b0b33338 100644 --- a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc +++ b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc @@ -375,6 +375,12 @@ TEST(GatherBlockQuantizedOpTest, UnsupportedTypes) { } #endif +// Runs GatherBlockQuantized with inputs that are expected to be rejected (no zero_points). +// Parameters: +// gather_axis - value for the "gather_axis" attribute +// quantize_axis - value for the "quantize_axis" attribute +// block_size - value for the "block_size" attribute +// bits - value for the "bits" attribute (default 4) template void Test_Fail_WithoutZeroPoints(int64_t gather_axis, int64_t quantize_axis, @@ -531,6 +537,15 @@ TEST(GatherBlockQuantizedOpTest, CudaInvalidBlockSizeZero) { Test_Fail_WithoutZeroPoints(0, 2, 0); Test_Fail_WithoutZeroPoints(0, 2, 0); } + +// uint8 data supports bits in {4, 8} and int4/uint4 supports bits == 4 on CUDA; other +// values are rejected by the constructor before the 8 / bits division that derives the +// uint8 packing factor. +TEST(GatherBlockQuantizedOpTest, CudaUnsupportedBits) { + Test_Fail_WithoutZeroPoints(0, 2, 16, 3); + Test_Fail_WithoutZeroPoints(0, 2, 16, 8); + Test_Fail_WithoutZeroPoints(0, 2, 16, 3); +} #endif template From e90181d3186ccb9902f324061f08496060135c3c Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:14:17 +0000 Subject: [PATCH 17/36] contrib: skip EmbedLayerNormalization error readback during CUDA graph capture The device kernel already validates word/position/segment ids and skips the embedding reads for any out-of-range id, so input handling stays in range regardless of the host-side readback. The readback only upgrades a skipped row into a clean INVALID_ARGUMENT status. cudaStreamSynchronize and device-to-host copies are not allowed on a stream that is capturing a CUDA graph, so make the readback conditional: query cudaStreamIsCapturing and perform the copy + synchronize + status check only when the stream is not capturing. This restores CUDA graph capture support for static-shape models (where EmbedLayerNormalization is typically the first node) while keeping ids in range device-side; error surfacing resumes on normal non-capturing runs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/embed_layer_norm.cc | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc index faf4cb805ed5a..81ba48f800dee 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc @@ -97,13 +97,22 @@ Status EmbedLayerNorm::ComputeInternal(OpKernelContext* context) const { segment_embedding_length, error_flag.get())); - // Surface any out-of-range id as a clean error instead of leaving an invalid output. - auto host_error_flag = AllocateBufferOnCPUPinned(1); - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(host_error_flag.get(), error_flag.get(), sizeof(int), - cudaMemcpyDeviceToHost, Stream(context))); - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(Stream(context))); - ORT_RETURN_IF(*host_error_flag.get() != 0, - "input id is out of range of the corresponding embedding table."); + // The kernel always validates ids device-side and skips the embedding reads for any + // out-of-range id, so input safety does not depend on the readback below; the readback only + // upgrades a skipped (silently zeroed) row into a clean error status. cudaStreamSynchronize and + // device-to-host copies are not permitted on a stream that is capturing a CUDA graph, so skip the + // status readback while the stream is capturing. Ids are still kept in range device-side, and + // error surfacing resumes on normal (non-capturing) runs. + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + CUDA_RETURN_IF_ERROR(cudaStreamIsCapturing(Stream(context), &capture_status)); + if (capture_status == cudaStreamCaptureStatusNone) { + auto host_error_flag = AllocateBufferOnCPUPinned(1); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(host_error_flag.get(), error_flag.get(), sizeof(int), + cudaMemcpyDeviceToHost, Stream(context))); + CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(Stream(context))); + ORT_RETURN_IF(*host_error_flag.get() != 0, + "input id is out of range of the corresponding embedding table."); + } return Status::OK(); } From 6dfadb6105be1dea60ec4d5f76ff9d3b265f741a Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:14:46 +0000 Subject: [PATCH 18/36] contrib: annotate GatherBlockQuantized CUDA test call arguments Use inline parameter-name comments for the positional arguments passed to Test_Fail_WithoutZeroPoints in the CUDA-scoped tests so the gather_axis, quantize_axis, block_size, and bits values are self-documenting at the call site. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/gather_block_quantized_op_test.cc | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc index 8e010b0b33338..3da208d8f9e45 100644 --- a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc +++ b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc @@ -533,18 +533,21 @@ TEST(GatherBlockQuantizedOpTest, CudaInvalidIndices) { // block_size == 0 is accepted by the constructor but is an invalid divisor for block // mapping, so ComputeInternal must reject it with a clean status. TEST(GatherBlockQuantizedOpTest, CudaInvalidBlockSizeZero) { - Test_Fail_WithoutZeroPoints(0, 2, 0); - Test_Fail_WithoutZeroPoints(0, 2, 0); - Test_Fail_WithoutZeroPoints(0, 2, 0); + Test_Fail_WithoutZeroPoints(/*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/0); + Test_Fail_WithoutZeroPoints(/*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/0); + Test_Fail_WithoutZeroPoints(/*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/0); } // uint8 data supports bits in {4, 8} and int4/uint4 supports bits == 4 on CUDA; other // values are rejected by the constructor before the 8 / bits division that derives the // uint8 packing factor. TEST(GatherBlockQuantizedOpTest, CudaUnsupportedBits) { - Test_Fail_WithoutZeroPoints(0, 2, 16, 3); - Test_Fail_WithoutZeroPoints(0, 2, 16, 8); - Test_Fail_WithoutZeroPoints(0, 2, 16, 3); + Test_Fail_WithoutZeroPoints( + /*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/16, /*bits=*/3); + Test_Fail_WithoutZeroPoints( + /*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/16, /*bits=*/8); + Test_Fail_WithoutZeroPoints( + /*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/16, /*bits=*/3); } #endif From 5ebd2ace89741864aaa095e29f5064319ecefc19 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:15:42 +0000 Subject: [PATCH 19/36] Restrict GridSample empty-spatial expect-failure test to CPU and CUDA EPs The zero-spatial-dimension validation guard currently lives only in the CPU and CUDA GridSample kernels. GetExecutionProviders() also adds the WebGPU and CoreML kernels on builds that enable them, and those separate kernels do not carry the guard, so the kExpectFailure substring assertion would falsely fail there. Add a CPU+CUDA-only provider helper and use it for the expect-failure case; positive tests are unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Agent-signed-off: Developer (cc69a935) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cpu/tensor/grid_sample_test_custom.cc | 22 ++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc b/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc index 9abb362baca22..43579d6217901 100644 --- a/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc +++ b/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc @@ -50,6 +50,24 @@ std::vector> GetExecutionProviders() { return execution_providers; } +// Only the CPU and CUDA GridSample kernels currently validate that input spatial dimensions +// are non-empty. WebGPU/CoreML use separate kernels without that guard, so the expect-failure +// test below must exclude them to avoid a false failure on those builds. +std::vector> GetCpuAndCudaExecutionProviders() { + std::vector> execution_providers; + + execution_providers.emplace_back(DefaultCpuExecutionProvider()); + +#ifdef USE_CUDA + execution_providers.emplace_back(DefaultCudaExecutionProvider()); +#ifdef ENABLE_CUDA_NHWC_OPS + execution_providers.push_back(DefaultCudaNHWCExecutionProvider()); +#endif +#endif + + return execution_providers; +} + template void RunTests(T& test, std::vector>&& execution_providers) { for (size_t idx = 0; idx < execution_providers.size(); ++idx) { @@ -610,8 +628,10 @@ TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_border_empty_spa test.AddInput("Grid", Grid_shape, Grid_data); test.AddOutput("Y", Y_shape, {TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f)}); + // The zero-spatial-dim guard currently covers the CPU and CUDA kernels; other EPs (WebGPU/CoreML) + // are out of scope for this test, so restrict the expect-failure run to CPU + CUDA only. RunTestsExpectFailure(test, "Input spatial dimensions must be non-empty for sampling", - GetExecutionProviders()); + GetCpuAndCudaExecutionProviders()); } } // namespace test From 5d1531236eb39563631fc7ec82a0a52f4058a212 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:16:18 +0000 Subject: [PATCH 20/36] contrib: skip GatherBlockQuantized CUDA host readback during graph capture The device-side index bounds check and kernel early-return already protect memory safety, so the host readback of the error flag only upgrades a silently incorrect result into a clean error. That readback performs a device-to-host copy and a stream synchronize, both of which are illegal while the stream is being captured for a CUDA graph. Query the capture status with cudaStreamIsCapturing and perform the readback only when the stream is not capturing, matching the idiom used by EmbedLayerNormalization. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../quantization/gather_block_quantized.cc | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc index b2e5cffa5d54e..65759c8ca0f27 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc @@ -205,12 +205,21 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) } auto host_index_out_of_bounds = AllocateBufferOnCPUPinned(1); - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(host_index_out_of_bounds.get(), index_out_of_bounds.get(), sizeof(int), - cudaMemcpyDeviceToHost, param.stream)); - CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(param.stream)); - ORT_RETURN_IF(*host_index_out_of_bounds.get() != 0, - "indices element out of data bounds. Each index must be within the inclusive range [", - -param.gather_axis_dim, ", ", param.gather_axis_dim - 1, "]."); + // The device-side bounds check and kernel early-return above always run, so memory + // safety does not depend on the host readback below. The readback only upgrades a + // silently incorrect result into a clean error, and it requires a device-to-host copy + // plus a stream synchronize. Both are illegal while the stream is being captured for a + // CUDA graph, so skip the readback during capture and let the device-side guard stand. + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + CUDA_RETURN_IF_ERROR(cudaStreamIsCapturing(param.stream, &capture_status)); + if (capture_status == cudaStreamCaptureStatusNone) { + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(host_index_out_of_bounds.get(), index_out_of_bounds.get(), sizeof(int), + cudaMemcpyDeviceToHost, param.stream)); + CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(param.stream)); + ORT_RETURN_IF(*host_index_out_of_bounds.get() != 0, + "indices element out of data bounds. Each index must be within the inclusive range [", + -param.gather_axis_dim, ", ", param.gather_axis_dim - 1, "]."); + } return Status::OK(); } From 5feef75554b66ebf9bd73e559580e7912b39b2e6 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:18:05 +0000 Subject: [PATCH 21/36] contrib: tidy EmbedLayerNormalization CUDA validation review minors Source-only follow-up to the EmbedLayerNormalization input-validation changes: - Widen the output offset to int64 to match the word/position/segment offsets; narrow explicitly at the LayerNorm call, which takes an int offset. - Reword the kernel and test comments to neutral "out of range of the embedding table" phrasing. - Trim the segment_embedding_length parameter comment to fit the column limit. - Document that the CUDA NHWC EP is intentionally left enabled for the negative-position-id test (it shares the same validated kernel and is skipped automatically if not registered). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/embed_layer_norm_impl.cu | 8 ++++---- .../contrib_ops/cuda/bert/embed_layer_norm_impl.h | 2 +- onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc | 9 +++++---- 3 files changed, 10 insertions(+), 9 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu index ea06a0da309f2..8de903d01a5ed 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu @@ -153,7 +153,7 @@ __global__ void EmbedLayerNormKernel( position_id = position_ids[sequence_position]; } - // Reject ids that fall outside their embedding tables instead of reading past the buffers. + // Flag any id that is out of range of its embedding table so it is rejected instead of indexed. is_valid_block = (word_id >= 0 && word_id < word_embedding_length) && (position_id >= 0 && position_id < position_embedding_length) && (nullptr == segment_embedding || @@ -175,7 +175,7 @@ __global__ void EmbedLayerNormKernel( const int64_t word_offset = static_cast(word_id) * hidden_size; const int64_t segment_offset = static_cast(segment_id) * hidden_size; // the output offset is given by b * (sequence_length * hidden_size) + s * hidden_size - const int output_offset = sequence_position * hidden_size; + const int64_t output_offset = static_cast(sequence_position) * hidden_size; cub::KeyValuePair thread_data(0.f, 0.f); @@ -197,8 +197,8 @@ __global__ void EmbedLayerNormKernel( thread_data = pair_sum(thread_data, cub::KeyValuePair(rldval, rldval * val_f)); } - // 3. layer norm on the sum - LayerNorm(thread_data, hidden_size, output_offset, beta, gamma, epsilon, output); + // 3. layer norm on the sum (LayerNorm takes an int offset; the widened offset fits the output) + LayerNorm(thread_data, hidden_size, static_cast(output_offset), beta, gamma, epsilon, output); } template diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h index b9d3bbc1d9304..bac7702161792 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h @@ -29,7 +29,7 @@ Status LaunchEmbedLayerNormKernel(cudaStream_t stream, const bool broadcast_position_ids, // Whether to broadcast position ids int word_embedding_length, // number of rows in word_embedding int position_embedding_length, // number of rows in position_embedding - int segment_embedding_length, // number of rows in segment_embedding (0 if none) + int segment_embedding_length, // rows in segment_embedding (0 if none) int* error_flag); // device flag set when an id is out of range } // namespace cuda diff --git a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc b/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc index 5ce54eda4157c..d50e344674c8f 100644 --- a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc +++ b/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc @@ -227,7 +227,7 @@ TEST(EmbedLayerNormTest, EmbedLayerNormBatch_Distill) { RunTest(embedlayernorm::EmbedLayerNormBatch_Distill()); } -// Regression test: negative position_ids must be rejected to not cause OOB read. +// Input validation test: a negative position id must be rejected rather than used to index the table. TEST(EmbedLayerNormTest, EmbedLayerNormNegativePositionIds) { int batch_size = 1; int sequence_size = 2; @@ -282,12 +282,13 @@ TEST(EmbedLayerNormTest, EmbedLayerNormNegativePositionIds) { tester.AddOutput("output", output_dims, std::vector(batch_size * sequence_size * hidden_size, 0.0f)); tester.AddOutput("mask_index", mask_index_dims, {0}); - // Both CPU and CUDA reject the out-of-range position id via input validation; other EPs are - // not exercised here. + // Both CPU and CUDA reject the out-of-range position id via input validation. The CUDA NHWC EP + // shares the same validated kernel, so it is intentionally left enabled (skipped automatically if + // the kernel is not registered for that EP). Only DML and OpenVINO are excluded here. tester.Run(OpTester::ExpectResult::kExpectFailure, "", {kDmlExecutionProvider, kOpenVINOExecutionProvider}); } -// An input id that points outside the word_embedding table must be rejected rather than read. +// An input id that points outside the word_embedding table must be rejected rather than indexed. // CPU validates per-index already; CUDA validates device-side and surfaces the same failure, so // this case runs on both EPs. TEST(EmbedLayerNormTest, EmbedLayerNormWordIdOutOfRange) { From f9b1eb0f60188e6c849a766f3ccd611a028944a4 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:20:37 +0000 Subject: [PATCH 22/36] contrib: make EmbedLayerNormalization CUDA output write index 64-bit For large batch * sequence_length * hidden_size (which can exceed 2^31 elements on high-memory GPUs), the EmbedLayerNormalization output write index computed as sequence_position * hidden_size overflowed 32-bit arithmetic and wrapped negative, producing an out-of-range write. Make the whole output write-index path 64-bit end-to-end: - Pass the already-widened int64 output offset straight to LayerNorm instead of narrowing it back to int. - Widen the shared LayerNorm helper's offset parameter and its internal index from int to int64_t so the offset is not truncated again. This is source-compatible with the only other caller (skip layer norm passes an int offset, which promotes cleanly). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu | 4 ++-- onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu index 8de903d01a5ed..cb6f5f39b4e09 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu @@ -197,8 +197,8 @@ __global__ void EmbedLayerNormKernel( thread_data = pair_sum(thread_data, cub::KeyValuePair(rldval, rldval * val_f)); } - // 3. layer norm on the sum (LayerNorm takes an int offset; the widened offset fits the output) - LayerNorm(thread_data, hidden_size, static_cast(output_offset), beta, gamma, epsilon, output); + // 3. layer norm on the sum (offset is int64 end-to-end to keep the output write index in range) + LayerNorm(thread_data, hidden_size, output_offset, beta, gamma, epsilon, output); } template diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index 0ed2c125405e3..ba119c0d0a5d0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -78,7 +78,7 @@ struct KeyValuePairSum { template __device__ inline void LayerNorm( - const cub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, + const cub::KeyValuePair& thread_data, const int ld, const int64_t offset, const T* beta, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation for mean/variance to avoid overflow in fp16/bf16. @@ -98,7 +98,7 @@ __device__ inline void LayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + const int64_t idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); const float b = (nullptr == beta) ? 0.f : static_cast(beta[i]); From 3ec828f677c03987aa32347a9fd501a0697700da Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:23:09 +0000 Subject: [PATCH 23/36] contrib: revert EmbedLayerNormalization output-offset widening from this change The output write-index 64-bit widening also requires touching the shared LayerNorm helper in layer_norm.cuh, which is used by other BERT kernels (skip layer norm) and needs its own cross-kernel review and regression testing. Keep that complete change as a separate follow-up so it does not gate this single-op input-validation change. Restore the output offset to its original 32-bit form and fully revert the shared layer_norm.cuh helper. The embedding-index validation and the read-side word/position/segment offset arithmetic are unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu | 4 ++-- onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu index cb6f5f39b4e09..3f202bd58858d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu @@ -175,7 +175,7 @@ __global__ void EmbedLayerNormKernel( const int64_t word_offset = static_cast(word_id) * hidden_size; const int64_t segment_offset = static_cast(segment_id) * hidden_size; // the output offset is given by b * (sequence_length * hidden_size) + s * hidden_size - const int64_t output_offset = static_cast(sequence_position) * hidden_size; + const int output_offset = sequence_position * hidden_size; cub::KeyValuePair thread_data(0.f, 0.f); @@ -197,7 +197,7 @@ __global__ void EmbedLayerNormKernel( thread_data = pair_sum(thread_data, cub::KeyValuePair(rldval, rldval * val_f)); } - // 3. layer norm on the sum (offset is int64 end-to-end to keep the output write index in range) + // 3. layer norm on the sum LayerNorm(thread_data, hidden_size, output_offset, beta, gamma, epsilon, output); } diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index ba119c0d0a5d0..0ed2c125405e3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -78,7 +78,7 @@ struct KeyValuePairSum { template __device__ inline void LayerNorm( - const cub::KeyValuePair& thread_data, const int ld, const int64_t offset, const T* beta, + const cub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation for mean/variance to avoid overflow in fp16/bf16. @@ -98,7 +98,7 @@ __device__ inline void LayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int64_t idx = offset + i; + const int idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); const float b = (nullptr == beta) ? 0.f : static_cast(beta[i]); From fca0005674c644ba8fc5e181c3af1d24f33687f1 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:26:06 +0000 Subject: [PATCH 24/36] contrib: use 64-bit element offsets in CUDA BERT LayerNorm indexing Support large tensors in the CUDA BERT LayerNorm device helpers by computing the global output element offset in 64-bit. For big batch * sequence_length * hidden_size (which exceeds 2^31 elements on high-memory GPUs), the previous 32-bit row offset arithmetic could not represent the index. - layer_norm.cuh: widen the global write offset / index of LayerNorm, SimplifiedLayerNorm, LayerNormSmall and SimplifiedLayerNormSmall to int64_t. The gamma/beta/bias indices (ld, i, threadIdx.x * ILP) stay int since they are bounded by hidden_size. - embed_layer_norm_impl.cu: compute output_offset in int64 and pass it to LayerNorm without narrowing. - skip_layer_norm_impl.cu: compute the per-block offset/index in int64 in both SkipLayerNormKernel and SkipLayerNormKernelSmall. Numerics are unchanged for normal tensor sizes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/embed_layer_norm_impl.cu | 4 ++-- onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh | 14 +++++++------- .../contrib_ops/cuda/bert/skip_layer_norm_impl.cu | 6 +++--- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu index 3f202bd58858d..53c98e74a2842 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu @@ -175,7 +175,7 @@ __global__ void EmbedLayerNormKernel( const int64_t word_offset = static_cast(word_id) * hidden_size; const int64_t segment_offset = static_cast(segment_id) * hidden_size; // the output offset is given by b * (sequence_length * hidden_size) + s * hidden_size - const int output_offset = sequence_position * hidden_size; + const int64_t output_offset = static_cast(sequence_position) * hidden_size; cub::KeyValuePair thread_data(0.f, 0.f); @@ -197,7 +197,7 @@ __global__ void EmbedLayerNormKernel( thread_data = pair_sum(thread_data, cub::KeyValuePair(rldval, rldval * val_f)); } - // 3. layer norm on the sum + // 3. layer norm on the sum (64-bit output offset to support large tensors) LayerNorm(thread_data, hidden_size, output_offset, beta, gamma, epsilon, output); } diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index 0ed2c125405e3..df7734f91a1a3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -78,7 +78,7 @@ struct KeyValuePairSum { template __device__ inline void LayerNorm( - const cub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, + const cub::KeyValuePair& thread_data, const int ld, const int64_t offset, const T* beta, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation for mean/variance to avoid overflow in fp16/bf16. @@ -98,7 +98,7 @@ __device__ inline void LayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + const int64_t idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); const float b = (nullptr == beta) ? 0.f : static_cast(beta[i]); @@ -108,7 +108,7 @@ __device__ inline void LayerNorm( template __device__ inline void SimplifiedLayerNorm( - const float& thread_data, const int ld, const int offset, const T* gamma, const float epsilon, T* output) { + const float& thread_data, const int ld, const int64_t offset, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation to avoid overflow in fp16/bf16. @@ -124,7 +124,7 @@ __device__ inline void SimplifiedLayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + const int64_t idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); output[idx] = static_cast(g * val * rsigma); @@ -133,7 +133,7 @@ __device__ inline void SimplifiedLayerNorm( template __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair& thread_data, - const int ld, const int idx, const T* beta, const T* gamma, + const int ld, const int64_t idx, const T* beta, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input @@ -182,8 +182,8 @@ __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair< } template -__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, const int idx, - const T* gamma, const float epsilon, T* output) { +__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, + const int64_t idx, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input // value is available in a register diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu index 7f5169639ff8d..c0e2de0a0a0c3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu @@ -75,7 +75,7 @@ __global__ void SkipLayerNormKernel( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, const int ld, int skip_size) { const float reverse_ld = 1.f / ld; - const int offset = blockIdx.x * ld; + const int64_t offset = static_cast(blockIdx.x) * ld; const bool has_bias = (bias != nullptr); // Reduce sum of x and x^2, and the results are divided by ld. @@ -84,7 +84,7 @@ __global__ void SkipLayerNormKernel( cub::KeyValuePair thread_data(0.f, 0.f); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + const int64_t idx = offset + i; T val = input[idx]; if (has_bias) { @@ -116,7 +116,7 @@ __global__ void SkipLayerNormKernelSmall( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, int ld, int skip_size) { const float rld = 1.f / ld; - const int idx = blockIdx.x * ld + threadIdx.x * ILP; + const int64_t idx = static_cast(blockIdx.x) * ld + threadIdx.x * ILP; using VecT = aligned_vector; T sum_v[ILP]; From 103957ac864ce29e9d8dc91dbebf87e596ccc358 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:34:52 +0000 Subject: [PATCH 25/36] contrib: move SkipLayerNorm 64-bit write-index out of the combined change The previous commit re-applied two logically separate changes together: the EmbedLayerNormalization output write-index widening (LayerNorm() offset/idx and output_offset) and the SkipLayerNormalization write-index widening. Split them so each lands as its own single-op change. This reverts only the SkipLayerNorm portion (SimplifiedLayerNorm, LayerNormSmall and SimplifiedLayerNormSmall offset/idx in layer_norm.cuh, plus both kernels in skip_layer_norm_impl.cu) so it can be re-applied as its own commit. The EmbedLayerNorm write-index widening (LayerNorm() and output_offset) is kept. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh | 10 +++++----- .../contrib_ops/cuda/bert/skip_layer_norm_impl.cu | 6 +++--- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index df7734f91a1a3..ba119c0d0a5d0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -108,7 +108,7 @@ __device__ inline void LayerNorm( template __device__ inline void SimplifiedLayerNorm( - const float& thread_data, const int ld, const int64_t offset, const T* gamma, const float epsilon, T* output) { + const float& thread_data, const int ld, const int offset, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation to avoid overflow in fp16/bf16. @@ -124,7 +124,7 @@ __device__ inline void SimplifiedLayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int64_t idx = offset + i; + const int idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); output[idx] = static_cast(g * val * rsigma); @@ -133,7 +133,7 @@ __device__ inline void SimplifiedLayerNorm( template __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair& thread_data, - const int ld, const int64_t idx, const T* beta, const T* gamma, + const int ld, const int idx, const T* beta, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input @@ -182,8 +182,8 @@ __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair< } template -__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, - const int64_t idx, const T* gamma, const float epsilon, T* output) { +__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, const int idx, + const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input // value is available in a register diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu index c0e2de0a0a0c3..7f5169639ff8d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu @@ -75,7 +75,7 @@ __global__ void SkipLayerNormKernel( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, const int ld, int skip_size) { const float reverse_ld = 1.f / ld; - const int64_t offset = static_cast(blockIdx.x) * ld; + const int offset = blockIdx.x * ld; const bool has_bias = (bias != nullptr); // Reduce sum of x and x^2, and the results are divided by ld. @@ -84,7 +84,7 @@ __global__ void SkipLayerNormKernel( cub::KeyValuePair thread_data(0.f, 0.f); for (int i = threadIdx.x; i < ld; i += TPB) { - const int64_t idx = offset + i; + const int idx = offset + i; T val = input[idx]; if (has_bias) { @@ -116,7 +116,7 @@ __global__ void SkipLayerNormKernelSmall( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, int ld, int skip_size) { const float rld = 1.f / ld; - const int64_t idx = static_cast(blockIdx.x) * ld + threadIdx.x * ILP; + const int idx = blockIdx.x * ld + threadIdx.x * ILP; using VecT = aligned_vector; T sum_v[ILP]; From 4f414762d317c82e1250985d933b730712628b52 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:35:29 +0000 Subject: [PATCH 26/36] contrib: use 64-bit element offsets in SkipLayerNorm CUDA write index Support large tensors in SkipLayerNormalization by computing the global output element offset in 64-bit. For big rows * hidden_size (rows = batch * sequence length), which can exceed 2^31 elements on high-memory GPUs, the previous 32-bit row offset arithmetic could not represent the index. - skip_layer_norm_impl.cu: compute the per-block offset/index in int64 in both SkipLayerNormKernel and SkipLayerNormKernelSmall. - layer_norm.cuh: widen the global write offset/index of SimplifiedLayerNorm, LayerNormSmall and SimplifiedLayerNormSmall (the helpers used only by SkipLayerNorm) to int64_t. The gamma/beta/bias indices (ld, i, threadIdx.x * ILP) stay int since they are bounded by hidden_size. LayerNorm() (used by EmbedLayerNorm) is widened separately in the EmbedLayerNorm write-index change. Numerics are unchanged for normal sizes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh | 10 +++++----- .../contrib_ops/cuda/bert/skip_layer_norm_impl.cu | 6 +++--- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index ba119c0d0a5d0..df7734f91a1a3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -108,7 +108,7 @@ __device__ inline void LayerNorm( template __device__ inline void SimplifiedLayerNorm( - const float& thread_data, const int ld, const int offset, const T* gamma, const float epsilon, T* output) { + const float& thread_data, const int ld, const int64_t offset, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation to avoid overflow in fp16/bf16. @@ -124,7 +124,7 @@ __device__ inline void SimplifiedLayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + const int64_t idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); output[idx] = static_cast(g * val * rsigma); @@ -133,7 +133,7 @@ __device__ inline void SimplifiedLayerNorm( template __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair& thread_data, - const int ld, const int idx, const T* beta, const T* gamma, + const int ld, const int64_t idx, const T* beta, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input @@ -182,8 +182,8 @@ __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair< } template -__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, const int idx, - const T* gamma, const float epsilon, T* output) { +__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, + const int64_t idx, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input // value is available in a register diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu index 7f5169639ff8d..c0e2de0a0a0c3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu @@ -75,7 +75,7 @@ __global__ void SkipLayerNormKernel( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, const int ld, int skip_size) { const float reverse_ld = 1.f / ld; - const int offset = blockIdx.x * ld; + const int64_t offset = static_cast(blockIdx.x) * ld; const bool has_bias = (bias != nullptr); // Reduce sum of x and x^2, and the results are divided by ld. @@ -84,7 +84,7 @@ __global__ void SkipLayerNormKernel( cub::KeyValuePair thread_data(0.f, 0.f); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + const int64_t idx = offset + i; T val = input[idx]; if (has_bias) { @@ -116,7 +116,7 @@ __global__ void SkipLayerNormKernelSmall( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, int ld, int skip_size) { const float rld = 1.f / ld; - const int idx = blockIdx.x * ld + threadIdx.x * ILP; + const int64_t idx = static_cast(blockIdx.x) * ld + threadIdx.x * ILP; using VecT = aligned_vector; T sum_v[ILP]; From e843cde65e0bbbb8ca6abad4d95d54cd26154ac4 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:37:03 +0000 Subject: [PATCH 27/36] contrib: guard throwing GatherBlockQuantized CUDA tests for no-exceptions builds The CudaShapeMismatch, CudaInvalidBlockSizeZero, and CudaUnsupportedBits cases reject their inputs by throwing: the first two via fail_shape_inference during Graph::Resolve and the last via an ORT_ENFORCE in the kernel constructor. A throw aborts under ORT_NO_EXCEPTIONS instead of surfacing as the failure that OpTester::ExpectResult::kExpectFailure expects, so wrap them in !defined(ORT_NO_EXCEPTIONS). CudaInvalidIndices stays unguarded because it fails through a device error flag turned into a Status, which does not throw. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../gather_block_quantized_op_test.cc | 27 +++++++++++++------ 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc index 3da208d8f9e45..5c9a3b1b56081 100644 --- a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc +++ b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc @@ -518,20 +518,30 @@ TEST(GatherBlockQuantizedOpTest, InvalidIndices) { // bounds on the device. These cases mirror the CPU coverage above but run on the CUDA // EP, which previously skipped these checks. They fall back to the CPU EP (which rejects // the same inputs) when no CUDA device is present. -TEST(GatherBlockQuantizedOpTest, CudaShapeMismatch) { - Test_ShapeMismatch_WithZeroPoints(); - Test_ShapeMismatch_WithZeroPoints(); - Test_ShapeMismatch_WithZeroPoints(); -} +// Out-of-bounds gathered indices are runtime data values, so they cannot be caught by +// shape inference; the device sets an error flag that ComputeInternal turns into a +// failed Status. This path does not throw, so it runs in all builds. TEST(GatherBlockQuantizedOpTest, CudaInvalidIndices) { Test_InvalidIndices_WithZeroPoints(); Test_InvalidIndices_WithZeroPoints(); Test_InvalidIndices_WithZeroPoints(); } -// block_size == 0 is accepted by the constructor but is an invalid divisor for block -// mapping, so ComputeInternal must reject it with a clean status. +#if !defined(ORT_NO_EXCEPTIONS) +// The following cases reject their inputs by throwing: the shape mismatch and the +// block_size == 0 case fail via fail_shape_inference during Graph::Resolve, and the +// unsupported-bits case fails via an ORT_ENFORCE in the kernel constructor. A throw +// aborts under ORT_NO_EXCEPTIONS, so these only run when exceptions are enabled, matching +// the convention OpTester::ExpectResult::kExpectFailure relies on for throwing failures. +TEST(GatherBlockQuantizedOpTest, CudaShapeMismatch) { + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); +} + +// block_size == 0 is an invalid divisor for relating data and scales shapes, so shape +// inference rejects it with fail_shape_inference during Graph::Resolve. TEST(GatherBlockQuantizedOpTest, CudaInvalidBlockSizeZero) { Test_Fail_WithoutZeroPoints(/*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/0); Test_Fail_WithoutZeroPoints(/*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/0); @@ -549,7 +559,8 @@ TEST(GatherBlockQuantizedOpTest, CudaUnsupportedBits) { Test_Fail_WithoutZeroPoints( /*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/16, /*bits=*/3); } -#endif +#endif // !defined(ORT_NO_EXCEPTIONS) +#endif // USE_CUDA template void Test_GatherAxis0_WithZeroPoints(int bits = 4) { From 0b432bf204a9099bd5b96c58df0e57e45612b553 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:41:37 +0000 Subject: [PATCH 28/36] Revert "contrib: guard throwing GatherBlockQuantized CUDA tests for no-exceptions builds" This reverts commit e843cde65e0bbbb8ca6abad4d95d54cd26154ac4. The GatherBlockQuantized test file has seven pre-existing throw-based expect-failure tests and no ORT_NO_EXCEPTIONS guards, so the provider test is not built with exceptions disabled. Guarding only the three new CUDA tests is inconsistent with that convention, so restore them to unguarded. The accurate comment on CudaInvalidBlockSizeZero (it is rejected by shape inference during Graph::Resolve, not by ComputeInternal) is kept. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../gather_block_quantized_op_test.cc | 23 +++++-------------- 1 file changed, 6 insertions(+), 17 deletions(-) diff --git a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc index 5c9a3b1b56081..b92d5640745b0 100644 --- a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc +++ b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc @@ -518,28 +518,18 @@ TEST(GatherBlockQuantizedOpTest, InvalidIndices) { // bounds on the device. These cases mirror the CPU coverage above but run on the CUDA // EP, which previously skipped these checks. They fall back to the CPU EP (which rejects // the same inputs) when no CUDA device is present. +TEST(GatherBlockQuantizedOpTest, CudaShapeMismatch) { + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); +} -// Out-of-bounds gathered indices are runtime data values, so they cannot be caught by -// shape inference; the device sets an error flag that ComputeInternal turns into a -// failed Status. This path does not throw, so it runs in all builds. TEST(GatherBlockQuantizedOpTest, CudaInvalidIndices) { Test_InvalidIndices_WithZeroPoints(); Test_InvalidIndices_WithZeroPoints(); Test_InvalidIndices_WithZeroPoints(); } -#if !defined(ORT_NO_EXCEPTIONS) -// The following cases reject their inputs by throwing: the shape mismatch and the -// block_size == 0 case fail via fail_shape_inference during Graph::Resolve, and the -// unsupported-bits case fails via an ORT_ENFORCE in the kernel constructor. A throw -// aborts under ORT_NO_EXCEPTIONS, so these only run when exceptions are enabled, matching -// the convention OpTester::ExpectResult::kExpectFailure relies on for throwing failures. -TEST(GatherBlockQuantizedOpTest, CudaShapeMismatch) { - Test_ShapeMismatch_WithZeroPoints(); - Test_ShapeMismatch_WithZeroPoints(); - Test_ShapeMismatch_WithZeroPoints(); -} - // block_size == 0 is an invalid divisor for relating data and scales shapes, so shape // inference rejects it with fail_shape_inference during Graph::Resolve. TEST(GatherBlockQuantizedOpTest, CudaInvalidBlockSizeZero) { @@ -559,8 +549,7 @@ TEST(GatherBlockQuantizedOpTest, CudaUnsupportedBits) { Test_Fail_WithoutZeroPoints( /*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/16, /*bits=*/3); } -#endif // !defined(ORT_NO_EXCEPTIONS) -#endif // USE_CUDA +#endif template void Test_GatherAxis0_WithZeroPoints(int bits = 4) { From 87139d35a2deaab0f5f0b1cf5826cc18fbff2794 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:44:14 +0000 Subject: [PATCH 29/36] Reapply "contrib: guard throwing GatherBlockQuantized CUDA tests for no-exceptions builds" This reverts commit 0b432bf204a9099bd5b96c58df0e57e45612b553. --- .../gather_block_quantized_op_test.cc | 23 ++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc index b92d5640745b0..5c9a3b1b56081 100644 --- a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc +++ b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc @@ -518,18 +518,28 @@ TEST(GatherBlockQuantizedOpTest, InvalidIndices) { // bounds on the device. These cases mirror the CPU coverage above but run on the CUDA // EP, which previously skipped these checks. They fall back to the CPU EP (which rejects // the same inputs) when no CUDA device is present. -TEST(GatherBlockQuantizedOpTest, CudaShapeMismatch) { - Test_ShapeMismatch_WithZeroPoints(); - Test_ShapeMismatch_WithZeroPoints(); - Test_ShapeMismatch_WithZeroPoints(); -} +// Out-of-bounds gathered indices are runtime data values, so they cannot be caught by +// shape inference; the device sets an error flag that ComputeInternal turns into a +// failed Status. This path does not throw, so it runs in all builds. TEST(GatherBlockQuantizedOpTest, CudaInvalidIndices) { Test_InvalidIndices_WithZeroPoints(); Test_InvalidIndices_WithZeroPoints(); Test_InvalidIndices_WithZeroPoints(); } +#if !defined(ORT_NO_EXCEPTIONS) +// The following cases reject their inputs by throwing: the shape mismatch and the +// block_size == 0 case fail via fail_shape_inference during Graph::Resolve, and the +// unsupported-bits case fails via an ORT_ENFORCE in the kernel constructor. A throw +// aborts under ORT_NO_EXCEPTIONS, so these only run when exceptions are enabled, matching +// the convention OpTester::ExpectResult::kExpectFailure relies on for throwing failures. +TEST(GatherBlockQuantizedOpTest, CudaShapeMismatch) { + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); +} + // block_size == 0 is an invalid divisor for relating data and scales shapes, so shape // inference rejects it with fail_shape_inference during Graph::Resolve. TEST(GatherBlockQuantizedOpTest, CudaInvalidBlockSizeZero) { @@ -549,7 +559,8 @@ TEST(GatherBlockQuantizedOpTest, CudaUnsupportedBits) { Test_Fail_WithoutZeroPoints( /*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/16, /*bits=*/3); } -#endif +#endif // !defined(ORT_NO_EXCEPTIONS) +#endif // USE_CUDA template void Test_GatherAxis0_WithZeroPoints(int bits = 4) { From 186cfda18f2a773f5dffd49eaf797f0d3c9301ae Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:46:18 +0000 Subject: [PATCH 30/36] contrib: consolidate scattered write-index widening for a single follow-up change Revert the write-index 64-bit widening currently spread across follow-up commits so it can re-land as one coherent change. This reverts only the write-offset int64 in the LayerNorm device helpers, the embed output_offset, and the SkipLayerNorm kernels; the EmbedLayerNormalization read-side offsets (word/position/segment) stay int64 as they are part of the input validation. No net behavior change once the consolidated widening is re-applied. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/embed_layer_norm_impl.cu | 6 +++--- onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh | 14 +++++++------- .../contrib_ops/cuda/bert/skip_layer_norm_impl.cu | 6 +++--- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu index 53c98e74a2842..37d0ebc0957b2 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu @@ -175,7 +175,7 @@ __global__ void EmbedLayerNormKernel( const int64_t word_offset = static_cast(word_id) * hidden_size; const int64_t segment_offset = static_cast(segment_id) * hidden_size; // the output offset is given by b * (sequence_length * hidden_size) + s * hidden_size - const int64_t output_offset = static_cast(sequence_position) * hidden_size; + const int output_offset = sequence_position * hidden_size; cub::KeyValuePair thread_data(0.f, 0.f); @@ -197,8 +197,8 @@ __global__ void EmbedLayerNormKernel( thread_data = pair_sum(thread_data, cub::KeyValuePair(rldval, rldval * val_f)); } - // 3. layer norm on the sum (64-bit output offset to support large tensors) - LayerNorm(thread_data, hidden_size, output_offset, beta, gamma, epsilon, output); + // 3. layer norm on the sum + LayerNorm(thread_data, hidden_size, static_cast(output_offset), beta, gamma, epsilon, output); } template diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index df7734f91a1a3..0ed2c125405e3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -78,7 +78,7 @@ struct KeyValuePairSum { template __device__ inline void LayerNorm( - const cub::KeyValuePair& thread_data, const int ld, const int64_t offset, const T* beta, + const cub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation for mean/variance to avoid overflow in fp16/bf16. @@ -98,7 +98,7 @@ __device__ inline void LayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int64_t idx = offset + i; + const int idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); const float b = (nullptr == beta) ? 0.f : static_cast(beta[i]); @@ -108,7 +108,7 @@ __device__ inline void LayerNorm( template __device__ inline void SimplifiedLayerNorm( - const float& thread_data, const int ld, const int64_t offset, const T* gamma, const float epsilon, T* output) { + const float& thread_data, const int ld, const int offset, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation to avoid overflow in fp16/bf16. @@ -124,7 +124,7 @@ __device__ inline void SimplifiedLayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int64_t idx = offset + i; + const int idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); output[idx] = static_cast(g * val * rsigma); @@ -133,7 +133,7 @@ __device__ inline void SimplifiedLayerNorm( template __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair& thread_data, - const int ld, const int64_t idx, const T* beta, const T* gamma, + const int ld, const int idx, const T* beta, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input @@ -182,8 +182,8 @@ __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair< } template -__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, - const int64_t idx, const T* gamma, const float epsilon, T* output) { +__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, const int idx, + const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input // value is available in a register diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu index c0e2de0a0a0c3..7f5169639ff8d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu @@ -75,7 +75,7 @@ __global__ void SkipLayerNormKernel( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, const int ld, int skip_size) { const float reverse_ld = 1.f / ld; - const int64_t offset = static_cast(blockIdx.x) * ld; + const int offset = blockIdx.x * ld; const bool has_bias = (bias != nullptr); // Reduce sum of x and x^2, and the results are divided by ld. @@ -84,7 +84,7 @@ __global__ void SkipLayerNormKernel( cub::KeyValuePair thread_data(0.f, 0.f); for (int i = threadIdx.x; i < ld; i += TPB) { - const int64_t idx = offset + i; + const int idx = offset + i; T val = input[idx]; if (has_bias) { @@ -116,7 +116,7 @@ __global__ void SkipLayerNormKernelSmall( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, int ld, int skip_size) { const float rld = 1.f / ld; - const int64_t idx = static_cast(blockIdx.x) * ld + threadIdx.x * ILP; + const int idx = blockIdx.x * ld + threadIdx.x * ILP; using VecT = aligned_vector; T sum_v[ILP]; From d724e4b96f2c9131443e8fe82e258128311e32d0 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:46:26 +0000 Subject: [PATCH 31/36] Revert "Reapply "contrib: guard throwing GatherBlockQuantized CUDA tests for no-exceptions builds"" This reverts commit 87139d35a2deaab0f5f0b1cf5826cc18fbff2794. --- .../gather_block_quantized_op_test.cc | 23 +++++-------------- 1 file changed, 6 insertions(+), 17 deletions(-) diff --git a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc index 5c9a3b1b56081..b92d5640745b0 100644 --- a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc +++ b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc @@ -518,28 +518,18 @@ TEST(GatherBlockQuantizedOpTest, InvalidIndices) { // bounds on the device. These cases mirror the CPU coverage above but run on the CUDA // EP, which previously skipped these checks. They fall back to the CPU EP (which rejects // the same inputs) when no CUDA device is present. +TEST(GatherBlockQuantizedOpTest, CudaShapeMismatch) { + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); +} -// Out-of-bounds gathered indices are runtime data values, so they cannot be caught by -// shape inference; the device sets an error flag that ComputeInternal turns into a -// failed Status. This path does not throw, so it runs in all builds. TEST(GatherBlockQuantizedOpTest, CudaInvalidIndices) { Test_InvalidIndices_WithZeroPoints(); Test_InvalidIndices_WithZeroPoints(); Test_InvalidIndices_WithZeroPoints(); } -#if !defined(ORT_NO_EXCEPTIONS) -// The following cases reject their inputs by throwing: the shape mismatch and the -// block_size == 0 case fail via fail_shape_inference during Graph::Resolve, and the -// unsupported-bits case fails via an ORT_ENFORCE in the kernel constructor. A throw -// aborts under ORT_NO_EXCEPTIONS, so these only run when exceptions are enabled, matching -// the convention OpTester::ExpectResult::kExpectFailure relies on for throwing failures. -TEST(GatherBlockQuantizedOpTest, CudaShapeMismatch) { - Test_ShapeMismatch_WithZeroPoints(); - Test_ShapeMismatch_WithZeroPoints(); - Test_ShapeMismatch_WithZeroPoints(); -} - // block_size == 0 is an invalid divisor for relating data and scales shapes, so shape // inference rejects it with fail_shape_inference during Graph::Resolve. TEST(GatherBlockQuantizedOpTest, CudaInvalidBlockSizeZero) { @@ -559,8 +549,7 @@ TEST(GatherBlockQuantizedOpTest, CudaUnsupportedBits) { Test_Fail_WithoutZeroPoints( /*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/16, /*bits=*/3); } -#endif // !defined(ORT_NO_EXCEPTIONS) -#endif // USE_CUDA +#endif template void Test_GatherAxis0_WithZeroPoints(int bits = 4) { From ece1b7a030ed4f91355cb2f8b7c6a4753c945bc3 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:47:03 +0000 Subject: [PATCH 32/36] contrib: use 64-bit element offsets in BERT LayerNorm write index Support large tensors in the CUDA BERT LayerNorm write path by computing the global output element offset in 64-bit. For big rows * hidden_size (rows = batch * sequence_length), which can exceed 2^31 elements on high-memory GPUs, the previous 32-bit row offset arithmetic could not represent the index. Three files, write-offset/index only (gamma/beta/bias indices ld, i, threadIdx.x * ILP stay int since they are bounded by hidden_size): - layer_norm.cuh: widen the global write offset/index of LayerNorm, SimplifiedLayerNorm, LayerNormSmall and SimplifiedLayerNormSmall to int64_t. All four have callers (EmbedLayerNorm via LayerNorm; SkipLayerNorm via all four) whose offset = row_index * hidden_size can exceed 2^31. - embed_layer_norm_impl.cu: compute output_offset in int64 and pass it to LayerNorm without narrowing. - skip_layer_norm_impl.cu: compute the per-block offset/index in int64 in both SkipLayerNormKernel and SkipLayerNormKernelSmall. Numerics are unchanged for normal tensor sizes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib_ops/cuda/bert/embed_layer_norm_impl.cu | 6 +++--- onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh | 14 +++++++------- .../contrib_ops/cuda/bert/skip_layer_norm_impl.cu | 6 +++--- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu index 37d0ebc0957b2..53c98e74a2842 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu @@ -175,7 +175,7 @@ __global__ void EmbedLayerNormKernel( const int64_t word_offset = static_cast(word_id) * hidden_size; const int64_t segment_offset = static_cast(segment_id) * hidden_size; // the output offset is given by b * (sequence_length * hidden_size) + s * hidden_size - const int output_offset = sequence_position * hidden_size; + const int64_t output_offset = static_cast(sequence_position) * hidden_size; cub::KeyValuePair thread_data(0.f, 0.f); @@ -197,8 +197,8 @@ __global__ void EmbedLayerNormKernel( thread_data = pair_sum(thread_data, cub::KeyValuePair(rldval, rldval * val_f)); } - // 3. layer norm on the sum - LayerNorm(thread_data, hidden_size, static_cast(output_offset), beta, gamma, epsilon, output); + // 3. layer norm on the sum (64-bit output offset to support large tensors) + LayerNorm(thread_data, hidden_size, output_offset, beta, gamma, epsilon, output); } template diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index 0ed2c125405e3..df7734f91a1a3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -78,7 +78,7 @@ struct KeyValuePairSum { template __device__ inline void LayerNorm( - const cub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, + const cub::KeyValuePair& thread_data, const int ld, const int64_t offset, const T* beta, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation for mean/variance to avoid overflow in fp16/bf16. @@ -98,7 +98,7 @@ __device__ inline void LayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + const int64_t idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); const float b = (nullptr == beta) ? 0.f : static_cast(beta[i]); @@ -108,7 +108,7 @@ __device__ inline void LayerNorm( template __device__ inline void SimplifiedLayerNorm( - const float& thread_data, const int ld, const int offset, const T* gamma, const float epsilon, T* output) { + const float& thread_data, const int ld, const int64_t offset, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation to avoid overflow in fp16/bf16. @@ -124,7 +124,7 @@ __device__ inline void SimplifiedLayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + const int64_t idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); output[idx] = static_cast(g * val * rsigma); @@ -133,7 +133,7 @@ __device__ inline void SimplifiedLayerNorm( template __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair& thread_data, - const int ld, const int idx, const T* beta, const T* gamma, + const int ld, const int64_t idx, const T* beta, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input @@ -182,8 +182,8 @@ __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair< } template -__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, const int idx, - const T* gamma, const float epsilon, T* output) { +__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, + const int64_t idx, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input // value is available in a register diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu index 7f5169639ff8d..c0e2de0a0a0c3 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu @@ -75,7 +75,7 @@ __global__ void SkipLayerNormKernel( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, const int ld, int skip_size) { const float reverse_ld = 1.f / ld; - const int offset = blockIdx.x * ld; + const int64_t offset = static_cast(blockIdx.x) * ld; const bool has_bias = (bias != nullptr); // Reduce sum of x and x^2, and the results are divided by ld. @@ -84,7 +84,7 @@ __global__ void SkipLayerNormKernel( cub::KeyValuePair thread_data(0.f, 0.f); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + const int64_t idx = offset + i; T val = input[idx]; if (has_bias) { @@ -116,7 +116,7 @@ __global__ void SkipLayerNormKernelSmall( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, int ld, int skip_size) { const float rld = 1.f / ld; - const int idx = blockIdx.x * ld + threadIdx.x * ILP; + const int64_t idx = static_cast(blockIdx.x) * ld + threadIdx.x * ILP; using VecT = aligned_vector; T sum_v[ILP]; From 86af9ca2dc0e403b040d04dfe8e132a496c5ccec Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:50:25 +0000 Subject: [PATCH 33/36] contrib: document 64-bit element offsets in BERT LayerNorm write index Add brief explanatory comments at each int64-widened global write offset/index in the CUDA BERT LayerNorm helpers and SkipLayerNorm kernels, noting that the offset (row_index * hidden_size = batch * sequence_length * hidden_size) can exceed 2^31 for large tensors and that gamma/beta indices intentionally stay int (bounded by hidden_size). This prevents a future reader from narrowing the index back to int and reintroducing the large-tensor limitation. Comment-only; no behavior change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh | 5 +++++ onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu | 2 ++ 2 files changed, 7 insertions(+) diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index df7734f91a1a3..a7bab0b10def4 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -98,6 +98,8 @@ __device__ inline void LayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { + // 64-bit global element offset: rows*ld (= batch*seq*hidden) can exceed 2^31 for large tensors. + // gamma/beta indices (i) stay int -- bounded by ld (hidden_size). const int64_t idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); @@ -124,6 +126,7 @@ __device__ inline void SimplifiedLayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { + // 64-bit element offset (see LayerNorm); large tensors can exceed 2^31. gamma index (i) stays int. const int64_t idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); @@ -135,6 +138,7 @@ template __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair& thread_data, const int ld, const int64_t idx, const T* beta, const T* gamma, const float epsilon, T* output) { + // idx is a 64-bit global element offset (caller passes row*ld = batch*seq*hidden, can exceed 2^31). // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input // value is available in a register @@ -184,6 +188,7 @@ __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair< template __device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, const int64_t idx, const T* gamma, const float epsilon, T* output) { + // idx is a 64-bit global element offset (caller passes row*ld = batch*seq*hidden, can exceed 2^31). // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input // value is available in a register diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu index c0e2de0a0a0c3..e1bab6b38da9b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu @@ -75,6 +75,7 @@ __global__ void SkipLayerNormKernel( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, const int ld, int skip_size) { const float reverse_ld = 1.f / ld; + // 64-bit global element offset: blockIdx.x*ld (= row*hidden) can exceed 2^31 for large tensors. const int64_t offset = static_cast(blockIdx.x) * ld; const bool has_bias = (bias != nullptr); @@ -116,6 +117,7 @@ __global__ void SkipLayerNormKernelSmall( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, int ld, int skip_size) { const float rld = 1.f / ld; + // 64-bit element offset (see SkipLayerNormKernel); large tensors can exceed 2^31. const int64_t idx = static_cast(blockIdx.x) * ld + threadIdx.x * ILP; using VecT = aligned_vector; From 6d63ef4f981c5949ccc1c566376c19fccf463fb7 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 18:38:20 +0000 Subject: [PATCH 34/36] Validate raw_data length and clamp empty ranges in contrib Range shape inference GetFirstElement read sizeof(T) bytes from a graph initializer's raw_data guarded only by the presence bit. Add a length check before the cast and fail shape inference cleanly when raw_data is shorter than the element size. Also clamp the computed element count to >= 0 in CalcRangeDim so shape inference matches the CPU kernel for empty/backward ranges. Add model-load regression tests in range_test.cc (guarded by #ifndef DISABLE_CONTRIB_OPS, Status-only assertions for no-exception builds). Agent-signed-off: Developer (b2fe149b) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../graph/contrib_ops/range_schema_defs.cc | 11 +- .../providers/cpu/generator/range_test.cc | 107 ++++++++++++++++++ 2 files changed, 117 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc b/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc index c2d853af86723..bd9848b9ee771 100644 --- a/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc @@ -60,6 +60,9 @@ static T GetFirstElement(const TensorProto* shapeInitializer) { if (utils::HasRawData(*shapeInitializer)) { const std::string& bytes = shapeInitializer->raw_data(); + if (bytes.size() < sizeof(T)) { + fail_shape_inference("Range: raw_data size is smaller than the element size for the given data type."); + } return *reinterpret_cast(bytes.c_str()); } return get_data(shapeInitializer); @@ -75,7 +78,13 @@ static int64_t CalcRangeDim(const TensorProto* startShapeInitializer, if (delta == 0) { fail_shape_inference("delta in Range operator can not be zero!"); } - return static_cast(ceil((1.0 * (limit - start)) / delta)); + // Mirror the CPU kernel (core/providers/cpu/generator/range.cc) which clamps empty or + // backward ranges to 0 so shape inference does not emit a negative dimension value. + int64_t n = static_cast(ceil((1.0 * (limit - start)) / delta)); + if (n <= 0) { + n = 0; + } + return n; } static int64_t CalcResultDim(const TensorProto* startShapeInitializer, diff --git a/onnxruntime/test/providers/cpu/generator/range_test.cc b/onnxruntime/test/providers/cpu/generator/range_test.cc index 04664628ecdd1..854efa3e07e3b 100644 --- a/onnxruntime/test/providers/cpu/generator/range_test.cc +++ b/onnxruntime/test/providers/cpu/generator/range_test.cc @@ -1,8 +1,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include "core/graph/constants.h" +#include "core/session/inference_session.h" #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "test/test_environment.h" namespace onnxruntime { namespace test { @@ -88,5 +91,109 @@ TEST(RangeTest, AlmostSameStartAndLimitHighDelta) { RunTest(2.0f, 2.01f, 1000000.0f, {1}, {2.0f}); } +#ifndef DISABLE_CONTRIB_OPS + +namespace { +// Adds a scalar double graph initializer whose raw_data is set to the provided bytes. +// The bytes length is intentionally controllable so tests can supply a truncated payload +// (shorter than sizeof(double)) to exercise input validation in shape inference. +void AddDoubleScalarInitializerWithRawData(ONNX_NAMESPACE::GraphProto* graph, + const std::string& name, + const std::string& raw_bytes) { + auto* initializer = graph->add_initializer(); + initializer->set_name(name); + initializer->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_DOUBLE); + initializer->set_raw_data(raw_bytes); +} + +std::string DoubleToRawData(double value) { + return std::string(reinterpret_cast(&value), sizeof(double)); +} + +// Builds a single com.microsoft Range node model whose start/limit/delta inputs are +// supplied as graph initializers, then loads and initializes it in a session. +common::Status BuildAndInitializeContribRangeModel(const std::string& start_raw_data, + const std::string& limit_raw_data, + const std::string& delta_raw_data, + InferenceSession& session_object) { + ONNX_NAMESPACE::ModelProto model; + model.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); + model.add_opset_import()->set_version(11); + auto* ms_opset = model.add_opset_import(); + ms_opset->set_domain(kMSDomain); + ms_opset->set_version(1); + + auto* graph = model.mutable_graph(); + graph->set_name("ContribRangeGraph"); + + AddDoubleScalarInitializerWithRawData(graph, "start", start_raw_data); + AddDoubleScalarInitializerWithRawData(graph, "limit", limit_raw_data); + AddDoubleScalarInitializerWithRawData(graph, "delta", delta_raw_data); + + auto* output = graph->add_output(); + output->set_name("Y"); + output->mutable_type()->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_DOUBLE); + + auto* range_node = graph->add_node(); + range_node->set_name("Range"); + range_node->set_op_type("Range"); + range_node->set_domain(kMSDomain); + range_node->add_input("start"); + range_node->add_input("limit"); + range_node->add_input("delta"); + range_node->add_output("Y"); + + std::string serialized_model; + if (!model.SerializeToString(&serialized_model)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to serialize test model."); + } + + std::stringstream model_stream(serialized_model); + ORT_RETURN_IF_ERROR(session_object.Load(model_stream)); + return session_object.Initialize(); +} +} // namespace + +// Verifies that Range shape inference rejects an initializer whose raw_data is shorter than +// the element size for its declared data type, returning a clean failure status instead of +// reading past the end of the buffer. The model must fail to load/initialize. +TEST(RangeTest, ContribOp_TruncatedRawData_FailsCleanly) { + // 'start' carries fewer bytes than sizeof(double); 'limit'/'delta' are well formed. + const std::string truncated_start(sizeof(double) / 2, '\0'); + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel(truncated_start, + DoubleToRawData(5.0), + DoubleToRawData(1.0), + session_object); + ASSERT_FALSE(status.IsOK()); +} + +// Verifies that shape inference clamps an empty/backward range to a zero-sized dimension, +// matching the CPU kernel behavior, instead of emitting a negative dimension value. +TEST(RangeTest, ContribOp_BackwardRange_InfersZeroDim) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_BackwardRange_InfersZeroDim"; + InferenceSession session_object{so, GetEnvironment()}; + // start > limit with a positive delta yields an empty range. + const auto status = BuildAndInitializeContribRangeModel(DoubleToRawData(5.0), + DoubleToRawData(0.0), + DoubleToRawData(1.0), + session_object); + ASSERT_STATUS_OK(status); + + const auto outputs = session_object.GetModelOutputs(); + ASSERT_STATUS_OK(outputs.first); + ASSERT_EQ(outputs.second->size(), 1u); + const auto* shape = (*outputs.second)[0]->Shape(); + ASSERT_NE(shape, nullptr); + ASSERT_EQ(shape->dim_size(), 1); + ASSERT_TRUE(shape->dim(0).has_dim_value()); + ASSERT_EQ(shape->dim(0).dim_value(), 0); +} + +#endif // DISABLE_CONTRIB_OPS + } // namespace test } // namespace onnxruntime From 19d740f77c1664f25cecad62986e0cf49867742b Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 18:51:36 +0000 Subject: [PATCH 35/36] Harden contrib Range shape inference and broaden regression coverage - GetFirstElement: read raw_data via std::memcpy into an aligned local after the length check, avoiding an unaligned access on strict-alignment targets. - CalcRangeDim and the CPU kernel ComputeRange: guard the element-count cast with std::isfinite so a non-finite computed count fails cleanly instead of an out-of-range cast (fail_shape_inference in the schema, a non-OK Status in the kernel). - range_test.cc: generalize the model-building helper over element type, dimensions and per-input bytes. Truncated raw_data is now modeled as dims=[0] with empty raw_data so the cases are tight regressions for the length check; add coverage for start/limit/delta positions, float and int64 element types, a zero-delta failure, and an exact-size success boundary. Throwing tests are additionally guarded by !defined(ORT_NO_EXCEPTIONS). Agent-signed-off: Developer (b2fe149b) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../graph/contrib_ops/range_schema_defs.cc | 12 +- .../core/providers/cpu/generator/range.cc | 7 +- .../providers/cpu/generator/range_test.cc | 200 +++++++++++++----- 3 files changed, 167 insertions(+), 52 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc b/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc index bd9848b9ee771..fd9ef7c6bd32c 100644 --- a/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc @@ -7,6 +7,7 @@ #include "core/graph/constants.h" #include "core/graph/op.h" #include +#include #include namespace onnxruntime { @@ -63,7 +64,10 @@ static T GetFirstElement(const TensorProto* shapeInitializer) { if (bytes.size() < sizeof(T)) { fail_shape_inference("Range: raw_data size is smaller than the element size for the given data type."); } - return *reinterpret_cast(bytes.c_str()); + // std::string data is only char-aligned, so copy the bytes into a properly aligned value. + T value; + std::memcpy(&value, bytes.data(), sizeof(T)); + return value; } return get_data(shapeInitializer); } @@ -80,7 +84,11 @@ static int64_t CalcRangeDim(const TensorProto* startShapeInitializer, } // Mirror the CPU kernel (core/providers/cpu/generator/range.cc) which clamps empty or // backward ranges to 0 so shape inference does not emit a negative dimension value. - int64_t n = static_cast(ceil((1.0 * (limit - start)) / delta)); + double count = ceil((1.0 * (limit - start)) / delta); + if (!std::isfinite(count)) { + fail_shape_inference("Range: the computed number of elements is not a finite value."); + } + int64_t n = static_cast(count); if (n <= 0) { n = 0; } diff --git a/onnxruntime/core/providers/cpu/generator/range.cc b/onnxruntime/core/providers/cpu/generator/range.cc index fcd4abcf72f87..997873b0f65a6 100644 --- a/onnxruntime/core/providers/cpu/generator/range.cc +++ b/onnxruntime/core/providers/cpu/generator/range.cc @@ -80,7 +80,12 @@ static Status ComputeRange( if (delta == T{0}) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "delta in Range operator can not be zero!"); } - int64_t n = static_cast(ceil((1.0 * (limit - start)) / delta)); + double count = ceil((1.0 * (limit - start)) / delta); + if (!std::isfinite(count)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Range: the computed number of elements is not a finite value."); + } + int64_t n = static_cast(count); if (n <= 0) n = 0; TensorShape shape = {n}; diff --git a/onnxruntime/test/providers/cpu/generator/range_test.cc b/onnxruntime/test/providers/cpu/generator/range_test.cc index 854efa3e07e3b..2d6173bdd9703 100644 --- a/onnxruntime/test/providers/cpu/generator/range_test.cc +++ b/onnxruntime/test/providers/cpu/generator/range_test.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "core/graph/constants.h" #include "core/session/inference_session.h" #include "gtest/gtest.h" @@ -94,28 +96,53 @@ TEST(RangeTest, AlmostSameStartAndLimitHighDelta) { #ifndef DISABLE_CONTRIB_OPS namespace { -// Adds a scalar double graph initializer whose raw_data is set to the provided bytes. -// The bytes length is intentionally controllable so tests can supply a truncated payload -// (shorter than sizeof(double)) to exercise input validation in shape inference. -void AddDoubleScalarInitializerWithRawData(ONNX_NAMESPACE::GraphProto* graph, - const std::string& name, - const std::string& raw_bytes) { - auto* initializer = graph->add_initializer(); - initializer->set_name(name); - initializer->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_DOUBLE); - initializer->set_raw_data(raw_bytes); +// Describes a single Range input supplied as a graph initializer: its declared dimensions +// and the exact raw_data bytes. Tests use this to craft both well-formed and truncated +// initializers (e.g. dims=[0] with empty raw_data) for the contrib Range schema. +struct RangeInputSpec { + std::vector dims; + std::string raw_data; +}; + +// Serializes a scalar value to its raw_data byte representation. +template +std::string ToRawData(T value) { + std::string bytes(sizeof(T), '\0'); + std::memcpy(bytes.data(), &value, sizeof(T)); + return bytes; } -std::string DoubleToRawData(double value) { - return std::string(reinterpret_cast(&value), sizeof(double)); +// A correctly-sized scalar initializer (no dims) holding a single value. +template +RangeInputSpec ScalarInput(T value) { + return RangeInputSpec{{}, ToRawData(value)}; +} + +// A zero-element initializer (dims=[0]) whose raw_data is empty. The declared size matches +// the (empty) payload, so initializer size validation accepts it; shape inference, however, +// still attempts to read the first element. +RangeInputSpec EmptyInput() { + return RangeInputSpec{{0}, std::string{}}; } -// Builds a single com.microsoft Range node model whose start/limit/delta inputs are -// supplied as graph initializers, then loads and initializes it in a session. -common::Status BuildAndInitializeContribRangeModel(const std::string& start_raw_data, - const std::string& limit_raw_data, - const std::string& delta_raw_data, - InferenceSession& session_object) { +void AddInitializer(ONNX_NAMESPACE::GraphProto* graph, const std::string& name, int data_type, + const RangeInputSpec& spec) { + auto* initializer = graph->add_initializer(); + initializer->set_name(name); + initializer->set_data_type(data_type); + for (int64_t dim : spec.dims) { + initializer->add_dims(dim); + } + initializer->set_raw_data(spec.raw_data); +} + +// Builds a single com.microsoft Range node model whose start/limit/delta inputs are supplied +// as graph initializers of the given element type, then loads and initializes it in a session. +common::Status BuildAndInitializeContribRangeModel(int data_type, + const RangeInputSpec& start, + const RangeInputSpec& limit, + const RangeInputSpec& delta, + InferenceSession& session_object) { ONNX_NAMESPACE::ModelProto model; model.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); model.add_opset_import()->set_version(11); @@ -126,13 +153,13 @@ common::Status BuildAndInitializeContribRangeModel(const std::string& start_raw_ auto* graph = model.mutable_graph(); graph->set_name("ContribRangeGraph"); - AddDoubleScalarInitializerWithRawData(graph, "start", start_raw_data); - AddDoubleScalarInitializerWithRawData(graph, "limit", limit_raw_data); - AddDoubleScalarInitializerWithRawData(graph, "delta", delta_raw_data); + AddInitializer(graph, "start", data_type, start); + AddInitializer(graph, "limit", data_type, limit); + AddInitializer(graph, "delta", data_type, delta); auto* output = graph->add_output(); output->set_name("Y"); - output->mutable_type()->mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_DOUBLE); + output->mutable_type()->mutable_tensor_type()->set_elem_type(data_type); auto* range_node = graph->add_node(); range_node->set_name("Range"); @@ -152,23 +179,20 @@ common::Status BuildAndInitializeContribRangeModel(const std::string& start_raw_ ORT_RETURN_IF_ERROR(session_object.Load(model_stream)); return session_object.Initialize(); } -} // namespace -// Verifies that Range shape inference rejects an initializer whose raw_data is shorter than -// the element size for its declared data type, returning a clean failure status instead of -// reading past the end of the buffer. The model must fail to load/initialize. -TEST(RangeTest, ContribOp_TruncatedRawData_FailsCleanly) { - // 'start' carries fewer bytes than sizeof(double); 'limit'/'delta' are well formed. - const std::string truncated_start(sizeof(double) / 2, '\0'); - SessionOptions so; - so.session_logid = "RangeTest.ContribOp_TruncatedRawData_FailsCleanly"; - InferenceSession session_object{so, GetEnvironment()}; - const auto status = BuildAndInitializeContribRangeModel(truncated_start, - DoubleToRawData(5.0), - DoubleToRawData(1.0), - session_object); - ASSERT_FALSE(status.IsOK()); +// Returns the inferred dim_value of output Y after a successful Initialize, or -1 if absent. +int64_t GetInferredOutputDim(const InferenceSession& session_object) { + const auto outputs = session_object.GetModelOutputs(); + if (!outputs.first.IsOK() || outputs.second->size() != 1u) { + return -1; + } + const auto* shape = (*outputs.second)[0]->Shape(); + if (shape == nullptr || shape->dim_size() != 1 || !shape->dim(0).has_dim_value()) { + return -1; + } + return shape->dim(0).dim_value(); } +} // namespace // Verifies that shape inference clamps an empty/backward range to a zero-sized dimension, // matching the CPU kernel behavior, instead of emitting a negative dimension value. @@ -177,22 +201,100 @@ TEST(RangeTest, ContribOp_BackwardRange_InfersZeroDim) { so.session_logid = "RangeTest.ContribOp_BackwardRange_InfersZeroDim"; InferenceSession session_object{so, GetEnvironment()}; // start > limit with a positive delta yields an empty range. - const auto status = BuildAndInitializeContribRangeModel(DoubleToRawData(5.0), - DoubleToRawData(0.0), - DoubleToRawData(1.0), - session_object); + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(5.0), ScalarInput(0.0), + ScalarInput(1.0), session_object); ASSERT_STATUS_OK(status); + ASSERT_EQ(GetInferredOutputDim(session_object), 0); +} - const auto outputs = session_object.GetModelOutputs(); - ASSERT_STATUS_OK(outputs.first); - ASSERT_EQ(outputs.second->size(), 1u); - const auto* shape = (*outputs.second)[0]->Shape(); - ASSERT_NE(shape, nullptr); - ASSERT_EQ(shape->dim_size(), 1); - ASSERT_TRUE(shape->dim(0).has_dim_value()); - ASSERT_EQ(shape->dim(0).dim_value(), 0); +// Verifies that a correctly-sized raw_data initializer (exactly sizeof(T) bytes) loads +// successfully and produces the expected inferred dimension, so the length check does not +// over-reject valid models. +TEST(RangeTest, ContribOp_ExactSizeRawData_LoadsSuccessfully) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_ExactSizeRawData_LoadsSuccessfully"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(5.0), + ScalarInput(1.0), session_object); + ASSERT_STATUS_OK(status); + ASSERT_EQ(GetInferredOutputDim(session_object), 5); } +// The following tests exercise failure paths that surface via fail_shape_inference, which +// reports the error by throwing an inference error. They are excluded from no-exception +// builds where such a throw would abort rather than yield a failure Status. +#if !defined(ORT_NO_EXCEPTIONS) + +// Verifies that Range shape inference rejects an initializer whose raw_data holds fewer bytes +// than its declared data type requires (here: a zero-element initializer with empty raw_data, +// which initializer size validation accepts), returning a clean failure status instead of +// reading more bytes than the initializer declares. One test per input position. +TEST(RangeTest, ContribOp_TruncatedStartRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedStartRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, EmptyInput(), ScalarInput(5.0), + ScalarInput(1.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + +TEST(RangeTest, ContribOp_TruncatedLimitRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedLimitRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), EmptyInput(), + ScalarInput(1.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + +TEST(RangeTest, ContribOp_TruncatedDeltaRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedDeltaRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(5.0), + EmptyInput(), session_object); + ASSERT_FALSE(status.IsOK()); +} + +// Exercises the length check for different sizeof(T) template instantiations (float, int64). +TEST(RangeTest, ContribOp_TruncatedFloatRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedFloatRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT, EmptyInput(), ScalarInput(5.0f), + ScalarInput(1.0f), session_object); + ASSERT_FALSE(status.IsOK()); +} + +TEST(RangeTest, ContribOp_TruncatedInt64RawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedInt64RawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_INT64, EmptyInput(), ScalarInput(5), + ScalarInput(1), session_object); + ASSERT_FALSE(status.IsOK()); +} + +// Verifies that a zero delta is rejected cleanly by shape inference. +TEST(RangeTest, ContribOp_ZeroDelta_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_ZeroDelta_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(10.0), + ScalarInput(0.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + +#endif // !defined(ORT_NO_EXCEPTIONS) + #endif // DISABLE_CONTRIB_OPS } // namespace test From f7f73c1c788283f2e72343081a45284224c8a2fb Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 18:59:47 +0000 Subject: [PATCH 36/36] Reject out-of-range element counts in contrib Range before the int64 cast The std::isfinite guard catches inf/NaN but not a finite count larger than the int64 range, where static_cast(count) is out of range. Handle the non-positive case before the cast (so a large-magnitude negative count cannot overflow it) and reject counts at or above 2^63 with a neutral message. Applied identically in the schema CalcRangeDim (fail_shape_inference) and the CPU kernel ComputeRange (non-OK Status), keeping the two sites consistent. Add a guarded regression test (start=0, limit=1e19, delta=1) asserting a clean non-OK status; verified it fails without the new guard. Agent-signed-off: Developer (b2fe149b) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/graph/contrib_ops/range_schema_defs.cc | 14 ++++++++++---- onnxruntime/core/providers/cpu/generator/range.cc | 15 ++++++++++++--- .../test/providers/cpu/generator/range_test.cc | 12 ++++++++++++ 3 files changed, 34 insertions(+), 7 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc b/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc index fd9ef7c6bd32c..2b2a565649175 100644 --- a/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc @@ -88,11 +88,17 @@ static int64_t CalcRangeDim(const TensorProto* startShapeInitializer, if (!std::isfinite(count)) { fail_shape_inference("Range: the computed number of elements is not a finite value."); } - int64_t n = static_cast(count); - if (n <= 0) { - n = 0; + // Empty or backward ranges clamp to 0; handle the non-positive case before the cast so a + // large-magnitude negative count can never reach (and overflow) the int64 conversion below. + if (count <= 0) { + return 0; } - return n; + // static_cast(INT64_MAX) rounds up to 2^63 (9223372036854775808.0), which is not + // representable as int64_t, so reject any count at or above that boundary before the cast. + if (count >= 9223372036854775808.0) { + fail_shape_inference("Range: the computed number of elements exceeds the supported range."); + } + return static_cast(count); } static int64_t CalcResultDim(const TensorProto* startShapeInitializer, diff --git a/onnxruntime/core/providers/cpu/generator/range.cc b/onnxruntime/core/providers/cpu/generator/range.cc index 997873b0f65a6..752ac2b5405ab 100644 --- a/onnxruntime/core/providers/cpu/generator/range.cc +++ b/onnxruntime/core/providers/cpu/generator/range.cc @@ -85,9 +85,18 @@ static Status ComputeRange( return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Range: the computed number of elements is not a finite value."); } - int64_t n = static_cast(count); - if (n <= 0) - n = 0; + // Empty or backward ranges clamp to 0; handle the non-positive case before the cast so a + // large-magnitude negative count can never reach (and overflow) the int64 conversion. + int64_t n = 0; + if (count > 0) { + // static_cast(INT64_MAX) rounds up to 2^63 (9223372036854775808.0), which is not + // representable as int64_t, so reject any count at or above that boundary before the cast. + if (count >= 9223372036854775808.0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Range: the computed number of elements exceeds the supported range."); + } + n = static_cast(count); + } TensorShape shape = {n}; T* y = ctx->Output(0, shape)->MutableData(); for (int64_t i = 0; i < n; ++i) { diff --git a/onnxruntime/test/providers/cpu/generator/range_test.cc b/onnxruntime/test/providers/cpu/generator/range_test.cc index 2d6173bdd9703..4c04766ebf359 100644 --- a/onnxruntime/test/providers/cpu/generator/range_test.cc +++ b/onnxruntime/test/providers/cpu/generator/range_test.cc @@ -293,6 +293,18 @@ TEST(RangeTest, ContribOp_ZeroDelta_FailsCleanly) { ASSERT_FALSE(status.IsOK()); } +// Verifies that a finite element count that exceeds the int64 range is rejected cleanly, +// rather than reaching an out-of-range conversion (start=0, limit=1e19, delta=1). +TEST(RangeTest, ContribOp_CountExceedsInt64Range_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_CountExceedsInt64Range_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(1e19), + ScalarInput(1.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + #endif // !defined(ORT_NO_EXCEPTIONS) #endif // DISABLE_CONTRIB_OPS