[extension] hotfix jit extension setup (#5402)

95c21e39 · Hongxin Liu · GitHub · 5d380a1a · 95c21e39
Unverified Commit 95c21e39 authored Feb 26, 2024 by Hongxin Liu Committed by GitHub Feb 26, 2024
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 5 deletions

setup.py setup.py +3 -5

No files found.
--- a/setup.py
+++ b/setup.py
@@ -121,8 +121,9 @@ setup(
            "tests",
            "scripts",
            "requirements",
+            "extensions",
            "*.egg-info",
-        )
+        ),
    ),
    description="An integrated large-scale model training system with efficient parallelization techniques",
    long_description=fetch_readme(),
@@ -153,10 +154,7 @@ setup(
    ],
    package_data={
        "colossalai": [
-            "_C/*.pyi",
-            "kernel/cuda_native/csrc/*",
-            "kernel/cuda_native/csrc/kernel/*",
-            "kernel/cuda_native/csrc/kernels/include/*",
+            "kernel/extensions/csrc/**/*",
        ]
    },
 )