This commit is contained in:
2026-06-03 16:49:38 +08:00
parent 319616b0a7
commit 3550efcd48
+13 -5
View File
@@ -28,7 +28,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE \
MINERU_MODEL_SOURCE=huggingface \
TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 \
HSA_ENABLE_SDMA=1 \
HSA_ENABLE_SDMA=1 \ # 启用系统 DMA 加速 GPU 显存传输
VLLM_TARGET_DEVICE=rocm
WORKDIR /opt
@@ -159,6 +159,13 @@ RUN set -ex && \
cd /opt/vllm && \
sed -i '109,121s/^/\/\/ /' csrc/mamba/mamba_ssm/selective_scan.h && \
echo "vllm mamba operator+ patch applied." && \
# cmake 别名兜底:ROCm 7.2 可能缺少 hiprand/hipblas cmake target
mkdir -p /opt/rocm/lib/cmake/hiprand && \
printf 'include(/opt/rocm/lib/cmake/rocrand/rocrand-config.cmake)\nif(TARGET roc::rocrand AND NOT TARGET hip::hiprand)\n add_library(hip::hiprand ALIAS roc::rocrand)\nendif()\n' \
> /opt/rocm/lib/cmake/hiprand/hiprand-config.cmake && \
mkdir -p /opt/rocm/lib/cmake/hipblas && \
printf 'include(/opt/rocm/lib/cmake/rocblas/rocblas-config.cmake)\nif(TARGET roc::rocblas AND NOT TARGET hip::hipblas)\n add_library(hip::hipblas ALIAS roc::rocblas)\nendif()\n' \
> /opt/rocm/lib/cmake/hipblas/hipblas-config.cmake && \
# cmake 配置
mkdir -p /opt/vllm_build && \
cmake -S /opt/vllm -B /opt/vllm_build -G Ninja \
@@ -177,13 +184,14 @@ RUN set -ex && \
cd /opt/vllm && ${VENV}/bin/pip install --no-cache-dir -e . --no-build-isolation && \
# 验证 PyTorch 没被 vllm 依赖覆盖
${VENV}/bin/python -c "import torch; v=torch.__version__; assert 'rocm' in v, f'PyTorch overwritten by vllm deps: {v}'; print('PyTorch OK:', v)" && \
# 清理可能的 CUDA triton 残余
${VENV}/bin/pip uninstall -y triton triton-rocm 2>/dev/null; \
# 先重装 ROCm PyTorch 覆盖可能的 CUDA 版,再清理 CUDA triton 元数据
# 顺序重要:pytorch-triton-rocm 和 triton 共享 triton/ 物理目录,必须先重装后卸载
${VENV}/bin/pip install --no-cache-dir --force-reinstall \
torch==2.11.0+rocm7.2 torchvision pytorch-triton-rocm \
--index-url ${TORCH_INDEX} && \
# 最终验证 vllm 平台检测
${VENV}/bin/python -c "from vllm.platforms import current_platform; print('Platform:', type(current_platform).__name__); print('is_rocm:', current_platform.is_rocm()); print('device_type:', current_platform.device_type); assert current_platform.is_rocm(), 'vllm ROCm detection failed!'; print('vllm OK')" && \
${VENV}/bin/pip uninstall -y triton triton-rocm 2>/dev/null; \
# vllm 平台模块导入验证(GPU 检测只能在运行时,容器构建时无 GPU 设备)
${VENV}/bin/python -c "from vllm.platforms import current_platform; print('Platform module:', type(current_platform).__name__); print('vllm import OK')" && \
# 清理构建目录(减小镜像体积,约 3-5GB)
rm -rf /opt/vllm_build