saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-06 04:28:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Beam search TopK improvement (#13594)

Browse source 
### Description
<!-- Describe your changes. -->

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
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
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.
This commit is contained in:
Yufeng Li 2022-11-22 21:24:27 -08:00 committed by GitHub
parent 7857f59d2b
commit c43ce64795
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 769 additions and 82 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h
View file

@ -32,6 +32,12 @@ struct BeamSearchState : public IBeamSearchState<T> {
this->next_indices = AllocateBuffer<int32_t>(allocator, next_indices_buffer_, SafeInt<size_t>(2) * batch_beam_size);
this->next_scores = AllocateBuffer<float>(allocator, next_scores_buffer_, SafeInt<size_t>(2) * batch_beam_size);
constexpr size_t max_parts_of_vocab = 128;
size_t topk_buffer_size = SafeInt<size_t>(batch_beam_size) * (max_parts_of_vocab + 1) * num_beams * 2 * 2;
this->topk_buffer = AllocateBuffer<float>(allocator, topk_temp_buffer_, topk_buffer_size);
if (use_position) {
this->next_positions = AllocateBuffer<int32_t>(allocator, next_positions_buffer_, batch_beam_size);
}
@ -50,9 +56,11 @@ struct BeamSearchState : public IBeamSearchState<T> {
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 {

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

@ -27,10 +27,20 @@ struct IBeamSearchState {
gsl::span<float> next_token_scores; // shape (batch_size, num_beams * vocab_size)
gsl::span<int32_t> next_tokens; // shape (batch_size, 2 * num_beams)
gsl::span<int32_t> next_indices; // shape (batch_size, 2 * num_beams)
gsl::span<float> next_scores; // shape (batch_size, 2 * num_beams)
gsl::span<int32_t> next_positions; // shape (batch_size, num_beams), empty for T5. Next position for position_ids.
gsl::span<float> beam_scores; // shape (batch_size, num_beams)
gsl::span<float> scores; // shape (max_length - sequence_length + 1, batch_size, num_beams * vocab_size)
gsl::span<float> remaining_scores; // portion of scores that is available for appending next token scores.
gsl::span<float> 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 {

446
onnxruntime/contrib_ops/cuda/transformers/beam_search_topk.cu Normal file
View file

@ -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 <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, int max_k>
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<T>::Min();
}
}
};
template <typename T, int max_k>
__device__ __forceinline__ TopK<T, max_k> reduce_topk_op(const TopK<T, max_k>& a, const TopK<T, max_k>& b) {
TopK<T, max_k> 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 <typename T, int max_k, int thread_block_size>
__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<T, max_k> 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<TopK<T, max_k>, thread_block_size> BlockReduce;
__shared__ typename BlockReduce::TempStorage temp_storage;
TopK<T, max_k> top_k_block = BlockReduce(temp_storage).Reduce(top_k_thread, reduce_topk_op<T, max_k>);
__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 <typename T, int max_k, int thread_block_size>
__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<T*>(shared_buf_extern);
int32_t* tokens_shared_buf =
reinterpret_cast<int32_t*>(shared_buf_extern + max_k * parts_per_beam * sizeof(int32_t));
typedef cub::BlockReduce<TopK<T, max_k>, 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<T, max_k> thread_topk;
for (int i = 0; i < max_k; ++i) {
thread_topk.key[i] = -1;
thread_topk.value[i] = NumericLimits<T>::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<T, max_k> topk_block = BlockReduce(temp_storage).Reduce(thread_topk, reduce_topk_op<T, max_k>);
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 <typename T, int max_k>
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<T, max_k, 32><<<batch_size * num_beams, 32, smem_stage2_size, stream>>>(
topk_values_tmp, topk_indices_tmp, K, vocab_size, parts_per_beam, output_values, output_indices);
return;
}
if (parts_per_beam <= 64) {
BeamSearchOnlineTopKStage2Kernel<T, max_k, 64><<<batch_size * num_beams, 64, smem_stage2_size, stream>>>(
topk_values_tmp, topk_indices_tmp, K, vocab_size, parts_per_beam, output_values, output_indices);
return;
}
BeamSearchOnlineTopKStage2Kernel<T, max_k, 128><<<batch_size * num_beams, 128, smem_stage2_size, stream>>>(
topk_values_tmp, topk_indices_tmp, K, vocab_size, parts_per_beam, output_values, output_indices);
return;
}
template <typename T, int max_k>
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<T, max_k, kThreadBlockSize>,
cudaFuncAttributePreferredSharedMemoryCarveout,
cudaSharedmemCarveoutMaxL1);
#endif // !USE_ROCM
BeamSearchOnlineTopKStage1Kernel<T, max_k, kThreadBlockSize>
<<<grid, kThreadBlockSize, 0, stream>>>(input, K, vocab_size, (vocab_size + voc_parts - 1) / voc_parts, output_values_tmp, output_indices_tmp);
LaunchBeamSearchOnlineTopKStage2Kernel<T, max_k>(
output_values_tmp,
output_indices_tmp,
batch_size,
num_beams,
vocab_size,
voc_parts,
K,
output_values,
output_indices,
stream);
}
template <typename T, typename I, int32_t max_k, int32_t thread_block_size>
__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<T, max_k> 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 <typename T, typename I>
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<T, I, K, 32><<<batch_size, 32, 0, stream>>>(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 <typename T>
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<T, K>(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

42
onnxruntime/contrib_ops/cuda/transformers/beam_search_topk.h Normal file
View file

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

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

@ -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<T>* 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<cudaStream_t>(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<float>();
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<Tensor>();
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<int32_t*>(topk_scores_1st_stage + batch_beam_size * max_parts_of_vocab * 2 * num_beams);
float* topk_scores_2nd_stage = reinterpret_cast<float*>(topk_tokens_1st_stage + batch_beam_size * max_parts_of_vocab * 2 * num_beams);
int32_t* topk_tokens_2nd_stage = reinterpret_cast<int32_t*>(topk_scores_2nd_stage + batch_beam_size * 2 * num_beams);
constexpr int axis = 1;
const unsigned top_k = static_cast<unsigned>(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<Tensor> topk_scores = Tensor::CreateDefault();
std::unique_ptr<Tensor> 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<float>();
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<Tensor>();
constexpr int axis = 1;
const unsigned top_k = static_cast<unsigned>(2 * num_beams);
constexpr bool largest = true;
constexpr bool sorted = true; // results returned in sorted order.
std::unique_ptr<Tensor> topk_scores = Tensor::CreateDefault();
std::unique_ptr<Tensor> 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<int64_t>();
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<float>();
#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<float>(),
topk_tokens->Data<int64_t>(),
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<const float> next_scores = gsl::make_span(
cpu_state->topk_scores.data(),
static_cast<typename gsl::span<float>::size_type>(topk_scores->Shape().Size()));
gsl::span<const float> next_scores(cpu_state->topk_scores.data(), beam_state->next_scores.size());
gsl::span<const int32_t> next_tokens(cpu_state->topk_tokens.data(), beam_state->next_tokens.size());
gsl::span<const int32_t> next_indices(cpu_state->topk_indices.data(), beam_state->next_indices.size());

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

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

42
onnxruntime/core/providers/cuda/math/topk_impl.cuh
View file

@ -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 <typename T>
struct NumericLimits {
static T Min() {
return std::numeric_limits<T>::lowest();
}
static T Max() {
return std::numeric_limits<T>::max();
}
};
template <>
struct NumericLimits<MLFloat16> {
static half Min() {
return -65504.0;
}
static half Max() {
return 65504.0;
}
};
template <>
struct NumericLimits<float> {
static float Min() {
return -INFINITY;
}
static float Max() {
return INFINITY;
}
};
template <>
struct NumericLimits<double> {
static double Min() {
return -HUGE_VAL;
}
static double Max() {
return HUGE_VAL;
}
};
#define BT GridDim::maxThreadsPerBlock
#define ALIGN(N) static_cast<int64_t>(pow(2, ceil(log2(static_cast<double>(N)))))
#define FROM(idx) (left_dim + (idx)*mid_dim + right_dim)

55
onnxruntime/core/providers/cuda/shared_inc/cuda_utils.h
View file

@ -6,11 +6,13 @@
#pragma once
#include <cuda_fp16.h>
#include <memory>
#include <type_traits>
#include <vector>
#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 <typename T, int32_t capacity = 8>
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<BitmaskElementType>::digits;
template <typename T>
struct NumericLimits {
__inline__ __host__ __device__ static T Min() {
return std::numeric_limits<T>::lowest();
}
__inline__ __host__ __device__ static T Max() {
return std::numeric_limits<T>::max();
}
};
template <>
struct NumericLimits<MLFloat16> {
__inline__ __host__ __device__ static half Min() {
return -65504.0;
}
__inline__ __host__ __device__ static half Max() {
return 65504.0;
}
};
template <>
struct NumericLimits<half> {
__inline__ __host__ __device__ static half Min() {
return -65504.0;
}
__inline__ __host__ __device__ static half Max() {
return 65504.0;
}
};
template <>
struct NumericLimits<float> {
__inline__ __host__ __device__ static float Min() {
return -INFINITY;
}
__inline__ __host__ __device__ static float Max() {
return INFINITY;
}
};
template <>
struct NumericLimits<double> {
__inline__ __host__ __device__ static double Min() {
return -HUGE_VAL;
}
__inline__ __host__ __device__ static double Max() {
return HUGE_VAL;
}
};
} // namespace cuda
} // namespace onnxruntime

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

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

135
onnxruntime/core/providers/cuda/test/beam_search_topk.cc Normal file
View file

@ -0,0 +1,135 @@
#ifndef NDEBUG
#include "contrib_ops/cuda/transformers/beam_search_topk.h"
#include <random>
#include <numeric>
#include <queue>
#include <algorithm>
#include <cuda_runtime.h>
namespace onnxruntime {
namespace cuda {
namespace test {
void FillAndShuffle(std::vector<float>& 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<float>& values,
std::vector<float>& top_k_values,
std::vector<int32_t>& top_k_tokens,
std::vector<int32_t>& 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<float, int32_t>;
for (int32_t b = 0; b < batch_size; b++) {
std::priority_queue<VK, std::vector<VK>, std::greater<VK>> 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<float> values(batch_x_beam_x_vocab);
FillAndShuffle(values, batch_size, beam_size, vocab_size);
std::vector<float> top_k_values_ref(batch_size * k);
std::vector<int32_t> top_k_tokens_ref(batch_size * k);
std::vector<int32_t> 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<float> top_k_values_host(batch_size * k);
std::vector<int32_t> top_k_token_host(batch_size * k);
std::vector<int32_t> 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