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

docs/source/llm-api/index.md

@@ -130,6 +130,50 @@ 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. Except for the isolation behavior described below, an individually
+configured environment variable 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 treat `DG_JIT_CACHE_DIR` equal to
+`<TRTLLM_CACHE_DIR>/deep_gemm` as unified-cache-managed and isolate it by rank
+and GPU, even when that value was set explicitly. Other explicit values are
+preserved. Set `TRTLLM_DEEP_GEMM_CACHE_PER_PROCESS=0` to preserve and share the
+unified path.
+
 ### FlashInfer JIT workspaces for MPI workers
 
 `trtllm-llmapi-launch` ranks and dynamically spawned `MpiPoolSession` workers
@@ -151,7 +195,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

@@ -15,7 +15,6 @@
 import gc
 import importlib
 import os
-import tempfile
 from functools import wraps
 from pathlib import Path
 from queue import Queue
@@ -27,7 +26,8 @@
 from tensorrt_llm._torch.utils import get_device_uuid
 from tensorrt_llm._torch.virtual_memory import (materialize_with_tag,
                                                 release_with_tag)
-from tensorrt_llm.executor.ray.utils import control_action_decorator
+from tensorrt_llm.executor.ray.utils import (_configure_deep_gemm_cache,
+                                             control_action_decorator)
 
 from ... import TorchLlmArgs
 from ...llmapi.llm_args import BaseLlmArgs, ExecutorMemoryType
@@ -73,9 +73,7 @@ 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}")
+        _configure_deep_gemm_cache(rank, self.gpu)
 
         torch.cuda.set_device(self.local_gpu)
 

tensorrt_llm/executor/ray/utils.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import functools
+import os
+import tempfile
 from contextlib import contextmanager
 from typing import Callable, Optional
 
@@ -22,6 +24,29 @@
     import tensorrt_llm.executor.ray.stub as ray
 
 
+def _configure_deep_gemm_cache(rank: int, gpu: int) -> None:
+    if os.environ.get("TRTLLM_DEEP_GEMM_CACHE_PER_PROCESS", "1") == "0":
+        return
+
+    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")):
+        from tensorrt_llm.logger import logger
+
+        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{gpu}")
+    else:
+        os.environ.setdefault(
+            "DG_JIT_CACHE_DIR",
+            os.path.join(tempfile.gettempdir(),
+                         f"deep_gemm_rank{rank}_gpu{gpu}"))
+
+
 @contextmanager
 def unwrap_ray_errors():
     try:

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/integration/test_lists/test-db/l0_cpu.yml

@@ -44,6 +44,7 @@ l0_cpu:
   - unittest/_torch/distributed
   - unittest/_torch/executor
   - unittest/test_startup.py
+  - unittest/test_unified_cache.py
   - unittest/_torch/disaggregation
   - unittest/_torch/peft
   - unittest/_torch/memory