Skip to content

Validate EmbedLayerNormalization embedding indices (CUDA) #29256

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
titaiwangms wants to merge 3 commits into from
Closed

Validate EmbedLayerNormalization embedding indices (CUDA) #29256

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3ab07d5
contrib: validate EmbedLayerNormalization CUDA embedding indices
titaiwangms Jun 25, 2026
23f5c61
contrib: skip EmbedLayerNormalization error readback during CUDA grap…
titaiwangms Jun 25, 2026
b2719b7
contrib: tidy EmbedLayerNormalization CUDA validation review minors
titaiwangms Jun 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,13 @@ inline Status CheckInputs(const TOpKernelContext* context, bool quantizedVersion
const Tensor* beta = context->template Input<Tensor>(6);
const Tensor* mask = context->template Input<Tensor>(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<Tensor>(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]) {
Expand Down Expand Up @@ -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) {
Expand Down
36 changes: 34 additions & 2 deletions onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,9 +61,18 @@ Status EmbedLayerNorm<T>::ComputeInternal(OpKernelContext* context) const {
int sequence_length = static_cast<int>(input_dims[1]);
size_t element_size = sizeof(T);

int word_embedding_length = static_cast<int>(word_embedding->Shape()[0]);
int position_embedding_length = static_cast<int>(position_embedding->Shape()[0]);
int segment_embedding_length =
(nullptr == segment_embedding) ? 0 : static_cast<int>(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<int>(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<T>(),
nullptr == mask_index ? nullptr : mask_index->MutableData<int32_t>(),
Expand All @@ -82,7 +91,30 @@ Status EmbedLayerNorm<T>::ComputeInternal(OpKernelContext* context) const {
element_size,
embedding_sum == nullptr ? nullptr : embedding_sum->MutableData<T>(),
position_ids == nullptr ? nullptr : position_ids->Data<int32_t>(),
broadcast_position_ids);
broadcast_position_ids,
word_embedding_length,
position_embedding_length,
segment_embedding_length,
error_flag.get()));

// 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<int>(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
Expand Down
48 changes: 36 additions & 12 deletions onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,8 @@ template <typename T, unsigned TPB>
__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
Expand All @@ -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;
Expand All @@ -150,17 +152,30 @@ __global__ void EmbedLayerNormKernel(
} else {
position_id = position_ids[sequence_position];
}

// 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 ||
(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<int64_t>(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<int64_t>(word_id) * hidden_size;
const int64_t segment_offset = static_cast<int64_t>(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<int64_t>(sequence_position) * hidden_size;

cub::KeyValuePair<float, float> thread_data(0.f, 0.f);

Expand All @@ -182,8 +197,8 @@ __global__ void EmbedLayerNormKernel(
thread_data = pair_sum(thread_data, cub::KeyValuePair<float, float>(rldval, rldval * val_f));
}

// 3. layer norm on the sum
LayerNorm<T, TPB>(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<T, TPB>(thread_data, hidden_size, static_cast<int>(output_offset), beta, gamma, epsilon, output);
}

template <typename T>
Expand All @@ -192,15 +207,18 @@ 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);

EmbedLayerNormKernel<T, tpb>
<<<grid, block, 0, stream>>>(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());
}
Expand All @@ -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));
Expand All @@ -241,15 +263,17 @@ Status LaunchEmbedLayerNormKernel(
reinterpret_cast<const half*>(word_embedding), reinterpret_cast<const half*>(position_embedding),
reinterpret_cast<const half*>(segment_embedding), epsilon,
reinterpret_cast<half*>(output), reinterpret_cast<half*>(embedding_sum), position_ids,
broadcast_position_ids);
broadcast_position_ids, word_embedding_length, position_embedding_length,
segment_embedding_length, error_flag);
} else {
return EmbedSkipLayerNorm<float>(
stream, hidden_size, batch_size, sequence_length, input_ids, segment_ids,
reinterpret_cast<const float*>(beta), reinterpret_cast<const float*>(gamma),
reinterpret_cast<const float*>(word_embedding), reinterpret_cast<const float*>(position_embedding),
reinterpret_cast<const float*>(segment_embedding), epsilon,
reinterpret_cast<float*>(output), reinterpret_cast<float*>(embedding_sum), position_ids,
broadcast_position_ids);
broadcast_position_ids, word_embedding_length, position_embedding_length,
segment_embedding_length, error_flag);
}
}

