[None][feat] Upgrade NIXL to v1.0.1 and UCX to 1.21

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,8 +1,8 @@
 #!/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"

jenkins/L0_Test.groovy

@@ -3344,7 +3344,24 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO
                     | tar -C "${llmSrc}/tensorrt_llm" -xv
             """
             withEnv(["MYPY_REQUIRE_BINDINGS=1"]) {
-                sh "cd ${llmSrc} && python3 -m pre_commit run type-check --all-files || (cat /root/.cache/pre-commit/pre-commit.log && /bin/false)"
+                // Strip the wheel's tensorrt_llm/libs (and its ucx/ subdir) out of
+                // LD_LIBRARY_PATH for this stage. The tar above populated
+                // ${llmSrc}/tensorrt_llm/{libs,bindings.*.so} from the wheel, so the
+                // source tree now has its own copies. With the wheel libs path on
+                // LD_LIBRARY_PATH, bindings.so's DT_NEEDED resolves libth_common.so
+                // to <wheel>/tensorrt_llm/libs/ while _common.py explicitly loads
+                // <src>/tensorrt_llm/libs/libth_common.so via torch.classes.load_library
+                // — two different absolute paths register the same Torch op twice
+                // and PyTorch aborts. Removing the wheel paths lets DT_NEEDED fall
+                // back to bindings.so's RUNPATH ($ORIGIN/libs = <src>/tensorrt_llm/libs/),
+                // matching the explicit load. This restores pre-#722cbdd071 behavior
+                // for type-check only; UCX/NIXL kv_cache_transceiver tests above
+                // still see the new LD_LIBRARY_PATH.
+                sh """
+                    TRTLLM_WHEEL_LIBS=\$(pip3 show tensorrt_llm | awk -F': ' '/^Location:/ { print \$2 }')/tensorrt_llm/libs
+                    export LD_LIBRARY_PATH=\$(echo "\$LD_LIBRARY_PATH" | tr ':' '\\n' | grep -vxF "\$TRTLLM_WHEEL_LIBS" | grep -vxF "\$TRTLLM_WHEEL_LIBS/ucx" | paste -sd:)
+                    cd ${llmSrc} && python3 -m pre_commit run type-check --all-files || (cat /root/.cache/pre-commit/pre-commit.log && /bin/false)
+                """
             }
 	}
     }
@@ -4410,6 +4427,20 @@ 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.
+                        // Use `pip3 show` (metadata only) instead of `import tensorrt_llm`,
+                        // because importing executes tensorrt_llm/__init__.py which prints a
+                        // version banner to stdout and would pollute the captured path.
+                        def trtllmLibsDir = sh(
+                            script: "pip3 show tensorrt_llm | grep \"Location\" | awk -F\":\" '{ gsub(/ /, \"\", \$2); print \$2\"/tensorrt_llm/libs\"}'",
+                            returnStdout: true,
+                        ).replaceAll("\\s","")
+                        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:pytorch-26.02-py3-x86_64-ubuntu24.04-trt10.15.1.29-skip-tritondevel-202605251043-14436
+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-202605251043-14436
+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-202605251043-14436
+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-202605251043-14436
+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-202605251043-14436

requirements-dev.txt

@@ -41,7 +41,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-cu13==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,