[None][feat] Upgrade NIXL to v1.0.1

docker/common/install_nixl.sh

@@ -4,7 +4,7 @@ set -ex
 GITHUB_URL="https://github.com"
 UCX_INSTALL_PATH="/usr/local/ucx/"
 CUDA_PATH="/usr/local/cuda"
-NIXL_VERSION="0.9.0"
+NIXL_VERSION="v1.0.1"
 NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
 OLD_LD_LIBRARY_PATH=$LD_LIBRARY_PATH
 

docker/common/install_ucx.sh

@@ -1,15 +1,20 @@
 #!/bin/bash
 set -ex
 
-UCX_VERSION="v1.20.x"
-UCX_COMMIT="f656dbdf93e72e60b5d6ca78b9e3d9e744e789bd"
+UCX_VERSION="v1.21.x"
+UCX_COMMIT="167a4c6a311d9a42e30a37dcc01b8a3e73ea2826"
 UCX_INSTALL_PATH="/usr/local/ucx/"
 CUDA_PATH="/usr/local/cuda"
 UCX_REPO="https://github.com/openucx/ucx.git"
 
 mkdir -p /third-party-source
 
-rm -rf ${UCX_INSTALL_PATH}
+# Drop the trailing slash so this also removes a pre-existing *symlink* at
+# ${UCX_INSTALL_PATH} (the NGC PyTorch base image ships /usr/local/ucx as a
+# symlink to /opt/hpcx/ucx). Without this, `make install` writes through the
+# symlink into /opt/hpcx/ucx, and the later /opt/hpcx/ucx -> /usr/local/ucx
+# replacement creates an infinite symlink loop.
+rm -rf "${UCX_INSTALL_PATH%/}"
 git clone -b ${UCX_VERSION} ${UCX_REPO}
 cd ucx
 git checkout ${UCX_COMMIT}
@@ -30,6 +35,25 @@ cd ucx
   --with-dm                       \
   --enable-mt
 make install -j$(nproc)
+# Replace any pre-existing UCX (e.g. HPC-X's UCX shipped in the NGC PyTorch
+# base image) with a symlink to our freshly-installed one, so every binary
+# in the container — PyTorch, NCCL, UCC, plus the C++/Python NIXL stack —
+# resolves the same UCX SONAMEs.
+for d in /opt/hpcx/ucx /opt/hpcx-*/ucx; do
+  # `[ -e ]` catches real dirs/files; `[ -L ]` separately catches symlinks
+  # (including broken ones, where -e returns false). Together they replace
+  # any pre-existing UCX — real dir, stale symlink to an old UCX, or our
+  # own symlink on re-run (idempotent).
+  if [ -e "$d" ] || [ -L "$d" ]; then
+    echo "Replacing pre-existing UCX at $d with symlink to ${UCX_INSTALL_PATH%/}"
+    rm -rf "$d"
+    ln -s "${UCX_INSTALL_PATH%/}" "$d"
+  fi
+done
+# Make /usr/local/ucx/lib known to the system dynamic linker so binaries
+# without an explicit RPATH (or with $ORIGIN-relative RPATH) still find it.
+echo "${UCX_INSTALL_PATH%/}/lib" > /etc/ld.so.conf.d/ucx.conf
+ldconfig
 cd ..
 rm -rf ucx  # Remove UCX source to save space
 echo "export LD_LIBRARY_PATH=${UCX_INSTALL_PATH}/lib:\$LD_LIBRARY_PATH" >> "${ENV}"

jenkins/L0_Test.groovy

@@ -4210,6 +4210,17 @@ def launchTestJobs(pipeline, testFilter)
                         }
                         echo "###### Run LLMAPI tests Start ######"
 
