[AMD] Add gfx1151 (Strix Halo / Ryzen AI MAX+) Docker image (#33939)

Co-authored-by: Hubert Lu <Hubert.Lu@amd.com>
2026-09-10 22:46:33 -07:00
parent 6182d45524
commit b9899b04c1
3 changed files with 309 additions and 0 deletions
@@ -0,0 +1,90 @@
+#!/bin/sh
+# Teach python/sglang/kernels/aot/setup_rocm.py to build for gfx1151 (Strix Halo).
+# Applied at image-build time by docker/rocm-gfx1151.Dockerfile; the repo files are
+# left untouched because gfx1151 is not a supported SGLang target.
+#
+# Two changes:
+#   1. Lift the {gfx942, gfx950, gfx1250} allowlist, which otherwise sys.exit(1)s.
+#      Current main gives non-gfx942 targets a 40KB TopK dynamic-LDS budget,
+#      which already fits gfx1151's 64KB workgroup limit.
+#   2. Force a single WARP_SIZE across the host and device compiler passes.
+#      include/utils.h resolves WARP_SIZE to 64 whenever __HIP_DEVICE_COMPILE__
+#      is undefined -- i.e. on the host pass -- and to 32 on a non-__GFX9__
+#      device pass. On CDNA both come out 64 and nothing is wrong, which is why
+#      upstream never sees this. On gfx1151 the two passes disagree, and the MoE
+#      TopK kernels use WARP_SIZE on both sides of the launch:
+#         moe_topk_softmax_kernels.cu  __launch_bounds__(WARPS_PER_CTA * WARP_SIZE)  -> device, 4*32 = 128
+#         moe_topk_softmax_kernels.cu  dim3 block_dim(WARP_SIZE, WARPS_PER_TB)       -> host,   64*4 = 256
+#      Launching 256 threads into a 128-thread bound fails with
+#      hipErrorLaunchFailure, poisons the queue, and typically surfaces as a
+#      page fault in whatever kernel runs next (moe_align_block_size_kernel),
+#      which makes it easy to misattribute. The same pattern is in
+#      moe_topk_sigmoid_kernels.cu. It also desynchronizes the launcher's
+#      TopkConstants math (ROWS_PER_WARP, VECs_PER_THREAD) from the kernel's.
+#      Only MoE models hit this; dense models never call these kernels.
+#      32 is simply correct here -- gfx1151 is a wave32 part -- so the override
+#      pins both passes to 32 rather than renaming the symbol per call site.
+#
+# Each edit is guarded: if the upstream line has changed, fail rather than
+# silently produce an image whose kernels were built with the wrong limits.
+set -e
+
+FILE="${1:?usage: sgl-kernel-gfx1151.sh <path to setup_rocm.py>}"
+UTILS="$(dirname "${FILE}")/include/utils.h"
+
+GATE_OLD='if amdgpu_target not in ["gfx942", "gfx950", "gfx1250"]:'
+GATE_NEW='if amdgpu_target not in ["gfx942", "gfx950", "gfx1250", "gfx1151"]:'
+
+FLAGS_OLD='    f"-DSGL_TOPK_DYNAMIC_SMEM_BYTES={topk_dynamic_smem_bytes}",'
+FLAGS_NEW='    f"-DSGL_TOPK_DYNAMIC_SMEM_BYTES={topk_dynamic_smem_bytes}",
+    # gfx1151 is wave32; pin both compiler passes to it (see utils.h below).
+    *(["-DSGL_ROCM_WARP_SIZE=32"] if amdgpu_target == "gfx1151" else []),'
+
+WARP_OLD='#if defined(__GFX9__) || !defined(__HIP_DEVICE_COMPILE__)
+#define WARP_SIZE 64'
+WARP_NEW='#if defined(SGL_ROCM_WARP_SIZE)
+#define WARP_SIZE SGL_ROCM_WARP_SIZE
+#elif defined(__GFX9__) || !defined(__HIP_DEVICE_COMPILE__)
+#define WARP_SIZE 64'
+
+for pattern in "${GATE_OLD}" "${FLAGS_OLD}"; do
+    if ! grep -qF "${pattern}" "${FILE}"; then
+        echo "ERROR: expected line not found in ${FILE}:" >&2
+        echo "  ${pattern}" >&2
+        echo "setup_rocm.py changed upstream; re-check this patch before building." >&2
+        exit 1
+    fi
+done
+
+if ! grep -qF "${WARP_OLD}" "${UTILS}"; then
+    echo "ERROR: expected WARP_SIZE block not found in ${UTILS}." >&2
+    echo "utils.h changed upstream; re-check the wave32 fix before building." >&2
+    exit 1
+fi
+
+python3 - "${UTILS}" "${WARP_OLD}" "${WARP_NEW}" <<'PY'
+import sys
+
+path, old, new = sys.argv[1:4]
+with open(path) as f:
+    src = f.read()
+with open(path, "w") as f:
+    f.write(src.replace(old, new, 1))
+PY
+
+python3 - "${FILE}" "${GATE_OLD}" "${GATE_NEW}" "${FLAGS_OLD}" "${FLAGS_NEW}" <<'PY'
+import sys
+
+path, gate_old, gate_new, flags_old, flags_new = sys.argv[1:6]
+with open(path) as f:
+    src = f.read()
+src = src.replace(gate_old, gate_new)
+src = src.replace(flags_old, flags_new, 1)
+with open(path, "w") as f:
+    f.write(src)
+PY
+
+echo "Patched ${FILE} for gfx1151:"
+grep -nF -e "${GATE_NEW}" -e "SGL_ROCM_WARP_SIZE" "${FILE}"
+echo "Patched ${UTILS} for wave32:"
+grep -nF "SGL_ROCM_WARP_SIZE" "${UTILS}"
@@ -0,0 +1,156 @@
+# SGLang for AMD Strix Halo / Ryzen AI MAX+ (gfx1151, RDNA3.5 iGPU).
+#
+# This is NOT a variant of docker/rocm.Dockerfile. That file targets CDNA
+# (gfx942/gfx950) and includes components which do not support gfx1151. This
+# image starts from AMD's stable ROCm/PyTorch image with native gfx1151 support.
+#
+# Build:
+#   docker build -f docker/rocm-gfx1151.Dockerfile -t sglang-rocm:gfx1151 .
+#
+# Run (Strix Halo has no discrete VRAM; the GPU carves out of system RAM):
+#   docker run -it --rm \
+#     --device=/dev/kfd --device=/dev/dri \
+#     --group-add video --group-add render \
+#     --security-opt seccomp=unconfined \
+#     --ipc=host --shm-size 16g \
+#     -p 30000:30000 \
+#     -v ~/.cache/huggingface:/root/.cache/huggingface \
+#     sglang-rocm:gfx1151 \
+#     python3 -m sglang.launch_server --model-path <model> \
+#       --attention-backend triton --host 0.0.0.0
+
+# ROCm 7.2.4 / PyTorch 2.9.1 is AMD's stable gfx1151-supported combination.
+# Pin the image digest so rebuilding cannot silently change the toolchain.
+ARG BASE_IMAGE="rocm/pytorch@sha256:7fe531fa185af260352fe7fbb3fa64ad749abe72adf0600a648c4692801b125a"
+
+# =============================================================================
+# Stage 1: stable ROCm + PyTorch for gfx1151.
+# Pullable and testable on its own:
+#   docker build --target rocm-torch -f docker/rocm-gfx1151.Dockerfile -t rocm-torch:gfx1151 .
+# =============================================================================
+FROM ${BASE_IMAGE} AS rocm-torch
+
+ARG GPU_ARCH=gfx1151
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        cmake \
+        libnuma-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+ENV PYTORCH_ROCM_ARCH=${GPU_ARCH}
+# ROCDXG requires this under WSL. It is inert when /dev/dxg is absent.
+ENV HSA_ENABLE_DXG_DETECTION=1
+
+# Fail loudly if the base image or its expected development toolchain changes.
+RUN python3 -c "import torch; print('torch', torch.__version__); assert torch.version.hip" \
+    && test -x /opt/rocm/bin/hipcc
+
+# =============================================================================
+# Stage 2: SGLang on top of the gfx1151 ROCm stack.
+# =============================================================================
+FROM rocm-torch AS sglang
+
+ARG GPU_ARCH=gfx1151
+# sgl-kernel's ROCm build (python/sglang/kernels/aot/setup_rocm.py) only accepts
+# gfx942/gfx950/gfx1250 and hard-exits on anything else; the patch below lifts
+# that gate. Set to 0 to skip the AOT kernels entirely and run Triton-only.
+ARG BUILD_SGL_KERNEL=1
+ARG MAX_JOBS=12
+
+WORKDIR /sgl-workspace
+
+# setuptools-rust builds the sglang-mm extension during the pip install below.
+ENV PATH="/root/.cargo/bin:${PATH}"
+RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal \
+    && rustc --version
+ENV CARGO_BUILD_JOBS=8
+
+COPY . /sgl-workspace/sglang
+
+# pyproject.toml pins the CUDA stack (torch, flashinfer[cu13], cuda-python,
+# ...). pyproject_other.toml carries the srt_hip extra, which is
+# torch-version-agnostic -- same swap docker/rocm.Dockerfile performs.
+RUN cd /sgl-workspace/sglang \
+    && rm -f python/pyproject.toml \
+    && mv python/pyproject_other.toml python/pyproject.toml
+
+# One problem in setup_rocm.py for this target, plus one in include/utils.h:
+# the arch gate sys.exit(1)s outside {gfx942, gfx950, gfx1250}, and WARP_SIZE
+# resolves to 64 on the host pass but 32 on the device pass for a wave32 part,
+# which mismatches the MoE TopK launch bounds. Current main already limits
+# non-gfx942 TopK dynamic LDS to 40KB, which fits gfx1151's 64KB limit. The
+# remaining two problems are fixed here rather than upstream:
+# gfx1151 is not a supported SGLang target, and the sources themselves compile
+# clean for it. Each edit greps for the expected text first, so a rewrite
+# upstream breaks the build loudly instead of silently misconfiguring kernels.
+COPY docker/patches/sgl-kernel-gfx1151.sh /tmp/sgl-kernel-gfx1151.sh
+
+RUN cd /sgl-workspace/sglang/python/sglang/kernels/aot \
+    && if [ "${BUILD_SGL_KERNEL}" = "1" ]; then \
+         rm -f pyproject.toml \
+         && mv pyproject_rocm.toml pyproject.toml \
+         && sh /tmp/sgl-kernel-gfx1151.sh setup_rocm.py \
+         && AMDGPU_TARGET=${GPU_ARCH} MAX_JOBS=${MAX_JOBS} python3 setup_rocm.py install; \
+       else \
+         echo "Skipping sgl-kernel build (BUILD_SGL_KERNEL=0)"; \
+       fi
+
+# Current main composes extras through self-references
+# (srt_hip -> sglang[runtime_common] -> sglang[runtime_base]). pip's resolver
+# recursively walks that cycle from an editable source checkout. Flatten those
+# three groups before installing the package itself without dependency solving.
+# Keep compressed-tensors at its last torch-2.9-compatible release.
+RUN cd /sgl-workspace/sglang \
+    && python3 - <<'PY'
+import subprocess
+import sys
+import tomllib
+from pathlib import Path
+
+project = tomllib.loads(Path("python/pyproject.toml").read_text())["project"]
+extras = project["optional-dependencies"]
+requirements = list(project["dependencies"])
+for group in ("runtime_base", "runtime_common", "srt_hip"):
+    requirements.extend(
+        "compressed-tensors==0.15.0"
+        if requirement == "compressed-tensors"
+        else requirement
+        for requirement in extras[group]
+        if not requirement.startswith("sglang[") and requirement != "torch"
+    )
+requirements = list(dict.fromkeys(requirements))
+subprocess.check_call(
+    [sys.executable, "-m", "pip", "install", "--no-cache-dir", *requirements]
+)
+PY
+RUN cd /sgl-workspace/sglang \
+    && pip install --no-cache-dir --no-deps -e python
+
+# aiter is not optional on ROCm despite being CDNA-oriented:
+# sglang/srt/layers/quantization/__init__.py imports quark, which imports
+# aiter.ops.triton at module scope, so `import sglang.srt.layers.activation`
+# fails outright without it. Installed WITHOUT PREBUILD_KERNELS -- that step
+# AOT-compiles the CDNA assembly kernels and is what actually fails on gfx1151.
+# In JIT mode aiter builds module_aiter_core for gfx1151 on demand instead.
+ARG AITER_REPO="https://github.com/ROCm/aiter.git"
+ARG AITER_COMMIT="c16d44b93a528b2a4bfd6d8d3409116d465872a9"
+
+RUN git clone --recursive ${AITER_REPO} /sgl-workspace/aiter \
+    && cd /sgl-workspace/aiter \
+    && git checkout ${AITER_COMMIT} \
+    && git submodule update --init --recursive \
+    && GPU_ARCHS=${GPU_ARCH} pip install --no-cache-dir --no-build-isolation \
+         --config-settings editable_mode=compat -e .
+
+# aiter's compiled attention/MoE kernels are CDNA-only; keep sglang on the
+# Triton paths. ServerArgs defaults to aiter on ROCm, so callers must pass
+# `--attention-backend triton` until the RDNA default is fixed upstream.
+# This is load-bearing beyond attention: aiter's RMSNorm uses v_pk_mul_f32,
+# a CDNA-only instruction, and its CK attention templates assume wave64.
+ENV SGLANG_USE_AITER=0
+
+WORKDIR /sgl-workspace/sglang
+CMD ["/bin/bash"]