[AMD] Update dsv3.2 AMD GPU docs and unify ROCm TileLang build (#17783)

Co-authored-by: wufann <715544327@qq.com>
2026-01-26 21:10:32 -08:00
parent a723d1c5ef
commit df42f4d386
2 changed files with 81 additions and 88 deletions
--- a/docker/rocm.Dockerfile
+++ b/docker/rocm.Dockerfile
@@ -1,8 +1,7 @@
 # Usage (to build SGLang ROCm docker image):
-#   docker build --build-arg SGL_BRANCH=v0.5.6.post2 --build-arg GPU_ARCH=gfx942 -t v0.5.6.post2-rocm630-mi30x -f rocm.Dockerfile .
-#   docker build --build-arg SGL_BRANCH=v0.5.6.post2 --build-arg GPU_ARCH=gfx942-rocm700 -t v0.5.6.post2-rocm700-mi30x -f rocm.Dockerfile .
-#   docker build --build-arg SGL_BRANCH=v0.5.6.post2 --build-arg GPU_ARCH=gfx950 -t v0.5.6.post2-rocm700-mi35x -f rocm.Dockerfile .
-
+#   docker build --build-arg SGL_BRANCH=v0.5.8 --build-arg GPU_ARCH=gfx942 -t v0.5.8-rocm630-mi30x -f rocm.Dockerfile .
+#   docker build --build-arg SGL_BRANCH=v0.5.8 --build-arg GPU_ARCH=gfx942-rocm700 -t v0.5.8-rocm700-mi30x -f rocm.Dockerfile .
+#   docker build --build-arg SGL_BRANCH=v0.5.8 --build-arg GPU_ARCH=gfx950 -t v0.5.8-rocm700-mi35x -f rocm.Dockerfile .

 # Default base images
 ARG BASE_IMAGE_942="rocm/sgl-dev:vllm20250114"
@@ -65,13 +64,8 @@ ARG LLVM_COMMIT="6520ace8227ffe2728148d5f3b9872a870b0a560"
 ARG MOONCAKE_REPO="https://github.com/kvcache-ai/Mooncake.git"
 ARG MOONCAKE_COMMIT="b6a841dc78c707ec655a563453277d969fb8f38d"

-ARG TILELANG_REPO="https://github.com/HaiShaw/tilelang.git"
-ARG TILELANG_BRANCH="dsv32-mi35x"
-ARG TILELANG_COMMIT="ae938cf885743f165a19656d1122ad42bb0e30b8"
-
-ARG TILELANG_GFX942_REPO="https://github.com/tile-ai/tilelang.git"
-ARG TILELANG_GFX942_BRANCH="main"
-ARG TILELANG_GFX942_COMMIT="2d8d3676eda18bd3d8e6fa783399ff96d3cd4ded"
+ARG TILELANG_REPO="https://github.com/tile-ai/tilelang.git"
+ARG TILELANG_COMMIT="ebf4a7cb8881432165ae8760e99d209d905c704a"

 ARG FHT_REPO="https://github.com/jeffdaily/fast-hadamard-transform.git"
 ARG FHT_BRANCH="rocm"
@@ -204,11 +198,12 @@ RUN find /sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/ \
 ENV PATH="/root/.cargo/bin:${PATH}"
 RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y \
    && rustc --version && cargo --version
+ENV CARGO_BUILD_JOBS=4

 # Build and install sgl-model-gateway
 RUN python3 -m pip install --no-cache-dir setuptools-rust \
    && cd /sgl-workspace/sglang/sgl-model-gateway/bindings/python \
-    && cargo build --release \
+    && /bin/bash -lc 'ulimit -n 8192 && cargo build --release' \
    && python3 -m pip install --no-cache-dir . \
    && rm -rf /root/.cache

@@ -219,82 +214,71 @@ ENV LIBGL_ALWAYS_INDIRECT=1
 RUN echo "LC_ALL=en_US.UTF-8" >> /etc/environment

 RUN /bin/bash -lc 'set -euo pipefail; \
-  # Build TileLang for gfx950 and gfx942-rocm700
-  if [ "${GPU_ARCH:-}" != "gfx950" ] && [ "${GPU_ARCH:-}" != "gfx942-rocm700" ]; then \
-    echo "[TileLang] Skipping (GPU_ARCH=${GPU_ARCH:-unset})"; \
-    exit 0; \
-  fi; \
  echo "[TileLang] Building TileLang for ${GPU_ARCH}"; \
-  if [ "$GPU_ARCH" = "gfx950" ]; then \
-    \
-    # System dependencies (NO llvm-dev to avoid llvm-config-16 shadowing)
-    apt-get update && apt-get install -y --no-install-recommends \
-        build-essential git wget curl ca-certificates gnupg \
-        libgtest-dev libgmock-dev \
-        libprotobuf-dev protobuf-compiler libgflags-dev libsqlite3-dev \
-        python3 python3-dev python3-setuptools python3-pip \
-        gcc libtinfo-dev zlib1g-dev libedit-dev libxml2-dev \
-        cmake ninja-build pkg-config libstdc++6 \
-    && rm -rf /var/lib/apt/lists/*; \
-    \
-    # Build GoogleTest static libs (Ubuntu package ships sources only)
-    cmake -S /usr/src/googletest -B /tmp/build-gtest -DBUILD_GTEST=ON -DBUILD_GMOCK=ON -DCMAKE_BUILD_TYPE=Release && \
-    cmake --build /tmp/build-gtest -j"$(nproc)" && \
-    cp -v /tmp/build-gtest/lib/*.a /usr/lib/x86_64-linux-gnu/ && \
-    rm -rf /tmp/build-gtest; \
-    \
-    # Keep setuptools < 80 (compat with base image)
-    python3 -m pip install --upgrade "setuptools>=77.0.3,<80" wheel cmake ninja && \
-    python3 -m pip cache purge || true; \
-    \
-    # Locate ROCm llvm-config; fallback to installing LLVM 18 if missing
-    LLVM_CONFIG_PATH=""; \
-    for p in /opt/rocm/llvm/bin/llvm-config /opt/rocm/llvm-*/bin/llvm-config /opt/rocm-*/llvm*/bin/llvm-config; do \
-      if [ -x "$p" ]; then LLVM_CONFIG_PATH="$p"; break; fi; \
-    done; \
-    if [ -z "$LLVM_CONFIG_PATH" ]; then \
-      echo "[TileLang] ROCm llvm-config not found; installing LLVM 18..."; \
-      curl -fsSL https://apt.llvm.org/llvm.sh -o /tmp/llvm.sh; \
-      chmod +x /tmp/llvm.sh; \
-      /tmp/llvm.sh 18; \
-      LLVM_CONFIG_PATH="$(command -v llvm-config-18)"; \
-      if [ -z "$LLVM_CONFIG_PATH" ]; then echo "ERROR: llvm-config-18 not found after install"; exit 1; fi; \
-    fi; \
-    echo "[TileLang] Using LLVM_CONFIG at: $LLVM_CONFIG_PATH"; \
-    export PATH="$(dirname "$LLVM_CONFIG_PATH"):/usr/local/bin:${PATH}"; \
-    export LLVM_CONFIG="$LLVM_CONFIG_PATH"; \
-    \
-    # Optional shim for tools that expect llvm-config-16
-    mkdir -p /usr/local/bin && \
-    printf "#!/usr/bin/env bash\nexec \"%s\" \"\$@\"\n" "$LLVM_CONFIG_PATH" > /usr/local/bin/llvm-config-16 && \
-    chmod +x /usr/local/bin/llvm-config-16; \
-    \
-    # TVM Python bits need Cython
-    python3 -m pip install --no-cache-dir "cython>=0.29.36,<3.0"; \
-    \
-    # Clone + pin TileLang (bundled TVM), then build
-    git clone --recursive --branch "${TILELANG_BRANCH}" "${TILELANG_REPO}" /opt/tilelang && \
-    cd /opt/tilelang && \
-    git fetch --depth=1 origin "${TILELANG_COMMIT}" || true && \
-    git checkout -f "${TILELANG_COMMIT}" && \
-    git submodule update --init --recursive && \
-    export CMAKE_ARGS="-DLLVM_CONFIG=${LLVM_CONFIG} ${CMAKE_ARGS:-}" && \
-    bash ./install_rocm.sh; \
-  else \
-    # Build GoogleTest static libs (Ubuntu package ships sources only)
-    apt-get install -y libgtest-dev libgmock-dev && \
-    cmake -S /usr/src/googletest -B /tmp/build-gtest -DBUILD_GTEST=ON -DBUILD_GMOCK=ON -DCMAKE_BUILD_TYPE=Release && \
-    cmake --build /tmp/build-gtest -j && \
-    cp -v /tmp/build-gtest/lib/*.a /usr/lib/x86_64-linux-gnu/ && \
-    rm -rf /tmp/build-gtest; \
-    # Build TileLang for gfx942-rocm700
-    git clone --branch "${TILELANG_GFX942_BRANCH}" "${TILELANG_GFX942_REPO}" /opt/tilelang && \
-    cd /opt/tilelang && \
-    git checkout -f "${TILELANG_GFX942_COMMIT}" && \
-    git submodule update --init --recursive && \
-    sed -i "/^[[:space:]]*\"torch/d" pyproject.toml && \
-    USE_ROCM=1 USE_CUDA=0 pip install -e . -v ; \
-  fi'
+  # System dependencies (NO llvm-dev to avoid llvm-config-16 shadowing)
+  apt-get update && apt-get install -y --no-install-recommends \
+      build-essential git wget curl ca-certificates gnupg \
+      libgtest-dev libgmock-dev \
+      libprotobuf-dev protobuf-compiler libgflags-dev libsqlite3-dev \
+      python3 python3-dev python3-setuptools python3-pip python3-apt \
+      gcc libtinfo-dev zlib1g-dev libedit-dev libxml2-dev \
+      cmake ninja-build pkg-config libstdc++6 software-properties-common \
+  && rm -rf /var/lib/apt/lists/*; \
+  \
+  # Prefer the container venv
+  VENV_PY="/opt/venv/bin/python"; \
+  VENV_PIP="/opt/venv/bin/pip"; \
+  if [ ! -x "$VENV_PY" ]; then VENV_PY="python3"; fi; \
+  if [ ! -x "$VENV_PIP" ]; then VENV_PIP="pip3"; fi; \
+  \
+  # Build GoogleTest static libs (Ubuntu package ships sources only)
+  cmake -S /usr/src/googletest -B /tmp/build-gtest -DBUILD_GTEST=ON -DBUILD_GMOCK=ON -DCMAKE_BUILD_TYPE=Release && \
+  cmake --build /tmp/build-gtest -j"$(nproc)" && \
+  cp -v /tmp/build-gtest/lib/*.a /usr/lib/x86_64-linux-gnu/ && \
+  rm -rf /tmp/build-gtest; \
+  \
+  # Keep setuptools < 80 (compat with base image)
+  "$VENV_PIP" install --upgrade "setuptools>=77.0.3,<80" wheel cmake ninja scikit-build-core && \
+  "$VENV_PIP" cache purge || true; \
+  \
+  # Locate ROCm llvm-config; fallback to installing LLVM 18 if missing
+  LLVM_CONFIG_PATH=""; \
+  for p in /opt/rocm/llvm/bin/llvm-config /opt/rocm/llvm-*/bin/llvm-config /opt/rocm-*/llvm*/bin/llvm-config; do \
+    if [ -x "$p" ]; then LLVM_CONFIG_PATH="$p"; break; fi; \
+  done; \
+  if [ -z "$LLVM_CONFIG_PATH" ]; then \
+    echo "[TileLang] ROCm llvm-config not found; installing LLVM 18..."; \
+    curl -fsSL https://apt.llvm.org/llvm-snapshot.gpg.key | gpg --dearmor -o /etc/apt/keyrings/llvm.gpg; \
+    echo "deb [signed-by=/etc/apt/keyrings/llvm.gpg] http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main" > /etc/apt/sources.list.d/llvm.list; \
+    apt-get update; \
+    apt-get install -y --no-install-recommends llvm-18; \
+    rm -rf /var/lib/apt/lists/*; \
+    LLVM_CONFIG_PATH="$(command -v llvm-config-18)"; \
+    if [ -z "$LLVM_CONFIG_PATH" ]; then echo "ERROR: llvm-config-18 not found after install"; exit 1; fi; \
+  fi; \
+  echo "[TileLang] Using LLVM_CONFIG at: $LLVM_CONFIG_PATH"; \
+  export PATH="$(dirname "$LLVM_CONFIG_PATH"):/usr/local/bin:${PATH}"; \
+  export LLVM_CONFIG="$LLVM_CONFIG_PATH"; \
+  \
+  # Optional shim for tools that expect llvm-config-16
+  mkdir -p /usr/local/bin && \
+  printf "#!/usr/bin/env bash\nexec \"%s\" \"\$@\"\n" "$LLVM_CONFIG_PATH" > /usr/local/bin/llvm-config-16 && \
+  chmod +x /usr/local/bin/llvm-config-16; \
+  \
+  # TVM Python bits need Cython + z3 before configure
+  "$VENV_PIP" install --no-cache-dir "cython>=0.29.36,<3.0" "apache-tvm-ffi>=0.1.6" "z3-solver>=4.13.0"; \
+  \
+  # Clone + pin TileLang (bundled TVM), then build
+  git clone --recursive "${TILELANG_REPO}" /opt/tilelang && \
+  cd /opt/tilelang && \
+  git fetch --depth=1 origin "${TILELANG_COMMIT}" || true && \
+  git checkout -f "${TILELANG_COMMIT}" && \
+  git submodule update --init --recursive && \
+  export CMAKE_ARGS="-DUSE_CUDA=OFF -DUSE_ROCM=ON -DROCM_PATH=/opt/rocm -DLLVM_CONFIG=${LLVM_CONFIG} -DSKBUILD_SABI_VERSION= ${CMAKE_ARGS:-}" && \
+  "$VENV_PIP" install -e . -v --no-build-isolation --no-deps; \
+  if [ -f pyproject.toml ]; then sed -i "/^[[:space:]]*\"torch/d" pyproject.toml || true; fi; \
+  "$VENV_PIP" cache purge || true; \
+  "$VENV_PY" -c "import tilelang; print(tilelang.__version__)"'

 # -----------------------
 # Hadamard-transform (HIP build)
--- a/docs/basic_usage/deepseek_v32.md
+++ b/docs/basic_usage/deepseek_v32.md
@@ -16,7 +16,13 @@ Note: This document is originally written for the usage of [DeepSeek-V3.2-Exp](h
 docker pull lmsysorg/sglang:latest

 # MI350/MI355
-docker pull lmsysorg/sglang:dsv32-rocm
+docker pull lmsysorg/sglang:v0.5.8-rocm700-mi35x
+
+# MI300
+# v0.5.8-rocm700-mi30x does not include PR #17504. Prefer the newest MI30x ROCm
+# image tag from Docker Hub when available, or build from source (below).
+docker pull lmsysorg/sglang:v0.5.8-rocm700-mi30x
+

 # NPUs
 docker pull lmsysorg/sglang:dsv32-a2
@@ -45,6 +51,9 @@ python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --ep

 # Launch with Pure TP
 python -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8
+
+# Launch with TP on MI30x/MI35x
+python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3.2-Exp --tp 8 --nsa-prefill-backend tilelang --nsa-decode-backend tilelang
 ```

 ### Configuration Tips