Sync from nvidia main branch to update cuda driver support

CMakeLists.txt

@@ -50,6 +50,9 @@ endif()
 set(USE_TRITONSERVER_DATATYPE "ON")
 message("-- Enable USE_TRITONSERVER_DATATYPE")
 
+# Required for Ubuntu 22.04/Triton 22.05+
+set(CMAKE_CUDA_FLAGS "-include stdio.h")
+
 set(CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/Modules)
 
 # Python.h needed by torch headers.

README.md

@@ -26,7 +26,7 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 -->
 
-**NOTE: Fastertransformer backend is currently undergoing restructuring. Build instructions are only tested with Triton container versions <= `23.04`**. 
+**NOTE: Fastertransformer backend is currently undergoing restructuring so might not work with all versions of Triton.** 
 
 # FasterTransformer Backend
 

docker/Dockerfile

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG TRITON_VERSION=22.12
+ARG TRITON_VERSION=23.05
 ARG BASE_IMAGE=nvcr.io/nvidia/tritonserver:${TRITON_VERSION}-py3
 FROM ${BASE_IMAGE}
 
-RUN apt-get update
-RUN apt-get install -y --no-install-recommends \
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
         autoconf \
         autogen \
         clangd \
@@ -35,34 +35,34 @@ RUN apt-get install -y --no-install-recommends \
         unzip \
         zstd \
         zip \
-        zsh
-RUN pip3 install torch==1.12.1+cu116 -f https://download.pytorch.org/whl/torch_stable.html && \
-    pip3 install --extra-index-url https://pypi.ngc.nvidia.com regex fire tritonclient[all] && \
-    pip3 install transformers huggingface_hub tokenizers SentencePiece sacrebleu datasets tqdm omegaconf rouge_score && \
-    pip3 install cmake==3.24.3
-
-RUN apt-get clean && \
+        zsh && \
+    apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
+RUN pip3 install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cu118 torch==2.0.1+cu118 && \
+    pip3 install --no-cache-dir --extra-index-url https://pypi.ngc.nvidia.com regex fire tritonclient[all] && \
+    pip3 install --no-cache-dir accelerate transformers huggingface_hub tokenizers SentencePiece sacrebleu datasets tqdm omegaconf rouge_score && \
+    pip3 install --no-cache-dir cmake==3.24.3
+
 # backend build
 ADD . /workspace/build/fastertransformer_backend
 RUN mkdir -p /workspace/build/fastertransformer_backend/build
 
 WORKDIR /workspace/build/fastertransformer_backend/build
 ARG FORCE_BACKEND_REBUILD=0
-RUN cmake \
+RUN CUDAFLAGS="-include stdio.h" cmake \
       -D CMAKE_EXPORT_COMPILE_COMMANDS=1 \
       -D CMAKE_BUILD_TYPE=Release \
       -D ENABLE_FP8=OFF \
       -D CMAKE_INSTALL_PREFIX=/opt/tritonserver \
       -D TRITON_COMMON_REPO_TAG="r${NVIDIA_TRITON_SERVER_VERSION}" \
       -D TRITON_CORE_REPO_TAG="r${NVIDIA_TRITON_SERVER_VERSION}" \
       -D TRITON_BACKEND_REPO_TAG="r${NVIDIA_TRITON_SERVER_VERSION}" \
-      ..
-RUN cd _deps/repo-ft-src/ && \
+      .. && \
+    cd _deps/repo-ft-src/ && \
     git log | head -n 3 2>&1 | tee /workspace/build/fastertransformer_backend/FT_version.txt && \
     cd /workspace/build/fastertransformer_backend/build && \
-    make -j"$(grep -c ^processor /proc/cpuinfo)" install && \
+    CUDAFLAGS="-include stdio.h" make -O -j"$(grep -c ^processor /proc/cpuinfo)" install && \
     rm /workspace/build/fastertransformer_backend/build/bin/*_example -rf && \
     rm /workspace/build/fastertransformer_backend/build/lib/lib*Backend.so -rf
 

docker/create_dockerfile_and_build.py

@@ -61,7 +61,7 @@ def create_dependencies(base_image):
 ARG BASE_IMAGE={base_image}
     '''.format(base_image=base_image)
     df += '''
-FROM ${BASE_IMAGE}
+FROM ${BASE_IMAGE} AS build_image
 RUN apt-get update && apt-get install -y --no-install-recommends \\
         autoconf \\
         autogen \\