saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

optimize topk for greedysearch (#14271)

Browse source 
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.
This commit is contained in:
Yufeng Li 2023-01-13 15:03:49 -08:00 committed by GitHub
parent 4f309f05ca
commit 8824f812e0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
16 changed files with 427 additions and 90 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
onnxruntime/contrib_ops/cpu/transformers/generate_impl_base.h
View file

@ -33,6 +33,40 @@ gsl::span<T> AllocateBuffer(AllocatorPtr allocator,
return span;
}
template <typename ElementType>
inline void AllocateTempBufferForGetGreedySearchTopOne(
int32_t batch_size,
AllocatorPtr allocator,
BufferUniquePtr& buffer,
gsl::span<ElementType>& stage_1_scores, // shape (batch_size, parts_of_vocab)
gsl::span<int32_t>& stage_1_tokens, // shape (batch_size, parts_of_vocab)
gsl::span<ElementType>& output_scores, // shape (batch_size)
gsl::span<int32_t>& 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<ElementType*>(topk_data);
stage_1_scores = gsl::make_span<ElementType>(stage_1_scores_data, stage_1_element_size);
int32_t* stage_1_token_data = reinterpret_cast<int32_t*>(
reinterpret_cast<float*>(stage_1_scores_data) + stage_1_element_size);
stage_1_tokens = gsl::make_span<int32_t>(stage_1_token_data, stage_1_element_size);
ElementType* output_score_data = reinterpret_cast<ElementType*>(stage_1_token_data + stage_1_element_size);
output_scores = gsl::make_span<ElementType>(output_score_data, output_element_size);
int32_t* output_token_data = reinterpret_cast<int32_t*>(
reinterpret_cast<float*>(output_score_data) + output_element_size);
output_tokens = gsl::make_span<int32_t>(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<Tensor>(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 <typename ParametersT>
@ -175,7 +209,6 @@ class GenerateBase {
return Status::OK();
}
protected:
bool IsCuda() const { return ort_stream_ != nullptr; }

8
onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc
View file

@ -504,11 +504,13 @@ Status GreedySearchProcessLogits(
#endif
gsl::span<const int64_t> next_token_indices = topk_indices.DataAsSpan<int64_t>();
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<int32_t>(next_token_indices[i]);
}
#ifdef DEBUG_GENERATION
gsl::span<const int64_t> next_tokens(greedy_state->next_tokens_cpu.data(),
greedy_state->next_tokens_cpu.size());
gsl::span<const int32_t> 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

19
onnxruntime/contrib_ops/cpu/transformers/generation_shared.h
View file

@ -57,13 +57,16 @@ struct IBeamSearchCpuState {
template <typename T>
struct IGreedySearchState {
gsl::span<int64_t> next_tokens_cpu; // shape (batch_size)
gsl::span<int32_t> sequences_space; // shape (2, batch_size, max_length)
gsl::span<int32_t> sequence_lengths; // shape (batch_size)
gsl::span<int32_t> next_positions; // shape (batch_size, num_beams). Next position value for position_ids.
gsl::span<bool> eos_meet; // shape (batch_size)
gsl::span<T> next_token_scores; // shape (batch_size, vocab_size)
gsl::span<int32_t> next_tokens; // shape (batch_size)
gsl::span<int32_t> sequences_space; // shape (2, batch_size, max_length)
gsl::span<int32_t> sequence_lengths; // shape (batch_size)
gsl::span<int32_t> next_positions; // shape (batch_size, num_beams). Next position value for position_ids.
gsl::span<bool> eos_meet; // shape (batch_size)
gsl::span<T> next_token_scores; // shape (batch_size, vocab_size)
gsl::span<int32_t> next_tokens; // shape (batch_size)
gsl::span<T> temp_topk_scores_buffer; // shape (batch_size, parts_of_vocab), temp buffer for topk stage 1 (GPU only)
gsl::span<int32_t> temp_topk_tokens_buffer; // shape (batch_size, parts_of_vocab), temp buffer for topk stage 1(GPU only)
gsl::span<T> topk_scores_buffer; // shape (batch_size), output buffer for topk stage 2 (GPU only)
gsl::span<int32_t> topk_tokens_buffer; // shape (batch_size), output buffer for topk stage 2 (GPU only)
};
template <typename T>
@ -77,7 +80,7 @@ struct ISamplingState {
gsl::span<float> h_softmaxed_score;
gsl::span<float> d_sampled;
gsl::span<float> h_sampled_all;
gsl::span<int64_t> d_indices;
gsl::span<int32_t> d_indices;
gsl::span<int> d_presence_mask;
BufferUniquePtr storage_buffer;

52
onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h
View file

@ -36,7 +36,7 @@ struct SamplingState : public ISamplingState<T> {
this->d_softmaxed_score = AllocateBuffer<float>(allocator, d_softmaxed_score_buffer_, SafeInt<size_t>(total_count));
this->d_sampled = AllocateBuffer<float>(allocator, d_sampled_buffer_, SafeInt<size_t>(batch_size));
this->h_sampled_all = AllocateBuffer<float>(cpu_allocator, h_sampled_all_buffer_, SafeInt<size_t>(batch_size * max_iter));
this->d_indices = AllocateBuffer<int64_t>(allocator, d_indices_buffer_, SafeInt<size_t>(batch_size));
this->d_indices = AllocateBuffer<int32_t>(allocator, d_indices_buffer_, SafeInt<size_t>(batch_size));
this->temp_storage_bytes = 0;
// TODO: Do not allocate this buffer if there's no presence_mask
this->d_presence_mask = AllocateBuffer<int>(allocator, d_presence_mask_buffer_, SafeInt<size_t>(total_count));
@ -79,7 +79,7 @@ struct GreedySearchState : public IGreedySearchState<T> {
int vocab_size,
int sequence_length,
int max_length,
bool /*is_cuda*/) {
bool is_cuda) {
// below buffers are on cpu
this->sequences_space = AllocateBuffer<int32_t>(cpu_allocator,
sequences_space_buffer_,
@ -91,15 +91,23 @@ struct GreedySearchState : public IGreedySearchState<T> {
this->eos_meet = AllocateBuffer<bool>(cpu_allocator, eos_meet_buffer_, batch_size);
memset(this->eos_meet.data(), 0, this->eos_meet.size_bytes());
this->next_tokens_cpu = AllocateBuffer<int64_t>(cpu_allocator,
next_tokens_cpu_buffer_,
SafeInt<size_t>(batch_size));
this->next_tokens = AllocateBuffer<int32_t>(cpu_allocator, next_tokens_buffer_, SafeInt<size_t>(batch_size));
// below buffers are on cpu or cuda
size_t next_token_size = SafeInt<size_t>(batch_size) * vocab_size;
this->next_token_scores = AllocateBuffer<T>(allocator, next_token_scores_buffer_, next_token_size);
this->next_positions = AllocateBuffer<int32_t>(allocator, next_positions_buffer_, batch_size);
if (is_cuda) {
AllocateTempBufferForGetGreedySearchTopOne<T>(
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<const int32_t> input_ids_in_cpu,
@ -109,8 +117,8 @@ struct GreedySearchState : public IGreedySearchState<T> {
gsl::span<int32_t> 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<gsl::index>(i) * max_length + j] = \
static_cast<int32_t>(input_ids_in_cpu[SafeInt<gsl::index>(i) * sequence_length + j]);
sequences_0[SafeInt<gsl::index>(i) * max_length + j] =
static_cast<int32_t>(input_ids_in_cpu[SafeInt<gsl::index>(i) * sequence_length + j]);
}
}
}
@ -120,9 +128,9 @@ struct GreedySearchState : public IGreedySearchState<T> {
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<T>& process_logits_func,
const GenerationDeviceHelper::DeviceCopyFunc<float>& 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_(&params),
process_logits_func_(process_logits_func) {
parameters_->ParseFromInputs(&context);
@ -186,11 +194,11 @@ Status GreedySearchBase<T, ParametersT>::CheckInputs(const OpKernelContextIntern
// input_ids : (batch_size, sequence_length)
// vocab_mask : (vocab_size) or nullptr
ORT_RETURN_IF_ERROR(this->CheckInputsImpl(parameters_,
context.Input<Tensor>(0), // input_ids
context.Input<Tensor>(4), // vocab_mask
context.Input<Tensor>(5), // prefix_vocab_mask
context.Input<Tensor>(6), // attention_mask
context.Input<Tensor>(7))); // presence_mask
context.Input<Tensor>(0), // input_ids
context.Input<Tensor>(4), // vocab_mask
context.Input<Tensor>(5), // prefix_vocab_mask
context.Input<Tensor>(6), // attention_mask
context.Input<Tensor>(7))); // presence_mask
return Status::OK();
}
@ -241,10 +249,6 @@ Status GreedySearchBase<T, ParametersT>::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<int32_t>(greedy_state.next_tokens_cpu[i]);
}
gsl::span<bool>& 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) {

6
onnxruntime/contrib_ops/cpu/transformers/sampling_cpu_helper.h
View file

@ -133,10 +133,10 @@ Status Sample(AllocatorPtr& allocator,
int64_t sampled_idx_dims[] = {static_cast<int64_t>(parameters->batch_size), 1};
TensorShape sampled_idx_shape(&sampled_idx_dims[0], 2);
gsl::span<int64_t>& next_token_idx = greedy_state->next_tokens_cpu;
gsl::span<int32_t>& next_token_idx = greedy_state->next_tokens;
OrtValue sampled_idx_ov;
Tensor::InitOrtValue(DataTypeImpl::GetType<int64_t>(),
Tensor::InitOrtValue(DataTypeImpl::GetType<int32_t>(),
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<int64_t>(allocatortemp,
ORT_RETURN_IF_ERROR(MultinomialComputeShared<int32_t>(allocatortemp,
input,
parameters->batch_size,
parameters->vocab_size,

4
onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.cu
View file

@ -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 <typename scalar_t, typename accscalar_t>
__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,

2
onnxruntime/contrib_ops/cuda/transformers/generation_cuda_impl.h
View file

@ -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,

55
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc
View file

@ -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<int64_t>(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<T>();
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<void*>(reinterpret_cast<const void*>(logits_data))
: reinterpret_cast<void*>(next_token_scores.data()),
allocator->Info(),
next_token_scores_value);
const Tensor& input = next_token_scores_value.Get<Tensor>();
constexpr int axis = 1;
constexpr unsigned top_k = static_cast<unsigned>(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<const CudaT*>(next_token_scores.data());
cuda::GreedySearchTopOne(
top_one_input,
batch_size,
is_reuse_logits_buffer ? padded_vocab_size : vocab_size,
reinterpret_cast<CudaT*>(greedy_state->temp_topk_scores_buffer.data()),
greedy_state->temp_topk_tokens_buffer.data(),
reinterpret_cast<CudaT*>(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<int64_t>();
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));

1
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h
View file

@ -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 {

169
onnxruntime/contrib_ops/cuda/transformers/greedy_search_top_one.cu Normal file
View file

@ -0,0 +1,169 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "greedy_search_top_one.h"
#include <cub/cub.cuh>
#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 <typename T>
struct TopOne {
int32_t key;
T value;
__device__ __host__ __forceinline__ TopOne(int32_t key = -1, T value = NumericLimits<T>::Min()) : key(key), value(value) {
}
__device__ __forceinline__ void Reduce(int32_t k, T v) {
if (value < v || key == -1) {
key = k;
value = v;
}
}
};
template <typename T>
__device__ __forceinline__ TopOne<T> ReduceTopOneOp(const TopOne<T>& a, const TopOne<T>& 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 <typename T, int thread_block_size>
__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<T> 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<TopOne<T>, thread_block_size> BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage;
TopOne<T> top_one_block = BlockReduce(temp_storage).Reduce(top_one_thread, ReduceTopOneOp<T>);
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 <typename T, int thread_block_size>
__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<T> 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<TopOne<T>, thread_block_size> BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage;
TopOne<T> top_one_block = BlockReduce(temp_storage).Reduce(thread_top_one, ReduceTopOneOp<T>);
if (thread_id == 0) {
output_values[batch_id] = top_one_block.value;
output_tokens[batch_id] = top_one_block.key;
}
}
template <typename T>
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<T, kThreadBlockSize><<<stage1_grid, kThreadBlockSize, 0, stream>>>(
input,
vocab_size,
(vocab_size + voc_parts - 1) / voc_parts,
tmp_values,
tmp_tokens);
constexpr int KThreadBlockSizeStage2 = 128;
GreedySearchTopOneStage2Kernel<T, KThreadBlockSizeStage2><<<batch_size, KThreadBlockSizeStage2, 0, stream>>>(
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

27
onnxruntime/contrib_ops/cuda/transformers/greedy_search_top_one.h Normal file
View file

@ -0,0 +1,27 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <stdint.h>
#include <cuda_runtime.h>
namespace onnxruntime {
namespace contrib {
namespace cuda {
template <typename T>
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

6
onnxruntime/contrib_ops/cuda/transformers/sampling_cuda_helper.h
View file

@ -150,7 +150,7 @@ Status Sample(AllocatorPtr& allocator,
dumper->Print("d_sampled", d_sampled.data(), parameters->batch_size, 1);
#endif
gsl::span<int64_t>& d_indices = sampling_state->d_indices;
gsl::span<int32_t>& d_indices = sampling_state->d_indices;
gsl::span<int>& 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));

4
onnxruntime/core/providers/cuda/cuda_provider_factory.cc
View file

@ -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;

1
onnxruntime/core/providers/cuda/test/all_tests.h
View file

@ -9,6 +9,7 @@ namespace test {
bool TestDeferredRelease();
bool TestDeferredReleaseWithoutArena();
bool TestBeamSearchTopK();
bool TestGreedySearchTopOne();
} // namespace test
} // namespace cuda

2
onnxruntime/core/providers/cuda/test/beam_search_topk.cc
View file

@ -127,6 +127,8 @@ bool TestBeamSearchTopK() {
}
}
CUDA_CALL_THROW(cudaFree(cuda_buffer));
return true;
}

106
onnxruntime/core/providers/cuda/test/greedy_search_top_one.cc Normal file
View file

@ -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 <algorithm>
#include <numeric>
#include <random>
#include <cuda_runtime.h>
namespace onnxruntime {
namespace cuda {
namespace test {
void FillAndShuffle(std::vector<float>& 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<float>& values,
std::vector<float>& top_k_values,
std::vector<int32_t>& top_k_tokens,
int32_t batch_size,
int32_t vocab_size) {
using VK = std::pair<float, int32_t>;
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<int32_t>(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<float> values(batch_x_vocab);
FillAndShuffle(values, batch_size, vocab_size);
std::vector<float> top_k_values_ref(batch_size);
std::vector<int32_t> 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<float*>(input_data_gpu + batch_x_vocab);
int32_t* stage_1_token_data = reinterpret_cast<int32_t*>(stage_1_scores_data + stage_1_element_size);
float* output_score_data = reinterpret_cast<float*>(stage_1_token_data + stage_1_element_size);
int32_t* output_token_data = reinterpret_cast<int32_t*>(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<float> top_k_values_host(batch_size);
std::vector<int32_t> 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