Skip to content

[TRTLLM-14475][chore] Self-sufficient transfer-agent dlopen and drop onnx/modelopt deps #16829

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Wanli-Jiang wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[TRTLLM-14475][chore] Self-sufficient transfer-agent dlopen and drop onnx/modelopt deps #16829

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions cpp/include/tensorrt_llm/executor/transferAgent.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -444,6 +444,14 @@ class BaseLoopbackAgent
virtual void executeLoopbackRequest(MemoryDescs const& memoryDescs, FileDescs const& fileDescs, bool isOffload) = 0;
};

/// @brief Promote the shared library containing this code (libtensorrt_llm.so) to the
/// process's global symbol scope. The KV cache transfer-agent wrapper libraries
/// (libtensorrt_llm_{nixl,ucx,mooncake}_wrapper.so) intentionally carry no DT_NEEDED on
/// libtensorrt_llm.so (the dependency would be circular) and resolve its symbols from the
/// global symbol table, while Python extension modules and their dependencies load with
/// RTLD_LOCAL. Idempotent; a no-op when the code is statically linked (e.g. unit tests).
void promoteHostLibraryToGlobalScope();

class DynLibLoader final
{
public:
Expand Down
2 changes: 2 additions & 0 deletions cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
#include "tensorrt_llm/executor/dataTransceiverState.h"
#include "tensorrt_llm/executor/serialization.h"
#include "tensorrt_llm/executor/serializeUtils.h"
#include "tensorrt_llm/executor/transferAgent.h"
#include "tensorrt_llm/runtime/utils/mpiUtils.h"
#include "tensorrt_llm/runtime/utils/pgUtils.h"
#include <algorithm>
Expand Down Expand Up @@ -501,6 +502,7 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa
if (backendType.value() == executor::CacheTransceiverConfig::BackendType::UCX)
{
std::lock_guard<std::mutex> lock(mDllMutex);
executor::kv_cache::promoteHostLibraryToGlobalScope();
mWrapperLibHandle = dllOpen(UCX_WRAPPER_LIB_NAME);
TLLM_CHECK_WITH_INFO(
mWrapperLibHandle != nullptr, "UCX wrapper library is not open correctly. error : %s", dlerror());
Expand Down
29 changes: 29 additions & 0 deletions cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,13 +22,41 @@
#include <algorithm>
#include <cuda.h>
#include <dlfcn.h>
#include <mutex>
#include <numeric>
#include <sstream>
#include <tuple>

namespace tensorrt_llm::executor::kv_cache
{

void promoteHostLibraryToGlobalScope()
{
static std::once_flag once;
std::call_once(once,
[]()
{
Dl_info info{};
if (dladdr(reinterpret_cast<void*>(&promoteHostLibraryToGlobalScope), &info) == 0
|| info.dli_fname == nullptr)
{
TLLM_LOG_DEBUG("dladdr failed; skipping global-scope promotion");
return;
}
// RTLD_NOLOAD promotes the visibility of the already-loaded library without
// loading a second copy. It fails when this code is statically linked into an
// executable (e.g. unit tests) — nothing to promote there.
void* handle = dlopen(info.dli_fname, RTLD_NOW | RTLD_GLOBAL | RTLD_NOLOAD);
if (handle == nullptr)
{
TLLM_LOG_DEBUG("global-scope promotion skipped for %s: %s", info.dli_fname, dlerror());
return;
}
TLLM_LOG_DEBUG("promoted %s to the global symbol scope", info.dli_fname);
// Keep the extra reference: the library must stay resident for the wrappers.
});
}

[[nodiscard]] DynLibLoader& DynLibLoader::getInstance()
{
static DynLibLoader instance;
Expand All @@ -37,6 +65,7 @@ namespace tensorrt_llm::executor::kv_cache

[[nodiscard]] void* DynLibLoader::getHandle(std::string const& name)
{
promoteHostLibraryToGlobalScope();
std::lock_guard<std::mutex> lock(mDllMutex);
auto it = mHandlers.find(name);
if (it != mHandlers.end())
Expand Down
3 changes: 0 additions & 3 deletions requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,6 @@ lark
lazy_loader~=0.5
mpi4py
numpy>=2.0.0,<2.4 # numba 0.63.1 requires numpy<2.4
onnx>=1.21.0
onnx_graphsurgeon>=0.5.2
graphviz
openai
polygraphy
Expand All @@ -25,7 +23,6 @@ sentencepiece>=0.1.99
# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-26-05.html#rel-26-05 uses 2.12.0a0.
torch>=2.11.0,<=2.13.0a0
torchvision
nvidia-modelopt[torch]~=0.37.0
# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-26-05.html#rel-26-05 uses 2.30.4
# torch 2.11.0+cu130 depends on nvidia-nccl-cu13==2.28.9
nvidia-nccl-cu13>=2.28.9,<=2.30.4
Expand Down
11 changes: 0 additions & 11 deletions tensorrt_llm/_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import ctypes
import os
import platform
import threading
Expand Down Expand Up @@ -45,16 +44,6 @@ def _init(log_level: object = None) -> None:

project_dir = str(Path(__file__).parent.absolute())

# Promote libtensorrt_llm.so symbols to the process's global scope. The KV
# cache transfer agent wrappers (libtensorrt_llm_nixl_wrapper.so,
# libtensorrt_llm_ucx_wrapper.so) are dlopen'ed at runtime without linking
# against libtensorrt_llm.so and resolve its symbols from the global symbol
# table, while Python extension modules and their dependencies load with
# RTLD_LOCAL. This promotion was previously a side effect of loading the
# TensorRT plugin library with RTLD_GLOBAL.
if platform.system() != "Windows":

@BowenFu BowenFu Jul 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Wanli-Jiang can you explain why we switch to implement this logic in C++? Is it that the python version no longer works?

@Wanli-Jiang Wanli-Jiang Jul 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Python preload wasn't broken — but it was a hack: on every import tensorrt_llm it force-loads the whole libtensorrt_llm.so with RTLD_GLOBAL via ctypes, purely to satisfy a C++ dlopen requirement (the nixl/ucx/mooncake wrappers carry no DT_NEEDED on it and need its symbols in the global scope). It also only kicks in when you enter through the Python import path.

The C++ version is simpler and cleaner: the library promotes its own symbols right at the wrapper-load site — promoteHostLibraryToGlobalScope() (RTLD_NOLOAD | RTLD_GLOBAL, call_once, no-op when statically linked). No Python involvement, works for any caller (C++ embeddings, executor, gtests), and the requirement lives where it belongs.

the Python preload was a hack (a whole-library RTLD_GLOBAL force-load on every import, just to satisfy a C++ dlopen need), and the C++ version is the simpler, cleaner home for it.

ctypes.CDLL(project_dir + "/libs/libtensorrt_llm.so", mode=ctypes.RTLD_GLOBAL)

# Load FT decoder layer and torch custom ops.
if platform.system() == "Windows":
ft_decoder_lib = project_dir + "/libs/th_common.dll"
Expand Down
6 changes: 0 additions & 6 deletions tensorrt_llm/_torch/auto_deploy/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,9 +37,3 @@
# import AutoDeploy LLM and LlmArgs (require TRT-LLM base classes)
from .llm import *
from .llm_args import *

try:
# This will overwrite the AutoModelForCausalLM.from_config to support modelopt quantization
import modelopt
except ImportError:
pass
6 changes: 2 additions & 4 deletions tensorrt_llm/_torch/auto_deploy/export/export.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,10 +32,8 @@
from ..utils.pipeline_cache_hooks import mark_pipeline_cache_hook
from .interface import apply_export_patches

try:
from modelopt.torch.quantization.utils import export_torch_mode as torch_export_context
except ImportError:
torch_export_context = nullcontext
# modelopt quantization support has been removed; use a null export context.
torch_export_context = nullcontext


# =====================================================================
Expand Down
37 changes: 0 additions & 37 deletions tensorrt_llm/_torch/auto_deploy/export/library/modelopt_context.py
View file
Open in desktop

This file was deleted.

16 changes: 3 additions & 13 deletions tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,19 +28,9 @@

from .logger import ad_logger

try:
# import modelopt to get quantize_op
from modelopt.torch.quantization import tensor_quant # noqa: F401

if hasattr(torch.ops, "tensorrt"):
modelopt_quantize_op = torch.ops.tensorrt.quantize_op
modelopt_dynamic_block_quantize_op = torch.ops.tensorrt.dynamic_block_quantize_op
else:
modelopt_quantize_op = None
modelopt_dynamic_block_quantize_op = None
except ImportError:
modelopt_quantize_op = None
modelopt_dynamic_block_quantize_op = None
# modelopt quantization support has been removed; these ops are no longer available.
modelopt_quantize_op = None

@BowenFu BowenFu Jul 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@nv-guomingz can you help review this part? Are we still using modelopt.torch.quantization.tensor_quant in torch path?

modelopt_dynamic_block_quantize_op = None

OpOrOverload = Union[OpOverloadPacket, OpOverload]
OperatorLike = Union[OpOrOverload, Callable]
Expand Down
Loading