+                        // Resolve the real tensorrt_llm install location after pip install,
+                        // and expose UCX shared libraries shipped inside the wheel
+                        // (tensorrt_llm/libs/ucx/*.so and libtensorrt_llm_ucx_wrapper.so)
+                        // so dlopen can find them at test runtime.
+                        def trtllmLibsDir = sh(
+                            script: "python3 -c 'import os, tensorrt_llm; print(os.path.join(os.path.dirname(tensorrt_llm.__file__), \"libs\"))'",
+                            returnStdout: true,
+                        ).trim()
+                        libEnv += ["LD_LIBRARY_PATH+trtllm_ucx=${trtllmLibsDir}/ucx"]
+                        libEnv += ["LD_LIBRARY_PATH+trtllm_libs=${trtllmLibsDir}"]
+
                         def config = VANILLA_CONFIG
                         if (cpu_arch == AARCH64_TRIPLE) {
                             config = LINUX_AARCH64_CONFIG

jenkins/current_image_tags.properties

@@ -13,8 +13,8 @@
 #     images are adopted from PostMerge pipelines, the abbreviated commit hash is used instead.
 IMAGE_NAME=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm
 
-LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.02-py3-x86_64-ubuntu24.04-trt10.15.1.29-skip-tritondevel-202605060827-13616
-LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-26.02-py3-sbsa-ubuntu24.04-trt10.15.1.29-skip-tritondevel-202605060827-13616
-LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-rocky8-x86_64-rocky8-py310-trt10.15.1.29-skip-tritondevel-202605060827-13616
-LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-rocky8-x86_64-rocky8-py312-trt10.15.1.29-skip-tritondevel-202605060827-13616
-LLM_SBSA_WHEEL_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:cuda-13.1.0-devel-ubuntu24.04-sbsa-ubuntu24.04-py312-trt10.15.1.29-skip-tritondevel-202605060827-13616
+LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm-staging/tritondevel:x86_64-tritondevel-torch_skip-d697cde-github-pr-13841-38849
+LLM_SBSA_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm-staging/tritondevel:sbsa-tritondevel-torch_skip-d697cde-github-pr-13841-38849
+LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm-staging/tritondevel:x86_64-rockylinux8-torch_skip-py310-d697cde-github-pr-13841-38849
+LLM_ROCKYLINUX8_PY312_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm-staging/tritondevel:x86_64-rockylinux8-torch_skip-py312-d697cde-github-pr-13841-38849
+LLM_SBSA_WHEEL_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm-staging/tritondevel:sbsa-ubuntu24-torch_skip-py312-d697cde-github-pr-13841-38849

requirements-dev.txt

@@ -40,7 +40,7 @@ opentelemetry-semantic-conventions-ai>=0.4.1
 fuzzywuzzy==0.18.0
 aiperf==0.6.0
 nanobind>=2.9.0
-nixl==0.9.0
+nixl==1.0.1
 cupti-python>=13.0,<13.2
 nvidia-cuda-cupti>=13.0,<13.2
 cxxfilt

tests/unittest/disaggregated/test_agent.py

@@ -55,9 +55,12 @@ class MemoryManager:
     allocated_memory: list[torch.Tensor] = field(default_factory=list)
 
     def allocate_memory(
-        self, size: int, name: str, memory_type=MemoryType.VRAM, device_id: int = 0
+        self, size: int, name: str, memory_type: str = "VRAM", device_id: int = 0
     ) -> RegMemoryDescs:
-        device = torch.device(f"cuda:{device_id}" if memory_type == MemoryType.VRAM else "cpu")
+        # `memory_type` is the string used by RegMemoryDescs (see base/agent.py:81);
+        # do not compare against `MemoryType.VRAM`, which is overridden to a C++ enum
+        # when the C++ binding is available.
+        device = torch.device(f"cuda:{device_id}" if memory_type == "VRAM" else "cpu")
 
         # Allocate memory block using torch.Tensor and track it
         block = torch.zeros(size, dtype=torch.uint8, device=device)

tests/unittest/others/test_kv_cache_transceiver.py

@@ -261,6 +261,7 @@ def create_hybrid_cache_manager(mapping,
         mamba_layer_mask=mamba_layer_mask,
         mamba_cache_dtype=mamba_conv_dtype,
         mamba_ssm_cache_dtype=mamba_ssm_dtype,
+        is_disagg=True,
         kv_cache_config=KvCacheConfig(
             max_tokens=256,
             enable_block_reuse=False,