From 3ab07d5041daa60cf0e0cac32b2c72a81eadd9eb Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:04:57 +0000 Subject: [PATCH 1/3] 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 23f5c61134c31665ca9636e0b164dc08113d98c7 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:14:17 +0000 Subject: [PATCH 2/3] 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 b2719b7a868c9ce64ef74fde50fc3559a2a78bc5 Mon Sep 17 00:00:00 2001 From: titaiwangms Date: Thu, 25 Jun 2026 01:18:05 +0000 Subject: [PATCH 3/3] 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) {