support NeoX-style rotary embedding (#14785)

### Description
<!-- Describe your changes. -->



### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

Co-authored-by: Ubuntu <wy@v100-2.0cdb2e52twzevn1i4fi45bylyg.jx.internal.cloudapp.net>

docs/ContribOperators.md

@@ -127,6 +127,8 @@ This version of the operator has been available since version 1 of the 'com.micr
 #### Attributes
 
 <dl>
+<dt><tt>do_rotary</tt> : int</dt>
+<dd>Whether to use rotary position embedding. Default value is 0.</dd>
 <dt><tt>mask_filter_value</tt> : float</dt>
 <dd>The value to be filled in the attention mask. Default value is -10000.0f</dd>
 <dt><tt>num_heads</tt> : int (required)</dt>
@@ -2588,6 +2590,8 @@ This version of the operator has been available since version 1 of the 'com.micr
 #### Attributes
 
 <dl>
+<dt><tt>do_rotary</tt> : int</dt>
+<dd>Whether to use rotary position embedding. Default value is 0.</dd>
 <dt><tt>mask_filter_value</tt> : float</dt>
 <dd>The value to be filled in the attention mask. Default value is -10000.0f</dd>
 <dt><tt>num_heads</tt> : int (required)</dt>

onnxruntime/contrib_ops/cpu/bert/attention_base.cc

@@ -125,6 +125,7 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
   }
 
   int64_t past_sequence_length = 0;
