wehub-resource-sync
229 lines
9.7 KiB
C++
229 lines
9.7 KiB
C++
// Copyright (c) 2024 PaddlePaddle Authors. 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.
|
|
|
|
#include "paddle/phi/kernels/top_p_sampling_kernel.h"
|
|
#include "xpu/refactor/customized_api.h"
|
|
|
|
#include "paddle/phi/backends/xpu/enforce_xpu.h"
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
|
|
#include "paddle/common/flags.h"
|
|
|
|
PHI_DEFINE_EXPORTED_bool(xpu_top_p_sampling_use_fp16,
|
|
false,
|
|
"use fp16 to improve the inference performance of "
|
|
"top_p_sampling xpu kernel");
|
|
PHI_DEFINE_EXPORTED_bool(xpu_use_rejection_top_p_sampling,
|
|
false,
|
|
"use Dual Pivot Rejection Sampling to improve the "
|
|
"inference performance of top_p_sampling. "
|
|
"The algorithm performs better when top_p is larger. "
|
|
"Note that top_p = 0 is not supported.");
|
|
PHI_DEFINE_EXPORTED_int32(
|
|
xpu_top_p_sampling_heuristic_threshold,
|
|
20,
|
|
"threshold of heuristic method used for xpu_top_p_sampling, default 20; if "
|
|
"heuristic_threshold = -1, xpu_top_p_sampling don't use heuristic method, "
|
|
"and will fallback to normal top_p_sampling; if heuristic_threshold > 0, "
|
|
"xpu_top_p_sampling will enable heuristic method");
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void TopPSamplingKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& ps,
|
|
const optional<DenseTensor>& threshold,
|
|
const optional<DenseTensor>& topp_seed,
|
|
int64_t random_seed,
|
|
int k,
|
|
const std::string& mode,
|
|
DenseTensor* out,
|
|
DenseTensor* ids,
|
|
DenseTensor* topk_scores,
|
|
DenseTensor* topk_ids) {
|
|
using XPUType = typename XPUTypeTrait<T>::Type;
|
|
|
|
const XPUType* x_ptr = reinterpret_cast<const XPUType*>(x.data<T>());
|
|
const XPUType* ps_ptr = reinterpret_cast<const XPUType*>(ps.data<T>());
|
|
XPUType* out_ptr = reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(out));
|
|
int64_t* ids_ptr = dev_ctx.template Alloc<int64_t>(ids);
|
|
auto x_dims = x.dims();
|
|
int64_t bs = x_dims[0];
|
|
int64_t vocab_size = x_dims[1];
|
|
|
|
XPUType* topk_scores_data = nullptr;
|
|
int64_t* topk_ids_data = nullptr;
|
|
if (k > 0) {
|
|
topk_scores_data =
|
|
reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(topk_scores));
|
|
topk_ids_data = dev_ctx.template Alloc<int64_t>(topk_ids);
|
|
int r = xpu::topk<XPUType, int64_t>(dev_ctx.x_context(),
|
|
x_ptr,
|
|
topk_scores_data,
|
|
topk_ids_data,
|
|
{bs, vocab_size},
|
|
k,
|
|
1,
|
|
true,
|
|
true);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu::topk");
|
|
}
|
|
std::vector<int64_t> infer_seed(bs, random_seed);
|
|
if (topp_seed.get_ptr() != nullptr) {
|
|
TensorToVector(*topp_seed, dev_ctx, &infer_seed);
|
|
}
|
|
|
|
std::uniform_real_distribution<float> dist(0.0, 1.0);
|
|
std::vector<float> rand_coeff_cpu;
|
|
for (int64_t i = 0; i < bs; i++) {
|
|
if (infer_seed[i] == -1) {
|
|
std::shared_ptr<std::mt19937_64> engine =
|
|
dev_ctx.GetGenerator()->GetCPUEngine();
|
|
rand_coeff_cpu.push_back(dist(*engine));
|
|
} else {
|
|
std::mt19937_64 engine(infer_seed[i]);
|
|
rand_coeff_cpu.push_back(dist(engine));
|
|
}
|
|
}
|
|
uint64_t seed_now = rand_coeff_cpu.empty() ? random_seed : rand_coeff_cpu[0];
|
|
uint64_t offset = 0;
|
|
xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
|
|
int* ids_int_ptr = RAII_GUARD.alloc<int>(ids->numel());
|
|
PADDLE_ENFORCE_EQ(
|
|
threshold.is_initialized(),
|
|
false,
|
|
errors::InvalidArgument(("threshold not supported in top_p_sampling")));
|
|
if (!FLAGS_xpu_use_rejection_top_p_sampling) {
|
|
float* rand_coeff_xpu = RAII_GUARD.alloc<float>(rand_coeff_cpu.size());
|
|
int r = xpu::do_host2device(dev_ctx.x_context(),
|
|
rand_coeff_cpu.data(),
|
|
rand_coeff_xpu,
|
|
rand_coeff_cpu.size() * sizeof(float));
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "do_host2device");
|
|
int heuristic_threshold = FLAGS_xpu_top_p_sampling_heuristic_threshold;
|
|
|
|
if ((!FLAGS_xpu_top_p_sampling_use_fp16) ||
|
|
std::is_same<T, phi::float16>::value) {
|
|
r = xpu::faster_top_p_sampling<XPUType, int>(dev_ctx.x_context(),
|
|
x_ptr,
|
|
ps_ptr,
|
|
rand_coeff_xpu,
|
|
ids_int_ptr,
|
|
bs,
|
|
vocab_size,
|
|
out_ptr,
|
|
nullptr,
|
|
heuristic_threshold);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "top_p_sampling");
|
|
} else {
|
|
using XPUTypeFP16 = typename XPUTypeTrait<phi::float16>::Type;
|
|
XPUTypeFP16* x_fp16_ptr = RAII_GUARD.alloc<XPUTypeFP16>(x.numel());
|
|
XPUTypeFP16* ps_fp16_ptr = RAII_GUARD.alloc<XPUTypeFP16>(ps.numel());
|
|
XPUTypeFP16* out_fp16_ptr = RAII_GUARD.alloc<XPUTypeFP16>(out->numel());
|
|
|
|
float fp16_scale = 32768.f; // experience value
|
|
r = xpu::scale_cast_fusion<XPUType, XPUTypeFP16>(
|
|
dev_ctx.x_context(), x_ptr, x_fp16_ptr, x.numel(), fp16_scale);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale_cast_fusion");
|
|
r = xpu::scale_cast_fusion<XPUType, XPUTypeFP16>(
|
|
dev_ctx.x_context(), ps_ptr, ps_fp16_ptr, ps.numel(), fp16_scale);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale_cast_fusion");
|
|
|
|
r = xpu::faster_top_p_sampling<XPUTypeFP16, int>(dev_ctx.x_context(),
|
|
x_fp16_ptr,
|
|
ps_fp16_ptr,
|
|
rand_coeff_xpu,
|
|
ids_int_ptr,
|
|
bs,
|
|
vocab_size,
|
|
out_fp16_ptr,
|
|
nullptr,
|
|
heuristic_threshold);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "top_p_sampling");
|
|
|
|
r = xpu::scale_cast_fusion<XPUTypeFP16, XPUType>(dev_ctx.x_context(),
|
|
out_fp16_ptr,
|
|
out_ptr,
|
|
out->numel(),
|
|
1.f / fp16_scale);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale_cast_fusion");
|
|
}
|
|
} else {
|
|
if ((!FLAGS_xpu_top_p_sampling_use_fp16) ||
|
|
std::is_same<T, phi::float16>::value) {
|
|
int r = xpu::top_k_top_p_sampling_from_probs<XPUType, int>(
|
|
dev_ctx.x_context(),
|
|
x_ptr,
|
|
nullptr,
|
|
ps_ptr,
|
|
nullptr,
|
|
ids_int_ptr,
|
|
vocab_size, // top_k
|
|
1.0f, // top_p
|
|
bs,
|
|
vocab_size,
|
|
true,
|
|
seed_now,
|
|
offset);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(
|
|
r, "xpu::top_k_top_p_sampling_from_probs<XPUType");
|
|
} else {
|
|
using XPUTypeFP16 = typename XPUTypeTrait<phi::float16>::Type;
|
|
XPUTypeFP16* x_fp16_ptr = RAII_GUARD.alloc<XPUTypeFP16>(x.numel());
|
|
XPUTypeFP16* ps_fp16_ptr = RAII_GUARD.alloc<XPUTypeFP16>(ps.numel());
|
|
|
|
float fp16_scale = 1.0f; // experience value
|
|
int r = xpu::cast<XPUType, XPUTypeFP16>(
|
|
dev_ctx.x_context(), x_ptr, x_fp16_ptr, x.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale_cast_fusion");
|
|
r = xpu::cast<XPUType, XPUTypeFP16>(
|
|
dev_ctx.x_context(), ps_ptr, ps_fp16_ptr, ps.numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale_cast_fusion");
|
|
r = xpu::top_k_top_p_sampling_from_probs<XPUTypeFP16, int>(
|
|
dev_ctx.x_context(),
|
|
x_fp16_ptr,
|
|
nullptr,
|
|
ps_fp16_ptr,
|
|
nullptr,
|
|
ids_int_ptr,
|
|
vocab_size, // top_k
|
|
1.0f, // top_p
|
|
bs,
|
|
vocab_size,
|
|
true,
|
|
seed_now,
|
|
offset);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(
|
|
r, "xpu::top_k_top_p_sampling_from_probs<XPUType");
|
|
}
|
|
int r = xpu::constant<XPUType>(dev_ctx.x_context(), out_ptr, bs, 0);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "xpu::constant");
|
|
}
|
|
int r = xpu::cast<int, int64_t>(
|
|
dev_ctx.x_context(), ids_int_ptr, ids_ptr, ids->numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(top_p_sampling,
|
|
XPU,
|
|
ALL_LAYOUT,
|
|
phi::TopPSamplingKernel,
|
|
float,
|
|
phi::float16) {}
|