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

[CPU] add c++ kernel to bind CPU cores and memory node (#7524)

Browse Source
This commit is contained in:
Chunyuan WU
2025-06-30 10:45:25 +08:00
committed by GitHub
parent 78700893ee
commit c5131f7a2f
7 changed files with 159 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
View File

@@ -227,6 +227,9 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding_cpu(
at::Tensor& cos_sin_cache,
bool is_neox);
// CPU and memory binding
std::string init_cpu_threads_env(const std::string& cpu_ids);
TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
// activation
m.def("silu_and_mul_cpu(Tensor input) -> Tensor");
@@ -353,6 +356,13 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
"rotary_embedding_cpu(Tensor positions, Tensor query, Tensor key, int head_size, Tensor cos_sin_cache, "
"bool is_neox) -> (Tensor, Tensor)");
m.impl("rotary_embedding_cpu", torch::kCPU, &rotary_embedding_cpu);
// CPU and memory binding
m.def("init_cpu_threads_env(str cpu_ids) -> str");
}
TORCH_LIBRARY_IMPL(sgl_kernel, CatchAll, m) {
m.impl("init_cpu_threads_env", init_cpu_threads_env);
}
REGISTER_EXTENSION(common_ops)