[AMD] Fix rocm sgl-kernel missing modules error (#4311)

Co-authored-by: yiakwy-xpu-ml-framework-team <leiwang2@amd.com>
2025-03-12 01:35:28 +08:00
parent 00f42707ea
commit 690e1f2371
3 changed files with 3 additions and 3 deletions
--- a/sgl-kernel/csrc/torch_extension_rocm.cc
+++ b/sgl-kernel/csrc/torch_extension_rocm.cc
@@ -63,4 +63,4 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
  m.impl("moe_align_block_size", torch::kCUDA, &moe_align_block_size);
 }

-REGISTER_EXTENSION(_kernels)
+REGISTER_EXTENSION(common_ops)
--- a/sgl-kernel/setup_rocm.py
+++ b/sgl-kernel/setup_rocm.py
@@ -79,7 +79,7 @@ ext_modules = [
 setup(
    name="sgl-kernel",
    version=_get_version(),
-    packages=find_packages(),
+    packages=find_packages(where="python"),
    package_dir={"": "python"},
    ext_modules=ext_modules,
    cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)},