From 36cc3ffdc705df06d007786fdd979f2e1493a07a Mon Sep 17 00:00:00 2001 From: Chunyuan WU Date: Thu, 3 Jul 2025 13:39:24 +0800 Subject: [PATCH] [CPU] [sgl-kernel] set dispatch key of initialize to CatchAll (#7734) --- sgl-kernel/csrc/cpu/torch_extension_cpu.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp index 4dc4f704f..44257dec5 100644 --- a/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp +++ b/sgl-kernel/csrc/cpu/torch_extension_cpu.cpp @@ -342,7 +342,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { // all reduce m.def("initialize(int size, int rank) -> ()"); - m.impl("initialize", torch::kCPU, &initialize); m.def("shm_allreduce(Tensor data, int reduce_op) -> ()"); m.impl("shm_allreduce", torch::kCPU, &shm_allreduce); m.def("shm_allgather(Tensor data, int dim) -> Tensor"); @@ -360,6 +359,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { TORCH_LIBRARY_IMPL(sgl_kernel, CatchAll, m) { m.impl("init_cpu_threads_env", init_cpu_threads_env); + m.impl("initialize", &initialize); } REGISTER_EXTENSION(common_ops)