build: add sccache-backed jit-cache builds and AOT diagnostics

.github/workflows/nightly-release.yml

@@ -176,6 +176,11 @@ jobs:
             -e FLASHINFER_DEV_RELEASE_SUFFIX="${FLASHINFER_DEV_RELEASE_SUFFIX}" \
             -e ARCH="${{ matrix.arch }}" \
             -e FLASHINFER_CUDA_ARCH_LIST="${FLASHINFER_CUDA_ARCH_LIST}" \
+            -e FLASHINFER_NVCC_THREADS="${FLASHINFER_NVCC_THREADS:-}" \
+            -e SCCACHE_BUCKET="${{ vars.SCCACHE_BUCKET }}" \
+            -e AWS_ACCESS_KEY_ID="${{ secrets.SCCACHE_AWS_ACCESS_KEY_ID }}" \
+            -e AWS_SECRET_ACCESS_KEY="${{ secrets.SCCACHE_AWS_SECRET_ACCESS_KEY }}" \
+            -e SCCACHE_REGION="${{ vars.SCCACHE_REGION }}" \
             -w /workspace \
             ${{ env.DOCKER_IMAGE }} \
             bash /workspace/scripts/build_flashinfer_jit_cache_whl.sh

.github/workflows/pr-test.yml

@@ -233,7 +233,18 @@ jobs:
           BUILD_NUMBER: ${{ github.run_number }}
 
       - name: Run Test
-        run: bash ci/bash.sh ${DOCKER_IMAGE} --no-gpu ./scripts/task_test_jit_cache_package_build_import.sh
+        env:
+          SCCACHE_BUCKET: ${{ vars.SCCACHE_BUCKET }}
+          SCCACHE_REGION: ${{ vars.SCCACHE_REGION }}
+          AWS_ACCESS_KEY_ID: ${{ secrets.SCCACHE_AWS_ACCESS_KEY_ID }}
+          AWS_SECRET_ACCESS_KEY: ${{ secrets.SCCACHE_AWS_SECRET_ACCESS_KEY }}
+        run: |
+          bash ci/bash.sh ${DOCKER_IMAGE} \
+            -e SCCACHE_BUCKET "${SCCACHE_BUCKET}" \
+            -e SCCACHE_REGION "${SCCACHE_REGION}" \
+            -e AWS_ACCESS_KEY_ID "${AWS_ACCESS_KEY_ID}" \
+            -e AWS_SECRET_ACCESS_KEY "${AWS_SECRET_ACCESS_KEY}" \
+            --no-gpu ./scripts/task_test_jit_cache_package_build_import.sh
 
   analyze-aot-failure:
     name: Analyze AOT Failure
@@ -314,7 +325,18 @@ jobs:
           BUILD_NUMBER: ${{ github.run_number }}
 
       - name: Run Test
-        run: bash ci/bash.sh ${DOCKER_IMAGE} --no-gpu ./scripts/task_test_jit_cache_package_build_import.sh
+        env:
+          SCCACHE_BUCKET: ${{ vars.SCCACHE_BUCKET }}
+          SCCACHE_REGION: ${{ vars.SCCACHE_REGION }}
+          AWS_ACCESS_KEY_ID: ${{ secrets.SCCACHE_AWS_ACCESS_KEY_ID }}
+          AWS_SECRET_ACCESS_KEY: ${{ secrets.SCCACHE_AWS_SECRET_ACCESS_KEY }}
+        run: |
+          bash ci/bash.sh ${DOCKER_IMAGE} \
+            -e SCCACHE_BUCKET "${SCCACHE_BUCKET}" \
+            -e SCCACHE_REGION "${SCCACHE_REGION}" \
+            -e AWS_ACCESS_KEY_ID "${AWS_ACCESS_KEY_ID}" \
+            -e AWS_SECRET_ACCESS_KEY "${AWS_SECRET_ACCESS_KEY}" \
+            --no-gpu ./scripts/task_test_jit_cache_package_build_import.sh
 
   # ---------------------------------------------------------------------------
   # GPU JIT Tests - SM86 (A10G) - Spot + On-Demand Rerun

.github/workflows/release.yml

@@ -203,6 +203,11 @@ jobs:
             -e FLASHINFER_LOCAL_VERSION="$FLASHINFER_LOCAL_VERSION" \
             -e ARCH="${{ matrix.arch }}" \
             -e FLASHINFER_CUDA_ARCH_LIST="${FLASHINFER_CUDA_ARCH_LIST}" \
