[TRTLLMINF-397][infra] Update dependencies to NGC PyTorch 26.08

ATTRIBUTIONS-Python.md

@@ -48450,7 +48450,7 @@ SOFTWARE.
 
 
 
-## polygraphy (0.49.26)
+## polygraphy (0.53.4)
 
 ### Licenses
 License: `Apache 2.0`

cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu

@@ -674,10 +674,16 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo
 // clang-format on
 
 // Generate explicit template instantiations.
+//
+// The parameter lists must repeat the definition's spelling verbatim, including `__restrict__` and any top-level
+// `const`. cudafe1 rewrites every kernel parameter into a reference when it emits `__wrapper__device_stub_*`, which
+// turns those otherwise-ignored qualifiers into part of the function type. Clang then rejects the instantiation;
+// GCC does not diagnose it.
 #define INST_DECODE(HD, KV_EB, STATE_EB, CR, NRW)                                                                      \
-    template __global__ void pagedKvCompressKernel<HD, KV_EB, STATE_EB, CR, NRW>(void const*, float const*, void*,     \
-        void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int, int, int,   \
-        int);
+    template __global__ void pagedKvCompressKernel<HD, KV_EB, STATE_EB, CR, NRW>(void const* __restrict__,             \
+        float const* __restrict__, void* __restrict__, void* __restrict__, int32_t const* __restrict__,                \
+        int32_t const* __restrict__, void* __restrict__, int32_t const* __restrict__, int32_t const* __restrict__,     \
+        int32_t const* __restrict__, int, int, int, int);
 FOREACH_DECODE_CONFIG(INST_DECODE)
 #undef INST_DECODE
 
