Add NGramRepeatBlock contrib op (#8078)

**Description**: 
Enforce no repetition of n-grams. Scores are set to `-inf` for tokens that form a repeated n-gram if added to the back of the input_ids.

**Motivation and Context**
Needed by transformer models in sequence generation algorithms (greedy search and beam search). This module has heavy impact on performance, and can be highly parallelized.

docs/ContribOperators.md

@@ -32,6 +32,7 @@ Do not modify directly.*
   * <a href="#com.microsoft.MaxpoolWithMask">com.microsoft.MaxpoolWithMask</a>
   * <a href="#com.microsoft.MulInteger">com.microsoft.MulInteger</a>
   * <a href="#com.microsoft.MurmurHash3">com.microsoft.MurmurHash3</a>
+  * <a href="#com.microsoft.NGramRepeatBlock">com.microsoft.NGramRepeatBlock</a>
   * <a href="#com.microsoft.NhwcMaxPool">com.microsoft.NhwcMaxPool</a>
   * <a href="#com.microsoft.Pad">com.microsoft.Pad</a>
   * <a href="#com.microsoft.QAttention">com.microsoft.QAttention</a>
@@ -1517,6 +1518,47 @@ This version of the operator has been available since version 1 of the 'com.micr
 </dl>
 
 
+### <a name="com.microsoft.NGramRepeatBlock"></a><a name="com.microsoft.ngramrepeatblock">**com.microsoft.NGramRepeatBlock**</a>
+
+  Enforce no repetition of n-grams. Scores are set to `-inf` for tokens that form a repeated n-gram if added to the back of the input_ids.
+
+#### Version
+
+This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
+
+#### Attributes
+
+<dl>
+<dt><tt>ngram_size</tt> : int (required)</dt>
+<dd>The NGram size.</dd>
+</dl>
+
+#### Inputs
+
+<dl>
+<dt><tt>input_ids</tt> : Tid</dt>
+<dd>2D input tensor with shape (batch_size, sequence_length)</dd>
+<dt><tt>scores</tt> : T</dt>
+<dd>2D input tensor with shape (batch_size, vocab_size)</dd>
+</dl>
+
+#### Outputs
+
+<dl>
+<dt><tt>scores_out</tt> : T</dt>
+<dd>2D output tensor with shape (batch_size, vocab_size)</dd>
+</dl>
+
+#### Type Constraints
+
+<dl>
+<dt><tt>Tid</tt> : tensor(int64)</dt>
+<dd>Constrain indices to integer types</dd>
+<dt><tt>T</tt> : tensor(float)</dt>
+<dd>Constrain scores input and output types to float tensors.</dd>
+</dl>
+
+
 ### <a name="com.microsoft.NhwcMaxPool"></a><a name="com.microsoft.nhwcmaxpool">**com.microsoft.NhwcMaxPool**</a>
 
 #### Version

docs/OperatorKernels.md

@@ -388,6 +388,7 @@ Do not modify directly.*
 |MatMulIntegerToFloat|*in* A:**T1**<br> *in* B:**T2**<br> *in* a_scale:**T3**<br> *in* b_scale:**T3**<br> *in* a_zero_point:**T1**<br> *in* b_zero_point:**T2**<br> *in* bias:**T3**<br> *out* Y:**T3**|1+|**T1** = tensor(uint8)<br/> **T2** = tensor(int8), tensor(uint8)<br/> **T3** = tensor(float)|
 |MaxpoolWithMask|*in* X:**T**<br> *in* M:**tensor(int32)**<br> *out* Y:**T**|1+|**X** = tensor(float)|
 |MurmurHash3|*in* X:**T1**<br> *out* Y:**T2**|1+|**T1** = tensor(double), tensor(float), tensor(int32), tensor(int64), tensor(string), tensor(uint32), tensor(uint64)<br/> **T2** = tensor(int32), tensor(uint32)|
+|NGramRepeatBlock|*in* input_ids:**Tid**<br> *in* scores:**T**<br> *out* scores_out:**T**|1+|**T** = tensor(float)<br/> **Tid** = tensor(int64)|
 |NhwcMaxPool|*in* x:**T**<br> *out* y:**T**|1+|**T** = tensor(uint8)|
 |Pad|*in* data:**T**<br> *in* pads:**tensor(int64)**<br> *in* value:**T**<br> *out* output:**T**|1+|**T** = tensor(float)|
 |QAttention|*in* input:**T1**<br> *in* weight:**T2**<br> *in* bias:**T3**<br> *in* input_scale:**T3**<br> *in* weight_scale:**T3**<br> *in* mask_index:**T4**<br> *in* input_zero_point:**T1**<br> *in* weight_zero_point:**T2**<br> *in* past:**T3**<br> *out* output:**T3**<br> *out* present:**T3**|1+|**T1** = tensor(uint8)<br/> **T2** = tensor(int8), tensor(uint8)<br/> **T3** = tensor(float)<br/> **T4** = tensor(int32)|
@@ -713,6 +714,7 @@ Do not modify directly.*
 |Inverse|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
 |Irfft|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
 |LongformerAttention|*in* input:**T**<br> *in* weight:**T**<br> *in* bias:**T**<br> *in* mask:**T**<br> *in* global_weight:**T**<br> *in* global_bias:**T**<br> *in* global:**G**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
+|NGramRepeatBlock|*in* input_ids:**Tid**<br> *in* scores:**T**<br> *out* scores_out:**T**|1+|**T** = tensor(float)<br/> **Tid** = tensor(int64)|
 |QAttention|*in* input:**T1**<br> *in* weight:**T2**<br> *in* bias:**T3**<br> *in* input_scale:**T3**<br> *in* weight_scale:**T3**<br> *in* mask_index:**T4**<br> *in* input_zero_point:**T1**<br> *in* weight_zero_point:**T2**<br> *in* past:**T3**<br> *out* output:**T3**<br> *out* present:**T3**|1+|**T1** = tensor(int8)<br/> **T2** = tensor(int8)<br/> **T3** = tensor(float), tensor(float16)<br/> **T4** = tensor(int32)|
 |QuantizeLinear|*in* x:**T1**<br> *in* y_scale:**T1**<br> *in* y_zero_point:**T2**<br> *out* y:**T2**|1+|**T1** = tensor(float16)<br/> **T2** = tensor(int8), tensor(uint8)|
 |Rfft|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|

onnxruntime/contrib_ops/cpu/bert/ngram_repeat_block.cc

@@ -0,0 +1,20 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "ngram_repeat_block.h"
+
+
+namespace onnxruntime {
+namespace contrib {
+
+ONNX_OPERATOR_KERNEL_EX(
+    NGramRepeatBlock,
+    kMSDomain,
+    1,
+    kCpuExecutionProvider,
+    KernelDefBuilder()
+        .TypeConstraint("Tid", DataTypeImpl::GetTensorType<int64_t>())
+        .TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
+    NGramRepeatBlock);
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cpu/bert/ngram_repeat_block.h

@@ -0,0 +1,85 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+#include "core/platform/threadpool.h"
+
+namespace onnxruntime {
+namespace contrib {
+class NGramRepeatBlock : public OpKernel {
+ public:
+  explicit NGramRepeatBlock(const OpKernelInfo& info) : OpKernel(info) {
+    ORT_ENFORCE(info.GetAttr<int64_t>("ngram_size", &ngram_size_).IsOK());
+    ORT_ENFORCE(ngram_size_ > 0);
+  }
+
+  Status Compute(OpKernelContext* context) const override {
+    const Tensor* input_ids = context->Input<Tensor>(0);
+    const Tensor* scores = context->Input<Tensor>(1);
+    Tensor* output = context->Output(0, scores->Shape());
+
+    const auto* scores_source = static_cast<const float*>(scores->DataRaw());
+    auto* scores_target = static_cast<float*>(output->MutableDataRaw());
+    if (scores_source != scores_target) {
+      memcpy(scores_target, scores_source, scores->Shape().Size() * sizeof(float));
+    }
+
+    const auto& input_ids_dims = input_ids->Shape().GetDims();
+    const auto& scores_dims = scores->Shape().GetDims();
+    ORT_ENFORCE(input_ids_dims.size() == 2);
+    ORT_ENFORCE(scores_dims.size() == 2);
+    int64_t batch_size = input_ids_dims[0];
+    int64_t cur_len = input_ids_dims[1];
+    ORT_ENFORCE(scores_dims[0] == batch_size);
+    int64_t vocab_size = scores_dims[1];
+
+    if (cur_len + 1 < ngram_size_) {
+      return Status::OK();
+    }
+
+    const auto* input_ids_data = static_cast<const int64_t*>(input_ids->DataRaw(input_ids->DataType()));
+
+    auto lambda = [&](int64_t b) {
+      for (int64_t i = 0; i < cur_len; ++i) {
+        if (i + ngram_size_ > cur_len) {
+          break;
+        }
+
+        bool is_banned = true;
+        for (int64_t j = 0; j < ngram_size_ - 1; ++j) {
+          auto token_at_tail = input_ids_data[b * cur_len + i + j];
+          auto token_to_cmp = input_ids_data[b * cur_len + cur_len + 1 - ngram_size_ + j];
+          if (token_at_tail != token_to_cmp) {
+            is_banned = false;
+            break;
+          }
+        }
+
+        if (is_banned) {
+          auto token_id = static_cast<int64_t>(input_ids_data[b * cur_len + i + ngram_size_ - 1]);
+          ORT_ENFORCE(token_id < vocab_size);
+          scores_target[b * vocab_size + token_id] = -std::numeric_limits<float>::infinity();
+        }
+      }
+    };
+
+    concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
+    concurrency::ThreadPool::TryParallelFor(
+      tp, batch_size, static_cast<double>(cur_len * ngram_size_),
+      [&lambda](ptrdiff_t first, ptrdiff_t last) {
+        for (auto b = static_cast<int64_t>(first), end = static_cast<int64_t>(last); b < end; ++b) {
+          lambda(b);
+        }
+      }
+    );
+
+    return Status::OK();
+  }
+ private:
+  int64_t ngram_size_;
+};
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc

@@ -33,6 +33,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1,
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Gelu);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, BiasGelu);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FastGelu);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, NGramRepeatBlock);
 
 #ifdef BUILD_MS_EXPERIMENTAL_OPS
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSExperimentalDomain, 1, DFT);
@@ -194,6 +195,7 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, BiasGelu)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Gelu)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FastGelu)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, NGramRepeatBlock)>,
 
 #ifdef BUILD_MS_EXPERIMENTAL_OPS
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSExperimentalDomain, 1, DFT)>,

onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block.cc

@@ -0,0 +1,71 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/cuda/cuda_common.h"
+#include "ngram_repeat_block.h"
+#include "ngram_repeat_block_impl.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+ONNX_OPERATOR_KERNEL_EX(
+    NGramRepeatBlock,
+    kMSDomain,
+    1,
+    kCudaExecutionProvider,
+    (*KernelDefBuilder::Create())
+        .TypeConstraint("Tid", DataTypeImpl::GetTensorType<int64_t>())
+        .TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
+    NGramRepeatBlock);
+
+using namespace ONNX_NAMESPACE;
+
+NGramRepeatBlock::NGramRepeatBlock(const OpKernelInfo& info) : CudaKernel(info) {
+  ORT_ENFORCE(info.GetAttr<int64_t>("ngram_size", &ngram_size_).IsOK());
+  ORT_ENFORCE(ngram_size_ > 0);
+}
+
+Status NGramRepeatBlock::ComputeInternal(OpKernelContext* context) const {
+  const Tensor* input_ids = context->Input<Tensor>(0);
+  const Tensor* scores = context->Input<Tensor>(1);
+  Tensor* output = context->Output(0, scores->Shape());
+
+  const auto* scores_source = static_cast<const float*>(scores->DataRaw());
+  auto* scores_target = static_cast<float*>(output->MutableDataRaw());
+  if (scores_source != scores_target) {
+    CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(scores_target, scores_source, scores->Shape().Size() * sizeof(float), cudaMemcpyDeviceToDevice, Stream()));
+  }
+
+  const auto& input_ids_dims = input_ids->Shape().GetDims();
+  const auto& scores_dims = scores->Shape().GetDims();
+  ORT_ENFORCE(input_ids_dims.size() == 2);
+  ORT_ENFORCE(scores_dims.size() == 2);
+  int64_t batch_size = input_ids_dims[0];
+  int64_t cur_len = input_ids_dims[1];
+  ORT_ENFORCE(scores_dims[0] == batch_size);
+  int64_t vocab_size = scores_dims[1];
+
+  if (cur_len + 1 < ngram_size_) {
+    return Status::OK();
+  }
+
+  const auto* input_ids_data = static_cast<const int64_t*>(input_ids->DataRaw(input_ids->DataType()));
+
+  NGramRepeatBlockImpl(
+      Stream(),
+      input_ids_data,
+      scores_target,
+      gsl::narrow_cast<int>(batch_size),
+      gsl::narrow_cast<int>(cur_len - 1),
+      gsl::narrow_cast<int>(cur_len),
+      gsl::narrow_cast<int>(vocab_size),
+      gsl::narrow_cast<int>(1),
+      gsl::narrow_cast<int>(ngram_size_));
+
+  return Status::OK();
+}
+
+}  //namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block.h