+            -e FLASHINFER_NVCC_THREADS="${FLASHINFER_NVCC_THREADS:-}" \
+            -e SCCACHE_BUCKET="${{ vars.SCCACHE_BUCKET }}" \
+            -e AWS_ACCESS_KEY_ID="${{ secrets.SCCACHE_AWS_ACCESS_KEY_ID }}" \
+            -e AWS_SECRET_ACCESS_KEY="${{ secrets.SCCACHE_AWS_SECRET_ACCESS_KEY }}" \
+            -e SCCACHE_REGION="${{ vars.SCCACHE_REGION }}" \
             -w /workspace \
             ${{ env.DOCKER_IMAGE }} \
             bash /workspace/scripts/build_flashinfer_jit_cache_whl.sh

flashinfer/jit/core.py

@@ -12,7 +12,7 @@
 
 from ..compilation_context import CompilationContext
 from . import env as jit_env
-from .cpp_ext import generate_ninja_build_for_op, run_ninja
+from .cpp_ext import generate_ninja_build_for_op, get_nvcc_parallelism_flags, run_ninja
 from .utils import write_if_different
 
 os.makedirs(jit_env.FLASHINFER_WORKSPACE_DIR, exist_ok=True)
@@ -298,9 +298,7 @@ def build(self, verbose: bool, need_lock: bool = True) -> None:
             FileLock(self.lock_path, thread_local=False) if need_lock else nullcontext()
         )
         with lock:
-            # Write ninja file if it doesn't exist (deferred case)
-            if not self.is_ninja_generated:
-                self.write_ninja()
+            self.write_ninja()
             run_ninja(self.build_dir, self.ninja_path, verbose)
 
     def load(self, so_path: Path):