@@ -1185,10 +1191,12 @@ __global__ void prefillReductionKernel(void const* __restrict__ kv_score_raw, fl
 // Explicit instantiations.  CR=128 uses four reduction groups to split the long
 // token loop; CR=4 stays single-warp because the reduction is too small to amortize
 // the merge overhead.
+// See the note above INST_DECODE: the parameter spelling must match the definition, `__restrict__` included.
 #define INST_PREFILL(HD, KV_EB, STATE_EB, CR, NRW)                                                                     \
-    template __global__ void prefillReductionKernel<HD, KV_EB, STATE_EB, CR, NRW>(void const*, float const*, void*,    \
-        void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int32_t const*,  \
-        int, int, int, int);
+    template __global__ void prefillReductionKernel<HD, KV_EB, STATE_EB, CR, NRW>(void const* __restrict__,            \
+        float const* __restrict__, void* __restrict__, void* __restrict__, int32_t const* __restrict__,                \
+        int32_t const* __restrict__, void* __restrict__, int32_t const* __restrict__, int32_t const* __restrict__,     \
+        int32_t const* __restrict__, int32_t const* __restrict__, int, int, int, int);
 
 #define INST_PREFILL_DTYPES(HD, CR, NRW)                                                                               \
     INST_PREFILL(HD, 2, 2, CR, NRW)                                                                                    \
@@ -1756,10 +1764,13 @@ __global__ void postProcessScatterKernel(void const* __restrict__ kv_comp, // [t
 // kNone supports bf16 (EB=2) and fp32 (EB=4) input types; the quantized
 // scale types only support bf16 input since the compressor output is bf16.
 // Each combination is instantiated with ROTATE_ACTIVATION=true and false.
+// See the note above INST_DECODE: the parameter spelling must match the definition, `__restrict__` included.
 #define INST_PPS(HD, EB, CST, AR)                                                                                      \
-    template __global__ void postProcessScatterKernel<HD, EB, CST, AR>(void const*, void*, void const*, float,         \
-        float const*, int32_t const*, int, int, void*, int32_t const*, int32_t const*, int32_t const*, int32_t const*, \
-        bool const*, int, int, int, int, int, int, void*, void*);
+    template __global__ void postProcessScatterKernel<HD, EB, CST, AR>(void const* __restrict__, void* __restrict__,   \
+        void const* __restrict__, float, float const* __restrict__, int32_t const* __restrict__, int, int,             \
+        void* __restrict__, int32_t const* __restrict__, int32_t const* __restrict__, int32_t const* __restrict__,     \
+        int32_t const* __restrict__, bool const* __restrict__, int, int, int, int, int, int, void* __restrict__,       \
+        void* __restrict__);
 
 #define INST_PPS_AR(HD, EB, CST)                                                                                       \
     INST_PPS(HD, EB, CST, true)                                                                                        \

cpp/tensorrt_llm/kernels/mhcKernels/mhcKernels.cu

@@ -322,13 +322,19 @@ __launch_bounds__(BLOCK_SIZE) __global__ void mhcBigFuseKernel(float const* __re
 #endif
 }
 
+// The parameter list below must repeat the definition's spelling verbatim, including `__restrict__` and any top-level
+// `const`. cudafe1 rewrites every kernel parameter into a reference when it emits `__wrapper__device_stub_*`, which
+// turns those otherwise-ignored qualifiers into part of the function type. Clang then rejects the instantiation;
+// GCC does not diagnose it.
 #define INST_BIGFUSE(NS, BS)                                                                                           \
-    template __global__ void mhcBigFuseKernel<NS, BS, /*kFuseNorm=*/false>(float const*, float const*,                 \
-        __nv_bfloat16 const*, float const*, float const*, float*, float*, __nv_bfloat16*, int, int, int, float, float, \
-        float, float, int, __nv_bfloat16 const*, float);                                                               \
-    template __global__ void mhcBigFuseKernel<NS, BS, /*kFuseNorm=*/true>(float const*, float const*,                  \
-        __nv_bfloat16 const*, float const*, float const*, float*, float*, __nv_bfloat16*, int, int, int, float, float, \
-        float, float, int, __nv_bfloat16 const*, float);
+    template __global__ void mhcBigFuseKernel<NS, BS, /*kFuseNorm=*/false>(float const* __restrict__,                  \
+        float const* __restrict__, __nv_bfloat16 const* __restrict__, float const* __restrict__,                       \
+        float const* __restrict__, float* __restrict__, float* __restrict__, __nv_bfloat16* __restrict__, int, int,    \
+        int, float, float, float, float, int, __nv_bfloat16 const* __restrict__, float);                               \
+    template __global__ void mhcBigFuseKernel<NS, BS, /*kFuseNorm=*/true>(float const* __restrict__,                   \
+        float const* __restrict__, __nv_bfloat16 const* __restrict__, float const* __restrict__,                       \
+        float const* __restrict__, float* __restrict__, float* __restrict__, __nv_bfloat16* __restrict__, int, int,    \
+        int, float, float, float, float, int, __nv_bfloat16 const* __restrict__, float);
 
 INST_BIGFUSE(1, 128)
 INST_BIGFUSE(1, 256)
@@ -535,9 +541,10 @@ __launch_bounds__(256) __global__ void mhcGemmSqrsumFmaKernel(__nv_bfloat16 cons
     }
 }
 
+// See the note above INST_BIGFUSE: the parameter spelling must match the definition, `__restrict__` included.
 #define INST_FMA(NPB)                                                                                                  \
-    template __global__ void mhcGemmSqrsumFmaKernel<NPB>(                                                              \
-        __nv_bfloat16 const*, float const*, float*, float*, int, int, int);
+    template __global__ void mhcGemmSqrsumFmaKernel<NPB>(__nv_bfloat16 const* __restrict__, float const* __restrict__, \
+        float* __restrict__, float* __restrict__, int, int, int);
 
 INST_FMA(1)
 INST_FMA(2)

cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.cpp

@@ -375,8 +375,8 @@ void* HostAccessibleDeviceAllocator::allocate(size_t memorySize)
         TLLM_CHECK_WITH_INFO(
             mAllowManagedFallback, "HostAccessibleDeviceAllocator is not supported on the current system.");
         TLLM_CUDA_CHECK(cudaMallocManaged(&devPtr, memorySize));
-        TLLM_CUDA_CHECK(cudaMemAdvise(
-            devPtr, memorySize, cudaMemAdviseSetPreferredLocation, {cudaMemLocationTypeDevice, currentDevId}));
+        TLLM_CUDA_CHECK(cudaMemAdvise(devPtr, memorySize, cudaMemAdviseSetPreferredLocation,
+            cudaMemLocation{cudaMemLocationTypeDevice, {currentDevId}}));
         hostPtr = devPtr;
     }
     recordAllocation(devPtr, memorySize, hostPtr, memDesc);

