add unittests without jit

Author: yzh119

flashinfer/aot.py

@@ -525,9 +525,6 @@ def gen_all_modules(
     # Add cuDNN FMHA module
     jit_specs.append(gen_cudnn_fmha_module())
 
-    # NOTE(Zihao): just for debugging, remove later
-    jit_specs = [gen_spdlog_module()]
-
     # dedup
     names = set()
     ret: List[JitSpec] = []

flashinfer/jit/cpp_ext.py

@@ -263,6 +263,10 @@ def _get_num_workers() -> Optional[int]:
 
 
 def run_ninja(workdir: Path, ninja_file: Path, verbose: bool) -> None:
+    if os.environ.get("FLASHINFER_DISABLE_JIT"):
+        raise RuntimeError(
+            "JIT compilation is disabled via FLASHINFER_DISABLE_JIT environment variable"
+        )
     workdir.mkdir(parents=True, exist_ok=True)
     command = [
         "ninja",

scripts/task_test_nightly_build.sh

@@ -23,6 +23,11 @@ pip install ${DIST_CUBIN_DIR}/*.whl
 echo "Installing flashinfer-jit-cache from ${DIST_JIT_CACHE_DIR}..."
 pip install ${DIST_JIT_CACHE_DIR}/*.whl
 
+# Disable JIT to verify that jit-cache package contains all necessary
+# precompiled modules for the test suite to pass without compilation
+echo "Disabling JIT compilation to test with precompiled cache only..."
+export FLASHINFER_DISABLE_JIT=1
+
 echo "Installing flashinfer-python from ${DIST_PYTHON_DIR}..."
 pip install ${DIST_PYTHON_DIR}/*.tar.gz