diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index ce59fa1337..adc13ee3ad 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -32,6 +32,7 @@ Do not modify directly.* * com.microsoft.MaxpoolWithMask * com.microsoft.MulInteger * com.microsoft.MurmurHash3 + * com.microsoft.NGramRepeatBlock * com.microsoft.NhwcMaxPool * com.microsoft.Pad * com.microsoft.QAttention @@ -1517,6 +1518,47 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.NGramRepeatBlock** + + 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 + +
+
ngram_size : int (required)
+
The NGram size.
+
+ +#### Inputs + +
+
input_ids : Tid
+
2D input tensor with shape (batch_size, sequence_length)
+
scores : T
+
2D input tensor with shape (batch_size, vocab_size)
+
+ +#### Outputs + +
+
scores_out : T
+
2D output tensor with shape (batch_size, vocab_size)
+
+ +#### Type Constraints + +
+
Tid : tensor(int64)
+
Constrain indices to integer types
+
T : tensor(float)
+
Constrain scores input and output types to float tensors.
+
+ + ### **com.microsoft.NhwcMaxPool** #### Version diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 12ea0af6ef..97935346e0 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -388,6 +388,7 @@ Do not modify directly.* |MatMulIntegerToFloat|*in* A:**T1**
*in* B:**T2**
*in* a_scale:**T3**
*in* b_scale:**T3**
*in* a_zero_point:**T1**
*in* b_zero_point:**T2**
*in* bias:**T3**
*out* Y:**T3**|1+|**T1** = tensor(uint8)
**T2** = tensor(int8), tensor(uint8)
**T3** = tensor(float)| |MaxpoolWithMask|*in* X:**T**
*in* M:**tensor(int32)**
*out* Y:**T**|1+|**X** = tensor(float)| |MurmurHash3|*in* X:**T1**
*out* Y:**T2**|1+|**T1** = tensor(double), tensor(float), tensor(int32), tensor(int64), tensor(string), tensor(uint32), tensor(uint64)
**T2** = tensor(int32), tensor(uint32)| +|NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)| |NhwcMaxPool|*in* x:**T**
*out* y:**T**|1+|**T** = tensor(uint8)| |Pad|*in* data:**T**
*in* pads:**tensor(int64)**
*in* value:**T**
*out* output:**T**|1+|**T** = tensor(float)| |QAttention|*in* input:**T1**
*in* weight:**T2**
*in* bias:**T3**
*in* input_scale:**T3**
*in* weight_scale:**T3**
*in* mask_index:**T4**
*in* input_zero_point:**T1**
*in* weight_zero_point:**T2**
*in* past:**T3**
*out* output:**T3**
*out* present:**T3**|1+|**T1** = tensor(uint8)
**T2** = tensor(int8), tensor(uint8)
**T3** = tensor(float)
**T4** = tensor(int32)| @@ -713,6 +714,7 @@ Do not modify directly.* |Inverse|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |Irfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |LongformerAttention|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* mask:**T**
*in* global_weight:**T**
*in* global_bias:**T**
*in* global:**G**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| +|NGramRepeatBlock|*in* input_ids:**Tid**
*in* scores:**T**
*out* scores_out:**T**|1+|**T** = tensor(float)
**Tid** = tensor(int64)| |QAttention|*in* input:**T1**
*in* weight:**T2**
*in* bias:**T3**
*in* input_scale:**T3**
*in* weight_scale:**T3**
*in* mask_index:**T4**
*in* input_zero_point:**T1**
*in* weight_zero_point:**T2**
*in* past:**T3**
*out* output:**T3**
*out* present:**T3**|1+|**T1** = tensor(int8)
**T2** = tensor(int8)
**T3** = tensor(float), tensor(float16)
**T4** = tensor(int32)| |QuantizeLinear|*in* x:**T1**
*in* y_scale:**T1**
*in* y_zero_point:**T2**
*out* y:**T2**|1+|**T1** = tensor(float16)
**T2** = tensor(int8), tensor(uint8)| |Rfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| diff --git a/onnxruntime/contrib_ops/cpu/bert/ngram_repeat_block.cc b/onnxruntime/contrib_ops/cpu/bert/ngram_repeat_block.cc new file mode 100644 index 0000000000..1e651861fe --- /dev/null +++ b/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()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + NGramRepeatBlock); +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/bert/ngram_repeat_block.h b/onnxruntime/contrib_ops/cpu/bert/ngram_repeat_block.h new file mode 100644 index 0000000000..509a35fc33 --- /dev/null +++ b/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("ngram_size", &ngram_size_).IsOK()); + ORT_ENFORCE(ngram_size_ > 0); + } + + Status Compute(OpKernelContext* context) const override { + const Tensor* input_ids = context->Input(0); + const Tensor* scores = context->Input(1); + Tensor* output = context->Output(0, scores->Shape()); + + const auto* scores_source = static_cast(scores->DataRaw()); + auto* scores_target = static_cast(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(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(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::infinity(); + } + } + }; + + concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); + concurrency::ThreadPool::TryParallelFor( + tp, batch_size, static_cast(cur_len * ngram_size_), + [&lambda](ptrdiff_t first, ptrdiff_t last) { + for (auto b = static_cast(first), end = static_cast(last); b < end; ++b) { + lambda(b); + } + } + ); + + return Status::OK(); + } + private: + int64_t ngram_size_; +}; +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc b/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc index 7933a06d26..0f9b2435a3 100644 --- a/onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc +++ b/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, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, #ifdef BUILD_MS_EXPERIMENTAL_OPS BuildKernelCreateInfo, diff --git a/onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block.cc b/onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block.cc new file mode 100644 index 0000000000..53fb3ffa9d --- /dev/null +++ b/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()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + NGramRepeatBlock); + +using namespace ONNX_NAMESPACE; + +NGramRepeatBlock::NGramRepeatBlock(const OpKernelInfo& info) : CudaKernel(info) { + ORT_ENFORCE(info.GetAttr("ngram_size", &ngram_size_).IsOK()); + ORT_ENFORCE(ngram_size_ > 0); +} + +Status NGramRepeatBlock::ComputeInternal(OpKernelContext* context) const { + const Tensor* input_ids = context->Input(0); + const Tensor* scores = context->Input(1); + Tensor* output = context->Output(0, scores->Shape()); + + const auto* scores_source = static_cast(scores->DataRaw()); + auto* scores_target = static_cast(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(input_ids->DataRaw(input_ids->DataType())); + + NGramRepeatBlockImpl( + Stream(), + input_ids_data, + scores_target, + gsl::narrow_cast(batch_size), + gsl::narrow_cast(cur_len - 1), + gsl::narrow_cast(cur_len), + gsl::narrow_cast(vocab_size), + gsl::narrow_cast(1), + gsl::narrow_cast(ngram_size_)); + + return Status::OK(); +} + +} //namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block.h b/onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block.h new file mode 100644 index 0000000000..eb219916ba --- /dev/null +++ b/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 diff --git a/onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block_impl.cu b/onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block_impl.cu new file mode 100644 index 0000000000..887f32721c --- /dev/null +++ b/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>>( + tokens_ptr, scores_ptr, max_predict_len, vocab_size, no_repeat_ngram_size); +} + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block_impl.h b/onnxruntime/contrib_ops/cuda/bert/ngram_repeat_block_impl.h new file mode 100644 index 0000000000..b96caf6a4d --- /dev/null +++ b/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 diff --git a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc b/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc index 64288cab64..86feca4f48 100644 --- a/onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc +++ b/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, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, // These ops were experimental ops in onnx domain which have been removed now. We add them here as // contrib ops to maintain backward compatibility diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index d276cea322..ab51a82523 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/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() { diff --git a/onnxruntime/test/contrib_ops/ngram_repeat_block_op_test.cc b/onnxruntime/test/contrib_ops/ngram_repeat_block_op_test.cc new file mode 100644 index 0000000000..8b062a09f6 --- /dev/null +++ b/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("input_ids", {3, 6}, + {0, 1, 2, 3, 1, 2, + 0, 0, 0, 1, 0, 0, + 0, 1, 0, 1, 0, 1}); + tester.AddInput("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("scores_out", {3, 4}, + {1.0f, 2.0f, 3.0f, -std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), -std::numeric_limits::infinity(), 7.0f, 8.0f, + -std::numeric_limits::infinity(), 10.0f, 11.0f, 12.0f}); + + if (HasCudaEnvironment(0)) { + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json b/onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json index 81706fb45f..6d79f1b5cc 100644 --- a/onnxruntime/test/testdata/kernel_def_hashes/contrib.cpu.json +++ b/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