cpp/tensorrt_llm/thop/CMakeLists.txt

@@ -31,6 +31,13 @@ endif()
 add_library(th_utils STATIC thUtils.cpp)
 set_property(TARGET th_utils PROPERTY POSITION_INDEPENDENT_CODE ON)
 set_property(TARGET th_utils PROPERTY CUDA_RESOLVE_DEVICE_SYMBOLS ON)
+# Newer Torch headers (e.g. ATen/native/cuda/Resize.h) compare ArrayRef against
+# OptionalArrayRef in the reverse order, relying on C++20's automatic generation
+# of the reversed comparison operator. th_utils/th_common only contain host .cpp
+# sources, so scope the bump to these targets instead of the repo-wide C++17
+# standard.
+set_target_properties(th_utils PROPERTIES CXX_STANDARD 20 CXX_STANDARD_REQUIRED
+                                                          ON)
 # Declare the dependency on the main shared library explicitly so consumers
 # (e.g. thUtilsTest) place it after th_utils on the link line; this was
 # previously satisfied transitively via the removed TensorRT plugin target.
@@ -146,6 +153,8 @@ add_library(
   mhcOp.cpp
   compressorOp.cpp)
 set_property(TARGET th_common PROPERTY POSITION_INDEPENDENT_CODE ON)
+set_target_properties(th_common PROPERTIES CXX_STANDARD 20 CXX_STANDARD_REQUIRED
+                                                           ON)
 target_link_libraries(
   th_common PRIVATE ${TORCH_LIBRARIES} th_utils ${Python3_LIBRARIES}
                     ${SHARED_TARGET} pg_utils)

docker/Dockerfile.multi

@@ -15,8 +15,8 @@
 # Multi-stage Dockerfile
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch
 ARG TRITON_IMAGE=nvcr.io/nvidia/tritonserver
-ARG BASE_TAG=26.05-py3
-ARG TRITON_BASE_TAG=26.05-py3
+ARG BASE_TAG=26.08-py3
+ARG TRITON_BASE_TAG=26.08-py3
 ARG DEVEL_IMAGE=devel
 
 FROM ${BASE_IMAGE}:${BASE_TAG} AS base
@@ -131,6 +131,12 @@ RUN --mount=type=bind,source=docker/common,target=/opt/docker/common \
     fi && \
     bash /opt/docker/common/install_ci.sh
 
+# Allow MPI/PRTE to run as root inside the container
+ENV OMPI_ALLOW_RUN_AS_ROOT=1 \
+    OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 \
+    PRTE_ALLOW_RUN_AS_ROOT=1 \
+    PRTE_ALLOW_RUN_AS_ROOT_CONFIRM=1
+
 FROM ${DEVEL_IMAGE} AS wheel
 WORKDIR /src/tensorrt_llm
 COPY cpp cpp

docker/Makefile

@@ -40,7 +40,13 @@ DOCKER_PROGRESS    ?= auto
 PLATFORM           ?= $(shell uname -m | grep -q 'aarch64' && echo "arm64" || echo "amd64")
 CUDA_ARCHS         ?= $(if $(filter arm64,$(PLATFORM)),'90-real;100-real;120-real',)
 BUILD_WHEEL_OPTS   ?=
