[TRTLLM-16212][feat] Unified caching via TRTLLM_CACHE_DIR environment variable

docs/source/llm-api/index.md

@@ -130,6 +130,47 @@ The following tips typically assist new LLM API users who are familiar with othe
 
   This limitation is applicable for multi-GPU inference only.
 
+### Unified compiled-artifact cache
+
+Set `TRTLLM_CACHE_DIR` before importing TensorRT-LLM to place its compiled
+artifact caches under a single root. TensorRT-LLM supplies the following
+defaults; an individually configured environment variable always takes
+precedence.
+
+| Environment variable | Default under `TRTLLM_CACHE_DIR` |
+|----------------------|------------------------------------|
+| `TLLM_AUTOTUNER_CACHE_PATH` | `autotuner/cache.json` |
+| `TORCHINDUCTOR_CACHE_DIR` | `inductor` |
+| `TRITON_CACHE_DIR` | `triton` |
+| `TORCH_EXTENSIONS_DIR` | `torch_extensions` |
+| `FLASHINFER_WORKSPACE_BASE` | `flashinfer` |
+| `CUTE_DSL_CACHE_DIR` | `cute_dsl` |
+| `DG_JIT_CACHE_DIR` | `deep_gemm` |
+| `TRTLLM_DG_CACHE_DIR` | `trtllm_deep_gemm` |
+| `CUDA_CACHE_PATH` | `cuda` |
+
+The directories are created by their respective consumers when needed.
+TensorRT-LLM does not add cache locking, cleanup, or per-rank isolation for
+these unified defaults (with the exceptions described below). The
+configured values are forwarded to dynamically spawned MPI workers.
+
+Configuring `TRTLLM_DG_CACHE_DIR`, either directly or through
+`TRTLLM_CACHE_DIR`, also defaults `TRTLLM_DG_JIT_DUMP_CUBIN` to `1` so its
+NVRTC-generated cubins are persisted. An explicitly configured dump setting
+still takes precedence. See the
+[DeepGEMM environment variables](source:examples/models/core/deepseek_v3/README.md#deepgemm)
+for details.
+
+Note that, when `TRTLLM_CACHE_DIR` is set, the FlashInfer workspace isolation
+mechanism described below still applies. This will prevent race conditions
+during cache read/writes, at the cost of limiting cache sharing. To disable
+workspace isolation, either set `TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0`,
+or set `FLASHINFER_WORKSPACE_BASE` to a different path from the one assigned
+under `TRTLLM_CACHE_DIR`.
+
+Ray workers similarly isolate the unified `DG_JIT_CACHE_DIR` by rank and GPU.
+Set `TRTLLM_DEEP_GEMM_CACHE_PER_PROCESS=0` to disable this behavior.
+
 ### FlashInfer JIT workspaces for MPI workers
 
 `trtllm-llmapi-launch` ranks and dynamically spawned `MpiPoolSession` workers
@@ -151,7 +192,7 @@ process-unique temporary workspace that is removed, artifacts included, when
 the process exits.
 
 An explicitly configured `FLASHINFER_WORKSPACE_BASE` takes precedence in both
-launch modes. An explicitly configured `FLASHINFER_CUBIN_DIR` is propagated
+launch modes, while setting `TRTLLM_CACHE_DIR` does not. An explicitly configured `FLASHINFER_CUBIN_DIR` is propagated
 unchanged to every rank and is the way to share downloaded artifacts between
 ranks; populate it before the multi-rank launch, either by installing the
 `flashinfer-cubin` package matching `flashinfer-python` or with a single-rank

examples/models/core/deepseek_v3/README.md

@@ -697,6 +697,7 @@ DeepGEMM-related behavior can be controlled by the following environment variabl
 | `TRTLLM_DG_JIT_DEBUG` | When set to `1`, enable JIT debugging. |
 | `TRTLLM_DG_JIT_USE_NVCC` | When set to `1`, use NVCC instead of NVRTC to compile the kernel, which has slightly better performance but requires CUDA Toolkit (>=12.3) and longer compilation time.|
 | `TRTLLM_DG_JIT_DUMP_CUBIN` | When set to `1`, dump the cubin file. This is only effective with NVRTC since NVCC will always dump the cubin file. NVRTC-based JIT will store the generated kernels in memory by default. If you want to persist the kernels across multiple runs, you can either use this variable or use NVCC. |
+| `TRTLLM_DG_CACHE_DIR` | Sets the root directory for TensorRT-LLM DeepGEMM JIT artifacts. Temporary files are stored under `tmp/`, and compiled kernels are stored under `cache/` when `TRTLLM_DG_JIT_DUMP_CUBIN=1` or NVCC is used. |
 
 #### MOE GEMM Optimization
 

tensorrt_llm/__init__.py

@@ -23,8 +23,9 @@
 
 from ._bootstrap import _init, _prepare_environment
 
-# Phase 1: DLL search path, Python-library preload and vendored triton_kernels
-# precedence. Must run before torch and before any TensorRT-LLM shared object.
+# Phase 1: cache environment, DLL search path, Python-library preload and
+# vendored triton_kernels precedence. Must run before torch and before any
+# TensorRT-LLM shared object.
 _prepare_environment()
 
 # The package's public surface is loaded lazily (PEP 562): importing

tensorrt_llm/_bootstrap.py

@@ -16,10 +16,10 @@
 
 Two phases, in this order, both driven from ``tensorrt_llm/__init__.py``:
 
-1. :func:`_prepare_environment` -- process environment defaults, DLL search
-   path, Python-library preload and vendored ``triton_kernels`` precedence.  It
-   must run *before* ``torch`` and before any TensorRT-LLM shared object is
-   loaded.
+1. :func:`_prepare_environment` -- process environment and compiled-artifact
+    cache defaults, DLL search path, Python-library preload and vendored
+    ``triton_kernels`` precedence. It must run *before* ``torch`` and before
+    any TensorRT-LLM shared object is loaded.
 2. :func:`_init` -- custom-op library loading and MPI initialization.  It runs
    after the package's own imports have completed.
 
@@ -45,6 +45,48 @@
 
 _inited = False
 
+_UNIFIED_CACHE_ROOT_ENV = "TRTLLM_CACHE_DIR"
+_TRTLLM_DG_CACHE_ENV = "TRTLLM_DG_CACHE_DIR"
+_TRTLLM_DG_DUMP_CUBIN_ENV = "TRTLLM_DG_JIT_DUMP_CUBIN"
+_UNIFIED_CACHE_ENV_VARS = (
+    "TLLM_AUTOTUNER_CACHE_PATH",
+    "TORCHINDUCTOR_CACHE_DIR",
+    "TRITON_CACHE_DIR",
+    "TORCH_EXTENSIONS_DIR",
+    "FLASHINFER_WORKSPACE_BASE",
+    "CUTE_DSL_CACHE_DIR",
+    "DG_JIT_CACHE_DIR",
+    _TRTLLM_DG_CACHE_ENV,
+    "CUDA_CACHE_PATH",
+)
+
+
+def _unified_cache_defaults(cache_root: str) -> dict[str, str]:
+    """Return cache environment defaults rooted at ``cache_root``."""
+    cache_root = os.path.expanduser(cache_root)
+    return {
+        "TLLM_AUTOTUNER_CACHE_PATH": os.path.join(cache_root, "autotuner", "cache.json"),
+        "TORCHINDUCTOR_CACHE_DIR": os.path.join(cache_root, "inductor"),
+        "TRITON_CACHE_DIR": os.path.join(cache_root, "triton"),
+        "TORCH_EXTENSIONS_DIR": os.path.join(cache_root, "torch_extensions"),
+        "FLASHINFER_WORKSPACE_BASE": os.path.join(cache_root, "flashinfer"),
+        "CUTE_DSL_CACHE_DIR": os.path.join(cache_root, "cute_dsl"),
+        "DG_JIT_CACHE_DIR": os.path.join(cache_root, "deep_gemm"),
+        _TRTLLM_DG_CACHE_ENV: os.path.join(cache_root, "trtllm_deep_gemm"),
+        "CUDA_CACHE_PATH": os.path.join(cache_root, "cuda"),
+    }
+
+
+def _setup_unified_cache() -> None:
+    """Configure compiled-artifact cache defaults before framework imports."""
+    cache_root = os.environ.get(_UNIFIED_CACHE_ROOT_ENV)
+    if cache_root:
+        for name, path in _unified_cache_defaults(cache_root).items():
+            os.environ.setdefault(name, path)
+
+    if os.environ.get(_TRTLLM_DG_CACHE_ENV):
+        os.environ.setdefault(_TRTLLM_DG_DUMP_CUBIN_ENV, "1")
+
 
 def _set_numexpr_thread_default() -> None:
     # TRT-LLM uses NumExpr only for scalar checkpoint formulas. Avoid creating
@@ -120,6 +162,7 @@ def _setup_vendored_triton_kernels():
 def _prepare_environment() -> None:
     """Phase 1: environment and library preparation, before the Torch import."""
     _set_numexpr_thread_default()
+    _setup_unified_cache()
     _add_trt_llm_dll_directory()
     _preload_python_lib()
     _setup_vendored_triton_kernels()

tensorrt_llm/executor/ray/gpu_worker.py

@@ -73,9 +73,24 @@ def __init__(self, worker_cls, worker_kwargs, world_size, rank):
         self.gpu = int(ray.get_gpu_ids()[0])
         self.local_gpu = self.physical_to_local_id(self.gpu)
 
-        # Per-worker DeepGemm JIT cache to avoid rename race across co-located workers
-        os.environ["DG_JIT_CACHE_DIR"] = os.path.join(
-            tempfile.gettempdir(), f"deep_gemm_rank{rank}_gpu{self.gpu}")
+        if os.environ.get("TRTLLM_DEEP_GEMM_CACHE_PER_PROCESS", "1") != "0":
+            cache_dir = os.environ.get("TRTLLM_CACHE_DIR")
+            if (cache_dir
+                    and os.environ.get("DG_JIT_CACHE_DIR") == os.path.join(
+                        os.path.expanduser(cache_dir), "deep_gemm")):
+                logger.warning_once(
+                    "TRTLLM_CACHE_DIR keeps DeepGEMM isolation enabled; set "
+                    "TRTLLM_DEEP_GEMM_CACHE_PER_PROCESS=0 for better cache "
+                    "reuse at the risk of concurrent writes.",
+                    key="deep_gemm_unified_cache_isolation")
+                os.environ["DG_JIT_CACHE_DIR"] = os.path.join(
+                    os.environ["DG_JIT_CACHE_DIR"],
+                    f"deep_gemm_rank{rank}_gpu{self.gpu}")
+            else:
+                os.environ.setdefault(
+                    "DG_JIT_CACHE_DIR",
+                    os.path.join(tempfile.gettempdir(),
+                                 f"deep_gemm_rank{rank}_gpu{self.gpu}"))
 
         torch.cuda.set_device(self.local_gpu)
 

tensorrt_llm/llmapi/mpi_session.py

@@ -17,6 +17,7 @@
 
 import zmq
 
+from tensorrt_llm._bootstrap import _UNIFIED_CACHE_ENV_VARS
 from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE
 from tensorrt_llm.logger import logger
 
@@ -549,27 +550,38 @@ def _start_mpi_pool(self):
         env = {
             key: value
             for key, value in os.environ.items()
-            if key.startswith("TRTLLM") or key.startswith("TLLM") or key in (
-                "FLASHINFER_WORKSPACE_BASE", "FLASHINFER_CUBIN_DIR",
-                "MASTER_ADDR", "MASTER_PORT")
+            if key.startswith("TRTLLM") or key.startswith("TLLM") or key in
+            _UNIFIED_CACHE_ENV_VARS or key in ("FLASHINFER_CUBIN_DIR",
+                                               "MASTER_ADDR", "MASTER_PORT")
         }
         workspace_managed = env.get(_FLASHINFER_WORKSPACE_MANAGED_ENV) == "1"
         env.update(self._env_overrides)
         explicit_workspace_override = (_FLASHINFER_WORKSPACE_ENV
                                        in self._env_overrides)
+        workspace_root = _FLASHINFER_WORKSPACE_ROOT
+        if cache_root := env.get("TRTLLM_CACHE_DIR"):
+            workspace_root = os.path.join(os.path.expanduser(cache_root),
+                                          "flashinfer")
+            workspace_managed |= env.get(
+                _FLASHINFER_WORKSPACE_ENV) == workspace_root
         isolate_workspace = (
             (self.n_workers > 1 or workspace_managed)
             and env.get("TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS", "1") != "0"
             and (_FLASHINFER_WORKSPACE_ENV not in env or
                  (workspace_managed and not explicit_workspace_override)))
         if isolate_workspace:
             env.pop(_FLASHINFER_WORKSPACE_ENV, None)
+            if cache_root:
+                logger.warning_once(
+                    "TRTLLM_CACHE_DIR keeps FlashInfer isolation enabled; "
+                    "set TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0 for better "
+                    "cache reuse at the risk of concurrent writes.",
+                    key="flashinfer_unified_cache_isolation")
         elif explicit_workspace_override:
             # The override is user-owned, including in any further nested pool.
             env.pop(_FLASHINFER_WORKSPACE_MANAGED_ENV, None)
-        python_args = ([
-            "-c", _FLASHINFER_WORKER_BOOTSTRAP, _FLASHINFER_WORKSPACE_ROOT
-        ] if isolate_workspace else None)
+        python_args = (["-c", _FLASHINFER_WORKER_BOOTSTRAP, workspace_root]
+                       if isolate_workspace else None)
         self.mpi_pool = MPIPoolExecutor(max_workers=self.n_workers,
                                         path=sys.path,
                                         env=env,

tensorrt_llm/llmapi/trtllm-llmapi-launch

@@ -84,15 +84,38 @@ PY
 # which include compiler inputs, inside the per-rank workspace. A caller-set
 # value is preserved. See docs/source/llm-api/index.md.
 function configure_flashinfer_workspace {
+    local workspace_root="${HOME:+$HOME/.cache/tensorrt_llm/flashinfer}"
+    if [ -n "$TRTLLM_CACHE_DIR" ]; then
+        local cache_root
+        if ! cache_root=$(python3 -S -c \
+            'import os, sys; print(os.path.expanduser(sys.argv[1]))' "$TRTLLM_CACHE_DIR"); then
+            log_stderr "Failed to expand TRTLLM_CACHE_DIR=$TRTLLM_CACHE_DIR"
+            use_unique_flashinfer_workspace
+            return
+        fi
+        workspace_root="${cache_root%/}/flashinfer"
+    fi
+
+    if [ -n "$FLASHINFER_WORKSPACE_BASE" ] && \
+        [ "$FLASHINFER_WORKSPACE_BASE" != "$workspace_root" ]; then
+        return
+    fi
+
     # PMI_RANK or PMI_ID still proves that this process is launcher-managed
     # when PMI_SIZE is unavailable. Isolate even rank 0 in that case so rank-0
     # processes from concurrent jobs cannot share FlashInfer's default workspace.
-    if [ -n "$FLASHINFER_WORKSPACE_BASE" ] || \
-        { [ "$tllm_mpi_size" -le 1 ] && [ -z "$PMI_RANK" ] && [ -z "$PMI_ID" ]; } || \
+    if { [ "$tllm_mpi_size" -le 1 ] && [ -z "$PMI_RANK" ] && [ -z "$PMI_ID" ]; } || \
         [ "${TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS:-1}" = "0" ]; then
+        if [ -n "$TRTLLM_CACHE_DIR" ]; then
+            export FLASHINFER_WORKSPACE_BASE="$workspace_root"
+        fi
         return
     fi
 
+    if [ -n "$TRTLLM_CACHE_DIR" ]; then
+        log_stderr "Warning: FlashInfer isolation remains enabled; set TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0 for better cache reuse at the risk of concurrent writes."
+    fi
+
     # Do not overwrite a descriptor inherited from the launcher caller. A
     # process-unique workspace still provides safe isolation in this rare case.
     if { true >&200; } 2>/dev/null; then
@@ -101,13 +124,11 @@ function configure_flashinfer_workspace {
         return
     fi
 
-    if [ -z "$HOME" ]; then
+    if [ -z "$workspace_root" ]; then
         log_stderr "HOME is unset; persistent FlashInfer JIT workspace isolation is unavailable"
         use_unique_flashinfer_workspace
         return
     fi
-
-    local workspace_root="$HOME/.cache/tensorrt_llm/flashinfer"
     if ! mkdir -p "$workspace_root"; then
         log_stderr "Failed to create the persistent FlashInfer JIT workspace root: $workspace_root"
         use_unique_flashinfer_workspace

tests/unittest/llmapi/test_mpi_session.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
+import shlex
 import subprocess  # nosec B404
 import sys
 import threading
@@ -283,14 +284,16 @@ def _launcher_env(tmp_path: Path, home: str) -> dict:
     """Environment for a rank-0 launcher-managed run of ``trtllm-llmapi-launch``.
 
     ``python3`` and ``openssl`` are stubbed so the launcher can hand out an IPC
-    address and an HMAC key without importing tensorrt_llm; the stubbed
-    ``python3 -S`` used for the workspace lock exits 0, so the persistent slot
-    is treated as acquired.
+    address and an HMAC key without importing tensorrt_llm. Stdlib-only
+    ``python3 -S`` calls use the test interpreter.
     """
     stub_bin = tmp_path / "bin"
     stub_bin.mkdir(exist_ok=True)
     python_stub = stub_bin / "python3"
     python_stub.write_text("#!/bin/sh\n"
+                           "if [ \"$1\" = \"-S\" ]; then\n"
+                           f"    exec {shlex.quote(sys.executable)} \"$@\"\n"
+                           "fi\n"
                            "if [ \"$1\" = \"-c\" ]; then\n"
                            "    echo ipc:///tmp/trtllm-pmi-workspace-test\n"
                            "fi\n")
@@ -378,6 +381,77 @@ def test_llmapi_launch_temporary_fallback_leaves_cubin_dir_unset(
     assert list(tmpdir.iterdir()) == []
 
 
+@pytest.mark.cpu_only
+def test_llmapi_launch_respects_unified_cache_root(tmp_path: Path) -> None:
+    stub_bin = tmp_path / "bin"
+    stub_bin.mkdir()
+    python_stub = stub_bin / "python3"
+    python_stub.write_text("#!/bin/sh\n"
+                           "if [ \"$1\" = \"-S\" ]; then\n"
+                           f"    exec {shlex.quote(sys.executable)} \"$@\"\n"
+                           "fi\n"
+                           "if [ \"$1\" = \"-c\" ]; then\n"
+                           "    echo ipc:///tmp/trtllm-unified-cache-test\n"
+                           "fi\n")
+    python_stub.chmod(0o755)
+    openssl_stub = stub_bin / "openssl"
+    openssl_stub.write_text("#!/bin/sh\nprintf '%064d\\n' 0\n")
+    openssl_stub.chmod(0o755)
+
+    env = os.environ.copy()
+    for name in (
+            "SLURM_NTASKS",
+            "SLURM_PROCID",
+            "OMPI_COMM_WORLD_SIZE",
+            "OMPI_COMM_WORLD_RANK",
+            "PMI_ID",
+            "FLASHINFER_WORKSPACE_BASE",
+            "FLASHINFER_CUBIN_DIR",
+            "TRTLLM_FLASHINFER_WORKSPACE_MANAGED",
+            "TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS",
+    ):
+        env.pop(name, None)
+    cache_root = tmp_path / "home" / "unified"
+    env["PMI_RANK"] = "0"
+    env["HOME"] = str(tmp_path / "home")
+    env["TRTLLM_CACHE_DIR"] = "~/unified"
+    env["FLASHINFER_WORKSPACE_BASE"] = str(cache_root / "flashinfer")
+    env["PATH"] = f"{stub_bin}{os.pathsep}{env['PATH']}"
+
+    launcher = Path(__file__).parents[
+        3] / "tensorrt_llm" / "llmapi" / "trtllm-llmapi-launch"
+    result = subprocess.run(  # nosec B603
+        ["bash", str(launcher), "/usr/bin/env"],
+        check=True,
+        capture_output=True,
+        env=env,
+        text=True,
+        timeout=10,
+    )
+
+    workspace = cache_root / "flashinfer" / "rank-0"
+    assert f"FLASHINFER_WORKSPACE_BASE={workspace}" in result.stdout
+    assert "TRTLLM_FLASHINFER_WORKSPACE_MANAGED=1" in result.stdout
+    assert "better cache reuse" in result.stderr
+
+
+@pytest.mark.cpu_only
+def test_llmapi_launch_matches_bootstrap_expansion_with_trailing_home(
+        tmp_path: Path) -> None:
+    home = tmp_path / "home"
+    home.mkdir()
+    cache_root = home / "unified"
+    env = _launcher_env(tmp_path, f"{home}{os.sep}")
+    env["TRTLLM_CACHE_DIR"] = "~/unified"
+    env["FLASHINFER_WORKSPACE_BASE"] = str(cache_root / "flashinfer")
+
+    result = _run_launcher_env(env)
+
+    workspace = cache_root / "flashinfer" / "rank-0"
+    assert f"FLASHINFER_WORKSPACE_BASE={workspace}" in result.stdout
+    assert "TRTLLM_FLASHINFER_WORKSPACE_MANAGED=1" in result.stdout
+
+
 @pytest.mark.cpu_only
 def test_llmapi_launch_aborts_when_no_workspace_is_available(
         tmp_path: Path) -> None: