diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 389450bc12c0..e6b255ff6232 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -219,6 +219,15 @@ 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. +Downloaded TRTLLM-Gen FMHA kernel cubins are the exception: FlashInfer only +loads them and never compiles them, so ranks with persistent slots share them +through a `cubins` directory next to the slots. A rank looks a kernel cubin up +in its own slot, then in the shared directory, and downloads it only when +neither has it, publishing the download to the shared directory with an atomic +rename. Every cubin read from the shared directory is checked against its +SHA-256. Headers and other compiler inputs are never shared. Set +`TRTLLM_FLASHINFER_SHARED_CUBINS=0` to keep these cubins per-rank as well. + 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 diff --git a/tensorrt_llm/_torch/flashinfer_cubin_pool.py b/tensorrt_llm/_torch/flashinfer_cubin_pool.py new file mode 100644 index 000000000000..45e37dbb62ae --- /dev/null +++ b/tensorrt_llm/_torch/flashinfer_cubin_pool.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Share FlashInfer runtime cubins between isolated MPI ranks. + +``trtllm-llmapi-launch`` and ``MpiPoolSession`` give every rank its own +FlashInfer workspace, and with ``FLASHINFER_CUBIN_DIR`` unset FlashInfer keeps +its downloaded artifacts inside that workspace too. That isolation is required +for compiler inputs: FlashInfer downloads trtllm-gen export headers into the +same directory and compiles against them, and a header replaced by a peer +rank's download mid-compilation breaks the build. + +Runtime cubins are different. A cubin is never handed to a compiler; FlashInfer +only reads its bytes and passes them to ``cuModuleLoadData``. Keeping them +per-rank means every rank downloads every kernel it uses separately, and a +TRTLLM-Gen FMHA kernel that warmup does not reach is downloaded on the serving +path the first time a request shape selects it. Under attention DP only the +rank that served a request has already fetched its kernels, so the other ranks +stall on the download in lockstep and the whole step waits for them. + +This module installs a read-through pool for the TRTLLM-Gen FMHA kernel cubins +only, the ``*.cubin`` files directly under FlashInfer's ``TRTLLM_GEN_FMHA`` +artifact directory. A lookup first tries the rank-local cache, then a pool +directory that sits next to the per-rank workspace slots, and only then +downloads. A rank that downloads a cubin publishes it to the pool with an +atomic rename, so peers see either a complete file or none. Peers read pooled +cubins in place, and every byte read from the pool is checked against the +artifact's SHA-256 before use. Headers, checksum files and every other +artifact take FlashInfer's own path unchanged and stay per-rank. +""" + +import hashlib +import os +import uuid +from pathlib import Path +from typing import Callable, Optional + +from ..logger import logger + +_CUBIN_SUFFIX = ".cubin" +_POOL_DIR_NAME = "cubins" +_MANAGED_ENV = "TRTLLM_FLASHINFER_WORKSPACE_MANAGED" +_WORKSPACE_ENV = "FLASHINFER_WORKSPACE_BASE" +_CUBIN_DIR_ENV = "FLASHINFER_CUBIN_DIR" +_DISABLE_ENV = "TRTLLM_FLASHINFER_SHARED_CUBINS" +_INSTALLED_MARKER = "_trtllm_shared_cubin_pool" + + +def resolve_shared_cubin_dir(environ: Optional[dict] = None) -> Optional[Path]: + """Return the shared cubin pool for this process, or None when not applicable. + + The pool exists only for TensorRT-LLM-managed per-rank workspaces: the + workspace is a ``rank-N`` slot and the pool is the ``cubins`` directory + beside it. A user-configured ``FLASHINFER_CUBIN_DIR`` is already shared and + is left alone, as is an unmanaged ``FLASHINFER_WORKSPACE_BASE``. + ``TRTLLM_FLASHINFER_SHARED_CUBINS=0`` disables the pool. + + Args: + environ: Environment to read; defaults to ``os.environ``. + """ + env = os.environ if environ is None else environ + if env.get(_DISABLE_ENV, "1") == "0": + return None + if env.get(_CUBIN_DIR_ENV): + return None + if env.get(_MANAGED_ENV) != "1": + return None + workspace = env.get(_WORKSPACE_ENV) + if not workspace: + return None + workspace_path = Path(workspace) + if not workspace_path.name.startswith("rank-"): + # Temporary fallback workspaces are process-unique and removed at exit, + # so there is no stable directory to share with peers. + return None + return workspace_path.parent / _POOL_DIR_NAME + + +def _read_verified(path: Path, sha256: str) -> bytes: + try: + data = path.read_bytes() + except OSError: + return b"" + if hashlib.sha256(data).hexdigest() != sha256: + return b"" + return data + + +def _is_pooled_artifact(file_name: str, pooled_prefix: str) -> bool: + if not file_name.endswith(_CUBIN_SUFFIX): + return False + prefix = pooled_prefix.rstrip("/") + "/" + if not file_name.startswith(prefix): + return False + # Only the kernel cubins directly under the artifact directory are pooled; + # FlashInfer joins the directory and the kernel name with a doubled slash. + return "/" not in file_name[len(prefix) :].lstrip("/") + + +def _write_atomically(path: Path, data: bytes) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temp_path = path.with_name(f"{path.name}.{uuid.uuid4().hex}.tmp") + try: + temp_path.write_bytes(data) + os.replace(temp_path, path) + finally: + if temp_path.exists(): + temp_path.unlink() + + +def make_pooled_get_artifact( + get_artifact: Callable[..., bytes], + local_cubin_dir: Path, + shared_cubin_dir: Path, + pooled_prefix: str, +) -> Callable[..., bytes]: + """Wrap FlashInfer's ``get_artifact`` with a shared pool for runtime cubins. + + Args: + get_artifact: FlashInfer's ``flashinfer.jit.cubin_loader.get_artifact``. + local_cubin_dir: This rank's FlashInfer cubin directory. + shared_cubin_dir: Pool directory shared by the job's ranks. + pooled_prefix: Artifact directory whose kernel cubins are pooled. + """ + + def pooled_get_artifact(file_name: str, sha256: str, *args, **kwargs) -> bytes: + if not _is_pooled_artifact(file_name, pooled_prefix): + return get_artifact(file_name, sha256, *args, **kwargs) + + data = _read_verified(local_cubin_dir / file_name, sha256) + if data: + return data + + shared_path = shared_cubin_dir / file_name + data = _read_verified(shared_path, sha256) + if data: + return data + + data = get_artifact(file_name, sha256, *args, **kwargs) + if data: + try: + _write_atomically(shared_path, data) + except OSError as error: + logger.debug(f"Could not share FlashInfer cubin {file_name}: {error}") + return data + + setattr(pooled_get_artifact, _INSTALLED_MARKER, True) + return pooled_get_artifact + + +def install_shared_cubin_pool() -> bool: + """Route FlashInfer's runtime cubin lookups through the shared pool. + + Returns True when the pool was installed. A no-op when the process does not + use a TensorRT-LLM-managed per-rank workspace, when the pool is disabled, or + when FlashInfer does not expose the expected loader. + """ + shared_cubin_dir = resolve_shared_cubin_dir() + if shared_cubin_dir is None: + return False + try: + from flashinfer.jit import cubin_loader + except ImportError: + return False + + try: + from flashinfer.artifacts import ArtifactPath + except ImportError: + return False + + get_artifact = getattr(cubin_loader, "get_artifact", None) + local_cubin_dir = getattr(cubin_loader, "FLASHINFER_CUBIN_DIR", None) + pooled_prefix = getattr(ArtifactPath, "TRTLLM_GEN_FMHA", None) + if not callable(get_artifact) or local_cubin_dir is None or not pooled_prefix: + return False + if getattr(get_artifact, _INSTALLED_MARKER, False): + return True + local_cubin_dir = Path(local_cubin_dir) + if local_cubin_dir.resolve() == shared_cubin_dir.resolve(): + return False + + pooled = make_pooled_get_artifact( + get_artifact, local_cubin_dir, shared_cubin_dir, pooled_prefix + ) + # The C++ cubin callback installed by setup_cubin_loader resolves + # get_artifact from this module's globals on every call, so rebinding the + # module attribute reaches the TRTLLM-Gen FMHA runtime loader. + cubin_loader.get_artifact = pooled + if getattr(cubin_loader, "get_cubin", None) is get_artifact: + cubin_loader.get_cubin = pooled + logger.info(f"FlashInfer runtime cubins are shared across ranks via {shared_cubin_dir}") + return True diff --git a/tensorrt_llm/_torch/flashinfer_utils.py b/tensorrt_llm/_torch/flashinfer_utils.py index a7eafaaf9282..7c478ff936aa 100644 --- a/tensorrt_llm/_torch/flashinfer_utils.py +++ b/tensorrt_llm/_torch/flashinfer_utils.py @@ -25,3 +25,8 @@ def get_env_enable_pdl() -> bool: print( "flashinfer is not installed properly, please try pip install or building from source codes" ) + +if IS_FLASHINFER_AVAILABLE: + from .flashinfer_cubin_pool import install_shared_cubin_pool + + install_shared_cubin_pool() diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index b7720072f057..808bb994fe69 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -81,6 +81,7 @@ l0_cpu: - unittest/_torch/speculative/hw_agnostic # l0_h100 selects with -m "not cpu_only", so this entry is what runs it. - unittest/_torch/speculative/test_spec_slot_pool_sizing.py + - unittest/_torch/flashinfer/test_flashinfer_cubin_pool.py - unittest/_torch/test_custom_config_registration.py - unittest/_torch/test_hf_quant_config.py - unittest/_torch/test_mmap_utils.py diff --git a/tests/unittest/_torch/flashinfer/test_flashinfer_cubin_pool.py b/tests/unittest/_torch/flashinfer/test_flashinfer_cubin_pool.py new file mode 100644 index 000000000000..d4f1d5f5966d --- /dev/null +++ b/tests/unittest/_torch/flashinfer/test_flashinfer_cubin_pool.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for sharing FlashInfer runtime cubins across isolated MPI ranks.""" + +import hashlib +import sys +import threading +import types +from pathlib import Path + +import pytest + +from tensorrt_llm._torch import flashinfer_cubin_pool as pool + +pytestmark = pytest.mark.cpu_only + +_FMHA = "deadbeef/fmha/trtllm-gen/" +# FlashInfer joins the artifact directory and the kernel name with a doubled slash. +_CUBIN = _FMHA + "/fmhaSm100fKernel_Test.cubin" +_HEADER = _FMHA + "/include/flashInferMetaInfo.h" +_GEMM_CUBIN = "deadbeef/gemm-0000000-0000000/Gemm_Test.cubin" + + +def _sha(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +class _Rank: + """One rank's view of FlashInfer: a private cubin dir and a recording downloader.""" + + def __init__(self, workspace_root: Path, rank: int, downloads: list): + self.workspace = workspace_root / f"rank-{rank}" + self.cubin_dir = self.workspace / ".cache" / "flashinfer" / "cubins" + self.rank = rank + self.downloads = downloads + self.env = { + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1", + "FLASHINFER_WORKSPACE_BASE": str(self.workspace), + } + + def download(self, file_name: str, sha256: str, session=None) -> bytes: + self.downloads.append((self.rank, file_name)) + data = file_name.encode() + path = self.cubin_dir / file_name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(data) + return data + + def pooled(self): + shared = pool.resolve_shared_cubin_dir(self.env) + assert shared is not None + return pool.make_pooled_get_artifact(self.download, self.cubin_dir, shared, _FMHA) + + +@pytest.fixture +def ranks(tmp_path: Path): + downloads: list = [] + root = tmp_path / "flashinfer" + return root, downloads, [_Rank(root, rank, downloads) for rank in range(4)] + + +def test_resolve_shared_dir_sits_beside_rank_slots(tmp_path: Path) -> None: + env = { + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1", + "FLASHINFER_WORKSPACE_BASE": str(tmp_path / "rank-3"), + } + assert pool.resolve_shared_cubin_dir(env) == tmp_path / "cubins" + + +@pytest.mark.parametrize( + "env", + [ + pytest.param({"FLASHINFER_WORKSPACE_BASE": "/x/rank-0"}, id="unmanaged-workspace"), + pytest.param({"TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1"}, id="no-workspace"), + pytest.param( + { + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1", + "FLASHINFER_WORKSPACE_BASE": "/tmp/trtllm-flashinfer-rank-0.abcd1234", + }, + id="temporary-fallback-workspace", + ), + pytest.param( + { + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1", + "FLASHINFER_WORKSPACE_BASE": "/x/rank-0", + "FLASHINFER_CUBIN_DIR": "/shared/cubins", + }, + id="explicit-cubin-dir", + ), + pytest.param( + { + "TRTLLM_FLASHINFER_WORKSPACE_MANAGED": "1", + "FLASHINFER_WORKSPACE_BASE": "/x/rank-0", + "TRTLLM_FLASHINFER_SHARED_CUBINS": "0", + }, + id="disabled", + ), + ], +) +def test_resolve_shared_dir_is_off_outside_managed_rank_slots(env: dict) -> None: + assert pool.resolve_shared_cubin_dir(env) is None + + +def test_cubin_downloaded_once_is_reused_by_peer_ranks(ranks) -> None: + root, downloads, (r0, r1, r2, _) = ranks + data = _CUBIN.encode() + + assert r0.pooled()(_CUBIN, _sha(data)) == data + assert r1.pooled()(_CUBIN, _sha(data)) == data + assert r2.pooled()(_CUBIN, _sha(data)) == data + + assert downloads == [(0, _CUBIN)] + assert (root / "cubins" / _CUBIN).read_bytes() == data + # Peers read the pooled cubin in place; nothing is written to their caches. + assert not (r1.cubin_dir / _CUBIN).exists() + assert not (r2.cubin_dir / _CUBIN).exists() + + +@pytest.mark.parametrize( + "file_name", + [ + pytest.param(_HEADER, id="fmha-header"), + pytest.param(_FMHA + "/checksums.txt", id="fmha-checksums"), + pytest.param(_FMHA + "/include/Test.cubin", id="fmha-subdirectory"), + pytest.param(_GEMM_CUBIN, id="other-artifact-cubin"), + ], +) +def test_only_fmha_kernel_cubins_are_shared(ranks, file_name: str) -> None: + root, downloads, (r0, r1, _, _) = ranks + data = file_name.encode() + + assert r0.pooled()(file_name, _sha(data)) == data + assert r1.pooled()(file_name, _sha(data)) == data + + assert downloads == [(0, file_name), (1, file_name)] + assert not (root / "cubins").exists() + + +def test_corrupt_pool_entry_is_ignored_and_replaced(ranks) -> None: + root, downloads, (r0, r1, _, _) = ranks + data = _CUBIN.encode() + shared = root / "cubins" / _CUBIN + shared.parent.mkdir(parents=True) + shared.write_bytes(b"truncated") + + assert r0.pooled()(_CUBIN, _sha(data)) == data + assert downloads == [(0, _CUBIN)] + assert shared.read_bytes() == data + + assert r1.pooled()(_CUBIN, _sha(data)) == data + assert downloads == [(0, _CUBIN)] + + +def test_failed_download_is_not_published(ranks) -> None: + root, _, (r0, _, _, _) = ranks + + def failing_download(file_name: str, sha256: str, session=None) -> bytes: + return b"" + + shared = pool.resolve_shared_cubin_dir(r0.env) + pooled = pool.make_pooled_get_artifact(failing_download, r0.cubin_dir, shared, _FMHA) + assert pooled(_CUBIN, _sha(_CUBIN.encode())) == b"" + assert not (root / "cubins" / _CUBIN).exists() + + +def test_concurrent_ranks_never_read_a_partial_cubin(tmp_path: Path) -> None: + downloads: list = [] + root = tmp_path / "flashinfer" + names = [_CUBIN.replace("Test", f"Test{i}") for i in range(4)] + errors: list = [] + + def run(rank: int) -> None: + pooled = _Rank(root, rank, downloads).pooled() + try: + for _ in range(20): + for name in names: + assert pooled(name, _sha(name.encode())) == name.encode() + except AssertionError as error: + errors.append(error) + + threads = [threading.Thread(target=run, args=(rank,)) for rank in range(16)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert not errors + assert not list((root / "cubins").rglob("*.tmp")) + + +def test_install_wraps_flashinfer_loader_once( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + workspace = tmp_path / "rank-0" + monkeypatch.setenv("TRTLLM_FLASHINFER_WORKSPACE_MANAGED", "1") + monkeypatch.setenv("FLASHINFER_WORKSPACE_BASE", str(workspace)) + monkeypatch.delenv("FLASHINFER_CUBIN_DIR", raising=False) + monkeypatch.delenv("TRTLLM_FLASHINFER_SHARED_CUBINS", raising=False) + + def get_artifact(file_name: str, sha256: str, session=None) -> bytes: + return b"" + + cubin_loader = types.SimpleNamespace( + get_artifact=get_artifact, + get_cubin=get_artifact, + FLASHINFER_CUBIN_DIR=workspace / ".cache" / "flashinfer" / "cubins", + ) + artifacts = types.SimpleNamespace(ArtifactPath=types.SimpleNamespace(TRTLLM_GEN_FMHA=_FMHA)) + jit = types.ModuleType("flashinfer.jit") + jit.cubin_loader = cubin_loader + flashinfer = types.ModuleType("flashinfer") + flashinfer.jit = jit + flashinfer.artifacts = artifacts + monkeypatch.setitem(sys.modules, "flashinfer", flashinfer) + monkeypatch.setitem(sys.modules, "flashinfer.jit", jit) + monkeypatch.setitem(sys.modules, "flashinfer.jit.cubin_loader", cubin_loader) + monkeypatch.setitem(sys.modules, "flashinfer.artifacts", artifacts) + + assert pool.install_shared_cubin_pool() + wrapped = cubin_loader.get_artifact + assert wrapped is not get_artifact + assert cubin_loader.get_cubin is wrapped + + assert pool.install_shared_cubin_pool() + assert cubin_loader.get_artifact is wrapped