-BUILD_WHEEL_ARGS   ?= $(shell grep '^ARG BUILD_WHEEL_ARGS=' Dockerfile.multi | grep -o '=.*' | tr -d '="')$(if $(CUDA_ARCHS), --cuda_architectures $(CUDA_ARCHS))$(if $(BUILD_WHEEL_OPTS), $(BUILD_WHEEL_OPTS))
+# The DLFW base image ships no libnvrtc_static.a, so cmake fails on a missing
+# CUDA::nvrtc_static target unless NVRTC is linked dynamically. Only the wheel and
+# release stages read BUILD_WHEEL_ARGS, and those always build FROM the DLFW devel
+# image; the rockylinux8/ubuntu24 targets stop at tritondevel and never see it, so
+# this needs no condition. Kept out of Dockerfile.multi's ARG default because the
+# `tr -d '="'` below would strip the '=' out of NVRTC_DYNAMIC_LINKING=ON.
+BUILD_WHEEL_ARGS   ?= $(shell grep '^ARG BUILD_WHEEL_ARGS=' Dockerfile.multi | grep -o '=.*' | tr -d '="') --extra-cmake-vars NVRTC_DYNAMIC_LINKING=ON$(if $(CUDA_ARCHS), --cuda_architectures $(CUDA_ARCHS))$(if $(BUILD_WHEEL_OPTS), $(BUILD_WHEEL_OPTS))
 BUILD_WHEEL_SCRIPT ?=
 TORCH_INSTALL_TYPE ?= skip
 CUDA_VERSION       ?=
@@ -208,21 +214,21 @@ jenkins-rockylinux8_%: PYTHON_VERSION_TAG_ID = $(if $(findstring 3.12,${PYTHON_V
 jenkins-rockylinux8_%: IMAGE_WITH_TAG = $(shell . ../jenkins/current_image_tags.properties && echo $$LLM_ROCKYLINUX8_${PYTHON_VERSION_TAG_ID}_DOCKER_IMAGE)
 jenkins-rockylinux8_%: STAGE = tritondevel
 jenkins-rockylinux8_%: BASE_IMAGE = nvcr.io/nvidia/cuda
-jenkins-rockylinux8_%: BASE_TAG = 13.2.1-devel-rockylinux8
+jenkins-rockylinux8_%: BASE_TAG = 13.3.1-devel-rockylinux8
 
 rockylinux8_%: STAGE = tritondevel
 rockylinux8_%: BASE_IMAGE = nvcr.io/nvidia/cuda
-rockylinux8_%: BASE_TAG = 13.2.1-devel-rockylinux8
+rockylinux8_%: BASE_TAG = 13.3.1-devel-rockylinux8
 
 # For x86_64
 ubuntu22_%: STAGE = tritondevel
 ubuntu22_%: BASE_IMAGE = nvcr.io/nvidia/cuda
-ubuntu22_%: BASE_TAG = 13.2.1-devel-ubuntu22.04
+ubuntu22_%: BASE_TAG = 13.3.1-devel-ubuntu22.04
 
 # For x86_64 and aarch64
 ubuntu24_%: STAGE = tritondevel
 ubuntu24_%: BASE_IMAGE = nvcr.io/nvidia/cuda
-ubuntu24_%: BASE_TAG = 13.2.1-devel-ubuntu24.04
+ubuntu24_%: BASE_TAG = 13.3.1-devel-ubuntu24.04
 
 trtllm_%: STAGE = release
 trtllm_%: PUSH_TO_STAGING := 0

docker/common/install_base.sh

@@ -114,6 +114,17 @@ init_ubuntu() {
   # the apt packages leaves pip3 functional.
 
   echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> "${ENV}"
+
+  # PRRTE cannot derive its install dirs on its own: the HPC-X Open MPI 5 in the
+  # NGC PyTorch base image is relocated away from the prefix it was built with,
+  # so prte_info resolves to /build-result/hpcx-*, a build-machine path absent
+  # from the container, and `prte --version` cannot find its help files. Open
+  # MPI itself is fine -- the base image already exports OPAL_PREFIX and puts
+  # /usr/local/mpi/bin on PATH. Only the DLFW base images carry /opt/hpcx.
+  if [ -d /opt/hpcx/ompi5 ]; then
+    echo 'export PRTE_PREFIX=/opt/hpcx/ompi5' >> "${ENV}"
+  fi
+
   # Remove previous TRT installation
   if [[ $(apt list --installed | grep libnvinfer) ]]; then
     apt-get remove --purge -y libnvinfer*