mirror of
https://github.com/PaddlePaddle/FastDeploy.git
synced 2025-10-05 00:33:03 +08:00
[OPs] Universal optimization and Fix early_stop cuda 700 (#3375)
Some checks failed
Deploy GitHub Pages / deploy (push) Has been cancelled
Some checks failed
Deploy GitHub Pages / deploy (push) Has been cancelled
* delete nonzero * delete setup_ops_base.py * check if * check gcp infer_seed.cpu() * fix repetition_early_stopper_kernel cuda 700
This commit is contained in:
@@ -776,6 +776,22 @@ void MergePrefillDecodeOutput(
|
||||
const int head_dim,
|
||||
const int max_token);
|
||||
|
||||
std::vector<paddle::Tensor> TopPSamplingReject(const paddle::Tensor &probs,
|
||||
const paddle::Tensor &top_p,
|
||||
const paddle::optional<paddle::Tensor> &top_k,
|
||||
int64_t seed);
|
||||
|
||||
std::vector<paddle::Tensor> TopKRenorm(const paddle::Tensor &probs,
|
||||
const paddle::Tensor &top_k);
|
||||
|
||||
std::vector<paddle::Tensor> MinPSamplingFromProbs(const paddle::Tensor &probs,
|
||||
const paddle::Tensor &min_p);
|
||||
|
||||
void SaveOutMmsgStatic(const paddle::Tensor& x,
|
||||
const paddle::Tensor& not_need_stop,
|
||||
int64_t rank_id,
|
||||
bool save_each_rank);
|
||||
|
||||
PYBIND11_MODULE(fastdeploy_ops, m) {
|
||||
|
||||
m.def("get_expert_token_num", &GetExpertTokenNum, py::arg("topk_ids"),
|
||||
@@ -1128,4 +1144,12 @@ PYBIND11_MODULE(fastdeploy_ops, m) {
|
||||
m.def("speculate_step_paddle",&SpeculateStepPaddle, "speculate_step_paddle function");
|
||||
|
||||
m.def("merge_prefill_decode_output", &MergePrefillDecodeOutput, "merge_prefill_decode_output function");
|
||||
|
||||
m.def("rejection_top_p_sampling", &TopPSamplingReject, "rejection_top_p_sampling function");
|
||||
|
||||
m.def("top_k_renorm_probs", &TopKRenorm, "top_k_renorm_probs function");
|
||||
|
||||
m.def("min_p_sampling", &MinPSamplingFromProbs, "min_p_sampling function");
|
||||
|
||||
m.def("save_output", &SaveOutMmsgStatic, "save_output function");
|
||||
}
|
||||
|
@@ -109,11 +109,11 @@ void GetOutputEp(const paddle::Tensor& x,
|
||||
return;
|
||||
}
|
||||
|
||||
void GetOutputStatic(const paddle::Tensor& x, int64_t rank_id, bool wait_flag) {
|
||||
void GetOutputEPStatic(const paddle::Tensor& x, int64_t rank_id, bool wait_flag) {
|
||||
GetOutputEp(x, rank_id, wait_flag, 1);
|
||||
}
|
||||
|
||||
void GetOutputDynamic(const paddle::Tensor& x,
|
||||
void GetOutputEPDynamic(const paddle::Tensor& x,
|
||||
int64_t rank_id,
|
||||
bool wait_flag,
|
||||
int msg_queue_id) {
|
||||
@@ -125,11 +125,11 @@ PD_BUILD_STATIC_OP(get_output_ep)
|
||||
.Attrs({"rank_id: int64_t", "wait_flag: bool"})
|
||||
.Outputs({"x_out"})
|
||||
.SetInplaceMap({{"x", "x_out"}})
|
||||
.SetKernelFn(PD_KERNEL(GetOutputStatic));
|
||||
.SetKernelFn(PD_KERNEL(GetOutputEPStatic));
|
||||
|
||||
PD_BUILD_STATIC_OP(get_output_ep_dynamic)
|
||||
.Inputs({"x"})
|
||||
.Attrs({"rank_id: int64_t", "wait_flag: bool", "msg_queue_id: int"})
|
||||
.Outputs({"x_out"})
|
||||
.SetInplaceMap({{"x", "x_out"}})
|
||||
.SetKernelFn(PD_KERNEL(GetOutputDynamic));
|
||||
.SetKernelFn(PD_KERNEL(GetOutputEPDynamic));
|
||||
|
@@ -19,7 +19,7 @@
|
||||
std::vector<paddle::Tensor> TopPSamplingReject(const paddle::Tensor &probs,
|
||||
const paddle::Tensor &top_p,
|
||||
const paddle::optional<paddle::Tensor> &top_k,
|
||||
int seed) {
|
||||
int64_t seed) {
|
||||
std::vector<int64_t> probs_shape = probs.shape();
|
||||
unsigned int batch_size = probs_shape[0];
|
||||
unsigned int vocab_size = probs_shape[1];
|
||||
@@ -82,7 +82,7 @@ TopPSamplingRejectInferDtype(const paddle::DataType &probs_dtype,
|
||||
PD_BUILD_STATIC_OP(rejection_top_p_sampling)
|
||||
.Inputs({"probs", "top_p", paddle::Optional("top_k")})
|
||||
.Outputs({"samples"})
|
||||
.Attrs({"seed: int"})
|
||||
.Attrs({"seed: int64_t"})
|
||||
.SetKernelFn(PD_KERNEL(TopPSamplingReject))
|
||||
.SetInferShapeFn(PD_INFER_SHAPE(TopPSamplingRejectInferShape))
|
||||
.SetInferDtypeFn(PD_INFER_DTYPE(TopPSamplingRejectInferDtype));
|
||||
|
Reference in New Issue
Block a user