Feat: Enable injection of extra include paths from env variable.

Author: kpoeppel

README.md

@@ -35,6 +35,21 @@ Your hardware needs to support CUDA Compute Capability $8.0$ or later. Make sure
 conda env create -n flashrnn -f environment_pt240cu124.yaml
 ```
 
+For all kinds of custom setups with torch and CUDA, keep in mind that versions have to match. Also, to make sure the correct CUDA libraries are included you can use the "FLASHRNN_EXTRA_INCLUDE_PATHS" environment variable now to inject different include paths, e.g.:
+
+```bash
+export FLASHRNN_EXTRA_INCLUDE_PATHS='/usr/local/include/cuda/:/usr/include/cuda/'
+```
+
+or within python:
+
+```python
+import os
+os.environ['FLASHRNN_EXTRA_INCLUDE_PATHS']='/usr/local/include/cuda/:/usr/include/cuda/'
+```
+
+
+
 ## Using FlashRNN
 
 FlashRNN employs a functional structure, none of the parameters are tied to the `flashrnn` function. To apply it simply use:

flashrnn/flashrnn/cuda_init.py

@@ -34,25 +34,33 @@ def defines_to_cflags(
 
 if torch.cuda.is_available():
     from packaging import version
+
     if version.parse(torch.__version__) >= version.parse("2.6.0"):
-        os.environ["CUDA_LIB"] = os.path.join(os.path.split(torch.utils.cpp_extension.include_paths(device_type="cuda")[-1])[0], "lib")
+        os.environ["CUDA_LIB"] = os.path.join(
+            os.path.split(torch.utils.cpp_extension.include_paths(device_type="cuda")[-1])[0], "lib"
+        )
     else:
-        os.environ["CUDA_LIB"] = os.path.join(os.path.split(torch.utils.cpp_extension.include_paths(cuda=True)[-1])[0], "lib")
+        os.environ["CUDA_LIB"] = os.path.join(
+            os.path.split(torch.utils.cpp_extension.include_paths(cuda=True)[-1])[0], "lib"
+        )
 
 
-EXTRA_INCLUDE_PATHS = ()
+EXTRA_INCLUDE_PATHS = () + (
+    tuple(os.environ["FLASHRNN_EXTRA_INCLUDE_PATHS"].split(":")) if "FLASHRNN_EXTRA_INCLUDE_PATHS" in os.environ else ()
+)
 if "CONDA_PREFIX" in os.environ:
     # This enforces adding the correct include directory from the CUDA installation via torch. If you use the system
     # installation, you might have to add the cflags yourself.
     from pathlib import Path
     from packaging import version
-    import sys
     import glob
+
     if version.parse(torch.__version__) >= version.parse("2.6.0"):
         matching_dirs = glob.glob(f"{os.environ['CONDA_PREFIX']}/targets/**", recursive=True)
-        EXTRA_INCLUDE_PATHS = tuple(map(str, (Path(os.environ["CONDA_PREFIX"]) / "targets").glob("**/include/")))[:1]
-
-
+        EXTRA_INCLUDE_PATHS = (
+            EXTRA_INCLUDE_PATHS
+            + tuple(map(str, (Path(os.environ["CONDA_PREFIX"]) / "targets").glob("**/include/")))[:1]
+        )
 
 
 def load(*, name, sources, extra_cflags=(), extra_cuda_cflags=(), **kwargs):