From 95c21e3950a86ad2f00e1b89690772bd87ae53ca Mon Sep 17 00:00:00 2001
From: Hongxin Liu <lhx0217@gmail.com>
Date: Mon, 26 Feb 2024 19:46:58 +0800
Subject: [PATCH] [extension] hotfix jit extension setup (#5402)

---
 setup.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/setup.py b/setup.py
index 1244bfff0..5f6d447dd 100644
--- 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/**/*",
         ]
     },
 )