diff --git a/.gitattributes b/.gitattributes index de0b56a7310..e72ba0fe7b7 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,7 +1,8 @@ *.a filter=lfs diff=lfs merge=lfs -text +*.dll filter=lfs diff=lfs merge=lfs -text *.lib filter=lfs diff=lfs merge=lfs -text *.so filter=lfs diff=lfs merge=lfs -text -*.dll filter=lfs diff=lfs merge=lfs -text +*.txz filter=lfs diff=lfs merge=lfs -text *.xz filter=lfs diff=lfs merge=lfs -text triton_backend/tools/gpt/input_data.json filter=lfs diff=lfs merge=lfs -text *cubin.cpp filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore index 3313ddd2752..a4c45ea6915 100644 --- a/.gitignore +++ b/.gitignore @@ -40,6 +40,9 @@ tensorrt_llm/libs tensorrt_llm/bindings.*.so tensorrt_llm/bindings.pyi tensorrt_llm/bindings/**/*.pyi +tensorrt_llm/deep_ep/ +tensorrt_llm/deep_ep_cpp_tllm.*.so +tensorrt_llm/deep_ep_cpp_tllm.pyi *docs/cpp_docs* *docs/source/_cpp_gen* docs/source/**/*.rst @@ -55,6 +58,7 @@ llm-test-workspace/ *.safetensors */tllm_debug/** *.patch +!cpp/tensorrt_llm/deep_ep/*.patch # Generated files cpp/include/tensorrt_llm/executor/version.h diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7f943c71141..00516b1afa7 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -27,6 +27,7 @@ repos: args: [--allow-multiple-documents] exclude: ".*/gitlab/.*.yml" - id: trailing-whitespace + exclude: '\.patch$' - id: check-toml - id: mixed-line-ending args: [--fix=lf] diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 5c2c3f0d8f7..f1a617b793f 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -297,4 +297,8 @@ if(BUILD_PYBIND) add_subdirectory(pybind) endif() +if(BUILD_DEEP_EP) + add_subdirectory(deep_ep) +endif() + add_subdirectory(plugins) diff --git a/cpp/tensorrt_llm/deep_ep/CMakeLists.txt b/cpp/tensorrt_llm/deep_ep/CMakeLists.txt new file mode 100644 index 00000000000..3e9874973b8 --- /dev/null +++ b/cpp/tensorrt_llm/deep_ep/CMakeLists.txt @@ -0,0 +1,207 @@ +set(DEEP_EP_COMMIT c381dadf43a85062f6a8947592017ee513abc70b) +set(NVSHMEM_URL_HASH + SHA256=eb2c8fb3b7084c2db86bd9fd905387909f1dfd483e7b45f7b3c3d5fcf5374b5a) + +add_custom_target(deep_ep) + +# CUDA architectures +# ================== + +# Filter CUDA arch >= 9.0 +set(DEEP_EP_CUDA_ARCHITECTURES "") +foreach(CUDA_ARCH IN LISTS CMAKE_CUDA_ARCHITECTURES) + string(REGEX MATCHALL "^([1-9][0-9]*)([0-9])[af]?(-real|-virtual)?$" MATCHES + ${CUDA_ARCH}) + if(NOT CMAKE_MATCH_0) + message(FATAL_ERROR "Invalid CUDA arch format: \"${CUDA_ARCH}\"") + endif() + set(CUDA_ARCH_MAJOR ${CMAKE_MATCH_1}) + set(CUDA_ARCH_MINOR ${CMAKE_MATCH_2}) + set(CUDA_ARCH_POSTFIX ${CMAKE_MATCH_3}) + if(${CUDA_ARCH_MAJOR} GREATER_EQUAL 9) + list(APPEND DEEP_EP_CUDA_ARCHITECTURES + "${CUDA_ARCH_MAJOR}${CUDA_ARCH_MINOR}${CUDA_ARCH_POSTFIX}") + endif() +endforeach() + +# Skip build if there is no suitable CUDA arch +if(WIN32) + set(DEEP_EP_CUDA_ARCHITECTURES "") +endif() +message( + STATUS "deep_ep DEEP_EP_CUDA_ARCHITECTURES: ${DEEP_EP_CUDA_ARCHITECTURES}") +file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/cuda_architectures.txt + "${DEEP_EP_CUDA_ARCHITECTURES}") +if(NOT DEEP_EP_CUDA_ARCHITECTURES) + return() +endif() + +# Prepare files +# ============= + +# Download DeepEP +include(FetchContent) +if(DEFINED ENV{GITHUB_MIRROR} AND NOT "$ENV{GITHUB_MIRROR}" STREQUAL "") + set(GITHUB_URL "$ENV{GITHUB_MIRROR}") +else() + set(GITHUB_URL "https://github.com") +endif() +set(DEEP_EP_URL + "${GITHUB_URL}/deepseek-ai/DeepEP/archive/${DEEP_EP_COMMIT}.tar.gz") +message(STATUS "deep_ep DEEP_EP_URL: ${DEEP_EP_URL}") +FetchContent_Declare(deep_ep_download URL ${DEEP_EP_URL}) +FetchContent_MakeAvailable(deep_ep_download) +set(DEEP_EP_SOURCE_DIR ${deep_ep_download_SOURCE_DIR}) + +# Copy and update python files +set(DEEP_EP_PYTHON_DEST ${CMAKE_CURRENT_BINARY_DIR}/python/deep_ep) +file(REMOVE_RECURSE ${DEEP_EP_PYTHON_DEST}) +file(MAKE_DIRECTORY ${DEEP_EP_PYTHON_DEST}) +configure_file(${DEEP_EP_SOURCE_DIR}/LICENSE ${DEEP_EP_PYTHON_DEST}/LICENSE + COPYONLY) +set(_files __init__.py buffer.py utils.py) +foreach(_f IN LISTS _files) + set(_src "${DEEP_EP_SOURCE_DIR}/deep_ep/${_f}") + set(_dst "${DEEP_EP_PYTHON_DEST}/${_f}") + file(READ "${_src}" _content) + string(REPLACE "deep_ep_cpp" "tensorrt_llm.deep_ep_cpp_tllm" _content + "${_content}") + string( + PREPEND + _content + "# Adapted from https://github.com/deepseek-ai/DeepEP/blob/${DEEP_EP_COMMIT}/deep_ep/${_f}\n" + ) + file(WRITE "${_dst}" "${_content}") + set_property( + DIRECTORY + APPEND + PROPERTY CMAKE_CONFIGURE_DEPENDS ${_src}) +endforeach() + +# Delete stale nvshmem on patch update +set(NVSHMEM_STAMP_FILE ${CMAKE_CURRENT_BINARY_DIR}/nvshmem_stamp.txt) +file(SHA256 ${DEEP_EP_SOURCE_DIR}/third-party/nvshmem.patch NVSHMEM_PATCH_HASH) +file(SHA256 ${CMAKE_CURRENT_SOURCE_DIR}/nvshmem_fast_build.patch + NVSHMEM_PATCH_2_HASH) +set(NVSHMEM_STAMP_CONTENT "${NVSHMEM_URL_HASH}") +string(APPEND NVSHMEM_STAMP_CONTENT " PATCH_COMMAND v1") +string(APPEND NVSHMEM_STAMP_CONTENT " ${NVSHMEM_PATCH_HASH}") +string(APPEND NVSHMEM_STAMP_CONTENT " 103") +string(APPEND NVSHMEM_STAMP_CONTENT " ${NVSHMEM_PATCH_2_HASH}") +set(OLD_NVSHMEM_STAMP_CONTENT "") +if(EXISTS ${NVSHMEM_STAMP_FILE}) + file(READ ${NVSHMEM_STAMP_FILE} OLD_NVSHMEM_STAMP_CONTENT) +endif() +if(NOT OLD_NVSHMEM_STAMP_CONTENT STREQUAL NVSHMEM_STAMP_CONTENT) + file(REMOVE_RECURSE ${CMAKE_CURRENT_BINARY_DIR}/nvshmem_project-prefix) + file(WRITE ${NVSHMEM_STAMP_FILE} "${NVSHMEM_STAMP_CONTENT}") +endif() +set_property( + DIRECTORY APPEND + PROPERTY CMAKE_CONFIGURE_DEPENDS + ${DEEP_EP_SOURCE_DIR}/third-party/nvshmem.patch + ${CMAKE_CURRENT_SOURCE_DIR}/nvshmem_fast_build.patch) + +# Add NVSHMEM +# =========== + +# NVSHMEM only works with GCC. Building NVSHMEM with Clang results in +# compilation errors. Using NVSHMEM with Clang results in slow builds and device +# link issues. +if(NOT CMAKE_CXX_COMPILER_ID STREQUAL "GNU") + set(CMAKE_C_COMPILER gcc) + set(CMAKE_CXX_COMPILER g++) + set(CMAKE_CUDA_HOST_COMPILER g++) +endif() + +# Add nvshmem external project +include(ExternalProject) +ExternalProject_Add( + nvshmem_project + URL file://${CMAKE_CURRENT_SOURCE_DIR}/nvshmem_src_3.2.5-1.txz + URL_HASH ${NVSHMEM_URL_HASH} + PATCH_COMMAND patch -p1 --forward --batch -i + ${DEEP_EP_SOURCE_DIR}/third-party/nvshmem.patch + COMMAND sed "s/TRANSPORT_VERSION_MAJOR 3/TRANSPORT_VERSION_MAJOR 103/" -i + src/CMakeLists.txt + COMMAND patch -p1 --forward --batch -i + ${CMAKE_CURRENT_SOURCE_DIR}/nvshmem_fast_build.patch + CMAKE_CACHE_ARGS + -DCMAKE_C_COMPILER:STRING=${CMAKE_C_COMPILER} + -DCMAKE_C_COMPILER_LAUNCHER:STRING=${CMAKE_C_COMPILER_LAUNCHER} + -DCMAKE_CXX_COMPILER:STRING=${CMAKE_CXX_COMPILER} + -DCMAKE_CXX_COMPILER_LAUNCHER:STRING=${CMAKE_CXX_COMPILER_LAUNCHER} + -DCMAKE_CUDA_ARCHITECTURES:STRING=${DEEP_EP_CUDA_ARCHITECTURES} + -DCMAKE_CUDA_HOST_COMPILER:STRING=${CMAKE_CUDA_HOST_COMPILER} + -DCMAKE_CUDA_COMPILER_LAUNCHER:STRING=${CMAKE_CUDA_COMPILER_LAUNCHER} + -DNVSHMEM_BUILD_EXAMPLES:BOOL=0 + -DNVSHMEM_BUILD_PACKAGES:BOOL=0 + -DNVSHMEM_BUILD_TESTS:BOOL=0 + -DNVSHMEM_IBGDA_SUPPORT:BOOL=1 + -DNVSHMEM_IBRC_SUPPORT:BOOL=0 + -DNVSHMEM_MPI_SUPPORT:BOOL=0 + -DNVSHMEM_PMIX_SUPPORT:BOOL=0 + -DNVSHMEM_SHMEM_SUPPORT:BOOL=0 + -DNVSHMEM_TIMEOUT_DEVICE_POLLING:BOOL=0 + -DNVSHMEM_UCX_SUPPORT:BOOL=0 + -DNVSHMEM_USE_GDRCOPY:BOOL=0 + -DNVSHMEM_USE_NCCL:BOOL=0 + INSTALL_COMMAND "" + BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/nvshmem-build + BUILD_BYPRODUCTS + ${CMAKE_CURRENT_BINARY_DIR}/nvshmem-build/src/lib/libnvshmem.a) +add_library(nvshmem_project::nvshmem STATIC IMPORTED) +add_dependencies(nvshmem_project::nvshmem nvshmem_project) +file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/nvshmem-build/src/include) +set_target_properties( + nvshmem_project::nvshmem + PROPERTIES IMPORTED_LOCATION + ${CMAKE_CURRENT_BINARY_DIR}/nvshmem-build/src/lib/libnvshmem.a + INTERFACE_INCLUDE_DIRECTORIES + ${CMAKE_CURRENT_BINARY_DIR}/nvshmem-build/src/include) + +# Add DeepEP cpp +# ============== + +# Let CMake generate `fatbinData` for CUDA separable compilation. Set to FALSE +# or TRUE are both OK, but it generates `code=lto_90a` rather than `code=sm_90a` +# for arch `90a-real` if set to TRUE. +set(CMAKE_INTERPROCEDURAL_OPTIMIZATION FALSE) + +# Find torch_python +find_library(TORCH_PYTHON_LIB torch_python REQUIRED + HINTS ${TORCH_INSTALL_PREFIX}/lib) + +# Add deep_ep_cpp_tllm +file(GLOB_RECURSE SRC_CPP ${DEEP_EP_SOURCE_DIR}/csrc/*.cpp) +file(GLOB_RECURSE SRC_CU ${DEEP_EP_SOURCE_DIR}/csrc/*.cu) +pybind11_add_module(deep_ep_cpp_tllm ${SRC_CPP} ${SRC_CU}) +set_target_properties( + deep_ep_cpp_tllm + PROPERTIES CXX_STANDARD_REQUIRED ON + CUDA_STANDARD_REQUIRED ON + CXX_STANDARD 17 + CUDA_STANDARD 17 + CUDA_SEPARABLE_COMPILATION ON + CUDA_ARCHITECTURES ${DEEP_EP_CUDA_ARCHITECTURES} + LINK_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/deep_ep_cpp_tllm.version + INSTALL_RPATH "$ORIGIN/libs/nvshmem;${TORCH_INSTALL_PREFIX}/lib" + BUILD_WITH_INSTALL_RPATH TRUE) +target_compile_options( + deep_ep_cpp_tllm + PRIVATE ${TORCH_CXX_FLAGS} -O3 $<$:-Xcompiler=-O3> + $<$:--ptxas-options=--register-usage-level=10>) +target_compile_definitions( + deep_ep_cpp_tllm PRIVATE DISABLE_AGGRESSIVE_PTX_INSTRS + TORCH_EXTENSION_NAME=deep_ep_cpp_tllm) +target_link_libraries( + deep_ep_cpp_tllm PRIVATE nvshmem_project::nvshmem ${TORCH_LIBRARIES} + ${TORCH_PYTHON_LIB}) +target_link_options( + deep_ep_cpp_tllm PRIVATE + -Wl,--version-script,${CMAKE_CURRENT_SOURCE_DIR}/deep_ep_cpp_tllm.version + -Wl,--no-undefined-version) + +# Set targets +# =========== +add_dependencies(deep_ep deep_ep_cpp_tllm nvshmem_project) diff --git a/cpp/tensorrt_llm/deep_ep/README.md b/cpp/tensorrt_llm/deep_ep/README.md new file mode 100644 index 00000000000..5c1ea0c797f --- /dev/null +++ b/cpp/tensorrt_llm/deep_ep/README.md @@ -0,0 +1,8 @@ +How to generate `nvshmem_fast_build.patch`? + +1. Build the project without applying the `nvshmem_fast_build.patch`. +2. Link NVSHMEM to DeepEP with one NVSHMEM object file omitted. +3. Repeat step 2 until no more object files can be omitted. +4. Remove the unused files from NVSHMEM's `CMakelists.txt`, and save the differences as `nvshmem_fast_build.patch`. + +The script `strip_nvshmem_helper.py` automatically performs steps 2 and 3. diff --git a/cpp/tensorrt_llm/deep_ep/deep_ep_cpp_tllm.version b/cpp/tensorrt_llm/deep_ep/deep_ep_cpp_tllm.version new file mode 100644 index 00000000000..f883136167b --- /dev/null +++ b/cpp/tensorrt_llm/deep_ep/deep_ep_cpp_tllm.version @@ -0,0 +1,4 @@ +{ + global: PyInit_deep_ep_cpp_tllm; + local: *; +}; diff --git a/cpp/tensorrt_llm/deep_ep/nvshmem_fast_build.patch b/cpp/tensorrt_llm/deep_ep/nvshmem_fast_build.patch new file mode 100644 index 00000000000..1b79c85841f --- /dev/null +++ b/cpp/tensorrt_llm/deep_ep/nvshmem_fast_build.patch @@ -0,0 +1,66 @@ +diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt +index cba899bba..c27337601 100644 +--- a/src/CMakeLists.txt ++++ b/src/CMakeLists.txt +@@ -264,48 +264,20 @@ set(NVSHMEM_HOST_SOURCES_NOMAXREGCOUNT + host/comm/rma.cu + host/stream/comm/quiet_on_stream.cu + host/stream/comm/cuda_interface_sync.cu +- host/stream/coll/alltoall/alltoall.cu + host/stream/coll/barrier/barrier.cu +- host/stream/coll/broadcast/broadcast.cu +- host/stream/coll/fcollect/fcollect.cu +- host/stream/coll/rdxn/reduce_and.cu +- host/stream/coll/rdxn/reduce_or.cu +- host/stream/coll/rdxn/reduce_xor.cu +- host/stream/coll/rdxn/reduce_min.cu + host/stream/coll/rdxn/reduce_max.cu +- host/stream/coll/rdxn/reduce_prod.cu +- host/stream/coll/rdxn/reduce_sum.cu + host/stream/coll/rdxn/reduce_team.cu +- host/stream/coll/reducescatter/reducescatter_and.cu +- host/stream/coll/reducescatter/reducescatter_or.cu +- host/stream/coll/reducescatter/reducescatter_xor.cu +- host/stream/coll/reducescatter/reducescatter_min.cu +- host/stream/coll/reducescatter/reducescatter_max.cu +- host/stream/coll/reducescatter/reducescatter_prod.cu +- host/stream/coll/reducescatter/reducescatter_sum.cu + ) + + set(NVSHMEM_HOST_SOURCES + host/bootstrap/bootstrap.cpp + host/bootstrap/bootstrap_loader.cpp + host/coll/cpu_coll.cpp +- host/coll/alltoall/alltoall.cpp +- host/coll/alltoall/alltoall_on_stream.cpp + host/coll/barrier/barrier.cpp + host/coll/barrier/barrier_on_stream.cpp +- host/coll/broadcast/broadcast.cpp +- host/coll/broadcast/broadcast_on_stream.cpp +- host/coll/fcollect/fcollect.cpp +- host/coll/fcollect/fcollect_on_stream.cpp +- host/coll/rdxn/rdxn.cpp +- host/coll/rdxn/rdxn_on_stream.cpp +- host/coll/reducescatter/reducescatter.cpp +- host/coll/reducescatter/reducescatter_on_stream.cpp + host/comm/putget.cpp +- host/comm/fence.cpp + host/comm/quiet.cpp + host/comm/sync.cpp +- host/comm/amo.cpp + host/proxy/proxy.cpp + host/transport/transport.cpp + host/transport/p2p/p2p.cpp +@@ -1006,3 +978,12 @@ set(CPACK_RPM_PACKAGE_REQUIRES_PREUN "/sbin/ldconfig") + + include(CPack) + # End Installation definitions ++ ++set_target_properties( ++ git_commit ++ nvshmem_device_project ++ nvshmem_bootstrap_pmi ++ nvshmem_bootstrap_pmi2 ++ nvshmem_host ++ nvshmem-info ++ PROPERTIES EXCLUDE_FROM_ALL TRUE) diff --git a/cpp/tensorrt_llm/deep_ep/nvshmem_src_3.2.5-1.txz b/cpp/tensorrt_llm/deep_ep/nvshmem_src_3.2.5-1.txz new file mode 100644 index 00000000000..76a5a2b78ed --- /dev/null +++ b/cpp/tensorrt_llm/deep_ep/nvshmem_src_3.2.5-1.txz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:eb2c8fb3b7084c2db86bd9fd905387909f1dfd483e7b45f7b3c3d5fcf5374b5a +size 618175 diff --git a/cpp/tensorrt_llm/deep_ep/strip_nvshmem_helper.py b/cpp/tensorrt_llm/deep_ep/strip_nvshmem_helper.py new file mode 100644 index 00000000000..e9901558970 --- /dev/null +++ b/cpp/tensorrt_llm/deep_ep/strip_nvshmem_helper.py @@ -0,0 +1,61 @@ +# A helper script to detect unused NVSHMEM object files. +# +# The script links NVSHMEM to DeepEP with one object file removed at a time and +# checks whether there are any undefined symbols. See README.md for details. +# This script is not tested or QA'ed, so you may need to update this script if +# the project structure changes or compilation options change. +import pathlib +import re +import subprocess + +project_dir = pathlib.Path(__file__).parent.parent.parent.parent + +# Run `find cpp/build | grep kernels/internode_ll.cu.o$` to get the directory +deep_ep_obj_dir = project_dir / "cpp/build/tensorrt_llm/deep_ep/CMakeFiles/deep_ep_cpp_tllm.dir/__/__/_deps/deep_ep_download-src/csrc" +assert deep_ep_obj_dir.is_dir() + +# Run `find cpp/build | grep host/bootstrap/bootstrap.cpp.o$` to get the directory +# Please set it to `nvshmem.dir` rather than `nvshmem_host.dir` +nvshmem_obj_dir = project_dir / "cpp/build/tensorrt_llm/deep_ep/nvshmem-build/src/CMakeFiles/nvshmem.dir" +assert nvshmem_obj_dir.is_dir() + +# Parse the `-gencode` arguments +with (project_dir / + "cpp/build/tensorrt_llm/deep_ep/cuda_architectures.txt").open() as f: + cuda_architectures = f.read() +pattern = re.compile(r'^([1-9][0-9]*[0-9][af]?)(-real|-virtual)?$') +gencode_args = [] +for cuda_arch in cuda_architectures.split(";"): + matches = re.match(pattern, cuda_arch) + assert matches is not None, f"Invalid cuda arch \"{cuda_arch}\"" + sm_version = matches.group(1) + postfix = matches.group(2) or "" + code = { + "": f"[compute_{sm_version},sm_{sm_version}]", + "-real": f"[sm_{sm_version}]", + "-virtual": f"[compute_{sm_version}]", + }[postfix] + gencode_args.append(f"-gencode=arch=compute_{sm_version},{code=:s}") + +temp_dir = project_dir / "cpp/build/tensorrt_llm/deep_ep/strip_nvshmem_helper" +temp_dir.mkdir(exist_ok=True) +ranlib = temp_dir / "liba.a" +if ranlib.exists(): + ranlib.unlink() + +deep_ep_obj_list = sorted(deep_ep_obj_dir.glob("kernels/**/*.o")) +nvshmem_obj_set = set(nvshmem_obj_dir.glob("**/*.o")) +for exclude_obj in sorted(nvshmem_obj_set): + # Create liba.a with one object file removed + subprocess.check_call( + ["ar", "rcs", ranlib, *(nvshmem_obj_set - {exclude_obj})]) + # Test whether there are undefined symbols + res = subprocess.call([ + "/usr/local/cuda/bin/nvcc", *gencode_args, "-Xlinker", "--no-undefined", + "-shared", *deep_ep_obj_list, ranlib, "-o", temp_dir / "a.out" + ]) + # If there are no undefined symbols, print "-" to indicate the file can be omitted + print("-" if res == 0 else "+", + str(exclude_obj.relative_to(nvshmem_obj_dir))[:-2]) + # Unlink the archive file because `ar` appends existing archives + ranlib.unlink() diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index 2fe4991fe8e..b456be578da 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -16,7 +16,6 @@ LABEL com.nvidia.ai-terms="https://www.nvidia.com/en-us/agreements/enterprise-so ENV BASH_ENV=${BASH_ENV:-/etc/bash.bashrc} ENV ENV=${ENV:-/etc/shinit_v2} ARG GITHUB_MIRROR="" -ENV GITHUB_MIRROR=$GITHUB_MIRROR RUN echo "Using GitHub mirror: $GITHUB_MIRROR" SHELL ["/bin/bash", "-c"] @@ -28,13 +27,13 @@ FROM base AS devel ARG PYTHON_VERSION="3.12.3" RUN echo "Using Python version: $PYTHON_VERSION" COPY docker/common/install_base.sh install_base.sh -RUN bash ./install_base.sh $PYTHON_VERSION && rm install_base.sh +RUN GITHUB_MIRROR=$GITHUB_MIRROR bash ./install_base.sh $PYTHON_VERSION && rm install_base.sh COPY docker/common/install_cmake.sh install_cmake.sh -RUN bash ./install_cmake.sh && rm install_cmake.sh +RUN GITHUB_MIRROR=$GITHUB_MIRROR bash ./install_cmake.sh && rm install_cmake.sh COPY docker/common/install_ccache.sh install_ccache.sh -RUN bash ./install_ccache.sh && rm install_ccache.sh +RUN GITHUB_MIRROR=$GITHUB_MIRROR bash ./install_ccache.sh && rm install_ccache.sh # Only take effect when the base image is Rocky Linux 8 with old CUDA version. COPY docker/common/install_cuda_toolkit.sh install_cuda_toolkit.sh @@ -61,7 +60,7 @@ RUN bash ./install_polygraphy.sh && rm install_polygraphy.sh # Install mpi4py COPY docker/common/install_mpi4py.sh install_mpi4py.sh -RUN bash ./install_mpi4py.sh && rm install_mpi4py.sh +RUN GITHUB_MIRROR=$GITHUB_MIRROR bash ./install_mpi4py.sh && rm install_mpi4py.sh # Install PyTorch ARG TORCH_INSTALL_TYPE="skip" @@ -72,10 +71,6 @@ RUN bash ./install_pytorch.sh $TORCH_INSTALL_TYPE && rm install_pytorch.sh RUN pip3 uninstall -y opencv && rm -rf /usr/local/lib/python3*/dist-packages/cv2/ RUN pip3 install opencv-python-headless --force-reinstall --no-deps --no-cache-dir -# Install DeepEP -COPY docker/common/install_deep_ep.sh install_deep_ep.sh -RUN bash ./install_deep_ep.sh && rm install_deep_ep.sh - # WARs against security issues inherited from pytorch:25.04 # * https://github.com/advisories/GHSA-vqfr-h8mv-ghfj # * https://github.com/advisories/GHSA-7cx3-6m66-7c5m @@ -120,9 +115,10 @@ COPY .gitmodules setup.py requirements.txt requirements-dev.txt constraints.txt RUN mkdir -p /root/.cache/pip /root/.cache/ccache ENV CCACHE_DIR=/root/.cache/ccache # Build the TRT-LLM wheel +ARG GITHUB_MIRROR="" ARG BUILD_WHEEL_ARGS="--clean --python_bindings --benchmarks" RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=${CCACHE_DIR} \ - python3 scripts/build_wheel.py ${BUILD_WHEEL_ARGS} + GITHUB_MIRROR=$GITHUB_MIRROR python3 scripts/build_wheel.py ${BUILD_WHEEL_ARGS} FROM ${DEVEL_IMAGE} AS release diff --git a/docker/common/install_base.sh b/docker/common/install_base.sh index 0b4eb91ca8a..d1c2f036d63 100644 --- a/docker/common/install_base.sh +++ b/docker/common/install_base.sh @@ -53,6 +53,7 @@ init_ubuntu() { llvm \ libclang-rt-dev \ libffi-dev \ + libibverbs-dev \ libnuma1 \ libnuma-dev \ python3-dev \ @@ -115,17 +116,19 @@ install_gcctoolset_rockylinux() { # https://catalog.ngc.nvidia.com/orgs/nvidia/containers/cuda echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> "${ENV}" dnf install \ - vim \ - wget \ - git-lfs \ - gcc-toolset-11 \ - libffi-devel \ - -y + patch \ + vim \ + wget \ + git-lfs \ + gcc-toolset-11 \ + libffi-devel \ + -y dnf install \ - openmpi \ - openmpi-devel \ - pigz \ - -y + openmpi \ + openmpi-devel \ + pigz \ + rdma-core-devel \ + -y echo "source scl_source enable gcc-toolset-11" >> "${ENV}" echo 'export PATH=/usr/lib64/openmpi/bin:$PATH' >> "${ENV}" } diff --git a/docker/common/install_deep_ep.sh b/docker/common/install_deep_ep.sh deleted file mode 100644 index c6c572eff9e..00000000000 --- a/docker/common/install_deep_ep.sh +++ /dev/null @@ -1,47 +0,0 @@ -#!/bin/bash - -set -euxo pipefail - -GITHUB_URL=${GITHUB_MIRROR:-https://github.com} -DEEP_EP_COMMIT=2b266cf6452134f993ab0fcb3ef2d5de7683c561 - -if [ "$(. /etc/os-release && echo $ID)" == "rocky" ]; then - echo "Skipping DeepEP installation in the Rocky distribution." - exit 0 -fi -libmlx5_dir=$(dirname $(ldconfig -p | grep libmlx5.so.1 | head -n1 | awk '{print $NF}')) - -export NVCC_APPEND_FLAGS="--threads 4" - -# Custom NVSHMEM -curl -fsSL https://developer.download.nvidia.com/compute/redist/nvshmem/3.2.5/source/nvshmem_src_3.2.5-1.txz | tar xz -pushd nvshmem_src -curl -fsSL $GITHUB_URL/deepseek-ai/DeepEP/raw/$DEEP_EP_COMMIT/third-party/nvshmem.patch | patch -p1 -sed "s/TRANSPORT_VERSION_MAJOR 3/TRANSPORT_VERSION_MAJOR 103/" -i src/CMakeLists.txt -ln -s libmlx5.so.1 "$libmlx5_dir/libmlx5.so" -cmake -S . -B build \ - -DCMAKE_INSTALL_PREFIX=/opt/custom_nvshmem \ - -DGDRCOPY_HOME=/usr/include \ - -DNVSHMEM_SHMEM_SUPPORT=0 \ - -DNVSHMEM_UCX_SUPPORT=0 \ - -DNVSHMEM_USE_NCCL=0 \ - -DNVSHMEM_MPI_SUPPORT=0 \ - -DNVSHMEM_IBGDA_SUPPORT=1 \ - -DNVSHMEM_PMIX_SUPPORT=0 \ - -DNVSHMEM_TIMEOUT_DEVICE_POLLING=0 \ - -DNVSHMEM_USE_GDRCOPY=1 \ - -DCMAKE_CUDA_ARCHITECTURES="90-real;100-real;120-real" \ - -DNVSHMEM_BUILD_TESTS=0 \ - -DNVSHMEM_BUILD_EXAMPLES=0 -cmake --build build -j`nproc` -make -C build install -popd - -# DeepEP -curl -fsSL $GITHUB_URL/deepseek-ai/DeepEP/archive/$DEEP_EP_COMMIT.tar.gz | tar xz -TORCH_CUDA_ARCH_LIST="9.0;10.0;12.0" NVSHMEM_DIR=/opt/custom_nvshmem pip install -v --no-cache-dir ./DeepEP-$DEEP_EP_COMMIT - -# Clean up -rm -r nvshmem_src -rm "$libmlx5_dir/libmlx5.so" -rm -r DeepEP-$DEEP_EP_COMMIT diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index 0479a1ab3b1..51cd760425a 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -592,6 +592,7 @@ pipeline { //Workspace normally is: /home/jenkins/agent/workspace/LLM/L0_MergeRequest@tmp/ HF_HOME="${env.WORKSPACE_TMP}/.cache/huggingface" CCACHE_DIR="${CCACHE_DIR}" + GITHUB_MIRROR="https://urm.nvidia.com/artifactory/github-go-remote" PIP_INDEX_URL="https://urm.nvidia.com/artifactory/api/pypi/pypi-remote/simple" // force datasets to be offline mode, to prevent CI jobs are downloading HF dataset causing test failures HF_DATASETS_OFFLINE=1 diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 5e1ebcb2ba9..a2481aeb99b 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -2220,6 +2220,7 @@ pipeline { //Workspace normally is: /home/jenkins/agent/workspace/LLM/L0_MergeRequest@tmp/ HF_HOME="${env.WORKSPACE_TMP}/.cache/huggingface" CCACHE_DIR="${CCACHE_DIR}" + GITHUB_MIRROR="https://urm.nvidia.com/artifactory/github-go-remote" PIP_INDEX_URL="https://urm.nvidia.com/artifactory/api/pypi/pypi-remote/simple" // force datasets to be offline mode, to prevent CI jobs are downloading HF dataset causing test failures HF_DATASETS_OFFLINE=1 diff --git a/jenkins/controlCCache.groovy b/jenkins/controlCCache.groovy index e8590554708..0384904d67e 100644 --- a/jenkins/controlCCache.groovy +++ b/jenkins/controlCCache.groovy @@ -1,7 +1,7 @@ import java.lang.InterruptedException -DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.05-py3-x86_64-ubuntu24.04-trt10.11.0.33-skip-tritondevel-202506271620-5539" +DOCKER_IMAGE = "urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.05-py3-x86_64-ubuntu24.04-trt10.11.0.33-skip-tritondevel-202507071100-5534" def createKubernetesPodConfig(image, arch = "amd64") { diff --git a/jenkins/current_image_tags.properties b/jenkins/current_image_tags.properties index 4ac0d51d3e7..cf86b5292b2 100644 --- a/jenkins/current_image_tags.properties +++ b/jenkins/current_image_tags.properties @@ -5,7 +5,7 @@ # for reuse in Dev Containers configuration. # Also, the file needs to be parseable by 'sh' for reuse by docker/Makefile. LLM_DOCKER_IMAGE_URI=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm -LLM_DOCKER_IMAGE_TAG_SUFFIX=-trt10.11.0.33-skip-tritondevel-202506271620-5539 +LLM_DOCKER_IMAGE_TAG_SUFFIX=-trt10.11.0.33-skip-tritondevel-202507071100-5534 LLM_DOCKER_IMAGE=${LLM_DOCKER_IMAGE_URI}:pytorch-25.05-py3-x86_64-ubuntu24.04${LLM_DOCKER_IMAGE_TAG_SUFFIX} LLM_SBSA_DOCKER_IMAGE=${LLM_DOCKER_IMAGE_URI}:pytorch-25.05-py3-aarch64-ubuntu24.04${LLM_DOCKER_IMAGE_TAG_SUFFIX} LLM_ROCKYLINUX8_PY310_DOCKER_IMAGE=${LLM_DOCKER_IMAGE_URI}:cuda-12.9.0-devel-rocky8-x86_64-rocky8-py310${LLM_DOCKER_IMAGE_TAG_SUFFIX} diff --git a/jenkins/license_cpp.json b/jenkins/license_cpp.json index e29d577d80e..d7cf0eea984 100644 --- a/jenkins/license_cpp.json +++ b/jenkins/license_cpp.json @@ -31,6 +31,7 @@ ], "skip": { "": "", + "tensorrt_llm/deep_ep/strip_nvshmem_helper.py": "py", "tensorrt_llm/kernels/selectiveScan/selectiveScan.h": "external", "tensorrt_llm/kernels/cutlass_kernels/python/generate_kernels.py": "py", "tensorrt_llm/kernels/cutlass_kernels/fp8_rowwise_gemm/fp8_rowwise_gemm_kernel_template_sm90.h": "dual license", diff --git a/scripts/build_wheel.py b/scripts/build_wheel.py index 3670596bcff..fd2d7f252f3 100755 --- a/scripts/build_wheel.py +++ b/scripts/build_wheel.py @@ -412,10 +412,12 @@ def main(*, if cpp_only: build_pyt = "OFF" build_pybind = "OFF" + build_deep_ep = "OFF" else: - targets.extend(["bindings", "th_common"]) + targets.extend(["th_common", "bindings", "deep_ep"]) build_pyt = "ON" build_pybind = "ON" + build_deep_ep = "ON" if benchmarks: targets.append("benchmarks") @@ -454,7 +456,7 @@ def main(*, ) cmake_def_args = " ".join(cmake_def_args) cmake_configure_command = ( - f'cmake -DCMAKE_BUILD_TYPE="{build_type}" -DBUILD_PYT="{build_pyt}" -DBUILD_PYBIND="{build_pybind}"' + f'cmake -DCMAKE_BUILD_TYPE="{build_type}" -DBUILD_PYT="{build_pyt}" -DBUILD_PYBIND="{build_pybind}" -DBUILD_DEEP_EP="{build_deep_ep}"' f' -DNVTX_DISABLE="{disable_nvtx}" -DBUILD_MICRO_BENCHMARKS={build_micro_benchmarks}' f' -DBUILD_WHEEL_TARGETS="{";".join(targets)}"' f' -DPython_EXECUTABLE={venv_python} -DPython3_EXECUTABLE={venv_python}' @@ -594,6 +596,13 @@ def symlink_remove_dst_tree(src, dst, dirs_exist_ok=True): "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/libdecoder_attention_1.so", lib_dir / "libdecoder_attention_1.so") + deep_ep_dir = pkg_dir / "deep_ep" + if deep_ep_dir.is_symlink(): + deep_ep_dir.unlink() + elif deep_ep_dir.is_dir(): + clear_folder(deep_ep_dir) + deep_ep_dir.rmdir() + bin_dir = pkg_dir / "bin" if bin_dir.exists(): clear_folder(bin_dir) @@ -605,19 +614,41 @@ def symlink_remove_dst_tree(src, dst, dirs_exist_ok=True): if not cpp_only: - def get_pybind_lib(): - pybind_build_dir = (build_dir / "tensorrt_llm" / "pybind") + def get_pybind_lib(subdirectory, name): + pybind_build_dir = (build_dir / "tensorrt_llm" / subdirectory) if on_windows: - pybind_lib = list(pybind_build_dir.glob("bindings.*.pyd")) + pybind_lib = list(pybind_build_dir.glob(f"{name}.*.pyd")) else: - pybind_lib = list(pybind_build_dir.glob("bindings.*.so")) + pybind_lib = list(pybind_build_dir.glob(f"{name}.*.so")) assert len( pybind_lib ) == 1, f"Exactly one pybind library should be present: {pybind_lib}" return pybind_lib[0] - install_file(get_pybind_lib(), pkg_dir) + install_file(get_pybind_lib("pybind", "bindings"), pkg_dir) + + with (build_dir / "tensorrt_llm" / "deep_ep" / + "cuda_architectures.txt").open() as f: + deep_ep_cuda_architectures = f.read().strip().strip(";") + if deep_ep_cuda_architectures: + install_file(get_pybind_lib("deep_ep", "deep_ep_cpp_tllm"), pkg_dir) + install_tree(build_dir / "tensorrt_llm" / "deep_ep" / "python" / + "deep_ep", + deep_ep_dir, + dirs_exist_ok=True) + (lib_dir / "nvshmem").mkdir(exist_ok=True) + install_file( + build_dir / "tensorrt_llm/deep_ep/nvshmem-build/License.txt", + lib_dir / "nvshmem") + install_file( + build_dir / + "tensorrt_llm/deep_ep/nvshmem-build/src/lib/nvshmem_bootstrap_uid.so.3", + lib_dir / "nvshmem") + install_file( + build_dir / + "tensorrt_llm/deep_ep/nvshmem-build/src/lib/nvshmem_transport_ibgda.so.103", + lib_dir / "nvshmem") if not skip_stubs: with working_directory(project_dir): build_run(f"\"{venv_python}\" -m pip install pybind11-stubgen") @@ -651,14 +682,13 @@ def get_pybind_lib(): if 'LD_LIBRARY_PATH' in env_ld: new_library_path += f":{env_ld['LD_LIBRARY_PATH']}" env_ld["LD_LIBRARY_PATH"] = new_library_path - try: + build_run( + f"\"{venv_python}\" -m pybind11_stubgen -o . bindings --exit-code", + env=env_ld) + if deep_ep_cuda_architectures: build_run( - f"\"{venv_python}\" -m pybind11_stubgen -o . bindings --exit-code", + f"\"{venv_python}\" -m pybind11_stubgen -o . deep_ep_cpp_tllm --exit-code", env=env_ld) - except CalledProcessError as ex: - print(f"Failed to build pybind11 stubgen: {ex}", - file=sys.stderr) - exit(1) if not skip_building_wheel: if dist_dir is None: diff --git a/setup.py b/setup.py index 1b06f0700fc..97a32ba1b96 100644 --- a/setup.py +++ b/setup.py @@ -104,7 +104,10 @@ def has_ext_modules(self): 'libs/libnvinfer_plugin_tensorrt_llm.so', 'libs/libtensorrt_llm_ucx_wrapper.so', 'libs/libdecoder_attention_0.so', 'libs/libtensorrt_llm_nixl_wrapper.so', - 'libs/libdecoder_attention_1.so', 'bindings.*.so', "include/**/*" + 'libs/libdecoder_attention_1.so', 'libs/nvshmem/License.txt', + 'libs/nvshmem/nvshmem_bootstrap_uid.so.3', + 'libs/nvshmem/nvshmem_transport_ibgda.so.103', 'bindings.*.so', + 'deep_ep/LICENSE', 'deep_ep_cpp_tllm.*.so', "include/**/*" ] package_data += [ diff --git a/tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py b/tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py index e0c7c67748f..f669bb64563 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py +++ b/tensorrt_llm/_torch/modules/fused_moe/deep_ep_utils.py @@ -5,11 +5,11 @@ import torch -from tensorrt_llm._utils import local_mpi_size, mpi_comm +from tensorrt_llm._utils import mpi_comm from tensorrt_llm.mapping import Mapping try: - from deep_ep import Buffer + from tensorrt_llm.deep_ep import Buffer deep_ep_installed = True except ModuleNotFoundError: deep_ep_installed = False @@ -54,7 +54,6 @@ def reserve(self, hidden_size: int, hidden_dtype: torch.dtype): self.buffer = Buffer(None, num_nvl_bytes, num_rdma_bytes, - num_nvl_peers=local_mpi_size(), comm=self.comm) def dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],