@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+#include "core/common/common.h"
+#include "core/providers/cuda/cuda_kernel.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+using namespace onnxruntime::cuda;
+
+class NGramRepeatBlock final : public CudaKernel {
+ public:
+  NGramRepeatBlock(const OpKernelInfo& op_kernel_info);
+  Status ComputeInternal(OpKernelContext* ctx) const override;
+ private:
+  int64_t ngram_size_;
+};
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block_impl.cu

@@ -0,0 +1,84 @@
+/*
+Copyright (c) Microsoft Corporation.
+Licensed under the MIT License.
+*/
+
+/*
+Kernel implementation for blocking repeated n-grams.
+*/
+
+#include "core/providers/cuda/cu_inc/common.cuh"
+#include "contrib_ops/cuda/bert/ngram_repeat_block_impl.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+using namespace onnxruntime::cuda;
+
+// Ban repeated ngrams of length = 'no_repeat_ngram_size'
+__global__ void banRepeatedTokens(const int64_t* __restrict__ tokens,
+                                  float* __restrict__ lprobs,
+                                  int max_predict_len, int vocab_size,
+                                  int no_repeat_ngram_size) {
+  auto row = blockIdx.x;
+  auto col = threadIdx.x;
+  auto start = row * (max_predict_len) + col;
+  // Each thread compares ngram starting from
+  // thread index with final ngram starting from
+  // step - no_repeat_ngram_size +2
+  auto check_start_pos = blockDim.x;
+  auto lprob_start = row * vocab_size;
+  bool is_banned = true;
+  extern __shared__ int64_t tokens_shm[];
+  tokens_shm[col] = tokens[start];
+  if (col == blockDim.x - 1) {
+    for (int i=1; i<no_repeat_ngram_size; i++){
+      if (col+i < max_predict_len){
+        tokens_shm[col + i] = tokens[start + i];
+      }
+    }
+  }
+  __syncthreads();
+
+  for (int k = 0; k < no_repeat_ngram_size - 1; k++) {
+    if (tokens_shm[col + k] != tokens_shm[check_start_pos + k]) {
+      is_banned = false;
+    }
+  }
+  if (is_banned == true) {
+    auto token_to_be_banned = tokens_shm[col + no_repeat_ngram_size - 1];
+    lprobs[lprob_start + token_to_be_banned] = -INFINITY;
+  }
+}
+
+// Allocate blocks and threads based on
+// batch size and sequence length and launch
+// kernel
+void NGramRepeatBlockImpl(
+    cudaStream_t stream,
+    const int64_t* tokens_ptr,
+    float* scores_ptr,
+    int bsz,
+    int step,
+    int max_predict_len,
+    int vocab_size,
+    int beam_size,
+    int no_repeat_ngram_size) {
+  int threads = step - no_repeat_ngram_size + 2;
+  if (threads <= 0) return;
+  int blocks = bsz * beam_size;
+  int shared_mem_size = (step + 1) * sizeof(int64_t);
+
+  // Launching N blocks where N is number of samples in a batch (beams*bsz)
+  // Launching T threads where T is number of previous ngrams in a sample
+  // Allocating shared mem per block for fastser access of input tokens since
+  // each token will be accessed N times to compare with current Ngram where
+  // N is Ngram size.
+  banRepeatedTokens<<<blocks, threads, shared_mem_size, stream>>>(
+      tokens_ptr, scores_ptr, max_predict_len, vocab_size, no_repeat_ngram_size);
+}
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block_impl.h

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+#include "core/providers/cuda/shared_inc/cuda_utils.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace cuda {
+
+using namespace onnxruntime::cuda;
+
+void NGramRepeatBlockImpl(
+    cudaStream_t stream,
+    const int64_t* tokens_ptr,
+    float* scores_ptr,
+    int bsz,
+    int step,
+    int max_predict_len,
+    int vocab_size,
+    int beam_size,
+    int no_repeat_ngram_size);
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc

