Add cutlass submodule for sgl-kernel (#2676)

2024-12-31 14:28:29 +08:00
parent 339c69a243
commit b4403985d0
4 changed files with 14 additions and 0 deletions
--- a/sgl-kernel/setup.py
+++ b/sgl-kernel/setup.py
@@ -58,6 +58,11 @@ def update_wheel_platform_tag():
    old_wheel.rename(new_wheel)


+cutlass = root / "3rdparty" / "cutlass"
+include_dirs = [
+    cutlass.resolve() / "include",
+    cutlass.resolve() / "tools" / "util" / "include",
+]
 nvcc_flags = [
    "-O3",
    "-Xcompiler",
@@ -82,6 +87,7 @@ ext_modules = [
            "src/sgl-kernel/csrc/moe_align_kernel.cu",
            "src/sgl-kernel/csrc/sgl_kernel_ops.cu",
        ],
+        include_dirs=include_dirs,
        extra_compile_args={
            "nvcc": nvcc_flags,
            "cxx": cxx_flags,