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

Support decoder masked self attention for greedy sampling (#15319)

Browse source
This commit is contained in:
Hariharan Seshadri 2023-04-05 23:08:43 -07:00 committed by GitHub
parent 71a4e7eb97
commit ca68ab6126
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 18 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
onnxruntime/contrib_ops/cpu/transformers/sampling.cc
View file

@ -140,9 +140,8 @@ Status Sampling::Compute(OpKernelContext* ctx) const {
init_greedy_state_func_ ? init_greedy_state_func_ : GenerationCpuDeviceHelper::InitGreedyState<float>,
device_copy_func_ ? device_copy_func_ : GenerationCpuDeviceHelper::DeviceCopy<float>,
update_gpt_feeds_func_ ? update_gpt_feeds_func_ : GenerationCpuDeviceHelper::UpdateGptFeeds<float>,
// TODO(hasesh): Add support for decoder masked multihead attention
nullptr,
0};
gpu_device_prop_,
gpu_device_arch_};
ORT_RETURN_IF_ERROR(impl.Initialize());
return impl.Execute(init_run_decoder_feeds_fetches_manager_, *decoder_feeds_fetches_manager_);
@ -165,9 +164,8 @@ Status Sampling::Compute(OpKernelContext* ctx) const {
init_greedy_state_fp16_func_,
device_copy_func_,
update_gpt_feeds_fp16_func_,
// TODO(hasesh): Add support for decoder masked multihead attention
nullptr,
0};
gpu_device_prop_,
gpu_device_arch_};
ORT_RETURN_IF_ERROR(impl.Initialize());

9
onnxruntime/contrib_ops/cpu/transformers/sampling.h
View file

@ -41,16 +41,16 @@ class Sampling : public IControlFlowKernel {
// device helpers that is same for both GPT and encoder-decoder models.
void SetDeviceHelpers(
const GenerationDeviceHelper::AddToFeedsFunc& add_to_feeds_func,
const GenerationDeviceHelper::ReorderPastStateFunc& reorder_past_state_func,
const GenerationDeviceHelper::AddToFeedsFunc& add_to_feeds_func,
const GenerationDeviceHelper::TopkFunc& topk_func,
const GenerationDeviceHelper::DeviceCopyFunc<float>& device_copy_func,
const GenerationDeviceHelper::GreedySearchProcessLogitsFunc<float>& process_logits_func,
const GenerationDeviceHelper::GreedySearchProcessLogitsFunc<MLFloat16>& process_logits_fp16_func,
const GenerationDeviceHelper::InitGreedyStateFunc<float>& init_greedy_state_func,
const GenerationDeviceHelper::InitGreedyStateFunc<MLFloat16>& init_greedy_state_fp16_func) {
add_to_feeds_func_ = add_to_feeds_func;
reorder_past_state_func_ = reorder_past_state_func;
add_to_feeds_func_ = add_to_feeds_func;
topk_func_ = topk_func;
device_copy_func_ = device_copy_func;
process_logits_func_ = process_logits_func;
@ -66,10 +66,13 @@ class Sampling : public IControlFlowKernel {
update_gpt_feeds_fp16_func_ = update_gpt_feeds_fp16_func;
}
const void* gpu_device_prop_ = nullptr;
int gpu_device_arch_ = 0;
private:
// Device specific functions
GenerationDeviceHelper::AddToFeedsFunc add_to_feeds_func_;
GenerationDeviceHelper::ReorderPastStateFunc reorder_past_state_func_;
GenerationDeviceHelper::AddToFeedsFunc add_to_feeds_func_;
GenerationDeviceHelper::TopkFunc topk_func_;
GenerationDeviceHelper::DeviceCopyFunc<float> device_copy_func_;

10
onnxruntime/contrib_ops/cuda/transformers/sampling.cc
View file

@ -32,8 +32,8 @@ transformers::CudaTensorConsoleDumper g_cuda_dumper_sampling;
Sampling::Sampling(const OpKernelInfo& info)
: onnxruntime::contrib::transformers::Sampling(info) {
SetDeviceHelpers(GenerationCudaDeviceHelper::AddToFeeds,
GenerationCudaDeviceHelper::ReorderPastState,
SetDeviceHelpers(GenerationCudaDeviceHelper::ReorderPastState,
GenerationCudaDeviceHelper::AddToFeeds,
GenerationCudaDeviceHelper::TopK,
GenerationCudaDeviceHelper::DeviceCopy<float>,
GenerationCudaDeviceHelper::GreedySearchProcessLogits<float>,
@ -45,6 +45,12 @@ Sampling::Sampling(const OpKernelInfo& info)
GenerationCudaDeviceHelper::UpdateGptFeeds<MLFloat16>);
SetConsoleDumper(&g_cuda_dumper_sampling);
gpu_device_prop_ = &reinterpret_cast<const CUDAExecutionProvider*>(info.GetExecutionProvider())->GetDeviceProp();
gpu_device_arch_ = static_cast<const cudaDeviceProp*>(gpu_device_prop_)->major * 100 +
static_cast<const cudaDeviceProp*>(gpu_device_prop_)->minor * 10;
}
Status Sampling::ComputeInternal(OpKernelContext* context) const {