diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
index bc24126443..9c3021514d 100644
--- a/.github/workflows/docker.yml
+++ b/.github/workflows/docker.yml
@@ -10,6 +10,10 @@ on:
       - "benchmark/**"
       - "tests/**"
       - "**/*.md"
+      - "autotest/**"
+      - "builder/**"
+      - "k8s/**"
+
     branches:
       - main
     tags:
@@ -31,9 +35,13 @@ jobs:
   publish_docker_image:
     runs-on: ubuntu-latest
     environment: 'prod'
+    strategy:
+      matrix:
+        cuda_version: [cu12, cu11]
     env:
+      CUDA_VERSION: ${{ matrix.cuda_version }}
       TAG_PREFIX: "openmmlab/lmdeploy"
-      TAG: "openmmlab/lmdeploy:latest"
+      TAG: "openmmlab/lmdeploy:latest-${{matrix.cuda_version}}"
     steps:
       - name: Checkout repository
         uses: actions/checkout@v3
@@ -64,18 +72,25 @@ jobs:
       - name: Update docker TAG from workflow input
         if: github.event_name == 'workflow_dispatch'
         run: |
-          export TAG=$TAG_PREFIX:${{github.event.inputs.image_tag}}
+          export TAG=$TAG_PREFIX:${{github.event.inputs.image_tag}}-${CUDA_VERSION}
           echo $TAG
           echo "TAG=${TAG}" >> $GITHUB_ENV
       - name: Build and push Docker image
         run: |
           echo $TAG
-          docker build . -f docker/Dockerfile -t ${TAG} --no-cache
+          docker build . -f docker/Dockerfile -t ${TAG} --build-arg CUDA_VERSION=${CUDA_VERSION}
           docker push $TAG
+      - name: Push docker image latest-cu12 as latest
+        if: endsWith(env.TAG, 'latest-cu12') == true
+        run: |
+          export latest_TAG=${TAG_PREFIX}:latest
+          echo $latest_TAG
+          docker tag $TAG $latest_TAG
+          docker push $latest_TAG
       - name: Push docker image with released tag
         if: startsWith(github.ref, 'refs/tags/') == true
         run: |
-          export RELEASE_TAG=${TAG_PREFIX}:${{github.ref_name}}
+          export RELEASE_TAG=${TAG_PREFIX}:${{github.ref_name}}-${CUDA_VERSION}
           echo $RELEASE_TAG
           docker tag $TAG $RELEASE_TAG
           docker push $RELEASE_TAG
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 76421c46ce..45f30d59a0 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,15 +1,37 @@
-FROM nvcr.io/nvidia/tritonserver:22.12-py3
+ARG CUDA_VERSION=cu12
 
-RUN rm /etc/apt/sources.list.d/cuda*.list && apt-get update && apt-get install -y --no-install-recommends \
-    rapidjson-dev libgoogle-glog-dev gdb python3.8-venv \
-    && rm -rf /var/lib/apt/lists/* && cd /opt && python3 -m venv py38
+FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 AS cu12
+ENV CUDA_VERSION_SHORT=cu121
 
-ENV PATH=/opt/py38/bin:$PATH
+FROM nvidia/cuda:11.8.0-devel-ubuntu22.04 AS cu11
+ENV CUDA_VERSION_SHORT=cu118
+
+FROM ${CUDA_VERSION} AS final
+
+ARG PYTHON_VERSION=3.10
+
+ARG TORCH_VERSION=2.3.0
+ARG TORCHVISION_VERSION=0.18.0
+
+RUN rm /etc/apt/sources.list.d/cuda*.list && apt-get update -y && apt-get install -y software-properties-common wget vim &&\
+    add-apt-repository ppa:deadsnakes/ppa -y && apt-get update -y && apt-get install -y --no-install-recommends \
+    ninja-build rapidjson-dev libgoogle-glog-dev gdb python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \
+    && apt-get clean -y && rm -rf /var/lib/apt/lists/* && cd /opt && python3 -m venv py3
+
+ENV PATH=/opt/py3/bin:$PATH
 
 RUN python3 -m pip install --no-cache-dir --upgrade pip setuptools==69.5.1 &&\
-    python3 -m pip install --no-cache-dir torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu118 &&\
+    python3 -m pip install --no-cache-dir torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} --index-url https://download.pytorch.org/whl/${CUDA_VERSION_SHORT} &&\
     python3 -m pip install --no-cache-dir cmake packaging wheel
 
+# install openmpi
+RUN wget https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-4.1.5.tar.gz &&\
+    tar xf openmpi-4.1.5.tar.gz && cd openmpi-4.1.5 && ./configure --prefix=/usr/local/openmpi &&\
+    make -j$(nproc) && make install && cd .. && rm -rf openmpi-4.1.5*
+
+ENV PATH=$PATH:/usr/local/openmpi/bin
+ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/openmpi/lib
+
 ENV NCCL_LAUNCH_MODE=GROUP
 
 # Should be in the lmdeploy root directory when building docker image
@@ -20,21 +42,14 @@ WORKDIR /opt/lmdeploy
 RUN cd /opt/lmdeploy &&\
     python3 -m pip install --no-cache-dir -r requirements.txt &&\
     mkdir -p build && cd build &&\
-    cmake .. \
-        -DCMAKE_BUILD_TYPE=RelWithDebInfo \
-        -DCMAKE_EXPORT_COMPILE_COMMANDS=1 \
-        -DCMAKE_INSTALL_PREFIX=/opt/tritonserver \
-        -DBUILD_PY_FFI=ON \
-        -DBUILD_MULTI_GPU=ON \
-        -DBUILD_CUTLASS_MOE=OFF \
-        -DBUILD_CUTLASS_MIXED_GEMM=OFF \
-        -DCMAKE_CUDA_FLAGS="-lineinfo" \
-        -DUSE_NVTX=ON &&\
-    make -j$(nproc) && make install &&\
+    sh ../generate.sh &&\
+    ninja -j$(nproc) && ninja install &&\
     cd .. &&\
     python3 -m pip install -e . &&\
     rm -rf build
 
-ENV LD_LIBRARY_PATH=/opt/tritonserver/lib:$LD_LIBRARY_PATH
+ENV LD_LIBRARY_PATH=/opt/lmdeploy/install/lib:$LD_LIBRARY_PATH
+ENV PATH=/opt/lmdeploy/install/bin:$PATH
+
 # explicitly set ptxas path for triton
 ENV TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
diff --git a/generate.sh b/generate.sh
index 6648d2e22a..3c6b5bceed 100755
--- a/generate.sh
+++ b/generate.sh
@@ -1,4 +1,5 @@
 #!/bin/sh
+WORKSPACE_PATH=$(dirname "$(readlink -f "$0")")
 
 builder="-G Ninja"
 
@@ -9,7 +10,7 @@ fi
 cmake ${builder} .. \
     -DCMAKE_BUILD_TYPE=RelWithDebInfo \
     -DCMAKE_EXPORT_COMPILE_COMMANDS=1 \
-    -DCMAKE_INSTALL_PREFIX=./install \
+    -DCMAKE_INSTALL_PREFIX=${WORKSPACE_PATH}/install \
     -DBUILD_PY_FFI=ON \
     -DBUILD_MULTI_GPU=ON \
     -DCMAKE_CUDA_FLAGS="-lineinfo" \