EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

fix bmm fp8 (#4926)

Browse Source
This commit is contained in:
Yineng Zhang
2025-03-30 12:15:20 -07:00
committed by GitHub
parent 9fccda3111
commit 195a09f57c
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
sgl-kernel/csrc/torch_extension.cc
View File

@@ -82,7 +82,10 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
/*
* From FlashInfer
*/
m.def("bmm_fp8", bmm_fp8);
m.def(
"bmm_fp8(Tensor A, Tensor B, Tensor! D, Tensor A_scale, Tensor B_scale, Tensor workspace_buffer, int "
"cublas_handle, int cuda_stream) -> ()");
m.impl("bmm_fp8", torch::kCUDA, &bmm_fp8);
m.def("min_p_sampling_from_probs", min_p_sampling_from_probs);
m.def("top_k_renorm_probs", top_k_renorm_probs);
m.def("top_p_renorm_probs", top_p_renorm_probs);