[https://nvbugs/6739553][fix] Isolate the FlashInfer cubin cache per rank

docs/source/llm-api/index.md

@@ -133,25 +133,30 @@ The following tips typically assist new LLM API users who are familiar with othe
 ### FlashInfer JIT workspaces for MPI workers
 
 `trtllm-llmapi-launch` ranks and dynamically spawned `MpiPoolSession` workers
-isolate FlashInfer JIT source-generation workspaces by default. Each process
-normally claims a locked, persistent cache slot under
-`~/.cache/tensorrt_llm/flashinfer`, while downloaded cubins remain in
-FlashInfer's shared cache. This prevents concurrent MPI processes from writing
-the same generated source files without forcing a cold JIT compilation on
-every launch. Set `TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0` before invoking
-the launcher or creating the LLM instance to disable this behavior.
+isolate FlashInfer JIT workspaces by default. Each process claims a locked,
+persistent cache slot under `~/.cache/tensorrt_llm/flashinfer` that holds its
+generated sources, compiled modules and downloaded FlashInfer artifacts, so
+concurrent MPI processes never write to each other's compiler inputs while JIT
+artifacts stay warm across launches. Set
+`TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0` before invoking the launcher or
+creating the LLM instance to disable this behavior.
 
 Persistent slots are not pruned automatically, so their count can grow with
 peak job concurrency. When no TensorRT-LLM processes are using the cache, the
 slots may be deleted safely; subsequent launches rebuild the removed JIT
 artifacts.
 
 If persistent workspace setup is unavailable, each process falls back to a
-process-unique temporary workspace that is removed when the process exits.
-Persistent cache reuse is not available for this fallback.
+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 also preserved.
+launch modes. 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
+warm-up run, and set `FLASHINFER_NO_DOWNLOAD=1` so a missing artifact fails
+instead of being downloaded concurrently into the shared directory.
 
 ### Cannot quit after generation
 

tensorrt_llm/llmapi/mpi_session.py

@@ -107,13 +107,9 @@
 
 os.environ["TRTLLM_FLASHINFER_WORKSPACE_MANAGED"] = "1"
 
-# Preserve FlashInfer's default cubin cache without importing
-# flashinfer.jit.env. The environment must be fully configured before that
-# module initializes its workspace constants.
-if "FLASHINFER_CUBIN_DIR" not in os.environ and os.environ.get("HOME"):
-    os.environ["FLASHINFER_CUBIN_DIR"] = str(
-        Path(os.environ["HOME"]) / ".cache" / "flashinfer" / "cubins"
-    )
+# FLASHINFER_CUBIN_DIR is left unset so FlashInfer keeps downloaded artifacts,
+# which include compiler inputs, inside the per-rank workspace. A caller-set
+# value is preserved. See docs/source/llm-api/index.md.
 
 from mpi4py.futures.server import main
 

tensorrt_llm/llmapi/trtllm-llmapi-launch

@@ -31,12 +31,6 @@ function mpi_world_size {
     fi
 }
 
-function set_shared_flashinfer_cubin_dir {
-    if [ -z "$FLASHINFER_CUBIN_DIR" ] && [ -n "$HOME" ]; then
-        export FLASHINFER_CUBIN_DIR="$HOME/.cache/flashinfer/cubins"
-    fi
-}
-
 function cleanup_flashinfer_temporary_workspace {
     local exit_status=$?
 
@@ -64,7 +58,6 @@ function use_unique_flashinfer_workspace {
     flashinfer_temporary_workspace="$workspace"
     export FLASHINFER_WORKSPACE_BASE="$workspace"
     export TRTLLM_FLASHINFER_WORKSPACE_MANAGED=1
-    set_shared_flashinfer_cubin_dir
     log_stderr "Using temporary FlashInfer JIT workspace after persistent workspace isolation failed: $FLASHINFER_WORKSPACE_BASE"
 }
 
@@ -87,6 +80,9 @@ except (ImportError, OSError, TypeError, ValueError) as error:
 PY
 }
 
+# FLASHINFER_CUBIN_DIR is left unset so FlashInfer keeps downloaded artifacts,
+# 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 {
     # 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
@@ -136,7 +132,6 @@ function configure_flashinfer_workspace {
             flashinfer_workspace_lock_held=1
             export FLASHINFER_WORKSPACE_BASE="$workspace"
             export TRTLLM_FLASHINFER_WORKSPACE_MANAGED=1
-            set_shared_flashinfer_cubin_dir
             log_stderr "FlashInfer JIT workspace: $FLASHINFER_WORKSPACE_BASE"
             return
         else

tests/unittest/llmapi/_run_mpi_comm_task.py

@@ -77,9 +77,10 @@ def main(
             assert all(
                 Path(workspace).parent == workspace_root
                 for workspace in workspaces)
-            assert cubin_dirs == {
-                str(Path.home() / ".cache" / "flashinfer" / "cubins")
-            }
+            # Unset means FlashInfer derives the artifact cache from each
+            # worker's isolated workspace, keeping downloaded compiler inputs
+            # per-rank.
+            assert cubin_dirs == {None}
 
             nested_session = MpiPoolSession(n_workers=2)
             try:
@@ -96,6 +97,11 @@ def main(
             assert all(
                 Path(workspace).parent == workspace_root
                 for workspace in nested_workspaces)
+            nested_cubin_dirs = {
+                cubin_dir
+                for _, cubin_dir in nested_worker_envs
+            }
+            assert nested_cubin_dirs == {None}
 
 
 if __name__ == "__main__":

tests/unittest/llmapi/test_flashinfer_workspace.py

@@ -70,9 +70,10 @@ def test_worker_bootstrap_reuses_rank_workspace_and_releases_lock(
     assert os.environ[_FLASHINFER_WORKSPACE_ENV] == str(workspace)
     assert os.environ[_FLASHINFER_MANAGED_ENV] == "1"
     assert workspace.is_dir()
-    assert os.environ[_FLASHINFER_CUBIN_ENV] == str(
-        Path.home() / ".cache" / "flashinfer" / "cubins"
-    )
+    # The artifact cache must stay unpinned so FlashInfer derives it from the
+    # isolated workspace; a shared directory lets a peer's download replace a
+    # trtllm-gen export header mid-compilation.
+    assert _FLASHINFER_CUBIN_ENV not in os.environ
 
     with (workspace / ".lock").open("a") as lock:
         fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
@@ -113,6 +114,7 @@ def test_worker_bootstrap_replaces_inherited_managed_workspace(
     rank = mpi_session.mpi4py.MPI.COMM_WORLD.Get_rank()
     assert os.environ[_FLASHINFER_WORKSPACE_ENV] == str(tmp_path / f"rank-{rank}")
     assert os.environ[_FLASHINFER_MANAGED_ENV] == "1"
+    assert _FLASHINFER_CUBIN_ENV not in os.environ
 
 
 def test_worker_bootstrap_falls_back_when_isolation_fails(
@@ -304,12 +306,21 @@ def __init__(self, **kwargs: object) -> None:
     assert env.get(_FLASHINFER_MANAGED_ENV) == expected_marker
 
 
-def test_mpi_pool_shares_cubins_but_isolates_workspaces(
-    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+@pytest.mark.parametrize("nested", [False, True])
+@pytest.mark.parametrize("explicit_cubin_dir", [False, True])
+def test_mpi_pool_isolates_workspaces_and_keeps_explicit_cubin_dir(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path, explicit_cubin_dir: bool, nested: bool
 ) -> None:
     monkeypatch.delenv(_FLASHINFER_WORKSPACE_ENV, raising=False)
     monkeypatch.delenv(_FLASHINFER_ISOLATION_ENV, raising=False)
-    monkeypatch.setenv(_FLASHINFER_CUBIN_ENV, str(tmp_path / "cubins"))
+    monkeypatch.delenv(_FLASHINFER_CUBIN_ENV, raising=False)
+    monkeypatch.delenv(_FLASHINFER_MANAGED_ENV, raising=False)
+    if nested:
+        # A pool created inside a TRT-LLM-managed worker replaces the inherited slot.
+        monkeypatch.setenv(_FLASHINFER_WORKSPACE_ENV, str(tmp_path / "parent-slot"))
+        monkeypatch.setenv(_FLASHINFER_MANAGED_ENV, "1")
+    if explicit_cubin_dir:
+        monkeypatch.setenv(_FLASHINFER_CUBIN_ENV, str(tmp_path / "cubins"))
     workspace_root = tmp_path / "workspaces"
     monkeypatch.setattr(mpi_session, "_FLASHINFER_WORKSPACE_ROOT", str(workspace_root))
     session = mpi_session.MpiPoolSession(n_workers=2)
@@ -322,7 +333,10 @@ def test_mpi_pool_shares_cubins_but_isolates_workspaces(
     assert len(workspaces) == 2
     assert all(Path(workspace).is_dir() for workspace in workspaces)
     assert all(Path(workspace).parent == workspace_root for workspace in workspaces)
-    assert all(cubin == str(tmp_path / "cubins") for _, cubin in worker_envs)
+    # Unset means FlashInfer derives the artifact cache from each isolated
+    # workspace, so downloaded compiler inputs stay per-rank.
+    expected_cubin = str(tmp_path / "cubins") if explicit_cubin_dir else None
+    assert all(cubin == expected_cubin for _, cubin in worker_envs)
 
 
 def test_mpi_pool_propagates_explicit_workspace(

tests/unittest/llmapi/test_mpi_session.py

@@ -196,10 +196,35 @@ def read_stream(stream, output_stream):
             raise subprocess.CalledProcessError(return_code, command)
 
 
-@pytest.mark.cpu_only
-def test_llmapi_launch_isolates_pmi_rank_without_size(tmp_path: Path) -> None:
+_LAUNCHER = (Path(__file__).parents[3] / "tensorrt_llm" / "llmapi" /
+             "trtllm-llmapi-launch")
+
+# Variables that would steer the launcher's workspace setup; each launcher test
+# starts from an environment without them and adds back only what it exercises.
+_LAUNCHER_ENV_SCRUB = (
+    "SLURM_NTASKS",
+    "SLURM_PROCID",
+    "OMPI_COMM_WORLD_SIZE",
+    "OMPI_COMM_WORLD_RANK",
+    "PMI_SIZE",
+    "PMI_ID",
+    "FLASHINFER_WORKSPACE_BASE",
+    "FLASHINFER_CUBIN_DIR",
+    "TRTLLM_FLASHINFER_WORKSPACE_MANAGED",
+    "TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS",
+)
+
+
+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.
+    """
     stub_bin = tmp_path / "bin"
-    stub_bin.mkdir()
+    stub_bin.mkdir(exist_ok=True)
     python_stub = stub_bin / "python3"
     python_stub.write_text("#!/bin/sh\n"
                            "if [ \"$1\" = \"-c\" ]; then\n"
@@ -210,66 +235,97 @@ def test_llmapi_launch_isolates_pmi_rank_without_size(tmp_path: Path) -> None:
     openssl_stub.write_text("#!/bin/sh\nprintf '%064d\\n' 0\n")
     openssl_stub.chmod(0o755)
 
-    home = tmp_path / "home"
-    home.mkdir()
     env = os.environ.copy()
-    for name in (
-            "SLURM_NTASKS",
-            "SLURM_PROCID",
-            "OMPI_COMM_WORLD_SIZE",
-            "OMPI_COMM_WORLD_RANK",
-            "PMI_SIZE",
-            "PMI_ID",
-            "FLASHINFER_WORKSPACE_BASE",
-            "FLASHINFER_CUBIN_DIR",
-            "TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS",
-    ):
+    for name in _LAUNCHER_ENV_SCRUB:
         env.pop(name, None)
     env["PMI_RANK"] = "0"
-    env["HOME"] = str(home)
+    env["HOME"] = home
     env["PATH"] = f"{stub_bin}{os.pathsep}{env['PATH']}"
+    return env
 
-    launcher = (Path(__file__).parents[3] / "tensorrt_llm" / "llmapi" /
-                "trtllm-llmapi-launch")
-    result = subprocess.run(  # nosec B603
-        ["bash", str(launcher), "/usr/bin/env"],
+
+def _run_launcher_env(env: dict) -> subprocess.CompletedProcess:
+    return subprocess.run(  # nosec B603
+        ["bash", str(_LAUNCHER), "/usr/bin/env"],
         check=True,
         capture_output=True,
         env=env,
         text=True,
         timeout=10,
     )
 
+
+@pytest.mark.cpu_only
+def test_llmapi_launch_isolates_pmi_rank_without_size(tmp_path: Path) -> None:
+    home = tmp_path / "home"
+    home.mkdir()
+    env = _launcher_env(tmp_path, str(home))
+
+    result = _run_launcher_env(env)
+
     workspace = home / ".cache" / "tensorrt_llm" / "flashinfer" / "rank-0"
     assert f"FLASHINFER_WORKSPACE_BASE={workspace}" in result.stdout
     assert "TRTLLM_FLASHINFER_WORKSPACE_MANAGED=1" in result.stdout
+    # The artifact cache must follow the per-rank workspace, so the launcher
+    # must not pin FLASHINFER_CUBIN_DIR to a shared directory.
+    assert "FLASHINFER_CUBIN_DIR=" not in result.stdout
+
+
+@pytest.mark.cpu_only
+def test_llmapi_launch_preserves_explicit_cubin_dir(tmp_path: Path) -> None:
+    home = tmp_path / "home"
+    home.mkdir()
+    shared = tmp_path / "shared-cubins"
+    env = _launcher_env(tmp_path, str(home))
+    env["FLASHINFER_CUBIN_DIR"] = str(shared)
+
+    result = _run_launcher_env(env)
+
+    workspace = home / ".cache" / "tensorrt_llm" / "flashinfer" / "rank-0"
+    assert f"FLASHINFER_WORKSPACE_BASE={workspace}" in result.stdout
+    lines = [
+        line for line in result.stdout.splitlines()
+        if line.startswith("FLASHINFER_CUBIN_DIR=")
+    ]
+    assert lines == [f"FLASHINFER_CUBIN_DIR={shared}"]
+
+
+@pytest.mark.cpu_only
+def test_llmapi_launch_temporary_fallback_leaves_cubin_dir_unset(
+        tmp_path: Path) -> None:
+    # A regular file as HOME makes the persistent workspace root impossible to
+    # create, which sends the launcher down the temporary-workspace fallback.
+    invalid_home = tmp_path / "home-file"
+    invalid_home.touch()
+    tmpdir = tmp_path / "tmp"
+    tmpdir.mkdir()
+    env = _launcher_env(tmp_path, str(invalid_home))
+    env["TMPDIR"] = str(tmpdir)
+
+    result = _run_launcher_env(env)
+
+    prefix = f"FLASHINFER_WORKSPACE_BASE={tmpdir}/trtllm-flashinfer-rank-0."
+    assert any(line.startswith(prefix)
+               for line in result.stdout.splitlines()), result.stdout
+    assert "TRTLLM_FLASHINFER_WORKSPACE_MANAGED=1" in result.stdout
+    assert "FLASHINFER_CUBIN_DIR=" not in result.stdout
+    assert "temporary FlashInfer JIT workspace" in result.stderr
+    # The fallback workspace, artifacts included, is removed at exit.
+    assert list(tmpdir.iterdir()) == []
 
 
 @pytest.mark.cpu_only
 def test_llmapi_launch_aborts_when_no_workspace_is_available(
         tmp_path: Path) -> None:
     env = os.environ.copy()
-    for name in (
-            "SLURM_NTASKS",
-            "SLURM_PROCID",
-            "OMPI_COMM_WORLD_SIZE",
-            "OMPI_COMM_WORLD_RANK",
-            "PMI_SIZE",
-            "PMI_ID",
-            "FLASHINFER_WORKSPACE_BASE",
-            "FLASHINFER_CUBIN_DIR",
-            "TRTLLM_FLASHINFER_WORKSPACE_MANAGED",
-            "TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS",
-    ):
+    for name in _LAUNCHER_ENV_SCRUB:
         env.pop(name, None)
     env["PMI_RANK"] = "0"
     env["HOME"] = ""
     env["TMPDIR"] = str(tmp_path / "missing")
 
-    launcher = (Path(__file__).parents[3] / "tensorrt_llm" / "llmapi" /
-                "trtllm-llmapi-launch")
     result = subprocess.run(  # nosec B603
-        ["/bin/bash", str(launcher), "/usr/bin/true"],
+        ["/bin/bash", str(_LAUNCHER), "/usr/bin/true"],
         check=False,
         capture_output=True,
         env=env,