Merge pull request #1064 from intel/sync_msft_26042026

Sync with Microsoft ONNX Runtime - 26042026

Author: ankitm3k

js/web/lib/wasm/jsep/webgpu/ops/attention.ts

@@ -809,14 +809,17 @@ export const applyAttention = (
 ) => {
   // Assumption is that presentKey/presentValue exists only if pastKey/pastValue exists.
   const outputCount = Math.min(context.outputCount, 1 + (pastKey ? 1 : 0) + (pastValue ? 1 : 0));
+  // When there are no present key/value outputs (outputCount <= 1), ignore past to match CPU EP semantics.
+  const effectivePastKey = outputCount > 1 ? pastKey : undefined;
+  const effectivePastValue = outputCount > 1 ? pastValue : undefined;
   const pastSequenceLength = outputCount > 1 ? parameters.pastSequenceLength : 0;
   const totalSequenceLength = pastSequenceLength + parameters.kvSequenceLength;
   const attentionBias =
     attentionBiasInput && ShapeUtil.size(attentionBiasInput.dims) > 0 ? attentionBiasInput : undefined;
 
   const inputsK = [q, k];
-  if (outputCount > 1 && pastKey && ShapeUtil.size(pastKey.dims) > 0) {
-    inputsK.push(pastKey);
+  if (effectivePastKey && ShapeUtil.size(effectivePastKey.dims) > 0) {
+    inputsK.push(effectivePastKey);
   }
   if (attentionBias) {
     inputsK.push(attentionBias);
@@ -833,7 +836,7 @@ export const applyAttention = (
       outputCount,
       q,
       k,
-      pastKey,
+      effectivePastKey,
       attentionBias,
       parameters,
       pastSequenceLength,
@@ -860,8 +863,8 @@ export const applyAttention = (
 
   // Run AttentionScore
   const inputsV = [probs, v];
-  if (outputCount > 1 && pastValue && ShapeUtil.size(pastValue.dims) > 0) {
-    inputsV.push(pastValue);
+  if (effectivePastValue && ShapeUtil.size(effectivePastValue.dims) > 0) {
+    inputsV.push(effectivePastValue);
   }
   if (seqLens) {
     inputsV.push(seqLens);
@@ -874,7 +877,7 @@ export const applyAttention = (
       outputCount,
       probs,
       v,
-      pastValue,
+      effectivePastValue,
       parameters,
       pastSequenceLength,
       seqLens,

onnxruntime/contrib_ops/cpu/attnlstm/uni_dir_attn_lstm.cc

@@ -216,7 +216,7 @@ void UniDirectionalAttnLstm<T>::LoadPeepholeWeights(const gsl::span<const T>& pe
 template <typename T>
 void UniDirectionalAttnLstm<T>::LoadBias(const gsl::span<const T>& WbRb_values) {
   const size_t hidden_size = CheckedToSizeT(hidden_size_);
-  const size_t Wb_to_Rb_offset = CheckedMulToSizeT(hidden_size_, 4);
+  const size_t Wb_to_Rb_offset = CheckedMulToSizeT(hidden_size, size_t{4});
 
   // add Wb and Rb
   auto copy_fused_bias = [&WbRb_values, hidden_size, Wb_to_Rb_offset](size_t offset, gsl::span<T>& out) {

onnxruntime/contrib_ops/cpu/bert/rotary_embedding.cc

@@ -3,12 +3,19 @@
 
 #include "contrib_ops/cpu/bert/rotary_embedding.h"
 #include "contrib_ops/cpu/bert/rotary_embedding_helper.h"
+#include "core/providers/cpu/llm/rotary_embedding_int32_utils.h"
+
+#include <algorithm>
+#include <limits>
 
 #include "core/mlas/inc/mlas.h"
 #include "core/platform/threadpool.h"
 
 using onnxruntime::concurrency::ThreadPool;
 using namespace onnxruntime::contrib::rotary_embedding_helper;
+using onnxruntime::rotary_embedding_int32_utils::CheckedAddToPtrdiff;
+
+using onnxruntime::rotary_embedding_int32_utils::CheckedPtrdiffMulToPtrdiff;
 
 namespace onnxruntime {
 namespace contrib {
@@ -32,8 +39,16 @@ REGISTER_KERNEL_TYPED(MLFloat16)
 template <typename T>
 RotaryEmbedding<T>::RotaryEmbedding(const OpKernelInfo& info) : OpKernel(info) {
   scale = info.GetAttrOrDefault<float>("scale", 1.0);
-  rotary_embedding_dim = static_cast<int>(info.GetAttrOrDefault<int64_t>("rotary_embedding_dim", 0));
-  num_heads = static_cast<int>(info.GetAttrOrDefault<int64_t>("num_heads", 0));
+  const int64_t rotary_embedding_dim_attr = info.GetAttrOrDefault<int64_t>("rotary_embedding_dim", 0);
+  const int64_t num_heads_attr = info.GetAttrOrDefault<int64_t>("num_heads", 0);
+  ORT_ENFORCE(rotary_embedding_dim_attr >= 0 && rotary_embedding_dim_attr <= std::numeric_limits<int>::max(),
+              "rotary_embedding_dim must be in range [0, ", std::numeric_limits<int>::max(),
+              "]. Actual value: ", rotary_embedding_dim_attr);
+  ORT_ENFORCE(num_heads_attr >= 0 && num_heads_attr <= std::numeric_limits<int>::max(),
+              "num_heads must be in range [0, ", std::numeric_limits<int>::max(),
+              "]. Actual value: ", num_heads_attr);
+  rotary_embedding_dim = static_cast<int>(rotary_embedding_dim_attr);
+  num_heads = static_cast<int>(num_heads_attr);
   interleaved = (info.GetAttrOrDefault<int64_t>("interleaved", 0) == 1);
   is_packed_batching = (info.GetAttrOrDefault<int64_t>("is_packed_batching", 0) == 1);
 
@@ -72,8 +87,12 @@ Status RunRotaryEmbedding(concurrency::ThreadPool* tp, RotaryParameters paramete
                              " exceeds cos/sin cache range [0, ", max_sequence_length, ")");
     }
   } else if (position_ids_format == 1) {
+    std::ptrdiff_t position_count = 0;
+    ORT_RETURN_IF_ERROR(onnxruntime::rotary_embedding_int32_utils::CheckedMulToPtrdiff(
+        batch_size, sequence_length, "position_ids element count", position_count));
+
     // Format 1: 2D array (batch_size, sequence_length)
-    for (int i = 0; i < batch_size * sequence_length; ++i) {
+    for (std::ptrdiff_t i = 0; i < position_count; ++i) {
       int64_t pos = position_ids[i];
       if (pos < 0 || pos >= static_cast<int64_t>(max_sequence_length)) {
         return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
@@ -84,7 +103,37 @@ Status RunRotaryEmbedding(concurrency::ThreadPool* tp, RotaryParameters paramete
   }
 
   // Parallel to calculate based on head_size
-  const int loop_len = batch_size * sequence_length * n_heads;
+  std::ptrdiff_t loop_len = 0;
+  ORT_RETURN_IF_ERROR(onnxruntime::rotary_embedding_int32_utils::CheckedMulToPtrdiff(
+      batch_size, sequence_length, n_heads, "total_elements", loop_len));
+
+  std::ptrdiff_t max_batch_offset = 0;
+  std::ptrdiff_t max_seq_offset = 0;
+  std::ptrdiff_t max_head_offset = 0;
+  std::ptrdiff_t max_block_offset = 0;
+  std::ptrdiff_t max_b_s_index = 0;
+  [[maybe_unused]] std::ptrdiff_t max_cache_offset = 0;
+  ORT_RETURN_IF_ERROR(onnxruntime::rotary_embedding_int32_utils::CheckedMulToPtrdiff(
+      std::max(batch_size - 1, 0), batch_stride, "max_batch_offset", max_batch_offset));
+  ORT_RETURN_IF_ERROR(onnxruntime::rotary_embedding_int32_utils::CheckedMulToPtrdiff(
+      std::max(sequence_length - 1, 0), seq_stride, "max_seq_offset", max_seq_offset));
+  ORT_RETURN_IF_ERROR(onnxruntime::rotary_embedding_int32_utils::CheckedMulToPtrdiff(
+      std::max(n_heads - 1, 0), head_stride, "max_head_offset", max_head_offset));
+  ORT_RETURN_IF_ERROR(onnxruntime::rotary_embedding_int32_utils::CheckedAddToPtrdiff(
+      max_batch_offset, max_seq_offset, "max_block_offset", max_block_offset));
+  ORT_RETURN_IF_ERROR(onnxruntime::rotary_embedding_int32_utils::CheckedAddToPtrdiff(
+      max_block_offset, max_head_offset, "max_block_offset", max_block_offset));
+  if (position_ids_format == 0) {
+    std::ptrdiff_t total_b_s_count = 0;
+    ORT_RETURN_IF_ERROR(onnxruntime::rotary_embedding_int32_utils::CheckedMulToPtrdiff(
+        batch_size, sequence_length, "total_b_s_count", total_b_s_count));
+    max_b_s_index = total_b_s_count > 0 ? total_b_s_count - 1 : 0;
+  } else {
+    max_b_s_index = std::max(max_sequence_length - 1, 0);
+  }
+  ORT_RETURN_IF_ERROR(onnxruntime::rotary_embedding_int32_utils::CheckedPtrdiffMulToPtrdiff(
+      max_b_s_index, half_rotary_emb_dim, "max_cache_offset", max_cache_offset));
+
   // The cost is calculated as:
   //   - head_size * sizeof(T) for reading input
   //   - head_size * sizeof(T) for writing output
@@ -97,16 +146,18 @@ Status RunRotaryEmbedding(concurrency::ThreadPool* tp, RotaryParameters paramete
       const int n = static_cast<int>(ptr % n_heads);
       // Identify the index of batch, sequence, and head (specific range) in the input/output tensor
       // for read/write
-      const int block_offset = b * batch_stride + s * seq_stride + n * head_stride;
+      const std::ptrdiff_t block_offset = static_cast<std::ptrdiff_t>(b) * batch_stride +
+                                          static_cast<std::ptrdiff_t>(s) * seq_stride +
+                                          static_cast<std::ptrdiff_t>(n) * head_stride;
 
       const T* input_data = input + block_offset;
       T* output_data = output + block_offset;
 
       // Cache is (M, H/2) or (M, rotary_embedding_dim/2)
-      const int position_id = (position_ids_format == 0)
-                                  ? static_cast<int>(position_ids[0]) + s
-                                  : static_cast<int>(position_ids[b * sequence_length + s]);
-      const int cache_offset = position_id * half_rotary_emb_dim;
+      const std::ptrdiff_t position_id = (position_ids_format == 0)
+                                             ? static_cast<std::ptrdiff_t>(position_ids[0]) + s
+                                             : static_cast<std::ptrdiff_t>(position_ids[static_cast<std::ptrdiff_t>(b) * sequence_length + s]);
+      const std::ptrdiff_t cache_offset = position_id * half_rotary_emb_dim;
       const T* cos_data = cos_cache + cache_offset;
       const T* sin_data = sin_cache + cache_offset;
 

onnxruntime/contrib_ops/cpu/bert/rotary_embedding_helper.h

@@ -2,13 +2,20 @@
 // Licensed under the MIT License.
 
 #pragma once
+
 #include "core/common/common.h"
 #include "core/providers/common.h"
+#include "core/providers/cpu/llm/rotary_embedding_int32_utils.h"
 
 namespace onnxruntime {
 namespace contrib {
 namespace rotary_embedding_helper {
 
+namespace detail {
+using onnxruntime::rotary_embedding_int32_utils::CheckedMulToInt32;
+using onnxruntime::rotary_embedding_int32_utils::NarrowNonNegativeToInt32;
+}  // namespace detail
+
 // Parameters deduced from node attributes and inputs/outputs.
 struct RotaryParameters {
   int batch_size;            // Batch size used by input
@@ -73,20 +80,52 @@ Status CheckInputs(const T* input,
   }
 
   // Get attributes from inputs
-  int batch_size = static_cast<int>(input_dims[0]);
-  int sequence_length = static_cast<int>(input_dims[1]);
-  int hidden_size = static_cast<int>(input_dims[2]);
+  int batch_size = 0;
+  int sequence_length = 0;
+  int hidden_size = 0;
+  int head_size = 0;
+
+  ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(input_dims[0], "batch_size", batch_size));
+  ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(input_dims[1], "sequence_length", sequence_length));
+  ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(input_dims[2], "hidden_size", hidden_size));
 
   bool transposed = false;
   if (input_dims.size() == 4) {
     // input is [batch, num_heads, seq, head_size]
-    sequence_length = static_cast<int>(input_dims[2]);
-    hidden_size = static_cast<int>(input_dims[1]) * static_cast<int>(input_dims[3]);
+    ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(input_dims[2], "sequence_length", sequence_length));
+    ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(input_dims[1], "num_heads", num_heads));
+    ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(input_dims[3], "head_size", head_size));
+    ORT_RETURN_IF_ERROR(detail::CheckedMulToInt32(num_heads, head_size, "hidden_size", hidden_size));
     transposed = true;
+  } else if (num_heads > 0) {
+    if (hidden_size % num_heads != 0) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "RotaryEmbedding: hidden_size=", hidden_size,
+                             " must be divisible by num_heads=", num_heads,
+                             " for rank-3 input");
+    }
+    head_size = static_cast<int>(hidden_size / num_heads);
+    if (batch_size > 0 && sequence_length > 0 && hidden_size > 0 && head_size <= 0) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "RotaryEmbedding: head_size must be greater than 0 for non-empty rank-3 input");
+    }
+  }
+  int max_sequence_length = 0;
+  ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(cos_cache_dims[0], "max_sequence_length", max_sequence_length));
+  if (rotary_embedding_dim == 0) {
+    int cache_width = 0;
+    ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(cos_cache_dims[1], "cache_width", cache_width));
+    if (head_size == 0) {
+      ORT_RETURN_IF_ERROR(detail::CheckedMulToInt32(cache_width, 2, "head_size", head_size));
+    }
+  } else {
+    if (!transposed) {
+      if (num_heads <= 0) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "RotaryEmbedding: num_heads must be greater than 0 for rank-3 input");
+      }
+    }
   }
-  int max_sequence_length = static_cast<int>(cos_cache_dims[0]);
-  int head_size = rotary_embedding_dim == 0 ? static_cast<int>(cos_cache_dims[1]) * 2
-                                            : static_cast<int>(hidden_size / num_heads);
   if (rotary_embedding_dim > 0 && rotary_embedding_dim > head_size) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "rotary_embedding_dim must be less than or equal to ",
                            "head_size");
@@ -96,11 +135,16 @@ Status CheckInputs(const T* input,
 
   // Check position_ids input shapes
   if (!onnxruntime::IsScalarOr1ElementVector(position_ids)) {
-    if (batch_size != static_cast<int>(position_ids_dims[0])) {
+    int position_ids_batch = 0;
+    int position_ids_sequence = 0;
+    ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(position_ids_dims[0], "position_ids_dim0", position_ids_batch));
+    ORT_RETURN_IF_ERROR(detail::NarrowNonNegativeToInt32(position_ids_dims[1], "position_ids_dim1", position_ids_sequence));
+
+    if (batch_size != position_ids_batch) {
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'position_ids' dimension 0 should be of size ",
                              "batch_size, got ", position_ids_dims[0]);
     }
-    if (sequence_length != static_cast<int>(position_ids_dims[1])) {
+    if (sequence_length != position_ids_sequence) {
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'position_ids' dimension 1 should be of size ",
                              "sequence_length, got ", position_ids_dims[1]);
     }
@@ -114,26 +158,44 @@ Status CheckInputs(const T* input,
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'cos_cache' dimension 0 should be same as ",
                            "max_sequence_length, got ", cos_cache_dims[0]);
   }
-  if ((head_size / 2) != static_cast<int>(cos_cache_dims[1]) && (rotary_embedding_dim > 0 && (rotary_embedding_dim / 2) != static_cast<int>(cos_cache_dims[1]))) {
+  if ((head_size / 2) != static_cast<int>(cos_cache_dims[1]) &&
+      (rotary_embedding_dim <= 0 || (rotary_embedding_dim / 2) != static_cast<int>(cos_cache_dims[1]))) {
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'cos_cache' dimension 1 should be same as ",
                            "head_size / 2 or rotary_embedding_dim / 2, got ", cos_cache_dims[1]);
   }
 
-  num_heads = num_heads > 0 ? num_heads : static_cast<int>(hidden_size / head_size);
+  if (num_heads <= 0) {
+    if (head_size == 0) {
+      if (batch_size == 0 || sequence_length == 0 || hidden_size == 0) {
+        num_heads = 0;
+      } else {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "RotaryEmbedding: head_size must be greater than 0 when inferring num_heads");
+      }
+    } else {
+      if (hidden_size % head_size != 0) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "RotaryEmbedding: hidden_size=", hidden_size,
+                               " must be divisible by inferred head_size=", head_size,
+                               " when inferring num_heads");
+      }
+      num_heads = static_cast<int>(hidden_size / head_size);
+    }
+  }
   // Calculate stride values
   int head_stride;
   int seq_stride;
   int batch_stride;
   if (transposed) {
     // Transposed input tensor shape is [batch, n_heads, seq_len, head_size]
     seq_stride = head_size;
-    head_stride = sequence_length * seq_stride;
-    batch_stride = num_heads * head_stride;
+    ORT_RETURN_IF_ERROR(detail::CheckedMulToInt32(sequence_length, seq_stride, "head_stride", head_stride));
+    ORT_RETURN_IF_ERROR(detail::CheckedMulToInt32(num_heads, head_stride, "batch_stride", batch_stride));
   } else {
     // Default input tensor shape is [batch, seq_len, hidden_size]
     head_stride = head_size;
-    seq_stride = num_heads * head_stride;
-    batch_stride = sequence_length * seq_stride;
+    ORT_RETURN_IF_ERROR(detail::CheckedMulToInt32(num_heads, head_stride, "seq_stride", seq_stride));
+    ORT_RETURN_IF_ERROR(detail::CheckedMulToInt32(sequence_length, seq_stride, "batch_stride", batch_stride));
   }
 
   // Set rotary parameters