@@ -431,7 +429,7 @@ def gen_jit_spec(
         cflags.insert(0, "-std=c++17")
 
     cuda_cflags = [
-        f"--threads={os.environ.get('FLASHINFER_NVCC_THREADS', '1')}",
+        *get_nvcc_parallelism_flags(),
         "-use_fast_math",
         "-DFLASHINFER_ENABLE_F16",
         "-DFLASHINFER_ENABLE_BF16",
@@ -446,7 +444,7 @@ def gen_jit_spec(
         cuda_cflags += [
             "-g",
             "-O0",
-            "-G",
+            "--device-debug",
             "-lineinfo",
             "--ptxas-options=-v",
             "-DCUTLASS_DEBUG_TRACE_LEVEL=2",
@@ -503,9 +501,8 @@ def build_jit_specs(
         if skip_prebuilt and spec.aot_path.exists():
             continue
         lines.append(f"subninja {spec.ninja_path}")
-        if not spec.is_ninja_generated:
-            with FileLock(spec.lock_path, thread_local=False):
-                spec.write_ninja()
+        with FileLock(spec.lock_path, thread_local=False):
+            spec.write_ninja()
     if not lines:
         return
 

flashinfer/jit/cpp_ext.py

@@ -4,6 +4,8 @@
 import logging
 import os
 import re
+import shlex
+import stat
 import subprocess
 import sys
 import sysconfig
@@ -91,6 +93,113 @@ def is_cuda_version_at_least(version_str: str) -> bool:
     return get_cuda_version() >= Version(version_str)
 
 
+def _get_positive_int_env(env_var_name: str, default: int) -> int:
+    value = os.environ.get(env_var_name)
+    if value is None:
+        return default
+
+    try:
+        int_value = int(value)
+    except ValueError:
+        logger.warning(
+            "Ignoring invalid %s=%r; using %s.", env_var_name, value, default
+        )
+        return default
+
+    if int_value < 1:
+        logger.warning("Ignoring %s=%r; value must be >= 1.", env_var_name, value)
+        return default
+
+    return int_value
+
+
+def get_nvcc_threads() -> int:
+    """Return the requested number of nvcc worker threads."""
+    return _get_positive_int_env("FLASHINFER_NVCC_THREADS", 1)
+
+
+def _extract_make_jobserver_auth(makeflags: Optional[str] = None) -> Optional[str]:
+    if makeflags is None:
+        makeflags = os.environ.get("MAKEFLAGS")
+    if not makeflags:
+        return None
+
+    matches = re.findall(r"(?:^|\s)--jobserver-(?:auth|fds)=([^\s]+)", makeflags)
+    if not matches:
+        return None
+    return matches[-1]
+
+
+def _has_fifo_make_jobserver(makeflags: Optional[str] = None) -> bool:
+    auth = _extract_make_jobserver_auth(makeflags)
+    if not auth or not auth.startswith("fifo:"):
+        return False
+
+    fifo_path = auth[len("fifo:") :]
+    if not fifo_path:
+        return False
+
+    try:
+        return stat.S_ISFIFO(os.stat(fifo_path).st_mode)
+    except OSError:
+        return False
+
+
+@functools.cache
+def _ninja_supports_make_jobserver() -> bool:
+    try:
+        output = subprocess.check_output(["ninja", "--version"], text=True).strip()
+    except (FileNotFoundError, subprocess.CalledProcessError):
+        return False
+
+    match = re.search(r"\d+(?:\.\d+)+", output)
+    if not match:
+        return False
+
+    return Version(match.group(0)) >= Version("1.13")
+
+
+def should_use_ninja_jobserver() -> bool:
+    """Return whether Ninja should inherit a GNU Make FIFO jobserver."""
+    return _has_fifo_make_jobserver() and _ninja_supports_make_jobserver()
+
+
+def _launcher_uses_sccache(launcher: Optional[str] = None) -> bool:
+    if launcher is None:
+        launcher = os.environ.get("FLASHINFER_NVCC_LAUNCHER", "")
+    if not launcher:
+        return False
+
+    try:
+        launcher_parts = shlex.split(launcher)
+    except ValueError:
+        launcher_parts = launcher.split()
+    if not launcher_parts:
+        return False
+
+    return Path(launcher_parts[0]).name == "sccache"
+
+
+def should_use_nvcc_jobserver(cuda_version: Optional[Version] = None) -> bool:
+    """Return whether nvcc can safely acquire slots from the inherited jobserver."""
+    if not sys.platform.startswith("linux"):
+        return False
+    if cuda_version is None:
+        cuda_version = get_cuda_version()
+    return cuda_version >= Version("13.0") and (
+        should_use_ninja_jobserver() or _launcher_uses_sccache()
+    )
+
+
+def get_nvcc_parallelism_flags(cuda_version: Optional[Version] = None) -> List[str]:
+    """Build nvcc flags controlled by FlashInfer parallelism environment variables."""
+    threads = get_nvcc_threads()
+    flags = [f"--threads={threads}"]
+    if threads > 1 and should_use_nvcc_jobserver(cuda_version):
+        flags.append("--jobserver")
+    return flags
+
+
 def join_multiline(vs: List[str]) -> str:
     return " $\n    ".join(vs)
 
@@ -245,13 +354,18 @@ def generate_ninja_build_for_op(
 
     cxx = os.environ.get("CXX", "c++")
     nvcc = os.environ.get("FLASHINFER_NVCC", "$cuda_home/bin/nvcc")
+    # Compiler launchers (e.g., sccache, ccache) — empty string when unset
+    cxx_launcher = os.environ.get("FLASHINFER_CXX_LAUNCHER", "")
+    nvcc_launcher = os.environ.get("FLASHINFER_NVCC_LAUNCHER", "")
 
     lines = [
         "ninja_required_version = 1.3",
         f"name = {name}",
         f"cuda_home = {cuda_home}",
         f"cxx = {cxx}",
         f"nvcc = {nvcc}",
+        f"cxx_launcher = {cxx_launcher}",
+        f"nvcc_launcher = {nvcc_launcher}",
         "",
         "common_cflags = " + join_multiline(common_cflags),
         "cflags = " + join_multiline(cflags),
@@ -261,12 +375,12 @@ def generate_ninja_build_for_op(
         "ldflags = " + join_multiline(ldflags),
         "",
         "rule compile",
-        "  command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags",
+        "  command = $cxx_launcher $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags",
         "  depfile = $out.d",
         "  deps = gcc",
         "",
         "rule cuda_compile",
-        "  command = $nvcc --generate-dependencies-with-compile --dependency-output $out.d $cuda_cflags -c $in -o $out $cuda_post_cflags",
+        "  command = $nvcc_launcher $nvcc --generate-dependencies-with-compile -MF $out.d $cuda_cflags -c $in -o $out $cuda_post_cflags",
         "  depfile = $out.d",
         "  deps = gcc",
         "",
@@ -333,7 +447,7 @@ def run_ninja(workdir: Path, ninja_file: Path, verbose: bool) -> None:
         str(ninja_file.resolve()),
     ]
     num_workers = _get_num_workers()
-    if num_workers is not None:
+    if num_workers is not None and not should_use_ninja_jobserver():
         command += ["-j", str(num_workers)]
 
     sys.stdout.flush()

scripts/build_flashinfer_jit_cache_whl.sh

@@ -8,18 +8,32 @@ echo "=========================================="
 echo "Building flashinfer-jit-cache wheel"
 echo "=========================================="
 
-# MAX_JOBS = min(nproc, max(1, MemAvailable_GB/4))
+# Parallelism: coordinate ninja jobs (MAX_JOBS) with nvcc internal threads (FLASHINFER_NVCC_THREADS).
+# Each nvcc invocation compiles multiple -gencode targets; --threads=N parallelizes them.
 MEM_AVAILABLE_GB=$(free -g | awk '/^Mem:/ {print $7}')
 NPROC=$(nproc)
-# MAX_JOBS=$(( MEM_AVAILABLE_GB / $([ "$(uname -m)" = "aarch64" ] && echo 8 || echo 4) ))
-MAX_JOBS=$(( MEM_AVAILABLE_GB / 8 ))
-if (( MAX_JOBS < 1 )); then
-  MAX_JOBS=1
-elif (( NPROC < MAX_JOBS )); then
-  MAX_JOBS=$NPROC
+
+# Set NVCC_THREADS to match number of gencode targets (capped at 8), unless already set
+NUM_ARCHS=$(echo "${FLASHINFER_CUDA_ARCH_LIST}" | wc -w)
+NVCC_THREADS=${FLASHINFER_NVCC_THREADS:-${NUM_ARCHS}}
+if (( NVCC_THREADS > 8 )); then NVCC_THREADS=8; fi
+if (( NVCC_THREADS > NPROC )); then NVCC_THREADS=${NPROC}; fi
+if (( NVCC_THREADS < 1 )); then NVCC_THREADS=1; fi
+
+# Memory budget: ~2GB per nvcc thread
+MEM_PER_JOB=$(( NVCC_THREADS * 2 ))
+MAX_JOBS=$(( MEM_AVAILABLE_GB / MEM_PER_JOB ))
+if (( MAX_JOBS < 1 )); then MAX_JOBS=1; fi
+
+# Cap total threads at available CPUs
+TOTAL_THREADS=$(( MAX_JOBS * NVCC_THREADS ))
+if (( TOTAL_THREADS > NPROC )); then
+  MAX_JOBS=$(( NPROC / NVCC_THREADS ))
+  if (( MAX_JOBS < 1 )); then MAX_JOBS=1; fi
 fi
 
 export MAX_JOBS
+export FLASHINFER_NVCC_THREADS="${NVCC_THREADS}"
 
 # Display build environment info
 echo "CUDA Version: ${CUDA_VERSION}"
@@ -30,6 +44,7 @@ echo "FlashInfer Local Version: ${FLASHINFER_LOCAL_VERSION}"
 echo "CUDA Architectures: ${FLASHINFER_CUDA_ARCH_LIST}"
 echo "Dev Release Suffix: ${FLASHINFER_DEV_RELEASE_SUFFIX}"
 echo "MAX_JOBS: ${MAX_JOBS}"
+echo "NVCC_THREADS: ${FLASHINFER_NVCC_THREADS}"
 echo "Python Version: $(python3 --version)"
 echo "Git commit: $(git rev-parse HEAD 2>/dev/null || echo 'unknown')"
 echo "Working directory: $(pwd)"
@@ -51,6 +66,32 @@ echo "::group::Install build system"
 pip install --upgrade build
 echo "::endgroup::"
 
+# Optional: set up sccache for compiler caching with S3 backend
+if [ -n "$SCCACHE_BUCKET" ]; then
+  echo "::group::Install sccache"
+  SCCACHE_VERSION="0.9.1"
+  SCCACHE_ARCH=$(uname -m)
+  curl -fsSL "https://github.com/mozilla/sccache/releases/download/v${SCCACHE_VERSION}/sccache-v${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz" | tar xz
+  mv "sccache-v${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache" /usr/local/bin/
+  rm -rf "sccache-v${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl"
+  chmod +x /usr/local/bin/sccache
+
+  # Namespace cache by CUDA version and CPU architecture
+  SCCACHE_SOURCE_ROOT=$(cd .. && pwd -P)
+  export SCCACHE_BASEDIRS="${SCCACHE_SOURCE_ROOT}${SCCACHE_BASEDIRS:+:${SCCACHE_BASEDIRS}}"
+  export SCCACHE_S3_KEY_PREFIX="cuda${CUDA_MAJOR}${CUDA_MINOR}-${SCCACHE_ARCH}"
+  export SCCACHE_IDLE_TIMEOUT=0
+  export FLASHINFER_NVCC_LAUNCHER="sccache"
+  export FLASHINFER_CXX_LAUNCHER="sccache"
+
+  sccache --start-server
+  echo "sccache version: $(sccache --version)"
+  echo "sccache bucket: ${SCCACHE_BUCKET}"
+  echo "sccache prefix: ${SCCACHE_S3_KEY_PREFIX}"
+  echo "sccache basedirs: ${SCCACHE_BASEDIRS}"
+  echo "::endgroup::"
+fi
+
 # Clean any previous builds
 echo "Cleaning previous builds..."
 rm -rf dist build *.egg-info
@@ -83,5 +124,12 @@ if [ -n "${OUTPUT_DIR}" ]; then
     cp -v dist/*.whl "${OUTPUT_DIR}/"
 fi
 
+# Print sccache stats if enabled
+if [ -n "$SCCACHE_BUCKET" ]; then
+  echo "::group::sccache stats"
+  sccache --show-stats
+  echo "::endgroup::"
+fi
+
 echo ""
 echo "Build process completed!"