@@ -35,6 +35,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1
 class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMulConj);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BiasSoftmax);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, BiasDropout);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, NGramRepeatBlock);
+
 
 // These ops were experimental ops in onnx domain which have been removed now. We add them here as
 // contrib ops to maintain backward compatibility
@@ -122,6 +124,7 @@ Status RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
     BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMul)>,
     BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ComplexMulConj)>,
     BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, MLFloat16, ComplexMulConj)>,
+    BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, NGramRepeatBlock)>,
 
       // These ops were experimental ops in onnx domain which have been removed now. We add them here as
       // contrib ops to maintain backward compatibility

onnxruntime/core/graph/contrib_ops/contrib_defs.cc

@@ -610,6 +610,29 @@ GELU (Gaussian Error Linear Unit) approximation: Y=0.5*X*(1+tanh(0.797885*X+0.03
       .TypeConstraint("T", {"tensor(float)", "tensor(float16)"}, "Constrain input and output types to float or half tensors.")
       .TypeConstraint("U", {"tensor(float)"}, "Constrain mean and inv_std_var to float tensors.")
       .TypeAndShapeInferenceFunction(ONNX_NAMESPACE::propagateShapeAndTypeFromFirstInput);
+
+
+  static const char* NGramRepeatBlock_ver1_doc = R"DOC(
+Enforce no repetition of n-grams. Scores are set to `-inf` for tokens that form a repeated n-gram if added to the back of the input_ids.
+)DOC";
+
+  ONNX_CONTRIB_OPERATOR_SCHEMA(NGramRepeatBlock)
+      .SetDomain(kMSDomain)
+      .SinceVersion(1)
+      .SetDoc(NGramRepeatBlock_ver1_doc)
+      .Attr("ngram_size", "The NGram size.", AttributeProto::INT)
+      .Input(0, "input_ids", "2D input tensor with shape (batch_size, sequence_length)", "Tid")
+      .Input(1, "scores", "2D input tensor with shape (batch_size, vocab_size)", "T")
+      .Output(0, "scores_out", "2D output tensor with shape (batch_size, vocab_size)", "T")
+      .TypeConstraint("Tid", {"tensor(int64)"}, "Constrain indices to integer types")
+      .TypeConstraint("T", {"tensor(float)"}, "Constrain scores input and output types to float tensors.")
+      .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
+        propagateElemTypeFromInputToOutput(ctx, 1, 0);
+        if (!hasInputShape(ctx, 1)) {
+          return;
+        }
+        propagateShapeFromInputToOutput(ctx, 1, 0);
+      });
 }
 
 void RegisterContribSchemas() {

onnxruntime/test/contrib_ops/ngram_repeat_block_op_test.cc

@@ -0,0 +1,41 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "gtest/gtest.h"
+#include "test/common/tensor_op_test_utils.h"
+#include "test/common/cuda_op_test_utils.h"
+#include "test/providers/provider_test_utils.h"
+
+namespace onnxruntime {
+namespace test {
+
+TEST(NGramRepeatBlockTest, NGramSize_3) {
+  OpTester tester("NGramRepeatBlock", 1, onnxruntime::kMSDomain);
+
+  tester.AddInput<int64_t>("input_ids", {3, 6},
+      {0, 1, 2, 3, 1, 2,
+       0, 0, 0, 1, 0, 0,
+       0, 1, 0, 1, 0, 1});
+  tester.AddInput<float>("scores", {3, 4},
+      {1.0f,  2.0f,  3.0f,  4.0f,
+       5.0f,  6.0f,  7.0f,  8.0f,
+       9.0f, 10.0f, 11.0f, 12.0f});
+  tester.AddAttribute("ngram_size", (int64_t)3);
+  tester.AddOutput<float>("scores_out", {3, 4},
+      {1.0f, 2.0f, 3.0f, -std::numeric_limits<float>::infinity(),
+       -std::numeric_limits<float>::infinity(), -std::numeric_limits<float>::infinity(), 7.0f, 8.0f,
+       -std::numeric_limits<float>::infinity(), 10.0f, 11.0f, 12.0f});
+
+  if (HasCudaEnvironment(0)) {
+    std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+    execution_providers.push_back(DefaultCudaExecutionProvider());
+    tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+  }
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCpuExecutionProvider());
+  tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+}
+
+}  // namespace test
+}  // namespace onnxruntime

onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json

@@ -151,6 +151,10 @@
         "NhwcMaxPool com.microsoft CPUExecutionProvider",
         8512357837341844248
     ],
+    [
+        "NGramRepeatBlock com.microsoft CPUExecutionProvider",
+        17162613206685017176
+    ],
     [
         "Pad com.microsoft CPUExecutionProvider",
         15076596470814458544