File size: 2,755 Bytes
57c3a10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | diff --git a/flashinfer/aot.py b/flashinfer/aot.py
index 2e2885a..981879f 100644
--- a/flashinfer/aot.py
+++ b/flashinfer/aot.py
@@ -378,15 +378,29 @@ def gen_all_modules(
jit_specs.append(gen_vllm_comm_module())
if add_misc:
- jit_specs += [
+ misc_modules = [
gen_cascade_module(),
gen_norm_module(),
- gen_nvshmem_module(),
+ ]
+
+ # Only add nvshmem module if nvidia.nvshmem is available or env paths are set
+ if (os.environ.get("NVSHMEM_INCLUDE_PATH") or
+ os.environ.get("NVSHMEM_LIBRARY_PATH")):
+ misc_modules.append(gen_nvshmem_module())
+ else:
+ try:
+ import nvidia.nvshmem
+ misc_modules.append(gen_nvshmem_module())
+ except ImportError:
+ pass # Skip nvshmem module if not available
+
+ misc_modules += [
gen_page_module(),
gen_quantization_module(),
gen_rope_module(),
gen_sampling_module(),
]
+ jit_specs += misc_modules
if has_sm90:
jit_specs.append(get_trtllm_utils_spec())
@@ -657,12 +671,12 @@ def main():
)
print("Total ops:", len(jit_specs))
- # Build
- build_jit_specs(jit_specs, verbose=True, skip_prebuilt=False)
+ # # Build
+ # build_jit_specs(jit_specs, verbose=True, skip_prebuilt=False)
- # Copy built kernels
- copy_built_kernels(jit_specs, out_dir)
- print("AOT kernels saved to:", out_dir)
+ # # Copy built kernels
+ # copy_built_kernels(jit_specs, out_dir)
+ # print("AOT kernels saved to:", out_dir)
if __name__ == "__main__":
diff --git a/flashinfer/jit/env.py b/flashinfer/jit/env.py
index 24104ee..5732a8f 100644
--- a/flashinfer/jit/env.py
+++ b/flashinfer/jit/env.py
@@ -73,10 +73,12 @@ def get_nvshmem_include_dirs():
if paths is not None:
return [pathlib.Path(p) for p in paths.split(os.pathsep) if p]
- import nvidia.nvshmem
-
- path = pathlib.Path(nvidia.nvshmem.__path__[0]) / "include"
- return [path]
+ try:
+ import nvidia.nvshmem
+ path = pathlib.Path(nvidia.nvshmem.__path__[0]) / "include"
+ return [path]
+ except ImportError:
+ return []
def get_nvshmem_lib_dirs():
@@ -84,7 +86,9 @@ def get_nvshmem_lib_dirs():
if paths is not None:
return [pathlib.Path(p) for p in paths.split(os.pathsep) if p]
- import nvidia.nvshmem
-
- path = pathlib.Path(nvidia.nvshmem.__path__[0]) / "lib"
- return [path]
+ try:
+ import nvidia.nvshmem
+ path = pathlib.Path(nvidia.nvshmem.__path__[0]) / "lib"
+ return [path]
+ except ImportError:
+ return []
|