skip importing torch and _kernels when building docs

docs/conf.py

@@ -11,6 +11,7 @@
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
 sys.path.insert(0, os.path.abspath("../python"))
+os.environ["BUILD_DOC"] = "1"
 
 import flashinfer
 

python/flashinfer/cascade.py

@@ -16,12 +16,19 @@
 import math
 from typing import Optional
 
-import torch
-
 try:
+    import torch
     from . import _kernels
-except ImportError:
-    _kernels = None
+except ImportError as e:
+    import os
+    import logging
+    if os.environ.get("BUILD_DOC", "0") == "1":
+        _kernels = None
+        logging.warning(
+            "Kernels are not loaded in documentation build mode."
+        )
+    else:
+        raise e
 
 from .decode import (
     batch_decode_with_padded_kv_cache_return_lse,

python/flashinfer/decode.py

@@ -16,12 +16,20 @@
 import math
 from typing import Optional, Union
 
-import torch
-
 try:
+    import torch
     from . import _kernels
-except ImportError:
-    _kernels = None
+except ImportError as e:
+    import os
+    import logging
+    if os.environ.get("BUILD_DOC", "0") == "1":
+        _kernels = None
+        logging.warning(
+            "Kernels are not loaded in documentation build mode."
+        )
+    else:
+        raise e
+
 
 from .utils import (
     RotaryMode,

python/flashinfer/page.py

@@ -16,11 +16,18 @@
 import torch
 
 try:
+    import torch
     from . import _kernels
-except ImportError:
-    _kernels = None
-from .utils import TensorLayout, check_kv_layout
-
+except ImportError as e:
+    import os
+    import logging
+    if os.environ.get("BUILD_DOC", "0") == "1":
+        _kernels = None
+        logging.warning(
+            "Kernels are not loaded in documentation build mode."
+        )
+    else:
+        raise e
 
 def append_paged_kv_cache(
     append_key: torch.Tensor,

python/flashinfer/prefill.py

@@ -16,12 +16,19 @@
 import math
 from typing import Optional
 
-import torch
-
 try:
+    import torch
     from . import _kernels
-except ImportError:
-    _kernels = None
+except ImportError as e:
+    import os
+    import logging
+    if os.environ.get("BUILD_DOC", "0") == "1":
+        _kernels = None
+        logging.warning(
+            "Kernels are not loaded in documentation build mode."
+        )
+    else:
+        raise e
 
 from .utils import (
     RotaryMode,