feat: use warp reduce as a simple example (#2304)

2024-12-01 22:43:50 +08:00
parent 5c18a03733
commit 47eb139f81
7 changed files with 193 additions and 17 deletions
--- a/sgl-kernel/setup.py
+++ b/sgl-kernel/setup.py
@@ -0,0 +1,20 @@
+from setuptools import find_packages, setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(
+    name="sgl-kernel",
+    version="0.0.2",
+    packages=find_packages(where="src"),
+    package_dir={"": "src"},
+    ext_modules=[
+        CUDAExtension(
+            "sgl_kernel.ops.warp_reduce_cuda",
+            [
+                "src/sgl-kernel/csrc/warp_reduce.cc",
+                "src/sgl-kernel/csrc/warp_reduce_kernel.cu",
+            ],
+        )
+    ],
+    cmdclass={"build_ext": BuildExtension},
+    install_requires=["torch"],
+)