From 8824f812e07603e36b2c9ba0d7f7593b008e5db4 Mon Sep 17 00:00:00 2001 From: Yufeng Li Date: Fri, 13 Jan 2023 15:03:49 -0800 Subject: [PATCH] optimize topk for greedysearch (#14271) Optimize top 1 computation in greedysearch. For vocabulary size 50k on A100, - batch size 1: from 220us to 10.4us. - batch size 4, from 230us to 11.5us. For generation of 50 tokens for example, it saves 50*0.2ms = 10ms. --- .../cpu/transformers/generate_impl_base.h | 55 ++++-- .../transformers/generation_device_helper.cc | 8 +- .../cpu/transformers/generation_shared.h | 19 +- .../transformers/greedy_search_impl_base.h | 52 +++--- .../cpu/transformers/sampling_cpu_helper.h | 6 +- .../cuda/transformers/generation_cuda_impl.cu | 4 +- .../cuda/transformers/generation_cuda_impl.h | 2 +- .../transformers/generation_device_helper.cc | 55 +++--- .../transformers/generation_device_helper.h | 1 + .../transformers/greedy_search_top_one.cu | 169 ++++++++++++++++++ .../cuda/transformers/greedy_search_top_one.h | 27 +++ .../cuda/transformers/sampling_cuda_helper.h | 6 +- .../providers/cuda/cuda_provider_factory.cc | 4 + .../core/providers/cuda/test/all_tests.h | 1 + .../providers/cuda/test/beam_search_topk.cc | 2 + .../cuda/test/greedy_search_top_one.cc | 106 +++++++++++ 16 files changed, 427 insertions(+), 90 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/transformers/greedy_search_top_one.cu create mode 100644 onnxruntime/contrib_ops/cuda/transformers/greedy_search_top_one.h create mode 100644 onnxruntime/core/providers/cuda/test/greedy_search_top_one.cc diff --git a/onnxruntime/contrib_ops/cpu/transformers/generate_impl_base.h b/onnxruntime/contrib_ops/cpu/transformers/generate_impl_base.h index 7f04f936a6..950ddf6d27 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/generate_impl_base.h +++ b/onnxruntime/contrib_ops/cpu/transformers/generate_impl_base.h @@ -33,6 +33,40 @@ gsl::span AllocateBuffer(AllocatorPtr allocator, return span; } +template +inline void AllocateTempBufferForGetGreedySearchTopOne( + int32_t batch_size, + AllocatorPtr allocator, + BufferUniquePtr& buffer, + gsl::span& stage_1_scores, // shape (batch_size, parts_of_vocab) + gsl::span& stage_1_tokens, // shape (batch_size, parts_of_vocab) + gsl::span& output_scores, // shape (batch_size) + gsl::span& output_tokens // shape (batch_size) +) { + constexpr size_t kMaxPartsPerVocab = 128; + const size_t stage_1_element_size = kMaxPartsPerVocab * batch_size; + const size_t output_element_size = batch_size; + + // Note: use float to allocate buffer for temporary value buffer to avoid unalignment + void* topk_data = allocator->Alloc((stage_1_element_size + output_element_size) * (sizeof(float) + sizeof(int32_t))); + BufferUniquePtr temp_buffer(topk_data, BufferDeleter(allocator)); + buffer = std::move(temp_buffer); + + ElementType* stage_1_scores_data = reinterpret_cast(topk_data); + stage_1_scores = gsl::make_span(stage_1_scores_data, stage_1_element_size); + + int32_t* stage_1_token_data = reinterpret_cast( + reinterpret_cast(stage_1_scores_data) + stage_1_element_size); + stage_1_tokens = gsl::make_span(stage_1_token_data, stage_1_element_size); + + ElementType* output_score_data = reinterpret_cast(stage_1_token_data + stage_1_element_size); + output_scores = gsl::make_span(output_score_data, output_element_size); + + int32_t* output_token_data = reinterpret_cast( + reinterpret_cast(output_score_data) + output_element_size); + output_tokens = gsl::make_span(output_token_data, output_element_size); +} + class GenerateBase { public: GenerateBase(OpKernelContextInternal& context, @@ -67,19 +101,19 @@ class GenerateBase { Status CheckScalarInput(const std::string& name, int index, bool required) const { auto* scalar_tensor = context_.Input(index); - if (scalar_tensor) { - if (!scalar_tensor->Shape().IsScalar()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, - FAIL, - "Node input ", name, " should be a scalar. Got shape of ", - scalar_tensor->Shape()); - } - } else if (required) { + if (scalar_tensor) { + if (!scalar_tensor->Shape().IsScalar()) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, - "Node input ", name, " is required"); + "Node input ", name, " should be a scalar. Got shape of ", + scalar_tensor->Shape()); } - return Status::OK(); + } else if (required) { + return ORT_MAKE_STATUS(ONNXRUNTIME, + FAIL, + "Node input ", name, " is required"); + } + return Status::OK(); } template @@ -175,7 +209,6 @@ class GenerateBase { return Status::OK(); } - protected: bool IsCuda() const { return ort_stream_ != nullptr; } diff --git a/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc b/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc index 518afa294a..17c7f1e6c6 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc +++ b/onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc @@ -504,11 +504,13 @@ Status GreedySearchProcessLogits( #endif gsl::span next_token_indices = topk_indices.DataAsSpan(); - gsl::copy(next_token_indices, greedy_state->next_tokens_cpu); + for (size_t i = 0; i < next_token_indices.size(); i++) { + greedy_state->next_tokens[i] = gsl::narrow_cast(next_token_indices[i]); + } #ifdef DEBUG_GENERATION - gsl::span next_tokens(greedy_state->next_tokens_cpu.data(), - greedy_state->next_tokens_cpu.size()); + gsl::span next_tokens(greedy_state->next_tokens.data(), + greedy_state->next_tokens.size()); dumper->Print("next_tokens before scorer", next_tokens.data(), batch_size, top_k); #endif diff --git a/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h b/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h index e7dca4d326..cf1d996885 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h +++ b/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h @@ -57,13 +57,16 @@ struct IBeamSearchCpuState { template struct IGreedySearchState { - gsl::span next_tokens_cpu; // shape (batch_size) - gsl::span sequences_space; // shape (2, batch_size, max_length) - gsl::span sequence_lengths; // shape (batch_size) - gsl::span next_positions; // shape (batch_size, num_beams). Next position value for position_ids. - gsl::span eos_meet; // shape (batch_size) - gsl::span next_token_scores; // shape (batch_size, vocab_size) - gsl::span next_tokens; // shape (batch_size) + gsl::span sequences_space; // shape (2, batch_size, max_length) + gsl::span sequence_lengths; // shape (batch_size) + gsl::span next_positions; // shape (batch_size, num_beams). Next position value for position_ids. + gsl::span eos_meet; // shape (batch_size) + gsl::span next_token_scores; // shape (batch_size, vocab_size) + gsl::span next_tokens; // shape (batch_size) + gsl::span temp_topk_scores_buffer; // shape (batch_size, parts_of_vocab), temp buffer for topk stage 1 (GPU only) + gsl::span temp_topk_tokens_buffer; // shape (batch_size, parts_of_vocab), temp buffer for topk stage 1(GPU only) + gsl::span topk_scores_buffer; // shape (batch_size), output buffer for topk stage 2 (GPU only) + gsl::span topk_tokens_buffer; // shape (batch_size), output buffer for topk stage 2 (GPU only) }; template @@ -77,7 +80,7 @@ struct ISamplingState { gsl::span h_softmaxed_score; gsl::span d_sampled; gsl::span h_sampled_all; - gsl::span d_indices; + gsl::span d_indices; gsl::span d_presence_mask; BufferUniquePtr storage_buffer; diff --git a/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h b/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h index 724db62219..bd8a5b97b9 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h +++ b/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h @@ -36,7 +36,7 @@ struct SamplingState : public ISamplingState { this->d_softmaxed_score = AllocateBuffer(allocator, d_softmaxed_score_buffer_, SafeInt(total_count)); this->d_sampled = AllocateBuffer(allocator, d_sampled_buffer_, SafeInt(batch_size)); this->h_sampled_all = AllocateBuffer(cpu_allocator, h_sampled_all_buffer_, SafeInt(batch_size * max_iter)); - this->d_indices = AllocateBuffer(allocator, d_indices_buffer_, SafeInt(batch_size)); + this->d_indices = AllocateBuffer(allocator, d_indices_buffer_, SafeInt(batch_size)); this->temp_storage_bytes = 0; // TODO: Do not allocate this buffer if there's no presence_mask this->d_presence_mask = AllocateBuffer(allocator, d_presence_mask_buffer_, SafeInt(total_count)); @@ -79,7 +79,7 @@ struct GreedySearchState : public IGreedySearchState { int vocab_size, int sequence_length, int max_length, - bool /*is_cuda*/) { + bool is_cuda) { // below buffers are on cpu this->sequences_space = AllocateBuffer(cpu_allocator, sequences_space_buffer_, @@ -91,15 +91,23 @@ struct GreedySearchState : public IGreedySearchState { this->eos_meet = AllocateBuffer(cpu_allocator, eos_meet_buffer_, batch_size); memset(this->eos_meet.data(), 0, this->eos_meet.size_bytes()); - this->next_tokens_cpu = AllocateBuffer(cpu_allocator, - next_tokens_cpu_buffer_, - SafeInt(batch_size)); this->next_tokens = AllocateBuffer(cpu_allocator, next_tokens_buffer_, SafeInt(batch_size)); // below buffers are on cpu or cuda size_t next_token_size = SafeInt(batch_size) * vocab_size; this->next_token_scores = AllocateBuffer(allocator, next_token_scores_buffer_, next_token_size); this->next_positions = AllocateBuffer(allocator, next_positions_buffer_, batch_size); + + if (is_cuda) { + AllocateTempBufferForGetGreedySearchTopOne( + batch_size, + allocator, + this->temp_topk_buffer_, + this->temp_topk_scores_buffer, + this->temp_topk_tokens_buffer, + this->topk_scores_buffer, + this->topk_tokens_buffer); + } } void SetSequence(gsl::span input_ids_in_cpu, @@ -109,8 +117,8 @@ struct GreedySearchState : public IGreedySearchState { gsl::span sequences_0 = this->sequences_space; for (size_t i = 0; i < batch_beam_size; i++) { for (int j = 0; j < sequence_length; j++) { - sequences_0[SafeInt(i) * max_length + j] = \ - static_cast(input_ids_in_cpu[SafeInt(i) * sequence_length + j]); + sequences_0[SafeInt(i) * max_length + j] = + static_cast(input_ids_in_cpu[SafeInt(i) * sequence_length + j]); } } } @@ -120,9 +128,9 @@ struct GreedySearchState : public IGreedySearchState { BufferUniquePtr sequence_lengths_buffer_; BufferUniquePtr next_token_scores_buffer_; BufferUniquePtr next_tokens_buffer_; - BufferUniquePtr next_tokens_cpu_buffer_; BufferUniquePtr next_positions_buffer_; BufferUniquePtr eos_meet_buffer_; + BufferUniquePtr temp_topk_buffer_; }; // Base class of gready search implementation that is common for both GPT-2 and Bart/T5. @@ -138,13 +146,13 @@ class GreedySearchBase : public GenerateBase { const GenerationDeviceHelper::TopkFunc& topk_func, const GenerationDeviceHelper::GreedySearchProcessLogitsFunc& process_logits_func, const GenerationDeviceHelper::DeviceCopyFunc& device_copy_func) - : GenerateBase(context, - decoder_session_state, - thread_pool, - ort_stream, - cuda_dumper, - topk_func, - device_copy_func), + : GenerateBase(context, + decoder_session_state, + thread_pool, + ort_stream, + cuda_dumper, + topk_func, + device_copy_func), parameters_(¶ms), process_logits_func_(process_logits_func) { parameters_->ParseFromInputs(&context); @@ -186,11 +194,11 @@ Status GreedySearchBase::CheckInputs(const OpKernelContextIntern // input_ids : (batch_size, sequence_length) // vocab_mask : (vocab_size) or nullptr ORT_RETURN_IF_ERROR(this->CheckInputsImpl(parameters_, - context.Input(0), // input_ids - context.Input(4), // vocab_mask - context.Input(5), // prefix_vocab_mask - context.Input(6), // attention_mask - context.Input(7))); // presence_mask + context.Input(0), // input_ids + context.Input(4), // vocab_mask + context.Input(5), // prefix_vocab_mask + context.Input(6), // attention_mask + context.Input(7))); // presence_mask return Status::OK(); } @@ -241,10 +249,6 @@ Status GreedySearchBase::GenerateNextToken( ORT_RETURN_IF_ERROR(ProcessLogits(logits, greedy_state, sampling_state, this->temp_space_allocator_, counter)); next_tokens = greedy_state.next_tokens; - for (size_t i = 0; i < next_tokens.size(); i++) { - next_tokens[i] = gsl::narrow_cast(greedy_state.next_tokens_cpu[i]); - } - gsl::span& eos_meet = greedy_state.eos_meet; for (size_t batch_id = 0; batch_id < next_tokens.size(); ++batch_id) { if (next_tokens[batch_id] == eos_token_id || eos_meet[batch_id] == true) { diff --git a/onnxruntime/contrib_ops/cpu/transformers/sampling_cpu_helper.h b/onnxruntime/contrib_ops/cpu/transformers/sampling_cpu_helper.h index 5caa7c0b27..5f955dd3d7 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/sampling_cpu_helper.h +++ b/onnxruntime/contrib_ops/cpu/transformers/sampling_cpu_helper.h @@ -133,10 +133,10 @@ Status Sample(AllocatorPtr& allocator, int64_t sampled_idx_dims[] = {static_cast(parameters->batch_size), 1}; TensorShape sampled_idx_shape(&sampled_idx_dims[0], 2); - gsl::span& next_token_idx = greedy_state->next_tokens_cpu; + gsl::span& next_token_idx = greedy_state->next_tokens; OrtValue sampled_idx_ov; - Tensor::InitOrtValue(DataTypeImpl::GetType(), + Tensor::InitOrtValue(DataTypeImpl::GetType(), sampled_idx_shape, next_token_idx.data(), allocator->Info(), @@ -145,7 +145,7 @@ Status Sample(AllocatorPtr& allocator, // Copy the allocator because MultinomialComputeShared() uses move(allocator) AllocatorPtr allocatortemp = allocator; - ORT_RETURN_IF_ERROR(MultinomialComputeShared(allocatortemp, + ORT_RETURN_IF_ERROR(MultinomialComputeShared(allocatortemp, input, parameters->batch_size, parameters->vocab_size, diff --git a/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu b/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu index 600eb50648..523603a550 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu +++ b/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu @@ -598,7 +598,7 @@ template void LaunchFilterLogitsKernel(float* d_sorted_logits_in, // Ref: https://github.com/pytorch/pytorch/blob/release/1.13/aten/src/ATen/native/cuda/MultinomialKernel.cu template -__global__ void sampleMultinomialOnce(int64_t* dest, +__global__ void sampleMultinomialOnce(int32_t* dest, int distributions, int categories, scalar_t* sampled, @@ -714,7 +714,7 @@ __global__ void sampleMultinomialOnce(int64_t* dest, // Only support n_sample = 1 void TorchMultinomialKernelLauncher(float* d_input, float* d_sampled, - int64_t* d_output, + int32_t* d_output, int batch_size, int vocab_size, int* d_presence_mask, diff --git a/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.h b/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.h index 0aa16a3b94..5209622823 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.h +++ b/onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.h @@ -103,7 +103,7 @@ void LaunchFilterLogitsKernel(float* d_sorted_logits_in, void TorchMultinomialKernelLauncher(float* d_input, float* d_sampled, - int64_t* d_output, + int32_t* d_output, int batch_size, int vocab_size, int* d_presence_mask, diff --git a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc index bf3f11a6b3..1a5a9ac5d9 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc +++ b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc @@ -16,6 +16,7 @@ #include "contrib_ops/cpu/transformers/subgraph_t5_decoder.h" #include "contrib_ops/cpu/transformers/subgraph_gpt.h" #include "contrib_ops/cuda/transformers/beam_search_topk.h" +#include "contrib_ops/cuda/transformers/greedy_search_top_one.h" #include "core/providers/cuda/nvtx_profile.h" #include "core/providers/cuda/nvtx_profile_context.h" #include "sampling_cuda_helper.h" @@ -525,6 +526,7 @@ Status GreedySearchProcessLogits( #endif ORT_UNUSED_PARAMETER(logits_processors); + ORT_UNUSED_PARAMETER(thread_pool); #ifndef DEBUG_GENERATION ORT_UNUSED_PARAMETER(dumper); #endif @@ -584,7 +586,7 @@ Status GreedySearchProcessLogits( #ifdef DEBUG_GENERATION dumper->Print("logits", logits); if (is_reuse_logits_buffer) { - //TODO: Handle padded logits in the logits buffer before printing its contents + // TODO: Handle padded logits in the logits buffer before printing its contents ORT_THROW("Dumping contents of logits buffer is not implemented yet"); } else { dumper->Print("next_token_scores", next_token_scores.data(), batch_size, vocab_size); @@ -639,7 +641,7 @@ Status GreedySearchProcessLogits( #ifdef DEBUG_GENERATION if (is_reuse_logits_buffer) { - //TODO: Handle padded logits in the logits buffer before printing its contents + // TODO: Handle padded logits in the logits buffer before printing its contents ORT_THROW("Dumping contents of logits buffer is not implemented yet"); } else { dumper->Print("next_token_scores after logits process", next_token_scores.data(), batch_size, vocab_size); @@ -662,43 +664,26 @@ Status GreedySearchProcessLogits( return Status::OK(); } - // next_tokens = torch.argmax(scores, dim=-1) - int64_t next_token_scores_dims[] = {static_cast(batch_size), - is_reuse_logits_buffer ? padded_vocab_size : vocab_size}; - TensorShape next_token_scores_shape(&next_token_scores_dims[0], 2); - auto element_type = DataTypeImpl::GetType(); - OrtValue next_token_scores_value; - - // TODO(hasesh): Same TODO as above about avoiding the const_cast here - Tensor::InitOrtValue(element_type, - next_token_scores_shape, - is_reuse_logits_buffer - ? const_cast(reinterpret_cast(logits_data)) - : reinterpret_cast(next_token_scores.data()), - allocator->Info(), - next_token_scores_value); - const Tensor& input = next_token_scores_value.Get(); - - constexpr int axis = 1; - constexpr unsigned top_k = static_cast(1); - constexpr bool largest = true; - constexpr bool sorted = false; - - auto topk_scores = Tensor::CreateDefault(); - auto topk_indices = Tensor::CreateDefault(); - ORT_RETURN_IF_ERROR(TopK(&input, axis, top_k, largest, sorted, allocator, stream, thread_pool, - *topk_scores, *topk_indices)); + const CudaT* top_one_input = is_reuse_logits_buffer ? logits_data + : reinterpret_cast(next_token_scores.data()); + cuda::GreedySearchTopOne( + top_one_input, + batch_size, + is_reuse_logits_buffer ? padded_vocab_size : vocab_size, + reinterpret_cast(greedy_state->temp_topk_scores_buffer.data()), + greedy_state->temp_topk_tokens_buffer.data(), + reinterpret_cast(greedy_state->topk_scores_buffer.data()), + greedy_state->topk_tokens_buffer.data(), + cuda_stream); #ifdef DEBUG_GENERATION - dumper->Print("topk_scores", *(topk_scores.get())); - dumper->Print("topk_indices", *(topk_indices.get())); + dumper->Print("topk_scores", greedy_state->topk_scores_buffer.data(), batch_size); + dumper->Print("topk_indices", greedy_state->topk_tokens_buffer.data(), batch_size); #endif - const int64_t* next_token_indices = topk_indices->Data(); - - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(greedy_state->next_tokens_cpu.data(), - next_token_indices, - greedy_state->next_tokens_cpu.size_bytes(), + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(greedy_state->next_tokens.data(), + greedy_state->topk_tokens_buffer.data(), + greedy_state->next_tokens.size_bytes(), cudaMemcpyDeviceToHost, cuda_stream)); CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream)); diff --git a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h index 755ac75e14..4233457bec 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h +++ b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h @@ -8,6 +8,7 @@ #include "core/providers/cuda/cuda_common.h" #include "core/common/gsl.h" +#include "contrib_ops/cpu/transformers/sequences.h" #include "contrib_ops/cpu/transformers/generation_shared.h" namespace onnxruntime { diff --git a/onnxruntime/contrib_ops/cuda/transformers/greedy_search_top_one.cu b/onnxruntime/contrib_ops/cuda/transformers/greedy_search_top_one.cu new file mode 100644 index 0000000000..68a2e16482 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/transformers/greedy_search_top_one.cu @@ -0,0 +1,169 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "greedy_search_top_one.h" + +#include + +#include "core/providers/cuda/shared_inc/cuda_utils.h" +#include "core/providers/cuda/cu_inc/common.cuh" + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +using namespace onnxruntime::cuda; + +template +struct TopOne { + int32_t key; + T value; + + __device__ __host__ __forceinline__ TopOne(int32_t key = -1, T value = NumericLimits::Min()) : key(key), value(value) { + } + + __device__ __forceinline__ void Reduce(int32_t k, T v) { + if (value < v || key == -1) { + key = k; + value = v; + } + } +}; + +template +__device__ __forceinline__ TopOne ReduceTopOneOp(const TopOne& a, const TopOne& b) { + if ((a.value > b.value) || (a.value == b.value && a.key != -1 && a.key < b.key)) { + return a; + } + + return b; +} + +// kernel to compute the top 1 on last axis for tensor with shape[batch, parts_of_vocab, vacab_part_size], +// and produce a tensor with shape [batch, parts_of_vocab] +// Its grid is [batch, parts_of_vocab] +template +__launch_bounds__(thread_block_size) __global__ void GreedySearchTopOneStage1Kernel( + const T* input, + int32_t vocab_size, + int32_t vocab_part_size, + T* output_values, + int32_t* output_token) { + TopOne top_one_thread; + + int batch = blockIdx.x; + int voc_part_id = blockIdx.y; + + int token_id_base = voc_part_id * vocab_part_size; + const T* input_block = input + batch * vocab_size; + // voc_part_size + for (int vocab_idx = threadIdx.x + token_id_base; + vocab_idx < vocab_part_size + token_id_base; + vocab_idx += blockDim.x) { + if (vocab_idx < vocab_size) { + top_one_thread.Reduce(vocab_idx, input_block[vocab_idx]); + } + } + + // reduce in thread block + typedef cub::BlockReduce, thread_block_size> BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + TopOne top_one_block = BlockReduce(temp_storage).Reduce(top_one_thread, ReduceTopOneOp); + if (threadIdx.x == 0) { + output_values[batch * gridDim.y + voc_part_id] = top_one_block.value; + output_token[batch * gridDim.y + voc_part_id] = top_one_block.key; + } +} + +// kernel to compute the top 1 on last axis for tensor with shape[batch, parts_of_vocab], +// and produce a tensor with shape [batch] +// Its grid is [batch] +template +__launch_bounds__(thread_block_size) __global__ void GreedySearchTopOneStage2Kernel( + const T* input_values, + const int32_t* input_tokens, + int32_t vocab_size, + int32_t vocab_parts, + T* output_values, + int32_t* output_tokens) { + const int batch_id = blockIdx.x; + const int thread_id = threadIdx.x; + + input_values += batch_id * vocab_parts; + input_tokens += batch_id * vocab_parts; + + TopOne thread_top_one; + for (int idx = thread_id; idx < vocab_parts; idx += thread_block_size) { + thread_top_one.Reduce(input_tokens[idx], input_values[idx]); + } + + typedef cub::BlockReduce, thread_block_size> BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + TopOne top_one_block = BlockReduce(temp_storage).Reduce(thread_top_one, ReduceTopOneOp); + if (thread_id == 0) { + output_values[batch_id] = top_one_block.value; + output_tokens[batch_id] = top_one_block.key; + } +} + +template +void GreedySearchTopOne( + const T* input, + int32_t batch_size, + int32_t vocab_size, + T* tmp_values, + int32_t* tmp_tokens, + T* output_values, + int32_t* output_tokens, + cudaStream_t stream) { + constexpr int kThreadBlockSize = GridDim::maxThreadsPerBlock; + + int voc_parts = 4; + if (batch_size < 256) { + voc_parts = (240 + batch_size - 1) / batch_size; + voc_parts = std::min(128, voc_parts); // we implement up to 128 + } + + dim3 stage1_grid(batch_size, voc_parts); + GreedySearchTopOneStage1Kernel<<>>( + input, + vocab_size, + (vocab_size + voc_parts - 1) / voc_parts, + tmp_values, + tmp_tokens); + + constexpr int KThreadBlockSizeStage2 = 128; + GreedySearchTopOneStage2Kernel<<>>( + tmp_values, + tmp_tokens, + vocab_size, + voc_parts, + output_values, + output_tokens); +} + +template void GreedySearchTopOne( + const float* input, + int32_t batch_size, + int32_t vocab_size, + float* tmp_values, + int32_t* tmp_tokens, + float* output_values, + int32_t* output_tokens, + cudaStream_t stream); + +template void GreedySearchTopOne( + const half* input, + int32_t batch_size, + int32_t vocab_size, + half* tmp_values, + int32_t* tmp_tokens, + half* output_values, + int32_t* output_tokens, + cudaStream_t stream); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/transformers/greedy_search_top_one.h b/onnxruntime/contrib_ops/cuda/transformers/greedy_search_top_one.h new file mode 100644 index 0000000000..009a410f0d --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/transformers/greedy_search_top_one.h @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template +void GreedySearchTopOne( + const T* input, + int32_t batch_size, + int32_t vocab_size, + T* tmp_values, + int32_t* tmp_tokens, + T* output_values, + int32_t* output_tokens, + cudaStream_t stream); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/transformers/sampling_cuda_helper.h b/onnxruntime/contrib_ops/cuda/transformers/sampling_cuda_helper.h index 092612e674..f9046ec049 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/sampling_cuda_helper.h +++ b/onnxruntime/contrib_ops/cuda/transformers/sampling_cuda_helper.h @@ -150,7 +150,7 @@ Status Sample(AllocatorPtr& allocator, dumper->Print("d_sampled", d_sampled.data(), parameters->batch_size, 1); #endif - gsl::span& d_indices = sampling_state->d_indices; + gsl::span& d_indices = sampling_state->d_indices; gsl::span& presence_mask = sampling_state->d_presence_mask; cuda::TorchMultinomialKernelLauncher(d_softmaxed_score.data(), d_sampled.data(), @@ -164,9 +164,9 @@ Status Sample(AllocatorPtr& allocator, dumper->Print("d_indices", d_indices.data(), parameters->batch_size, 1); #endif - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(greedy_state->next_tokens_cpu.data(), + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(greedy_state->next_tokens.data(), sampling_state->d_indices.data(), - greedy_state->next_tokens_cpu.size_bytes(), + greedy_state->next_tokens.size_bytes(), cudaMemcpyDeviceToHost, cuda_stream)); diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index 30cbbf32ac..d7858fd1c4 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -202,6 +202,10 @@ struct ProviderInfo_CUDA_Impl : ProviderInfo_CUDA { return false; } + if (!onnxruntime::cuda::test::TestGreedySearchTopOne()) { + return false; + } + // TODO(wechi): brings disabled tests in onnxruntime/test/providers/cuda/* // back alive here. return true; diff --git a/onnxruntime/core/providers/cuda/test/all_tests.h b/onnxruntime/core/providers/cuda/test/all_tests.h index 3d38fe8e02..9f36474ba2 100644 --- a/onnxruntime/core/providers/cuda/test/all_tests.h +++ b/onnxruntime/core/providers/cuda/test/all_tests.h @@ -9,6 +9,7 @@ namespace test { bool TestDeferredRelease(); bool TestDeferredReleaseWithoutArena(); bool TestBeamSearchTopK(); +bool TestGreedySearchTopOne(); } // namespace test } // namespace cuda diff --git a/onnxruntime/core/providers/cuda/test/beam_search_topk.cc b/onnxruntime/core/providers/cuda/test/beam_search_topk.cc index 5431ce1495..443dd63f58 100644 --- a/onnxruntime/core/providers/cuda/test/beam_search_topk.cc +++ b/onnxruntime/core/providers/cuda/test/beam_search_topk.cc @@ -127,6 +127,8 @@ bool TestBeamSearchTopK() { } } + CUDA_CALL_THROW(cudaFree(cuda_buffer)); + return true; } diff --git a/onnxruntime/core/providers/cuda/test/greedy_search_top_one.cc b/onnxruntime/core/providers/cuda/test/greedy_search_top_one.cc new file mode 100644 index 0000000000..9b79c45661 --- /dev/null +++ b/onnxruntime/core/providers/cuda/test/greedy_search_top_one.cc @@ -0,0 +1,106 @@ +#ifndef NDEBUG + +#include "contrib_ops/cuda/transformers/greedy_search_top_one.h" +#include "core/providers/cuda/shared_inc/cuda_call.h" + +#include +#include +#include + +#include + +namespace onnxruntime { +namespace cuda { +namespace test { + +void FillAndShuffle(std::vector& values, int32_t batch_size, int32_t vocab_size) { + std::random_device rd; + std::mt19937 generator(rd()); + for (int32_t batch = 0; batch < batch_size; batch++) { + for (int32_t vocab = 0; vocab < vocab_size; vocab++) { + values[batch * vocab_size + vocab] = vocab; + } + std::shuffle(values.begin() + batch * vocab_size, + values.begin() + batch * vocab_size + vocab_size, + generator); + } +} + +void ComputeTop1Reference(const std::vector& values, + std::vector& top_k_values, + std::vector& top_k_tokens, + int32_t batch_size, + int32_t vocab_size) { + using VK = std::pair; + + for (int32_t b = 0; b < batch_size; b++) { + int32_t base_idx = b * vocab_size; + + auto max_itr = std::max_element(values.begin() + base_idx, values.begin() + base_idx + vocab_size); + top_k_values[b] = *max_itr; + top_k_tokens[b] = static_cast(std::distance(values.begin() + base_idx, max_itr)); + } +} + +bool TestGreedySearchTopOne() { + int32_t batch_size = 4; + int32_t vocab_size = 50257; + int32_t batch_x_vocab = batch_size * vocab_size; + std::vector values(batch_x_vocab); + FillAndShuffle(values, batch_size, vocab_size); + + std::vector top_k_values_ref(batch_size); + std::vector top_k_tokens_ref(batch_size); + ComputeTop1Reference(values, top_k_values_ref, top_k_tokens_ref, batch_size, vocab_size); + + constexpr size_t kMaxPartsPerVocab = 128; + const size_t stage_1_element_size = kMaxPartsPerVocab * batch_size; + const size_t output_element_size = batch_size; + + size_t input_buffer_size = batch_x_vocab * sizeof(float); + size_t top1_temp_buffer = stage_1_element_size * (sizeof(float) + sizeof(int32_t)); + size_t output_buffer = output_element_size * (sizeof(float) + sizeof(int32_t)); + size_t buffer_size = input_buffer_size + top1_temp_buffer + output_buffer; + + void* topk_data = nullptr; + CUDA_CALL_THROW(cudaMalloc(&topk_data, buffer_size)); + + float* input_data_gpu = (float*)topk_data; + float* stage_1_scores_data = reinterpret_cast(input_data_gpu + batch_x_vocab); + int32_t* stage_1_token_data = reinterpret_cast(stage_1_scores_data + stage_1_element_size); + + float* output_score_data = reinterpret_cast(stage_1_token_data + stage_1_element_size); + int32_t* output_token_data = reinterpret_cast(output_score_data + output_element_size); + + CUDA_CALL_THROW(cudaMemcpy(input_data_gpu, values.data(), input_buffer_size, cudaMemcpyHostToDevice)); + + contrib::cuda::GreedySearchTopOne( + input_data_gpu, + batch_size, + vocab_size, + stage_1_scores_data, + stage_1_token_data, + output_score_data, + output_token_data, + NULL /*stream*/); + + std::vector top_k_values_host(batch_size); + std::vector top_k_token_host(batch_size); + CUDA_CALL_THROW(cudaMemcpy(top_k_values_host.data(), output_score_data, batch_size * sizeof(float), cudaMemcpyDeviceToHost)); + CUDA_CALL_THROW(cudaMemcpy(top_k_token_host.data(), output_token_data, batch_size * sizeof(float), cudaMemcpyDeviceToHost)); + for (int32_t i = 0; i < batch_size; i++) { + if (top_k_values_ref[i] != top_k_values_host[i] || + top_k_tokens_ref[i] != top_k_token_host[i]) { + return false; + } + } + + CUDA_CALL_THROW(cudaFree(topk_data)); + + return true; +} + +} // namespace test +} // namespace cuda +} // namespace onnxruntime +#endif