From c43ce64795df8d0284c8b612693ff237439f25d3 Mon Sep 17 00:00:00 2001 From: Yufeng Li Date: Tue, 22 Nov 2022 21:24:27 -0800 Subject: [PATCH] Beam search TopK improvement (#13594) ### Description TopK in BeamSearch retrieves top 2*beam next tokens based on logit score, specifically computing top [batch, 2*beam] tokens based on score [batch, beam, vocab_size]. ### Motivation and Context Current implementation use batch as the grid and each thread block compute top 2*beam from [beam, vocab_size]. It is inefficient because: 1. batch size is usually small( <32) and can not fully leverage GPU's SMs; 2. vocab_size is usually more than 50k. It is inefficient to compute 50k * beam in one thread block. This PR split the topk computation into multiple stages: - for small beam size, split [batch, beam, vocab_size] to [batch, beam, parts_of_vocab, vocab_size_per_part] - 1st stage, each thread block compute top 2*beam from vocab_sizer_per_part and gets [batch, beam, parts_of_vocab, 2*beam] - 2nd stage, each thread block compute top 2*beam from parts_of_vocab *(2*beam} and gets [batch, beam, 2*beam] - last stage, compute [batch, 2*beam] from [batch, beam, 2*beam] - for large beam size, 1st stage computes [batch, beam, 2*beam] from [batch, beam, vocab_size] and 2nd stage computes [batch, 2*beam] from [batch, beam, 2*beam]. With the change, performance improves a lot, it reduces ~100us from 2ms for batch:4, beam:4, vocab_size:~50k. --- .../cpu/transformers/beam_search_impl_base.h | 8 + .../cpu/transformers/generation_shared.h | 10 + .../cuda/transformers/beam_search_topk.cu | 446 ++++++++++++++++++ .../cuda/transformers/beam_search_topk.h | 42 ++ .../transformers/generation_device_helper.cc | 108 +++-- .../providers/cuda/cuda_provider_factory.cc | 4 + .../core/providers/cuda/math/topk_impl.cuh | 42 +- .../providers/cuda/shared_inc/cuda_utils.h | 55 ++- .../core/providers/cuda/test/all_tests.h | 1 + .../providers/cuda/test/beam_search_topk.cc | 135 ++++++ 10 files changed, 769 insertions(+), 82 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/transformers/beam_search_topk.cu create mode 100644 onnxruntime/contrib_ops/cuda/transformers/beam_search_topk.h create mode 100644 onnxruntime/core/providers/cuda/test/beam_search_topk.cc diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h index 847b916cf2..174d78af4f 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h +++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h @@ -32,6 +32,12 @@ struct BeamSearchState : public IBeamSearchState { this->next_indices = AllocateBuffer(allocator, next_indices_buffer_, SafeInt(2) * batch_beam_size); + this->next_scores = AllocateBuffer(allocator, next_scores_buffer_, SafeInt(2) * batch_beam_size); + + constexpr size_t max_parts_of_vocab = 128; + size_t topk_buffer_size = SafeInt(batch_beam_size) * (max_parts_of_vocab + 1) * num_beams * 2 * 2; + this->topk_buffer = AllocateBuffer(allocator, topk_temp_buffer_, topk_buffer_size); + if (use_position) { this->next_positions = AllocateBuffer(allocator, next_positions_buffer_, batch_beam_size); } @@ -50,9 +56,11 @@ struct BeamSearchState : public IBeamSearchState { BufferUniquePtr next_token_scores_buffer_; BufferUniquePtr next_tokens_buffer_; BufferUniquePtr next_indices_buffer_; + BufferUniquePtr next_scores_buffer_; BufferUniquePtr next_positions_buffer_; BufferUniquePtr beam_scores_buffer_; BufferUniquePtr scores_buffer_; + BufferUniquePtr topk_temp_buffer_; }; struct BeamSearchCpuState : public IBeamSearchCpuState { diff --git a/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h b/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h index a5794e5f4e..edbebcc81a 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h +++ b/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h @@ -27,10 +27,20 @@ struct IBeamSearchState { gsl::span next_token_scores; // shape (batch_size, num_beams * vocab_size) gsl::span next_tokens; // shape (batch_size, 2 * num_beams) gsl::span next_indices; // shape (batch_size, 2 * num_beams) + gsl::span next_scores; // shape (batch_size, 2 * num_beams) gsl::span next_positions; // shape (batch_size, num_beams), empty for T5. Next position for position_ids. gsl::span beam_scores; // shape (batch_size, num_beams) gsl::span scores; // shape (max_length - sequence_length + 1, batch_size, num_beams * vocab_size) gsl::span remaining_scores; // portion of scores that is available for appending next token scores. + gsl::span topk_buffer; // temp buffer for topk computation, including: + // 1st stage needs: + // temp score: (batch_size * num_beams * parts_vocab, 2 * num_beams) + // temp token: (batch_size * num_beams * parts_vocab, 2 * num_beams) + // 2nd stage needs: + // temp score: (batch_size * num_beams, 2 * num_beams) + // temp token: (batch_size * num_beams, 2 * num_beams) + // in total, it will be: + // 2 * (batch_size * num_beams * (parts_vocab + 1), 2 * num_beams) }; struct IBeamSearchCpuState { diff --git a/onnxruntime/contrib_ops/cuda/transformers/beam_search_topk.cu b/onnxruntime/contrib_ops/cuda/transformers/beam_search_topk.cu new file mode 100644 index 0000000000..5c54c03a05 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/transformers/beam_search_topk.cu @@ -0,0 +1,446 @@ +/* + * Copyright (c) 2020-2021, NVIDIA CORPORATION. All rights reserved. + * + * 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 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* Modifications Copyright (c) Microsoft. */ + +#include "beam_search_topk.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 TopK { + int32_t key[max_k]; + T value[max_k]; + + __device__ __forceinline__ void Insert(T elem, int elem_id) { + T v = value[max_k - 1]; + if (v < elem || + (key[max_k - 1] == -1) || + ((elem == value[max_k - 1]) && (elem_id < key[max_k - 1]))) { + value[max_k - 1] = elem; + key[max_k - 1] = elem_id; + } + + for (int k = max_k - 2; k >= 0; --k) { + if (value[k + 1] > value[k] || + key[k] == -1 || + ((value[k + 1] == value[k]) && (key[k + 1] < key[k]))) { + T u2 = value[k]; + int p2 = key[k]; + value[k] = value[k + 1]; + key[k] = key[k + 1]; + value[k + 1] = u2; + key[k + 1] = p2; + } + } + } + + __device__ __forceinline__ void Init() { + for (int i = 0; i < max_k; i++) { + key[i] = -1; + value[i] = NumericLimits::Min(); + } + } +}; + +template +__device__ __forceinline__ TopK reduce_topk_op(const TopK& a, const TopK& b) { + TopK res = a; + for (int i = 0; i < max_k; ++i) + res.Insert(b.value[i], b.key[i]); + return res; +} + +// kernel to compute the top k on last axis for tensor with shape: [batch, beam_size, parts_of_vocab, vacab_part_size] +// Its grid is [batch * beam_size, parts_of_vocab] +template +__launch_bounds__(thread_block_size) __global__ void BeamSearchOnlineTopKStage1Kernel( + const T* input, + int32_t k, + int32_t vocab_size, + int32_t vocab_part_size, + T* output_values, + int32_t* output_token) { + TopK top_k_thread; + top_k_thread.Init(); + + int batch_beam = blockIdx.x; + int voc_part_id = blockIdx.y; + + int token_id_base = voc_part_id * vocab_part_size; + const T* input_block = input + batch_beam * vocab_size; + // voc_part_size + for (int i = threadIdx.x + token_id_base; i < vocab_part_size + token_id_base; i += blockDim.x) { + if (i < vocab_size) { + top_k_thread.Insert(input_block[i], i); + } + } + + // reduce in thread block + typedef cub::BlockReduce, thread_block_size> BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + TopK top_k_block = BlockReduce(temp_storage).Reduce(top_k_thread, reduce_topk_op); + __syncthreads(); + + output_values += batch_beam * gridDim.y * k + voc_part_id * k; + output_token += batch_beam * gridDim.y * k + voc_part_id * k; + if (threadIdx.x == 0) { + for (int i = 0; i < k; i++) { + output_values[i] = top_k_block.value[i]; + output_token[i] = top_k_block.key[i]; + } + } +} + +template +__launch_bounds__(thread_block_size) __global__ void BeamSearchOnlineTopKStage2Kernel( + const T* input_values, + const int32_t* input_tokens, + int32_t k, + int32_t vocab_size, + int32_t parts_per_beam, + T* output_values, + int32_t* output_indices) { + const int vector_id = blockIdx.x; + const int thread_id = threadIdx.x; + + extern __shared__ char shared_buf_extern[]; + T* value_shared_buf = reinterpret_cast(shared_buf_extern); + int32_t* tokens_shared_buf = + reinterpret_cast(shared_buf_extern + max_k * parts_per_beam * sizeof(int32_t)); + + typedef cub::BlockReduce, thread_block_size> BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + input_values += vector_id * k * parts_per_beam; + input_tokens += vector_id * k * parts_per_beam; + + TopK thread_topk; + for (int i = 0; i < max_k; ++i) { + thread_topk.key[i] = -1; + thread_topk.value[i] = NumericLimits::Min(); + } + + for (int idx = thread_id; idx < k * parts_per_beam; idx += thread_block_size) { + value_shared_buf[idx] = input_values[idx]; + tokens_shared_buf[idx] = input_tokens[idx]; + } + __syncthreads(); + + if (thread_id < parts_per_beam) { + T* b_v = value_shared_buf + thread_id * k; + int32_t* b_i = tokens_shared_buf + thread_id * k; + for (int i = 0; i < k; i++) { + thread_topk.Insert(b_v[i], b_i[i]); + } + } + + TopK topk_block = BlockReduce(temp_storage).Reduce(thread_topk, reduce_topk_op); + + if (thread_id == 0) { + output_values += vector_id * k; + output_indices += vector_id * k; + + for (int i = 0; i < k; ++i) { + if (i < k) { + output_values[i] = topk_block.value[i]; + output_indices[i] = topk_block.key[i]; + } + } + } +} + +template +void LaunchBeamSearchOnlineTopKStage2Kernel( + const T* topk_values_tmp, + const int32_t* topk_indices_tmp, + int32_t batch_size, + int32_t num_beams, + int32_t vocab_size, + int32_t parts_per_beam, + int32_t K, + T* output_values, + int32_t* output_indices, + cudaStream_t stream) { + ORT_ENFORCE(parts_per_beam <= 128, "Parts per beam should not be greater than 128"); + + int smem_stage2_size = parts_per_beam * max_k * 2 * sizeof(int32_t); + + if (parts_per_beam <= 32) { + BeamSearchOnlineTopKStage2Kernel<<>>( + topk_values_tmp, topk_indices_tmp, K, vocab_size, parts_per_beam, output_values, output_indices); + return; + } + + if (parts_per_beam <= 64) { + BeamSearchOnlineTopKStage2Kernel<<>>( + topk_values_tmp, topk_indices_tmp, K, vocab_size, parts_per_beam, output_values, output_indices); + return; + } + + BeamSearchOnlineTopKStage2Kernel<<>>( + topk_values_tmp, topk_indices_tmp, K, vocab_size, parts_per_beam, output_values, output_indices); + return; +} + +template +void TopKLauncherMaxK( + const T* input, + int batch_size, + int num_beams, + int vocab_size, + int K, + T* output_values, + int32_t* output_indices, + T* output_values_tmp, + int32_t* output_indices_tmp, + cudaStream_t stream) { + constexpr int kThreadBlockSize = (max_k < 16) ? (max_k < 8) ? 256 : 128 : 64; + + int voc_parts = 4; + if (batch_size * num_beams < 256) { + // volta has 80 SMs, so we aim for three waves + voc_parts = (240 + batch_size * num_beams - 1) / (batch_size * num_beams); + voc_parts = std::min(128, voc_parts); // we implement up to 128 + } + + dim3 grid(batch_size * num_beams, voc_parts); + +#ifndef USE_ROCM + cudaFuncSetAttribute(BeamSearchOnlineTopKStage1Kernel, + cudaFuncAttributePreferredSharedMemoryCarveout, + cudaSharedmemCarveoutMaxL1); +#endif // !USE_ROCM + + BeamSearchOnlineTopKStage1Kernel + <<>>(input, K, vocab_size, (vocab_size + voc_parts - 1) / voc_parts, output_values_tmp, output_indices_tmp); + + LaunchBeamSearchOnlineTopKStage2Kernel( + output_values_tmp, + output_indices_tmp, + batch_size, + num_beams, + vocab_size, + voc_parts, + K, + output_values, + output_indices, + stream); +} + +template +__launch_bounds__(thread_block_size) __global__ void BatchTopKKernel( + const T* topk_scores, + const I* topk_tokens, + int32_t* next_indices, + int32_t* next_tokens, + T* next_scores, + int32_t batch_size, + int32_t num_beams, + int32_t k) { + int thread_id = threadIdx.x; + int block_id = blockIdx.x; + TopK thread_topk; + if (thread_id == 0) { + thread_topk.Init(); + + int index_block = block_id * num_beams * k; + for (int32_t i = 0; i < num_beams * k; i++) { + thread_topk.Insert(topk_scores[index_block + i], index_block + i); + } + + int index_next = block_id * k; + for (int i = 0; i < k; i++) { + next_tokens[index_next + i] = topk_tokens[thread_topk.key[i]]; + next_indices[index_next + i] = (thread_topk.key[i] - index_block) / k; + next_scores[index_next + i] = thread_topk.value[i]; + } + } +} + +template +void LaunchBatchTopKKernel(const T* topk_scores, + const I* topk_tokens, + int32_t* next_indices, + int32_t* next_tokens, + T* next_scores, + int32_t batch_size, + int32_t num_beams, + int32_t k, + cudaStream_t stream) { + ORT_ENFORCE(k <= 256, "LaunchBatchTopKKernel doesn't support k >= 256"); + +#define BatchTopKKernelLauncher(K) \ + BatchTopKKernel<<>>(topk_scores, \ + topk_tokens, \ + next_indices, \ + next_tokens, \ + next_scores, \ + batch_size, \ + num_beams, \ + k); + + if (k <= 4) { + BatchTopKKernelLauncher(4); + } else if (k <= 8) { + BatchTopKKernelLauncher(8); + } else if (k <= 16) { + BatchTopKKernelLauncher(16); + } else if (k <= 32) { + BatchTopKKernelLauncher(32); + } else if (k <= 64) { + BatchTopKKernelLauncher(64); + } else if (k <= 128) { + BatchTopKKernelLauncher(128); + } else { + BatchTopKKernelLauncher(256); + } +} + +template void LaunchBatchTopKKernel(const float* topk_scores, + const int32_t* topk_tokens, + int32_t* next_indices, + int32_t* next_tokens, + float* next_scores, + int32_t batch_size, + int32_t num_beams, + int32_t k, + cudaStream_t stream); + +template void LaunchBatchTopKKernel(const float* topk_scores, + const int64_t* topk_tokens, + int32_t* next_indices, + int32_t* next_tokens, + float* next_scores, + int32_t batch_size, + int32_t num_beams, + int32_t k, + cudaStream_t stream); + +template void LaunchBatchTopKKernel(const half* topk_scores, + const int32_t* topk_tokens, + int32_t* next_indices, + int32_t* next_tokens, + half* next_scores, + int32_t batch_size, + int32_t num_beams, + int32_t k, + cudaStream_t stream); + +template void LaunchBatchTopKKernel(const half* topk_scores, + const int64_t* topk_tokens, + int32_t* next_indices, + int32_t* next_tokens, + half* next_scores, + int32_t batch_size, + int32_t num_beams, + int32_t k, + cudaStream_t stream); + +template +void BeamSearchTopK( + const T* input, + int32_t batch_size, + int32_t num_beams, + int32_t vocab_size, + int32_t k, + T* tmp_values_1st_stage, + int32_t* tmp_indices_1st_stage, + T* tmp_values_2nd_stage, + int32_t* tmp_indices_2nd_stage, + T* output_values, + int32_t* output_tokens, + int32_t* output_indices, + cudaStream_t stream) { + ORT_ENFORCE(k <= 64, "BeamSearchTopK doesn't support k > 64"); + +#define TopKLauncher(K) \ + TopKLauncherMaxK(input, \ + batch_size, \ + num_beams, \ + vocab_size, \ + k, tmp_values_2nd_stage, \ + tmp_indices_2nd_stage, \ + tmp_values_1st_stage, \ + tmp_indices_1st_stage, \ + stream); + + if (k <= 4) { + TopKLauncher(4) + } else if (k <= 8) { + TopKLauncher(8) + } else if (k <= 16) { + TopKLauncher(16) + } else if (k <= 32) { + TopKLauncher(32) + } else { + TopKLauncher(64) + } + + LaunchBatchTopKKernel(tmp_values_2nd_stage, + tmp_indices_2nd_stage, + output_indices, + output_tokens, + output_values, + batch_size, + num_beams, + k, + stream); +} + +template void BeamSearchTopK( + const float* input, + int32_t batch_size, + int32_t num_beams, + int32_t vocab_size, + int32_t k, + float* tmp_values_1st_stage, + int32_t* tmp_indices_1st_stage, + float* tmp_values_2st_stage, + int32_t* tmp_indices_2st_stage, + float* output_values, + int32_t* output_tokens, + int32_t* output_indices, + cudaStream_t stream); + +template void BeamSearchTopK( + const half* input, + int32_t batch_size, + int32_t num_beams, + int32_t vocab_size, + int32_t k, + half* tmp_values_1st_stage, + int32_t* tmp_indices_1st_stage, + half* tmp_values_2st_stage, + int32_t* tmp_indices_2st_stage, + half* output_values, + int32_t* output_tokens, + int32_t* output_indices, + cudaStream_t stream); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/transformers/beam_search_topk.h b/onnxruntime/contrib_ops/cuda/transformers/beam_search_topk.h new file mode 100644 index 0000000000..6877481eb6 --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/transformers/beam_search_topk.h @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include + +namespace onnxruntime { +namespace contrib { +namespace cuda { + +template +void LaunchBatchTopKKernel( + const T* topk_scores, + const I* topk_indices, + int32_t* next_indices, + int32_t* next_tokens, + T* next_scores, + int32_t batch_size, + int32_t num_beams, + int32_t k, + cudaStream_t stream); + +template +void BeamSearchTopK( + const T* input, + int32_t batch_size, + int32_t num_beams, + int32_t vocab_size, + int32_t k, + T* tmp_values_1st_stage, + int32_t* tmp_indices_1st_stage, + T* tmp_values_2st_stage, + int32_t* tmp_indices_2st_stage, + T* output_values, + int32_t* output_tokens, + int32_t* output_indices, + cudaStream_t stream); + +} // namespace cuda +} // namespace contrib +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc index 589928d4a2..11dabc4e97 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc +++ b/onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc @@ -14,6 +14,7 @@ #include "contrib_ops/cuda/transformers/dump_cuda_tensor.h" #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" namespace onnxruntime { namespace concurrency { @@ -176,6 +177,8 @@ void InitBeamState(transformers::IBeamSearchState* beam_state, cudaMemsetAsync(beam_state->next_token_scores.data(), 0, beam_state->next_token_scores.size_bytes(), cuda_stream); cudaMemsetAsync(beam_state->next_tokens.data(), 0, beam_state->next_tokens.size_bytes(), cuda_stream); cudaMemsetAsync(beam_state->next_indices.data(), 0, beam_state->next_indices.size_bytes(), cuda_stream); + cudaMemsetAsync(beam_state->next_scores.data(), 0, beam_state->next_scores.size_bytes(), cuda_stream); + cudaMemsetAsync(beam_state->topk_buffer.data(), 0, beam_state->topk_buffer.size_bytes(), cuda_stream); // Initialize score of first beam of each group with 0 and the rest with -1e9. cuda::LaunchInitKernel(beam_state->beam_scores.data(), batch_size, num_beams, reinterpret_cast(stream)); @@ -215,6 +218,7 @@ Status ProcessLogits(const OrtValue& logits, // const transformers::IConsoleDumper* dumper) { // tensor dumper ORT_UNUSED_PARAMETER(logits_processors); + ORT_UNUSED_PARAMETER(thread_pool); #ifndef DEBUG_GENERATION ORT_UNUSED_PARAMETER(dumper); @@ -327,7 +331,6 @@ Status ProcessLogits(const OrtValue& logits, // #ifdef DEBUG_GENERATION dumper->Print("next_token_scores after logits process", next_token_scores.data(), batch_size, num_beams, vocab_size); #endif - // Add beam score to next token scores. Corresponding python code is like: // next_token_scores = next_token_scores + beam_scores[:, None].expand_as(next_token_scores) cuda::LaunchAddProbsKernel(next_token_scores.data(), beam_state->beam_scores.data(), @@ -347,50 +350,77 @@ Status ProcessLogits(const OrtValue& logits, // beam_state->remaining_scores = beam_state->remaining_scores.subspan(next_token_scores.size()); } - // Apply top-k selection like the following: - // next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size) - // next_token_scores, next_tokens = torch.topk(next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True) - int64_t next_token_scores_dims[] = {batch_size, num_beams * vocab_size}; - TensorShape next_token_scores_shape(&next_token_scores_dims[0], 2); - auto element_type = DataTypeImpl::GetType(); - OrtValue next_token_scores_value; - Tensor::InitOrtValue(element_type, next_token_scores_shape, next_token_scores.data(), allocator->Info(), - next_token_scores_value); - const Tensor& input = next_token_scores_value.Get(); + if (num_beams <= 32) { + constexpr size_t max_parts_of_vocab = 128; + float* topk_tmp_buffer = beam_state->topk_buffer.data(); + float* topk_scores_1st_stage = topk_tmp_buffer; + int32_t* topk_tokens_1st_stage = reinterpret_cast(topk_scores_1st_stage + batch_beam_size * max_parts_of_vocab * 2 * num_beams); + float* topk_scores_2nd_stage = reinterpret_cast(topk_tokens_1st_stage + batch_beam_size * max_parts_of_vocab * 2 * num_beams); + int32_t* topk_tokens_2nd_stage = reinterpret_cast(topk_scores_2nd_stage + batch_beam_size * 2 * num_beams); - constexpr int axis = 1; - const unsigned top_k = static_cast(2 * num_beams); - constexpr bool largest = true; - constexpr bool sorted = true; // results returned in sorted order. + cuda::BeamSearchTopK(next_token_scores.data(), + batch_size, + num_beams, + vocab_size, + 2 * num_beams, + topk_scores_1st_stage, + topk_tokens_1st_stage, + topk_scores_2nd_stage, + topk_tokens_2nd_stage, + beam_state->next_scores.data(), + beam_state->next_tokens.data(), + beam_state->next_indices.data(), + cuda_stream); - std::unique_ptr topk_scores = Tensor::CreateDefault(); - std::unique_ptr topk_indices = Tensor::CreateDefault(); - ORT_RETURN_IF_ERROR(TopK(&input, axis, top_k, largest, sorted, allocator, stream, thread_pool, - *topk_scores, *topk_indices)); + // Select [batch_size, 2 * num_beams] from [batch_size * num_beams, 2 * num_beams] +#ifdef DEBUG_GENERATION + dumper->Print("next_tokens before scorer", beam_state->next_tokens.data(), batch_size, 2 * num_beams); + dumper->Print("next_indices before scorer", beam_state->next_indices.data(), batch_size, 2 * num_beams); + dumper->Print("next_scores before scorer", beam_state->next_scores.data(), batch_size, 2 * num_beams); +#endif + } else { + // Apply top-k selection like the following: + // next_token_scores = next_token_scores.view(batch_size, num_beams * vocab_size) + // next_token_scores, next_tokens = torch.topk(next_token_scores, 2 * num_beams, dim=1, largest=True, sorted=True) + // int64_t next_token_scores_dims[] = {batch_size, num_beams * vocab_size}; + int64_t next_token_scores_dims[] = {batch_size * num_beams, vocab_size}; + + TensorShape next_token_scores_shape(&next_token_scores_dims[0], 2); + auto element_type = DataTypeImpl::GetType(); + OrtValue next_token_scores_value; + Tensor::InitOrtValue(element_type, next_token_scores_shape, next_token_scores.data(), allocator->Info(), + next_token_scores_value); + const Tensor& input = next_token_scores_value.Get(); + + constexpr int axis = 1; + const unsigned top_k = static_cast(2 * num_beams); + constexpr bool largest = true; + constexpr bool sorted = true; // results returned in sorted order. + + std::unique_ptr topk_scores = Tensor::CreateDefault(); + std::unique_ptr topk_tokens = Tensor::CreateDefault(); + ORT_RETURN_IF_ERROR(TopK(&input, axis, top_k, largest, sorted, allocator, stream, thread_pool, + *topk_scores, *topk_tokens)); #ifdef DEBUG_GENERATION - dumper->Print("topk_scores", *(topk_scores.get())); - dumper->Print("topk_indices", *(topk_indices.get())); + dumper->Print("topk_scores", *(topk_scores.get())); + dumper->Print("topk_tokens", *(topk_tokens.get())); #endif - // Convert indices in range [0, num_beams * vocab_size) to token ID of range [0, vocab_size) like the following: - // next_indices = (next_tokens / vocab_size).long() - // next_tokens = next_tokens % vocab_size - const int64_t* next_token_indices = topk_indices->Data(); - cuda::LaunchNextTokenKernel(next_token_indices, beam_state->next_indices.data(), beam_state->next_tokens.data(), - batch_size, top_k, vocab_size, cuda_stream); - - const float* data = topk_scores->Data(); - -#ifdef DEBUG_GENERATION - dumper->Print("next_scores before scorer", data, batch_size, top_k); - dumper->Print("next_tokens before scorer", beam_state->next_tokens.data(), batch_size, top_k); - dumper->Print("next_indices before scorer", beam_state->next_indices.data(), batch_size, top_k); -#endif + cuda::LaunchBatchTopKKernel(topk_scores->Data(), + topk_tokens->Data(), + beam_state->next_indices.data(), + beam_state->next_tokens.data(), + beam_state->next_scores.data(), + batch_size, + num_beams, + 2 * num_beams, + cuda_stream); + } CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(cpu_state->topk_scores.data(), - data, - topk_scores->Shape().Size() * sizeof(float), + beam_state->next_scores.data(), + beam_state->next_scores.size_bytes(), cudaMemcpyDeviceToHost, cuda_stream)); CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(cpu_state->topk_tokens.data(), @@ -405,9 +435,7 @@ Status ProcessLogits(const OrtValue& logits, // cuda_stream)); CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(cuda_stream)); - gsl::span next_scores = gsl::make_span( - cpu_state->topk_scores.data(), - static_cast::size_type>(topk_scores->Shape().Size())); + gsl::span next_scores(cpu_state->topk_scores.data(), beam_state->next_scores.size()); gsl::span next_tokens(cpu_state->topk_tokens.data(), beam_state->next_tokens.size()); gsl::span next_indices(cpu_state->topk_indices.data(), beam_state->next_indices.size()); diff --git a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc index d91dcb9434..a83c0e4f12 100644 --- a/onnxruntime/core/providers/cuda/cuda_provider_factory.cc +++ b/onnxruntime/core/providers/cuda/cuda_provider_factory.cc @@ -198,6 +198,10 @@ struct ProviderInfo_CUDA_Impl : ProviderInfo_CUDA { return false; } + if (!onnxruntime::cuda::test::TestBeamSearchTopK()) { + return false; + } + // TODO(wechi): brings disabled tests in onnxruntime/test/providers/cuda/* // back alive here. return true; diff --git a/onnxruntime/core/providers/cuda/math/topk_impl.cuh b/onnxruntime/core/providers/cuda/math/topk_impl.cuh index 7434d33ef8..a353f58155 100644 --- a/onnxruntime/core/providers/cuda/math/topk_impl.cuh +++ b/onnxruntime/core/providers/cuda/math/topk_impl.cuh @@ -1,8 +1,10 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#pragma once #include "topk_impl.h" #include "core/providers/cuda/cu_inc/common.cuh" +#include "core/providers/cuda/shared_inc/cuda_utils.h" #include "device_atomic_functions.h" #include "cub/cub.cuh" #include "cub/util_type.cuh" @@ -24,46 +26,6 @@ struct KV { int64_t val; }; -template -struct NumericLimits { - static T Min() { - return std::numeric_limits::lowest(); - } - static T Max() { - return std::numeric_limits::max(); - } -}; - -template <> -struct NumericLimits { - static half Min() { - return -65504.0; - } - static half Max() { - return 65504.0; - } -}; - -template <> -struct NumericLimits { - static float Min() { - return -INFINITY; - } - static float Max() { - return INFINITY; - } -}; - -template <> -struct NumericLimits { - static double Min() { - return -HUGE_VAL; - } - static double Max() { - return HUGE_VAL; - } -}; - #define BT GridDim::maxThreadsPerBlock #define ALIGN(N) static_cast(pow(2, ceil(log2(static_cast(N))))) #define FROM(idx) (left_dim + (idx)*mid_dim + right_dim) diff --git a/onnxruntime/core/providers/cuda/shared_inc/cuda_utils.h b/onnxruntime/core/providers/cuda/shared_inc/cuda_utils.h index e7137648e7..fa987866c0 100644 --- a/onnxruntime/core/providers/cuda/shared_inc/cuda_utils.h +++ b/onnxruntime/core/providers/cuda/shared_inc/cuda_utils.h @@ -6,11 +6,13 @@ #pragma once +#include #include #include #include -#include "core/common/gsl.h" +#include "core/common/gsl.h" +#include "core/framework/float16.h" #include "core/providers/cuda/shared_inc/fast_divmod.h" namespace onnxruntime { @@ -50,7 +52,6 @@ void Fill(cudaStream_t stream, T* output, T value, int64_t count); */ template struct TArray { - #if defined(USE_ROCM) #define TARRAY_CONSTRUCTOR_SPECIFIERS __host__ __device__ #else @@ -117,5 +118,55 @@ struct TArray { using BitmaskElementType = uint32_t; constexpr int kNumBitsPerBitmaskElement = std::numeric_limits::digits; +template +struct NumericLimits { + __inline__ __host__ __device__ static T Min() { + return std::numeric_limits::lowest(); + } + __inline__ __host__ __device__ static T Max() { + return std::numeric_limits::max(); + } +}; + +template <> +struct NumericLimits { + __inline__ __host__ __device__ static half Min() { + return -65504.0; + } + __inline__ __host__ __device__ static half Max() { + return 65504.0; + } +}; + +template <> +struct NumericLimits { + __inline__ __host__ __device__ static half Min() { + return -65504.0; + } + __inline__ __host__ __device__ static half Max() { + return 65504.0; + } +}; + +template <> +struct NumericLimits { + __inline__ __host__ __device__ static float Min() { + return -INFINITY; + } + __inline__ __host__ __device__ static float Max() { + return INFINITY; + } +}; + +template <> +struct NumericLimits { + __inline__ __host__ __device__ static double Min() { + return -HUGE_VAL; + } + __inline__ __host__ __device__ static double Max() { + return HUGE_VAL; + } +}; + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/test/all_tests.h b/onnxruntime/core/providers/cuda/test/all_tests.h index d8ef8486f9..3d38fe8e02 100644 --- a/onnxruntime/core/providers/cuda/test/all_tests.h +++ b/onnxruntime/core/providers/cuda/test/all_tests.h @@ -8,6 +8,7 @@ namespace test { // Test header provides function declarations in EP-side bridge. bool TestDeferredRelease(); bool TestDeferredReleaseWithoutArena(); +bool TestBeamSearchTopK(); } // 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 new file mode 100644 index 0000000000..03c46ec1fe --- /dev/null +++ b/onnxruntime/core/providers/cuda/test/beam_search_topk.cc @@ -0,0 +1,135 @@ +#ifndef NDEBUG + +#include "contrib_ops/cuda/transformers/beam_search_topk.h" +#include +#include +#include +#include +#include + +namespace onnxruntime { +namespace cuda { +namespace test { + +void FillAndShuffle(std::vector& values, int32_t batch_size, int32_t beam_size, int32_t vocab_size) { + std::random_device rd; + std::mt19937 generator(rd()); + for (int32_t batch = 0; batch < batch_size; batch++) { + int32_t batch_base_idx = batch * beam_size * vocab_size; + for (int32_t beam = 0; beam < beam_size; beam++) { + int32_t value = beam; + int32_t beam_base_idx = beam * vocab_size; + for (int32_t vocab = 0; vocab < vocab_size; vocab++) { + values[batch_base_idx + beam_base_idx + vocab] = (float)(value); + value += beam_size; + } + std::shuffle(values.begin() + batch_base_idx + beam_base_idx, + values.begin() + batch_base_idx + beam_base_idx + vocab_size, + generator); + } + } +} + +void ComputeTopKReference(const std::vector& values, + std::vector& top_k_values, + std::vector& top_k_tokens, + std::vector& top_k_indices, + int32_t batch_size, + int32_t beam_size, + int32_t vocab_size, + int32_t k) { + auto value_compare = [&values](const int32_t idx_1, const int32_t idx_2) { return values[idx_1] > values[idx_2]; }; + + using VK = std::pair; + + for (int32_t b = 0; b < batch_size; b++) { + std::priority_queue, std::greater> queue; + + int32_t base_idx = b * beam_size * vocab_size; + + // initialize queue with k elements + for (int32_t i = 0; i < k; i++) { + queue.push({values[base_idx + i], i}); + } + for (int32_t i = k; i < beam_size * vocab_size; i++) { + if (values[base_idx + i] > queue.top().first) { + queue.pop(); + queue.push({values[base_idx + i], i}); + } + } + + int32_t top_k_base_idx = b * k; + for (int32_t i = k - 1; i >= 0; i--) { + top_k_values[top_k_base_idx + i] = queue.top().first; + top_k_tokens[top_k_base_idx + i] = queue.top().second % vocab_size; + top_k_indices[top_k_base_idx + i] = queue.top().second / vocab_size; + queue.pop(); + } + } +} + +bool TestBeamSearchTopK() { + int32_t batch_size = 4; + int32_t beam_size = 4; + int32_t vocab_size = 50257; + int32_t k = 2 * beam_size; + int32_t batch_x_beam_x_vocab = batch_size * beam_size * vocab_size; + std::vector values(batch_x_beam_x_vocab); + FillAndShuffle(values, batch_size, beam_size, vocab_size); + + std::vector top_k_values_ref(batch_size * k); + std::vector top_k_tokens_ref(batch_size * k); + std::vector top_k_indices_ref(batch_size * k); + ComputeTopKReference(values, top_k_values_ref, top_k_tokens_ref, top_k_indices_ref, batch_size, beam_size, vocab_size, k); + + const int32_t max_vocab_parts = 128; + size_t buffer_size = batch_x_beam_x_vocab * 4 // input + + batch_size * beam_size * k * (max_vocab_parts + 1) * 2 * 4 // tmp + + batch_size * k * 3 * 4; // output size + void* cuda_buffer = nullptr; + cudaMalloc(&cuda_buffer, buffer_size); + float* values_device = (float*)cuda_buffer; + float* top_k_1st_values_tmp = (float*)(values_device + batch_x_beam_x_vocab); + int32_t* top_k_1st_tokens_tmp = (int32_t*)(top_k_1st_values_tmp + batch_size * beam_size * k * max_vocab_parts); + float* top_k_2nd_values_tmp = (float*)(top_k_1st_tokens_tmp + batch_size * beam_size * k * max_vocab_parts); + int32_t* top_k_2nd_tokens_tmp = (int32_t*)(top_k_2nd_values_tmp + batch_size * beam_size * k); + float* top_k_value = (float*)(top_k_2nd_tokens_tmp + batch_size * beam_size * k); + int32_t* top_k_token = (int32_t*)(top_k_value + batch_size * k); + int32_t* top_k_indices = (int32_t*)(top_k_token + batch_size * k); + cudaMemcpy(values_device, values.data(), batch_x_beam_x_vocab * 4, cudaMemcpyHostToDevice); + + contrib::cuda::BeamSearchTopK(values_device, + batch_size, + beam_size, + vocab_size, + k, + top_k_1st_values_tmp, + top_k_1st_tokens_tmp, + top_k_2nd_values_tmp, + top_k_2nd_tokens_tmp, + top_k_value, + top_k_token, + top_k_indices, + NULL /*stream*/); + + std::vector top_k_values_host(batch_size * k); + std::vector top_k_token_host(batch_size * k); + std::vector top_k_indices_host(batch_size * k); + cudaMemcpy(top_k_values_host.data(), top_k_value, batch_size * k * 4, cudaMemcpyDeviceToHost); + cudaMemcpy(top_k_token_host.data(), top_k_token, batch_size * k * 4, cudaMemcpyDeviceToHost); + cudaMemcpy(top_k_indices_host.data(), top_k_indices, batch_size * k * 4, cudaMemcpyDeviceToHost); + for (int32_t i = 0; i < batch_size * k; i++) { + if (top_k_values_ref[i] != top_k_values_host[i] || + top_k_tokens_ref[i] != top_k_token_host[i] || + top_k_indices_ref[i] != top_k_indices_host[i]) { + return false; + } + } + + return true; +} + +} // namespace test +} // namespace cuda +} // namespace onnxruntime +#endif \ No newline at end of file