+  int64_t original_past_sequence_length = 0;
   if (past != nullptr) {  // past is optional
     if (k_hidden_size != v_hidden_size) {
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'past' expect k_hidden_size == v_hidden_size");
@@ -157,6 +158,7 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
 
     if (!past_present_share_buffer_) {
       past_sequence_length = past_dims[3];
+      original_past_sequence_length = past_sequence_length;
     } else {
       if (past_seq_len == nullptr || !onnxruntime::IsScalarOr1ElementVector(past_seq_len)) {
         return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
@@ -237,6 +239,7 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
     output_parameters->batch_size = static_cast<int>(batch_size);
     output_parameters->sequence_length = static_cast<int>(sequence_length);
     output_parameters->past_sequence_length = static_cast<int>(past_sequence_length);
+    output_parameters->original_past_sequence_length = static_cast<int>(original_past_sequence_length);
     output_parameters->kv_sequence_length = static_cast<int>(kv_sequence_length);
     output_parameters->total_sequence_length = static_cast<int>(total_sequence_length);
     output_parameters->max_sequence_length = static_cast<int>(max_sequence_length);
@@ -248,6 +251,7 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape,
     output_parameters->num_heads = num_heads_;
     output_parameters->is_unidirectional = is_unidirectional_;
     output_parameters->past_present_share_buffer = (past_present_share_buffer_ != 0 && past != nullptr);
+    output_parameters->do_rotary = do_rotary_;
     output_parameters->mask_filter_value = mask_filter_value_;
     output_parameters->scale = scale_;
     output_parameters->mask_type = mask_type;

onnxruntime/contrib_ops/cpu/bert/attention_base.h

@@ -37,6 +37,7 @@ class AttentionBase {
     num_heads_ = static_cast<int>(num_heads);
 
     is_unidirectional_ = info.GetAttrOrDefault<int64_t>("unidirectional", 0) == 1;
+    do_rotary_ = info.GetAttrOrDefault<int64_t>("do_rotary", 0) == 1;
     mask_filter_value_ = info.GetAttrOrDefault<float>("mask_filter_value", -10000.0f);
     scale_ = info.GetAttrOrDefault<float>("scale", 0.0f);
 
@@ -70,6 +71,7 @@ class AttentionBase {
   std::vector<int64_t> qkv_hidden_sizes_;  // Q, K, V hidden sizes parsed from the qkv_hidden_sizes attribute.
   bool require_same_hidden_size_;          // whether the implementation supports different hidden sizes of Q/K/V.
   bool past_present_share_buffer_;         // whether or not the past (if used) and present tensor share the same buffer
+  bool do_rotary_;                         // whether or not to use rotary embeddings
   float mask_filter_value_;                // the value to be used for filtered out positions
   float scale_;                            // the scale to be used for softmax
 };

onnxruntime/contrib_ops/cpu/bert/attention_common.h

@@ -38,18 +38,20 @@ enum AttentionKernelType {
 struct AttentionParameters {
   int batch_size;
   int sequence_length;
-  int kv_sequence_length;     // input sequence length of K or V
-  int past_sequence_length;   // sequence length in past state of K or V
-  int total_sequence_length;  // total sequence length of K or V
-  int max_sequence_length;    // max sequence length from 4D mask
-  int input_hidden_size;      // first dimension of weights for input projection
-  int hidden_size;            // hidden size of Q or K
-  int head_size;              // hidden size per head of Q or K
-  int v_hidden_size;          // hidden size of V
-  int v_head_size;            // hidden size per head of V
+  int kv_sequence_length;            // input sequence length of K or V
+  int past_sequence_length;          // sequence length in past state of K or V
+  int original_past_sequence_length; // original sequence length in past state of K or V
+  int total_sequence_length;         // total sequence length of K or V
+  int max_sequence_length;           // max sequence length from 4D mask
+  int input_hidden_size;             // first dimension of weights for input projection
+  int hidden_size;                   // hidden size of Q or K
+  int head_size;                     // hidden size per head of Q or K
+  int v_hidden_size;                 // hidden size of V
+  int v_head_size;                   // hidden size per head of V
   int num_heads;
   bool is_unidirectional;
   bool past_present_share_buffer;
+  bool do_rotary;
   float mask_filter_value;
   float scale;
   AttentionMaskType mask_type;

onnxruntime/contrib_ops/cuda/bert/add_bias_transpose.cu

@@ -3,32 +3,7 @@
 #include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cuda/cu_inc/common.cuh"
 #include "contrib_ops/cuda/bert/add_bias_transpose.h"
-
-namespace onnxruntime {
-namespace cuda {
-
-struct __align__(8) Half4 {
-  half2 x;
-  half2 y;
-};
-
-__device__ __forceinline__ Half4 operator+(const Half4& a, const Half4& b) {
-  Half4 r;
-  r.x = a.x + b.x;
-  r.y = a.y + b.y;
-  return r;
-}
-
-__device__ __forceinline__ float2 operator+(const float2& a, const float2& b) {
-  return make_float2(a.x + b.x, a.y + b.y);
-}
-
-__device__ __forceinline__ float4 operator+(const float4& a, const float4& b) {
-  return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w);
-}
-
-}  // namespace cuda
-}  // namespace onnxruntime
+#include "contrib_ops/cuda/bert/rotary_embedding_util.h"
 
 using namespace onnxruntime::cuda;
 
@@ -240,6 +215,151 @@ __global__ void AddBiasTransposeQKV(int M, const T* input, const T* biases, T* o
   }
 }
 
+#ifndef USE_ROCM
+template <typename T>
+__global__ void AddBiasTransposeQKV(int M, const T* input, const T* biases, T* output, T* qkv_add_bias,
+                                    const int rotary_embedding_dim, const int head_size, const int step,
+                                    const int format) {
+  // AddBiasTransposeQKV with rotary embedding
+  // Format 1 for unfused attention, or fused causal attention
+  //     Input:  BxSxMxNxH
+  //     Output: MxBxNxSxH
+  //     qkv_add_bias: BxSxMxNxH
+  // Format 2 for fused TRT attention
+  //     Input:  BxSxMxNxH
+  //     Output: BxSxNxMxH
+  //     qkv_add_bias: BxSxMxNxH
+  // Format 3 for cutlass memory efficient attention
+  //     Input:  BxSxMxNxH
+  //     Output: MxBxSxNxH
+  // B is batch_size, S is sequence_length, M is number of matrices, N is num_heads, H is head_size
+  int n = blockIdx.y;
+  int s = blockIdx.x;
+  int b = blockIdx.z;
+
+  const int seq_len = (gridDim.x == step) ? s : step;
+
+  const int num_heads = gridDim.y;
+
+  const int sequence_length = gridDim.x;
+  const int batch_size = gridDim.z;
+  const int H = head_size;
+  const int NH = num_heads * head_size;
+  const int NHS = NH * sequence_length;
+
+  constexpr int vec_size = Vec_t<T>::size;
+  using Vec_t = typename Vec_t<T>::Type;
+
+  extern __shared__ __align__(sizeof(float2)) char smem_[];
+
+  int tidx = threadIdx.x;
+  const int head_idx = tidx * vec_size;
+
+  if (head_idx < head_size) {
+    const bool is_masked = head_idx >= head_size;
+
+    const int input_offset_base = n * head_size + (s * M) * NH + b * NHS * M;
+    const int src_q_idx = input_offset_base + head_idx;
+    const int src_k_idx = input_offset_base + NH + head_idx;
+    const int src_v_idx = input_offset_base + 2 * NH + head_idx;
+
+    Vec_t q, k, v;
+    Vec_t q_bias, k_bias, v_bias;
+
+    if (!is_masked) {
+      q = *reinterpret_cast<const Vec_t*>(&input[src_q_idx]);
+      k = *reinterpret_cast<const Vec_t*>(&input[src_k_idx]);
+      v = *reinterpret_cast<const Vec_t*>(&input[src_v_idx]);
+
+      q_bias = *reinterpret_cast<const Vec_t*>(&biases[n * H + head_idx]);
+      k_bias = *reinterpret_cast<const Vec_t*>(&biases[NH + n * H + head_idx]);
+      v_bias = *reinterpret_cast<const Vec_t*>(&biases[2 * NH + n * H + head_idx]);
+    }
+
+    q = add(q, q_bias);
+    k = add(k, k_bias);
+    v = add(v, v_bias);
+
+    const bool do_rotary = !is_masked && vec_size * tidx < rotary_embedding_dim;
+
+    T* q_smem = reinterpret_cast<T*>(smem_);
+    T* k_smem = q_smem + rotary_embedding_dim;
+
+    const int half_rotary_dim = rotary_embedding_dim / 2;
+    const int half_idx        = (head_idx) / half_rotary_dim;
+    const int intra_half_idx  = (head_idx) % half_rotary_dim;
+    const int smem_pitch      = half_rotary_dim;
+
+    if (do_rotary) {
+      *reinterpret_cast<Vec_t*>(q_smem + half_idx * smem_pitch + intra_half_idx) = q;
+      *reinterpret_cast<Vec_t*>(k_smem + half_idx * smem_pitch + intra_half_idx) = k;
+    }
+
+    __syncthreads();
+
+    const int transpose_idx = half_idx * (half_rotary_dim / 2) + intra_half_idx / 2;
+    constexpr int tidx_factor = vec_size / 2;
+
+    if (do_rotary) {
+      vec_from_smem_transpose(q, q_smem, transpose_idx, smem_pitch);
+      vec_from_smem_transpose(k, k_smem, transpose_idx, smem_pitch);
+
+      apply_rotary_embedding(q, k, transpose_idx / tidx_factor, rotary_embedding_dim, seq_len);
+
+      write_smem_transpose(q, q_smem, transpose_idx, smem_pitch);
+      write_smem_transpose(k, k_smem, transpose_idx, smem_pitch);
+    }
+
+    __syncthreads();
+
+    if (do_rotary) {
+      q = *reinterpret_cast<Vec_t*>(q_smem + half_idx * smem_pitch + intra_half_idx);
+      k = *reinterpret_cast<Vec_t*>(k_smem + half_idx * smem_pitch + intra_half_idx);
+    }
+
+    int dest_q_idx;
+    int dest_k_idx;
+    int dest_v_idx;
+
+    // Format 1
+    if (format == 1) {
+      const int output_offset_base = s * head_size + n * sequence_length * H + b * NHS;
+      dest_q_idx = output_offset_base + head_idx;
+      dest_k_idx = output_offset_base + NHS * batch_size + head_idx;
+      dest_v_idx = output_offset_base + 2 * NHS * batch_size + head_idx;
+    }
+
+    // Format 2
+    if (format == 2) {
+      const int output_offset_base = M * (b * NHS + s * NH + n * H);
+      dest_q_idx = output_offset_base + head_idx;
+      dest_k_idx = output_offset_base + H + head_idx;
+      dest_v_idx = output_offset_base + 2 * H + head_idx;
+    }
+
+    // Format 3
+    if (format == 3) {
+      const int output_offset_base = n * H + s * NH + b * NHS;
+      dest_q_idx = output_offset_base + head_idx;
+      dest_k_idx = output_offset_base + NHS * batch_size + head_idx;
+      dest_v_idx = output_offset_base + 2 * NHS * batch_size + head_idx;
+    }
+
+    if (!is_masked) {
+      *reinterpret_cast<Vec_t*>(&output[dest_q_idx]) = q;
+      *reinterpret_cast<Vec_t*>(&output[dest_k_idx]) = k;
+      *reinterpret_cast<Vec_t*>(&output[dest_v_idx]) = v;
+
+      if (nullptr != qkv_add_bias) {
+        *reinterpret_cast<Vec_t*>(&qkv_add_bias[src_q_idx]) = q;
+        *reinterpret_cast<Vec_t*>(&qkv_add_bias[src_k_idx]) = k;
+        *reinterpret_cast<Vec_t*>(&qkv_add_bias[src_v_idx]) = v;
+      }
+    }
+  }
+}
+#endif
+
 // this suppose 3 matrix in total
 template <typename T>
 __global__ void AddBiasTransposeQKV(const T* input, const T* biases, T* output, int v_head_size) {
@@ -320,6 +440,7 @@ __global__ void AddBiasTransposeQKVLarge(const int head_size, const T* input, co
   }
 }
 
+
 template <typename T>
 __global__ void AddBiasTransposeCutlass(const T* input, const T* biases, T* output, int v_head_size) {
   // Format 3 for cutlass memory efficient attention
@@ -518,8 +639,35 @@ template <typename T>
 void InvokeAddBiasTranspose(
     cudaStream_t stream, const int num_matrices, const int format, const int max_threads_per_block,
     const int batch_size, const int sequence_length, const int num_heads, const int qk_head_size,
-    const T* input, const T* biases, T* output, T* qkv_add_bias, const int v_head_size, int total_matrix_count) {
+    const T* input, const T* biases, T* output, T* qkv_add_bias, const int v_head_size, int total_matrix_count,
+    bool do_rotary = false, int original_past_sequence_length = 0) {
   assert(num_heads <= max_threads_per_block);
+
+  if (do_rotary) {
+#ifdef USE_ROCM
+    ORT_THROW("Rotary Attention is not supported on ROCm");
+#elif !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 530
+    if (format != 1 && format != 2 && format != 3) {
+      ORT_THROW("format must be 1, 2 or 3 for rotary attention");
+    }
+    if (v_head_size != -1 && qk_head_size != v_head_size) {
+      ORT_THROW("qk_head_size must be equal to v_head_size for rotary attention");
+    }
+
+    const int step = original_past_sequence_length == 0 ? sequence_length : original_past_sequence_length;
+    size_t smem_size = 2 * qk_head_size * sizeof(T);
+
+    const dim3 grid(sequence_length, num_heads, batch_size);
+    const dim3 block((qk_head_size / 2 + 31) / 32 * 32, 1, 1);
+    AddBiasTransposeQKV<T><<<grid, block, smem_size, stream>>>(total_matrix_count, input, biases, output,
+                                                               qkv_add_bias, qk_head_size, qk_head_size,
+                                                               step, format);
+#else
+    ORT_THROW("Rotary Attention is supported on sm >= 530. Current sm is", __CUDA_ARCH__);
+#endif
+    return;
+  }
+
   const dim3 grid(sequence_length, batch_size, num_matrices);
   if (qk_head_size * num_heads <= max_threads_per_block) {
     const dim3 block(qk_head_size, num_heads, 1);
@@ -575,10 +723,10 @@ template <>
 void LaunchAddBiasTranspose(
     cudaStream_t stream, const int num_matrices, const int format, const int max_threads_per_block,
     const int batch_size, const int sequence_length, const int num_heads, const int qk_head_size,
-    const half* input, const half* biases, half* output,
-    bool enable_half4, const int v_head_size, half* qkv_add_bias, int total_matrix_count) {
+    const half* input, const half* biases, half* output, bool enable_half4, const int v_head_size,
+    half* qkv_add_bias, int total_matrix_count, bool do_rotary, int original_past_sequence_length) {
   total_matrix_count = std::max(num_matrices, total_matrix_count);
-  if (enable_half4 && 0 == (qk_head_size % 4) && (v_head_size == -1 || 0 == (v_head_size % 4))) {
+  if (enable_half4 && 0 == (qk_head_size % 4) && (v_head_size == -1 || 0 == (v_head_size % 4)) && !do_rotary) {
     const int H = qk_head_size / 4;
     const int H_v = v_head_size / 4;
     const Half4* input2 = reinterpret_cast<const Half4*>(input);
@@ -588,7 +736,7 @@ void LaunchAddBiasTranspose(
     InvokeAddBiasTranspose<Half4>(stream, num_matrices, format, max_threads_per_block,
                                   batch_size, sequence_length, num_heads, H, input2, biases2, output2,
                                   qkv_add_bias2, H_v, total_matrix_count);
-  } else if (0 == (qk_head_size & 1) && (v_head_size == -1 || 0 == (v_head_size & 1))) {
+  } else if (0 == (qk_head_size & 1) && (v_head_size == -1 || 0 == (v_head_size & 1)) && !do_rotary) {
     const int H = qk_head_size / 2;
     const int H_v = v_head_size / 2;
     const half2* input2 = reinterpret_cast<const half2*>(input);
@@ -602,7 +750,7 @@ void LaunchAddBiasTranspose(
     InvokeAddBiasTranspose<half>(
         stream, num_matrices, format, max_threads_per_block,
         batch_size, sequence_length, num_heads, qk_head_size, input, biases, output,
-        qkv_add_bias, v_head_size, total_matrix_count);
+        qkv_add_bias, v_head_size, total_matrix_count, do_rotary, original_past_sequence_length);
   }
 }
 
@@ -610,10 +758,11 @@ template <>
 void LaunchAddBiasTranspose(
     cudaStream_t stream, const int num_matrices, const int format, const int max_threads_per_block,
     const int batch_size, const int sequence_length, const int num_heads, const int qk_head_size,
-    const float* input, const float* biases, float* output,
-    bool /*enable_half4*/, const int v_head_size, float* qkv_add_bias, int total_matrix_count) {
+    const float* input, const float* biases, float* output, bool /*enable_half4*/,
+    const int v_head_size, float* qkv_add_bias, int total_matrix_count, bool do_rotary,
+    int original_past_sequence_length) {
   total_matrix_count = std::max(num_matrices, total_matrix_count);
-  if (0 == (qk_head_size % 4) && (v_head_size == -1 || 0 == (v_head_size % 4))) {
+  if (0 == (qk_head_size % 4) && (v_head_size == -1 || 0 == (v_head_size % 4)) && !do_rotary) {
     const int H = qk_head_size / 4;
     const float4* input2 = reinterpret_cast<const float4*>(input);
     const float4* biases2 = reinterpret_cast<const float4*>(biases);
@@ -623,7 +772,7 @@ void LaunchAddBiasTranspose(
         stream, num_matrices, format, max_threads_per_block,
         batch_size, sequence_length, num_heads, H, input2, biases2, output2,
         qkv_add_bias2, v_head_size / 4, total_matrix_count);
-  } else if (0 == (qk_head_size & 1) && (v_head_size == -1 || 0 == (v_head_size & 1))) {
+  } else if (0 == (qk_head_size & 1) && (v_head_size == -1 || 0 == (v_head_size & 1)) && !do_rotary) {
     const int H = qk_head_size / 2;
     const float2* input2 = reinterpret_cast<const float2*>(input);
     const float2* biases2 = reinterpret_cast<const float2*>(biases);
@@ -637,7 +786,7 @@ void LaunchAddBiasTranspose(
     InvokeAddBiasTranspose<float>(
         stream, num_matrices, format, max_threads_per_block,
         batch_size, sequence_length, num_heads, qk_head_size, input, biases, output,
-        qkv_add_bias, v_head_size, total_matrix_count);
+        qkv_add_bias, v_head_size, total_matrix_count, do_rotary, original_past_sequence_length);
   }
 }
 

onnxruntime/contrib_ops/cuda/bert/add_bias_transpose.h

@@ -33,7 +33,7 @@ void LaunchAddBiasTranspose(
     cudaStream_t stream, const int num_matrices, const int format, const int max_threads_per_block,
     const int batch_size, const int sequence_length, const int num_heads, const int qk_head_size,
     const T* input, const T* biases, T* output, bool enable_half4, const int v_head_size, T* qkv_add_bias = nullptr,
-    int total_matrix_count = -1);
+    int total_matrix_count = -1, bool do_rotary = false, int original_past_sequence_length = 0);
 
 // Add (bias) and Transpose for separated inputs of Q, K and V, and output Trt format.
 // For self attention:

onnxruntime/contrib_ops/cuda/bert/attention_impl.cu

@@ -330,8 +330,8 @@ Status PrepareQkv(contrib::AttentionParameters& parameters,
       // format 2: BxSx(NH + NH + NH) => BxSxNx(H + H + H)
       LaunchAddBiasTranspose(stream, matrix_to_transpose, format, max_threads_per_block,
                              batch_size, sequence_length, num_heads, qk_head_size,
-                             data.gemm_buffer, data.bias, qkv,
-                             true, v_head_size, qkv_add_bias, 3);
+                             data.gemm_buffer, data.bias, qkv, true, v_head_size, qkv_add_bias,
+                             3, parameters.do_rotary, parameters.original_past_sequence_length);
     }
   } else if (data.key == nullptr) {  // gemm_buffer == nullptr and packed qkv
     assert(data.bias == nullptr);

onnxruntime/core/graph/contrib_ops/bert_defs.cc

@@ -231,6 +231,10 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
               "(2, batch_size, num_heads, max_sequence_length, head_size)",
               AttributeProto::INT,
               OPTIONAL_VALUE)
+        .Attr("do_rotary",
+              "Whether to use rotary position embedding. Default value is 0.",
+              AttributeProto::INT,
+              OPTIONAL_VALUE)
         .Attr("mask_filter_value",
               "The value to be filled in the attention mask. Default value is -10000.0f",
               AttributeProto::FLOAT,

onnxruntime/core/graph/contrib_ops/quantization_defs.cc

@@ -949,6 +949,8 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
         .Attr("num_heads", "Number of attention heads", AttributeProto::INT)
         .Attr("unidirectional", "Whether every token can only attend to previous tokens. Default value is 0.",
               AttributeProto::INT, static_cast<int64_t>(0))
+        .Attr("do_rotary", "Whether to use rotary position embedding. Default value is 0.",
+              AttributeProto::INT, OPTIONAL_VALUE)
         .Attr("past_present_share_buffer", "Corresponding past and present are same tensor, its shape is "
               "(2, batch_size, num_heads, max_sequence_length, head_size)",
               AttributeProto::INT, OPTIONAL_VALUE)

onnxruntime/test/contrib_ops/attention_op_test.cc

@@ -62,7 +62,8 @@ static void RunAttentionTest(
     const std::vector<float>& relative_position_bias_data = {},
     int kv_sequence_length = 0,
     bool past_present_share_buffer = false,
-    bool use_scale = false) {
+    bool use_scale = false,
+    bool do_neox_rotary = false) {
   input_hidden_size = (input_hidden_size == 0 ? hidden_size : input_hidden_size);  // By default, no pruning.
   kv_sequence_length = (kv_sequence_length == 0 ? sequence_length : kv_sequence_length);
   past_present_share_buffer = past_present_share_buffer && use_past_state;
@@ -82,6 +83,9 @@ static void RunAttentionTest(
     if (use_scale && !enable_rocm) {
       tester.AddAttribute<float>("scale", static_cast<float>(1.f / sqrt(head_size)));
     }
+    if (do_neox_rotary && !enable_rocm) {
+      tester.AddAttribute<int64_t>("do_rotary", static_cast<int64_t>(do_neox_rotary ? 1 : 0));
+    }
 
     int32_t qkv_hidden_size_sum;
     int32_t v_hidden_size;
@@ -267,19 +271,20 @@ static void RunAttentionTest(
     const std::vector<float>& relative_position_bias_data = {},
     int kv_sequence_length = 0,
     bool past_present_share_buffer = false,
-    bool use_scale = false) {
+    bool use_scale = false,
+    bool do_neox_rotary = false) {
   RunAttentionTest(input_data, weights_data, false, bias_data, mask_index_data, output_data,
                    batch_size, sequence_length, hidden_size, number_of_heads,
                    use_float16, is_unidirectional, use_past_state, past_sequence_length,
                    past_data, present_data, mask_type, input_hidden_size, max_sequence_length,
                    disable_cpu, disable_cuda, disable_rocm, qkv_sizes, relative_position_bias_data,
-                   kv_sequence_length, past_present_share_buffer, use_scale);
+                   kv_sequence_length, past_present_share_buffer, use_scale, do_neox_rotary);
   RunAttentionTest(input_data, weights_data, true, bias_data, mask_index_data, output_data,
                    batch_size, sequence_length, hidden_size, number_of_heads,
                    use_float16, is_unidirectional, use_past_state, past_sequence_length,
                    past_data, present_data, mask_type, input_hidden_size, max_sequence_length,
                    disable_cpu, disable_cuda, disable_rocm, qkv_sizes, relative_position_bias_data,
-                   kv_sequence_length, past_present_share_buffer, use_scale);
+                   kv_sequence_length, past_present_share_buffer, use_scale, do_neox_rotary);
 }
 
 TEST(AttentionTest, AttentionBatch1) {
@@ -1716,6 +1721,51 @@ TEST(AttentionTest, AttentionWithNormFactor) {
                    true /*use_scale*/);
 }
 
+TEST(AttentionTest, AttentionWithNeoXRotaryEmbedding) {
+  int batch_size = 2;
+  int sequence_length = 2;
+  int hidden_size = 4;
+  int number_of_heads = 2;
+
+  std::vector<float> input_data = {
+      0.5f, 0.2f, 0.3f, -0.6f,
+      0.8f, -0.5f, 0.0f, 1.f,
+      0.8f, -0.5f, 0.0f, 1.f,
+      0.5f, 0.2f, 0.3f, -0.6f};
+
+  std::vector<float> weight_data = {
+      0.1f, -0.2f, 0.3f, 1.0f, 1.1f, 0.3f, 0.5f, 0.2f, 0.3f, -0.6f, 1.5f, 2.0f,
+      0.5f, 0.1f, 0.4f, 1.6f, 1.0f, 2.0f, 0.4f, 0.8f, 0.9f, 0.1f, -1.3f, 0.7f,
+      0.3f, 0.2f, 4.0f, 2.2f, 1.6f, 1.1f, 0.7f, 0.2f, 0.4f, 1.0f, 1.2f, 0.5f,
+      0.2f, 0.1f, 0.4f, 1.6f, 2.4f, 3.3f, 2.1f, 4.2f, 8.4f, 0.0f, 2.1f, 3.2f};
+
+  std::vector<float> bias_data = {
+      -0.5f, 0.6f, 1.2f, 2.1f, 0.5f, 0.7f, 0.2f, 1.2f, 0.5f, 0.4f, 0.3f, 1.2f};
+
+  // Test mask start position > 0.
+  std::vector<int32_t> mask_index_data = {0, 1, 1, 1};
+
+  std::vector<float> output_data = {
+      3.0146f, 0.1142f, 3.9834f, 5.3394f,
+      8.69f, -0.13f, 4.25f, 5.65f,
+      8.69f, -0.13f, 4.25f, 5.65f,
+      -1.4697f, 0.3071f, 4.25f, 5.65f};
+
+  bool use_float16 = true;
+  bool is_unidirectional = true;
+  bool use_past_state = false;
+  int past_sequence_length = 0;
+  const std::vector<float>* past_data = nullptr;
+  const std::vector<float>* present_data = nullptr;
+  RunAttentionTest(input_data, weight_data, bias_data, mask_index_data, output_data,
+                   batch_size, sequence_length, hidden_size, number_of_heads,
+                   use_float16, is_unidirectional, use_past_state, past_sequence_length, past_data, present_data,
+                   AttentionMaskType::MASK_2D_KEY_PADDING, 0 /*input_hidden_size*/, 0 /*max_sequence_length*/,
+                   true /*disable_cpu*/, false /*disable_cuda*/, true /*disable_rocm*/, {} /*qkv_sizes*/,
+                   {} /*relative_position_bias_data*/, 0 /*kv_sequence_length*/, false /*past_present_share_buffer*/,
+                   true /*use_scale*/, true /*use_neox_rotary_embedding*/);
+}
+
 TEST(AttentionTest, AttentionMask1DEndNoWord) {
   int batch_size = 2;
   int sequence_length = 2;

onnxruntime/test/python/transformers/test_parity_neox_attention.py

@@ -0,0 +1,310 @@
+# --------------------------------------------------------------------------
+# Copyright 2020 The HuggingFace Inc. team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+# --------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.  See License.txt in the project root for
+# license information.
+# -------------------------------------------------------------------------
+
+import numpy as np
+import torch
+from torch import nn
+
+torch.set_printoptions(threshold=10000)
+
+
+def create_neox_attention_graph(
+    batch_size,
+    seq_len,
+    hidden_size,
+    qkv_weight,
+    qkv_bias,
+    num_heads,
+    use_rotary,
+):
+    from onnx import TensorProto, helper
+
+    nodes = [
+        helper.make_node(
+            "Attention",
+            [
+                "input",
+                "weight",
+                "bias",
+            ],
+            ["output"],
+            "NeoXAttention_0",
+            num_heads=num_heads,
+            unidirectional=1,
+            do_rotary=(use_rotary is True),
+            domain="com.microsoft",
+        ),
+    ]
+
+    initializers = [
+        helper.make_tensor("weight", TensorProto.FLOAT, [hidden_size, 3 * hidden_size], qkv_weight.flatten().tolist()),
+        helper.make_tensor("bias", TensorProto.FLOAT, [3 * hidden_size], qkv_bias.flatten().tolist()),
+    ]
+
+    graph = helper.make_graph(
+        nodes,
+        "NeoXAttention_Graph",
+        [
+            helper.make_tensor_value_info("input", TensorProto.FLOAT, [batch_size, seq_len, hidden_size]),
+        ],
+        [
+            helper.make_tensor_value_info("output", TensorProto.FLOAT, [batch_size, seq_len, hidden_size]),
+        ],
+        initializers,
+    )
+
+    model = helper.make_model(graph)
+    return model.SerializeToString()
+
+
+class RotaryEmbedding(torch.nn.Module):
+    def __init__(self, dim, max_position_embeddings, base=10000, device=None):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+        # Build here to make `torch.jit.trace` work.
+        self.max_seq_len_cached = max_position_embeddings
+        t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.cos_cached = emb.cos()[None, None, :, :]
+        self.sin_cached = emb.sin()[None, None, :, :]
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        # This `if` block is unlikely to be run after we build sin/cos in `__init__`. Keep the logic here just in case.
+        if seq_len > self.max_seq_len_cached:
+            self.max_seq_len_cached = seq_len
+            t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            # Different from paper, but it uses a different permutation in order to obtain the same calculation
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+            self.cos_cached = emb.cos()[None, None, :, :]
+            self.sin_cached = emb.sin()[None, None, :, :]
+        return self.cos_cached[:seq_len, ...].to(x.device), self.sin_cached[:seq_len, ...].to(x.device)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, offset: int = 0):
+    cos = cos[..., offset : q.shape[-2] + offset, :]
+    sin = sin[..., offset : q.shape[-2] + offset, :]
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class GPTNeoXAttention(nn.Module):
+    def __init__(self, batch_size, seq_len, num_head, hidden_size, use_rotary):
+        super().__init__()
+        self.use_rotary = use_rotary
+        self.num_attention_heads = num_head
+        self.hidden_size = hidden_size
+        self.head_size = self.hidden_size // self.num_attention_heads
+        self.rotary_ndims = int(self.head_size)
+        max_positions = 2048
+        self.register_buffer(
+            "bias",
+            torch.tril(torch.ones((max_positions, max_positions), dtype=torch.uint8)).view(
+                1, 1, max_positions, max_positions
+            ),
+        )
+        self.register_buffer("masked_bias", torch.tensor(-1e9))
+        self.rotary_emb = RotaryEmbedding(self.rotary_ndims, 2048, 10000)
+        self.norm_factor = torch.sqrt(torch.tensor(self.head_size, dtype=torch.float32)).to(torch.get_default_dtype())
+        self.query_key_value = nn.Linear(hidden_size, 3 * hidden_size)
+
+        # self.query_key_value.weight.data.copy_(torch.tensor(np.ones((3 * hidden_size, hidden_size))))
+        # self.query_key_value.bias.data.copy_(torch.tensor(np.zeros((3 * hidden_size))))
+
+        self.onnx_graph = create_neox_attention_graph(
+            batch_size,
+            seq_len,
+            self.hidden_size,
+            self.query_key_value.weight.reshape(self.num_attention_heads, 3, -1)
+            .transpose(0, 1)
+            .reshape(3 * self.hidden_size, -1)
+            .transpose(0, 1),
+            self.query_key_value.bias.reshape(self.num_attention_heads, 3, -1).transpose(0, 1).reshape(-1),
+            self.num_attention_heads,
+            self.use_rotary,
+        )
+
+    @classmethod
+    def _merge_heads(cls, tensor, num_attention_heads, attn_head_size):
+        """
+        Merges attn_head_size dim and num_attn_heads dim into hidden dim
+        """
+        # tensor [bs, num_attention_heads, seq_len, attn_head_size]
+        tensor = tensor.permute(0, 2, 1, 3).contiguous()
+        # -> [bs, seq_len, num_attention_heads, attn_head_size]
+        tensor = tensor.view(tensor.size(0), tensor.size(1), num_attention_heads * attn_head_size)
+        # -> [bs, seq_len, hidden_size]
+        return tensor
+
+    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
+        # q, k, v: [bs, num_attention_heads, seq_len, attn_head_size]
+        # compute causal mask from causal mask buffer
+        batch_size, num_attention_heads, query_length, attn_head_size = query.size()
+        key_length = key.size(-2)
+
+        causal_mask = self.bias[:, :, key_length - query_length : key_length, :key_length].bool()
+
+        query = query.view(batch_size * num_attention_heads, query_length, attn_head_size)
+        key = key.view(batch_size * num_attention_heads, key_length, attn_head_size)
+        attn_scores = torch.zeros(
+            batch_size * num_attention_heads,
+            query_length,
+            key_length,
+            dtype=query.dtype,
+            device=key.device,
+        )
+        attn_scores = torch.baddbmm(
+            attn_scores,
+            query,
+            key.transpose(1, 2),
+            beta=1.0,
+            alpha=(torch.tensor(1.0, dtype=self.norm_factor.dtype, device=self.norm_factor.device) / self.norm_factor),
+        )
+        attn_scores = attn_scores.view(batch_size, num_attention_heads, query_length, key_length)
+
+        mask_value = torch.finfo(attn_scores.dtype).min
+        # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
+        # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
+        mask_value = torch.tensor(mask_value, dtype=attn_scores.dtype).to(attn_scores.device)
+        attn_scores = torch.where(causal_mask, attn_scores, mask_value)
+
+        if attention_mask is not None:
+            # Apply the attention mask
+            attn_scores = attn_scores + attention_mask
+
+        attn_weights = nn.functional.softmax(attn_scores, dim=-1)
+        attn_weights = attn_weights.to(value.dtype)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attn_weights = attn_weights * head_mask
+
+        attn_output = torch.matmul(attn_weights, value)
+        return attn_output, attn_weights
+
+    def onnx_forward(
+        self,
+        hidden_states,
+    ):
+        ort_inputs = {
+            "input": np.ascontiguousarray(hidden_states.cpu().numpy()),
+        }
+
+        from onnxruntime import InferenceSession, SessionOptions
+
+        sess_options = SessionOptions()
+        ort_session = InferenceSession(self.onnx_graph, sess_options, providers=["CUDAExecutionProvider"])
+        ort_output = ort_session.run(None, ort_inputs)
+
+        output = torch.tensor(ort_output)
+
+        return output
+
+    def torch_forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        layer_past=None,
+        use_cache=False,
+        output_attentions=False,
+    ):
+        has_layer_past = layer_past is not None
+
+        # Compute QKV
+        # Attention heads [batch, seq_len, hidden_size]
+        #   --> [batch, seq_len, (np * 3 * head_size)]
+        qkv = self.query_key_value(hidden_states)
+
+        # [batch, seq_len, (num_heads * 3 * head_size)]
+        #   --> [batch, seq_len, num_heads, 3 * head_size]
+        new_qkv_shape = qkv.size()[:-1] + (self.num_attention_heads, 3 * self.head_size)
+        qkv = qkv.view(*new_qkv_shape)
+
+        # [batch, seq_len, num_attention_heads, 3 * head_size] --> 3 [batch, num_attention_heads, seq_len, head_size]
+        query = qkv[..., : self.head_size].permute(0, 2, 1, 3)
+        key = qkv[..., self.head_size : 2 * self.head_size].permute(0, 2, 1, 3)
+        value = qkv[..., 2 * self.head_size :].permute(0, 2, 1, 3)
+
+        if self.use_rotary:
+            # Compute rotary embeddings on rotary_ndims
+            query_rot = query[..., : self.rotary_ndims]
+            query_pass = query[..., self.rotary_ndims :]
+            key_rot = key[..., : self.rotary_ndims]
+            key_pass = key[..., self.rotary_ndims :]
+
+            # Compute token offset for rotary embeddings (when decoding)
+            seq_len = key.shape[-2]
+            offset = 0
+            if has_layer_past:
+                offset = layer_past[0].shape[-2]
+                seq_len += offset
+            cos, sin = self.rotary_emb(value, seq_len=seq_len)
+            query, key = apply_rotary_pos_emb(query_rot, key_rot, cos, sin, offset=offset)
+            query = torch.cat((query, query_pass), dim=-1)
+            key = torch.cat((key, key_pass), dim=-1)
+            # print("query", query.shape, query)
+            # print("key", key.shape, key)
+            # print("value", value.shape, value)
+
+        # Cache QKV values
+        if has_layer_past:
+            past_key = layer_past[0]
+            past_value = layer_past[1]
+            key = torch.cat((past_key, key), dim=-2)
+            value = torch.cat((past_value, value), dim=-2)
+        present = (key, value) if use_cache else None
+
+        # Compute attention
+        attn_output, _ = self._attn(query, key, value, attention_mask, head_mask)
+
+        # Reshape outputs
+        attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_size)
+
+        return attn_output
+
+
+if __name__ == "__main__":
+    for batch_size in [1, 2, 4, 8]:
+        for seq_len in [32, 128, 512, 1024, 2048]:
+            for num_head in [12]:
+                for hidden_size in [768]:
+                    attn = GPTNeoXAttention(batch_size, seq_len, num_head, hidden_size, use_rotary=True)
+
+                    hidden_states = torch.normal(mean=0.5, std=0.1, size=(batch_size, seq_len, hidden_size)).to(
+                        torch.float32
+                    )
+
+                    torch_output = attn.torch_forward(hidden_states)
+                    ort_output = attn.onnx_forward(hidden_states)
+                    print(
+                        "Parity check with shape BNSH = ({},{},{},{})".format(
+                            batch_size, seq_len, num_head, hidden_size
+                        )
+                    )
+                    if torch.allclose(torch_output, ort_output, atol=1e-6):
+                        print("Success!")
+                    else:
+                        print("Failure!")