Isolate CI torch extension cache

Author: Qubitium

.github/workflows/unit_tests.yml

@@ -244,8 +244,19 @@ jobs:
             --shell)"
           eval "$runtime_output"
 
+          test_cache_key="${{ matrix.test_script }}"
+          test_cache_key="${test_cache_key//\//_}"
+          test_cache_key="${test_cache_key//./_}"
+          test_cache_key="${test_cache_key//-/_}"
+          export GPTQMODEL_TORCH_EXTENSIONS_DIR="/tmp/gptqmodel/torch_extensions/${{ github.run_id }}/${{ github.run_attempt }}/${ENV_NAME}/${test_cache_key}"
+          export TORCH_EXTENSIONS_DIR="$GPTQMODEL_TORCH_EXTENSIONS_DIR"
+          mkdir -p "$GPTQMODEL_TORCH_EXTENSIONS_DIR"
+          echo "GPTQMODEL_TORCH_EXTENSIONS_DIR=$GPTQMODEL_TORCH_EXTENSIONS_DIR" >> "$GITHUB_ENV"
+          echo "TORCH_EXTENSIONS_DIR=$TORCH_EXTENSIONS_DIR" >> "$GITHUB_ENV"
+
           echo "-- setting up env --"
           echo "env_name: $ENV_NAME"
+          echo "torch extensions dir: $GPTQMODEL_TORCH_EXTENSIONS_DIR"
           # will clean later
 
           echo "-- ls -ahl /opt/uv/venvs before --"

gptqmodel/exllamav3/ext.py

@@ -11,7 +11,6 @@
 import os
 import sys
 from pathlib import Path
-from typing import Optional
 
 import torch
 
@@ -112,20 +111,6 @@ def _exllamav3_include_paths() -> list[str]:
     )
 
 
-def _legacy_build_root() -> Optional[Path]:
-    build_root = os.environ.get("GPTQMODEL_EXT_BUILD")
-    if not build_root:
-        return None
-    return Path(build_root) / extension_name
-
-
-def _default_build_root() -> Path:
-    legacy_root = _legacy_build_root()
-    if legacy_root is not None:
-        return legacy_root
-    return default_torch_ops_build_root("exllamav3")
-
-
 def _extra_cflags() -> list[str]:
     if windows:
         flags = ["/O2", "/std:c++17"]
@@ -184,7 +169,7 @@ def _prepare_build_env() -> None:
     ),
     sources=_source_files,
     build_root_env="GPTQMODEL_EXLLAMAV3_BUILD_ROOT",
-    default_build_root=_default_build_root,
+    default_build_root=lambda: default_torch_ops_build_root("exllamav3"),
     display_name="ExLlamaV3",
     extra_cflags=_extra_cflags,
     extra_cuda_cflags=_extra_cuda_cflags,

gptqmodel/utils/cpp.py

@@ -60,6 +60,7 @@
 # the best overall tradeoff across Marlin, AWQ, QQQ, ExLlama, and ParoQuant.
 _DEFAULT_NVCC_THREADS = "8"
 _GLOBAL_KERNEL_REBUILD_ENV = "GPTQMODEL_KERNEL_REBUILD"
+_TORCH_OPS_BUILD_ROOT_ENV = "GPTQMODEL_TORCH_EXTENSIONS_DIR"
 
 
 def _nvcc_path() -> Optional[str]:
@@ -271,6 +272,9 @@ def close(self, *, succeeded: bool, elapsed_seconds: Optional[float] = None) ->
 def default_torch_ops_build_root(subdir: str) -> Path:
     """Return the default on-disk cache root for torch.ops JIT extensions."""
 
+    override_root = os.getenv(_TORCH_OPS_BUILD_ROOT_ENV)
+    if override_root:
+        return Path(override_root).expanduser() / subdir
     return Path.home() / ".cache" / "gptqmodel" / "torch_extensions" / subdir
 
 

tests/test_torch_ops_jit_extension.py

@@ -226,6 +226,24 @@ def test_default_torch_ops_build_root_ignores_removed_global_override(monkeypatc
     )
 
 
+def test_default_torch_ops_build_root_respects_ci_override(monkeypatch):
+    monkeypatch.setenv("GPTQMODEL_TORCH_EXTENSIONS_DIR", "/tmp/gptqmodel-ci")
+
+    assert cpp_module.default_torch_ops_build_root("marlin") == Path("/tmp/gptqmodel-ci") / "marlin"
+
+
+def test_torch_ops_jit_extension_prefers_explicit_build_root_over_global_default(monkeypatch, tmp_path):
+    loader = _make_loader(
+        tmp_path,
+        default_build_root=lambda: cpp_module.default_torch_ops_build_root("unit_test_ops"),
+    )
+
+    monkeypatch.setenv("GPTQMODEL_TORCH_EXTENSIONS_DIR", "/tmp/gptqmodel-ci")
+    monkeypatch.setenv("UNIT_TEST_BUILD_ROOT", "/tmp/unit-test-override")
+
+    assert loader.base_build_root() == Path("/tmp/unit-test-override")
+
+
 def test_torch_ops_jit_extension_prefers_cached_binary(monkeypatch, tmp_path):
     """Guard cache reuse so startup skips expensive JIT rebuilds when ops are already built."""