[MUSA]: Add flashinfer sampling backend (#24978)
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: zhiguo.qin <zhiguo.qin@mthreads.com> Co-authored-by: R0CKSTAR <yeahdongcn@gmail.com>
This commit is contained in:
co-authored by
github-actions[bot]
zhiguo.qin
R0CKSTAR
parent
d9fa84b25b
commit
897587b03a
@@ -278,6 +278,16 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
|
||||
m.def("top_p_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_p_arr, float top_p_val) -> ()");
|
||||
m.impl("top_p_renorm_probs", torch::kMUSA, &top_p_renorm_probs);
|
||||
|
||||
m.def(
|
||||
"min_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? maybe_min_p_arr, float "
|
||||
"min_p_val, bool deterministic, Generator? gen) -> ()");
|
||||
m.impl("min_p_sampling_from_probs", torch::kMUSA, &min_p_sampling_from_probs);
|
||||
|
||||
m.def(
|
||||
"top_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? maybe_top_p_arr, "
|
||||
"float top_p_val, bool deterministic, Generator? gen) -> ()");
|
||||
m.impl("top_p_sampling_from_probs", torch::kMUSA, &top_p_sampling_from_probs);
|
||||
|
||||
/*
|
||||
* From csrc/musa
|
||||
*/
|
||||
|
||||
Reference in New Issue
Block a user