generated.patch

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 []