Expand Down
6 changes: 5 additions & 1 deletion onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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, // rows in segment_embedding (0 if none)
int* error_flag); // device flag set when an id is out of range

} // namespace cuda
} // namespace contrib
Expand Down
101 changes: 98 additions & 3 deletions onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -282,8 +282,103 @@ TEST(EmbedLayerNormTest, EmbedLayerNormNegativePositionIds) {
tester.AddOutput<float>("output", output_dims, std::vector<float>(batch_size * sequence_size * hidden_size, 0.0f));
tester.AddOutput<int32_t>("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. 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 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) {
int batch_size = 1;
int sequence_size = 2;
int hidden_size = 4;

std::vector<int64_t> input_ids_dims = {batch_size, sequence_size};
std::vector<int64_t> word_embedding_dims = {6, hidden_size};
std::vector<int64_t> position_embedding_dims = {3, hidden_size};
std::vector<int64_t> segment_embedding_dims = {2, hidden_size};
std::vector<int64_t> gamma_dims = {hidden_size};
std::vector<int64_t> beta_dims = {hidden_size};
std::vector<int64_t> output_dims = {batch_size, sequence_size, hidden_size};
std::vector<int64_t> mask_index_dims = {batch_size};

OpTester tester("EmbedLayerNormalization", 1, onnxruntime::kMSDomain);
// Second id (99) is outside the 6-row word_embedding table.
tester.AddInput<int32_t>("input_ids", input_ids_dims, {1, 99});
tester.AddInput<int32_t>("segment_ids", input_ids_dims, {0, 1});
tester.AddInput<float>("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<float>("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<float>("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<float>("gamma", gamma_dims, {0.25f, 0.15f, 0.45f, -0.66f}, /*is_initializer=*/true);
tester.AddInput<float>("beta", beta_dims, {0.6f, 0.2f, 0.5f, -0.6f}, /*is_initializer=*/true);
tester.AddAttribute("epsilon", embedlayernorm::kEpsilon);

tester.AddOutput<float>("output", output_dims, std::vector<float>(batch_size * sequence_size * hidden_size, 0.0f));
tester.AddOutput<int32_t>("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<int64_t> input_ids_dims = {batch_size, sequence_size};
std::vector<int64_t> word_embedding_dims = {6, hidden_size};
// Only 1 row, but sequence_size is 2 and no position_ids are supplied.
std::vector<int64_t> position_embedding_dims = {1, hidden_size};
std::vector<int64_t> segment_embedding_dims = {2, hidden_size};
std::vector<int64_t> gamma_dims = {hidden_size};
std::vector<int64_t> beta_dims = {hidden_size};
std::vector<int64_t> output_dims = {batch_size, sequence_size, hidden_size};
std::vector<int64_t> mask_index_dims = {batch_size};

OpTester tester("EmbedLayerNormalization", 1, onnxruntime::kMSDomain);
tester.AddInput<int32_t>("input_ids", input_ids_dims, {1, 3});
tester.AddInput<int32_t>("segment_ids", input_ids_dims, {0, 1});
tester.AddInput<float>("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<float>("position_embedding", position_embedding_dims,
{0.1f, 0.1f, 0.4f, 0.6f},
/*is_initializer=*/true);
tester.AddInput<float>("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<float>("gamma", gamma_dims, {0.25f, 0.15f, 0.45f, -0.66f}, /*is_initializer=*/true);
tester.AddInput<float>("beta", beta_dims, {0.6f, 0.2f, 0.5f, -0.6f}, /*is_initializer=*/true);
tester.AddAttribute("epsilon", embedlayernorm::kEpsilon);

tester.AddOutput<float>("output", output_dims, std::vector<float>(batch_size * sequence_size * hidden_size, 0.0f));
tester.AddOutput<int32_t>("mask_index", mask_index_dims, {0});

tester.Run(OpTester::ExpectResult::kExpectFailure, "", {kDmlExecutionProvider, kOpenVINOExecutionProvider});
}

} // namespace test
Expand Down
Loading