diff --git a/CMakeLists.txt b/CMakeLists.txt index 3790db9263..54374c2d01 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -49,7 +49,7 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS) set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14") # Supported AMD GPU architectures. -set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") +set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1250;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201") # ROCm installation prefix. Default to /opt/rocm but allow override via # -DROCM_PATH=/your/rocm/path when invoking cmake. @@ -1422,6 +1422,38 @@ if(VLLM_GPU_LANG STREQUAL "HIP") "csrc/rocm/skinny_gemms.cu" "csrc/rocm/skinny_gemms_int4.cu" "csrc/rocm/attention.cu") + set(VLLM_ROCM_EXT_FLAGS ${VLLM_GPU_FLAGS}) + + # skinny_gemms*.cu are built on gfx9/gfx11 ISA (MFMA, dot2/dot4, legacy + # s_waitcnt asm) that gfx1250 (gfx12) does not provide. Exclude them from the + # gfx1250 build and disable their op registrations (VLLM_SKIP_SKINNY_GEMMS); + # vLLM falls back to default/Triton GEMM for those ops on gfx1250. + list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms.cu") + list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms_int4.cu") + set(VLLM_SKINNY_ARCHES ${VLLM_GPU_ARCHES}) + list(FILTER VLLM_SKINNY_ARCHES EXCLUDE REGEX "gfx1250") + if(VLLM_SKINNY_ARCHES) + message(STATUS "Building skinny_gemms for archs: ${VLLM_SKINNY_ARCHES}") + set(VLLM_SKINNY_SRC + "csrc/rocm/skinny_gemms.cu" + "csrc/rocm/skinny_gemms_int4.cu") + hipify_sources_target(VLLM_SKINNY_HIP_SRCS _rocm_C_skinny "${VLLM_SKINNY_SRC}") + unset(_VLLM_LAST_HIPIFY_TARGET) + add_library(_rocm_C_skinny OBJECT ${VLLM_SKINNY_HIP_SRCS}) + add_dependencies(_rocm_C_skinny hipify_all) + set_source_files_properties(${VLLM_SKINNY_HIP_SRCS} PROPERTIES LANGUAGE ${VLLM_GPU_LANG}) + set_target_properties(_rocm_C_skinny PROPERTIES + ${VLLM_GPU_LANG}_ARCHITECTURES "${VLLM_SKINNY_ARCHES}" + POSITION_INDEPENDENT_CODE ON) + target_include_directories(_rocm_C_skinny PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/csrc) + target_compile_options(_rocm_C_skinny PRIVATE + $<$:${VLLM_ROCM_EXT_FLAGS}>) + target_compile_definitions(_rocm_C_skinny PRIVATE "-DTORCH_EXTENSION_NAME=_rocm_C") + target_link_libraries(_rocm_C_skinny PRIVATE torch) + else() + message(STATUS "Only gfx1250, skipping skinny_gemms") + list(APPEND VLLM_ROCM_EXT_FLAGS "-DVLLM_SKIP_SKINNY_GEMMS") + endif() set(VLLM_ROCM_HAS_GFX1100 OFF) if(VLLM_GPU_ARCHES MATCHES "gfx1100") @@ -1437,10 +1469,16 @@ if(VLLM_GPU_LANG STREQUAL "HIP") DESTINATION vllm LANGUAGE ${VLLM_GPU_LANG} SOURCES ${VLLM_ROCM_EXT_SRC} - COMPILE_FLAGS ${VLLM_GPU_FLAGS} + COMPILE_FLAGS ${VLLM_ROCM_EXT_FLAGS} ARCHITECTURES ${VLLM_GPU_ARCHES} USE_SABI 3 WITH_SOABI) + + if(TARGET _rocm_C_skinny) + target_link_libraries(_rocm_C PRIVATE _rocm_C_skinny) + else() + target_compile_definitions(_rocm_C PRIVATE VLLM_SKIP_SKINNY_GEMMS) + endif() if(VLLM_ROCM_HAS_GFX1100) target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100) diff --git a/csrc/quickreduce/base.h b/csrc/quickreduce/base.h index 6c3456d06f..f330ed401e 100644 --- a/csrc/quickreduce/base.h +++ b/csrc/quickreduce/base.h @@ -79,6 +79,7 @@ union BufferResource { }; }; +#if !defined(__gfx1250__) __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32"); @@ -86,6 +87,16 @@ __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( __quickreduce_device_inline__ static void buffer_store_dwordx4( int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32"); +#else +__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4( + int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) {} + +__quickreduce_device_inline__ static void buffer_store_dwordx4(int32x4_t data, + int32x4_t srsrc, + int32_t voffset, + int32_t soffset, + int32_t aux) {} +#endif __quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) { #if defined(__gfx942__) diff --git a/csrc/rocm/attention.cu b/csrc/rocm/attention.cu index 4ac255d0a7..5ca55fe3b9 100644 --- a/csrc/rocm/attention.cu +++ b/csrc/rocm/attention.cu @@ -2405,6 +2405,16 @@ template __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, const bit16x8& inpB, const floatx8& inpC) { + #if defined(__gfx1250__) + // gfx1250 (gfx12 family) does not provide the gfx12 WMMA variant used by + // gfx1200/1201 (needs wmma-128b-insts). This custom-attention WMMA path is + // unsupported on gfx1250; trap if ever launched (fail loud, not + // silent-wrong). + (void)inpA; + (void)inpB; + __builtin_trap(); + return inpC; + #else if constexpr (std::is_same::value) { return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(inpA, inpB, inpC); } else if constexpr (std::is_same::value) { @@ -2412,6 +2422,7 @@ __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, } else { static_assert(false, "unsupported 16b dtype"); } + #endif } template diff --git a/csrc/rocm/torch_bindings.cpp b/csrc/rocm/torch_bindings.cpp index 021359cb63..92dbed7f02 100644 --- a/csrc/rocm/torch_bindings.cpp +++ b/csrc/rocm/torch_bindings.cpp @@ -14,6 +14,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { // vLLM custom ops for rocm +// skinny_gemms.cu (LLMM1/wvSplitK/wvSplitKrc/wvSplitKQ) is excluded on gfx1250 +// (gfx9/gfx11 ISA, unsupported there); skip these registrations to avoid +// undefined symbols. vLLM uses default/Triton GEMM for these ops on gfx1250. +#ifndef VLLM_SKIP_SKINNY_GEMMS // Custom gemm op for matrix-vector multiplication rocm_ops.def( "LLMM1(Tensor in_a, Tensor in_b, int rows_per_block) -> " @@ -46,6 +50,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) { "Tensor scale_a, " " Tensor scale_b, int CuCount) -> ()"); rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ); +#endif // VLLM_SKIP_SKINNY_GEMMS #ifdef VLLM_ROCM_GFX1100 // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100). diff --git a/docker/Dockerfile.rocm_base_gfx1250 b/docker/Dockerfile.rocm_base_gfx1250 new file mode 100644 index 0000000000..3fad6800ab --- /dev/null +++ b/docker/Dockerfile.rocm_base_gfx1250 @@ -0,0 +1,325 @@ +ARG BASE_IMAGE=ubuntu:24.04 +ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/ +ARG ROCM_SDK_VERSION=7.14.0a20260623 +ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623 +ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623 +ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623 +ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623 +ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623 + + +ARG FA_BRANCH="jpvillam/gfx1250_wip" +ARG FA_REPO="https://github.com/jpvillam-amd/flash-attention.git" +ARG AITER_BRANCH="main" +ARG AITER_REPO="https://github.com/ROCm/aiter.git" +ARG MORI_BRANCH="v1.1.0" +ARG MORI_REPO="https://github.com/ROCm/mori.git" + +# Sccache configuration (only used in release pipeline) +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME=vllm-build-sccache +ARG SCCACHE_REGION_NAME=us-west-2 +ARG SCCACHE_S3_NO_CREDENTIALS=0 + +FROM ${BASE_IMAGE} AS base + +ARG PYTORCH_ROCM_ARCH=gfx1250 +ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV AITER_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV MORI_GPU_ARCHS=gfx942;gfx950 +ENV FA_GPU_ARCHS=gfx942;gfx950;gfx1250 + +# TODO: Unset these when support is available for gfx1250 +ENV ENABLE_CK=0 +ENV SKIP_CK_BUILD="TRUE" +ENV FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE" +ARG PREBUILD_KERNELS=0 +ARG TRITON_KERNEL_BRANCH="padroute" +ARG TRITON_KERNEL_REPO="https://github.com/jpvillam-amd/triton.git" + +# Required for RCCL in ROCm7.1 +ENV HSA_NO_SCRATCH_RECLAIM=1 + +ARG PYTHON_VERSION=3.12 +ENV PYTHON_VERSION=${PYTHON_VERSION} + +RUN mkdir -p /app +WORKDIR /app +ENV DEBIAN_FRONTEND=noninteractive + +# Install Python and other dependencies +RUN apt-get update -y \ + && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev libnuma-dev libdrm-dev pkg-config g++ \ + && for i in 1 2 3; do \ + add-apt-repository -y ppa:deadsnakes/ppa && break || \ + { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \ + done \ + && apt-get update -y \ + && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \ + python${PYTHON_VERSION}-lib2to3 python-is-python3 \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \ + && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \ + && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \ + && python3 --version + +ENV VIRTUAL_ENV=/opt/venv +RUN python${PYTHON_VERSION} -m venv "${VIRTUAL_ENV}" && \ + "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML +ENV PATH=${VIRTUAL_ENV}/bin:$PATH + +RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython +RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/* + +# Install sccache if USE_SCCACHE is enabled (for release builds) +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME +ARG SCCACHE_REGION_NAME +ARG SCCACHE_S3_NO_CREDENTIALS +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ + fi + +## +## Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index +## +ARG ROCM_WHEEL_INDEX +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION +ARG ROCM_SDK_VERSION +ARG APEX_VERSION +ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages +ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel +ENV ROCM_HOME=${ROCM_PATH} +ENV ROCM_SOURCE_DIR=${ROCM_PATH} +ENV ROCM_BIN=${ROCM_PATH}/bin +ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake +ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode +ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH +ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib +ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake +ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi + +# torch/torchvision/torchaudio must be pinned to mutually-consistent builds +# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm +# sdk version is derived from torch's own dependency pin unless overridden, +# which keeps the set consistent and avoids pip backtracking. +RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \ + --extra-index-url https://pypi.org/simple \ + "torch[device-all]==${TORCH_VERSION}" \ + "torchvision==${TORCHVISION_VERSION}" \ + "torchaudio==${TORCHAUDIO_VERSION}" \ + "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \ + rocm-sdk init + +# Torch runtime deps that may not be published on the ROCm wheel index; +# install them from PyPI afterwards. +RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \ + "networkx>=2.5.1" jinja2 "fsspec>=0.8.5" + + + +# Expose the rocm-sdk wheel as a conventional /opt/rocm install so downstream +# builds (Dockerfile.rocm: vLLM csrc, RIXL/UCX, ROCShmem/DeepEP) keep working. +RUN ln -sfn "${ROCM_PATH}" /opt/rocm; + +RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \ + sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \ + "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \ + fi + +# The ROCm SDK wheel ships a broken CMake export for hsakmt +# This patch includes the right paths for the numa build target +RUN <<'EOF' +set -eu +TARGETS="/opt/rocm/lib/cmake/hsakmt/hsakmtTargets.cmake" +[ -f "$TARGETS" ] || exit 0 # nothing to patch +grep -q NUMA_LIBRARY "$TARGETS" && exit 0 # already patched + +# 1. Point libdrm's -L at the copy bundled in the wheel, not the builder path. +sed -i 's|-L/__w/[^;"]*|-L${_IMPORT_PREFIX}/lib/rocm_sysdeps/lib|g' "$TARGETS" + +# 2. Drop the nonexistent RHEL libc path (libc is linked implicitly anyway). +sed -i 's|/usr/lib64/libc.so;||g' "$TARGETS" + +# 3. Define the numa::numa target the export references but forgot to create. +cat >> "$TARGETS" <<'CMAKE' + +if(NOT TARGET numa::numa) + find_library(NUMA_LIBRARY NAMES numa REQUIRED) + add_library(numa::numa UNKNOWN IMPORTED) + set_target_properties(numa::numa PROPERTIES IMPORTED_LOCATION "${NUMA_LIBRARY}") +endif() +CMAKE +EOF + +# Clone custom triton_kernels for install later +RUN mkdir -p /app/patched_triton_kernels; \ + cd /app/patched_triton_kernels \ + && git clone ${TRITON_KERNEL_REPO} \ + && cd triton \ + && git checkout ${TRITON_KERNEL_BRANCH} \ + && git submodule update --init --recursive +ENV TRITON_KERNELS_SRC_DIR="/app/patched_triton_kernels/triton/python/triton_kernels/triton_kernels/" + + +# Setup sccache for HIP compilation via HIP_CLANG_PATH +# This creates wrapper scripts in a separate directory and points HIP to use them +# This avoids modifying the original ROCm binaries which can break detection +# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm) +# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1 +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + echo "Setting up sccache wrappers for HIP compilation..." \ + && mkdir -p /opt/sccache-wrappers \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \ + && chmod +x /opt/sccache-wrappers/clang++ \ + && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \ + && chmod +x /opt/sccache-wrappers/clang \ + && echo "sccache wrappers created in /opt/sccache-wrappers"; \ + fi + +# Set sccache environment variables only when USE_SCCACHE=1 +# This prevents S3 config from leaking into images when sccache is not used +ARG USE_SCCACHE +ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} +ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} +ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} +ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} + + +### +### AMD SMI Build +### +FROM base AS build_amdsmi +RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \ + && pip wheel . --wheel-dir=dist +RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install + + +### +### MORI Build TODO(Build needs fixing) +### +FROM base AS build_mori +ARG MORI_BRANCH +ARG MORI_REPO +ARG MORI_GPU_ARCHS +RUN mkdir -p /app/install; \ + git clone ${MORI_REPO} \ + && cd mori \ + && git checkout ${MORI_BRANCH} \ + && git submodule update --init --recursive \ + && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \ + && cp /app/mori/dist/*.whl /app/install; + + +### +### FlashAttention Build +### +# Remove && git submodule update --init \ for CK +FROM base AS build_fa +ARG FA_BRANCH +ARG FA_REPO +ARG USE_SCCACHE +RUN mkdir -p /app/install; \ + git clone ${FA_REPO} \ + && cd flash-attention \ + && git checkout ${FA_BRANCH} \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && GPU_ARCHS=$(echo ${FA_GPU_ARCHS} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && cp dist/*.whl /app/install; + + +### +### AITER Build +### +FROM base AS build_aiter +ARG AITER_BRANCH +ARG AITER_REPO +ARG USE_SCCACHE +RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO} +RUN cd aiter \ + && git submodule update --init --recursive \ + && pip install -r requirements.txt +RUN pip install pyyaml && cd aiter \ + && if [ "$USE_SCCACHE" = "1" ]; then \ + export HIP_CLANG_PATH=/opt/sccache-wrappers \ + && sccache --show-stats; \ + fi \ + && AITER_USE_SYSTEM_TRITON=1 PREBUILD_KERNELS=${PREBUILD_KERNELS} GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \ + && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \ + && ls /app/aiter/dist/*.whl +RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install + + +### +### Final Build +### + +# Wheel release stage - +# only includes dependencies used by wheel release pipeline +FROM base AS debs_wheel_release +RUN mkdir /app/debs +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi +RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs + +# Full debs stage - includes Mori (used by Docker releases) +FROM base AS debs +RUN mkdir /app/debs +RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi +RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \ + cp /install/*.whl /app/debs +RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \ + if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi + +FROM base AS final +RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \ + pip install /install/*.whl + +ARG BASE_IMAGE +ARG ROCM_WHEEL_INDEX +ARG ROCM_SDK_VERSION +ARG TORCH_VERSION +ARG TORCHVISION_VERSION +ARG TORCHAUDIO_VERSION +ARG FA_BRANCH +ARG FA_REPO +ARG AITER_BRANCH +ARG AITER_REPO +ARG MORI_BRANCH +ARG MORI_REPO +RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \ + && echo "ROCM_WHEEL_INDEX: ${ROCM_WHEEL_INDEX}" >> /app/versions.txt \ + && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \ + && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \ + && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \ + && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \ + && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \ + && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \ + && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \ + && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \ + && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \ + && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250 new file mode 100644 index 0000000000..4c0008f6ec --- /dev/null +++ b/docker/Dockerfile.rocm_gfx1250 @@ -0,0 +1,729 @@ +# default base image +ARG REMOTE_VLLM="0" +ARG COMMON_WORKDIR=/app +ARG BASE_IMAGE=rocm/vllm-dev:base +ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base +# NIC backend for MoRI RDMA support. +# By default (all), drivers and userspace libraries for all supported NIC types +# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. +# To install drivers for a single NIC type only, set NIC_BACKEND explicitly: +# --build-arg NIC_BACKEND=ainic # AMD AINIC (Pensando) only +# --build-arg NIC_BACKEND=bnxt # Broadcom Thor-2 only +# --build-arg NIC_BACKEND=none # Install nothing. +ARG NIC_BACKEND=all +# AMD AINIC apt repo settings +# Users can specify a custom version compatible with their host drivers. +# The default version has been tested with ioinic-dkms=25.11.1.001 +ARG AINIC_VERSION=1.117.3-hydra +ARG UBUNTU_CODENAME=jammy + +# Sccache configuration. Release builds use this today; CI can opt in when a +# shared S3-compatible cache backend is available. +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME=vllm-build-sccache +ARG SCCACHE_REGION_NAME=us-west-2 +ARG SCCACHE_S3_NO_CREDENTIALS=0 + +FROM ${BASE_IMAGE} AS base + +ARG ARG_PYTORCH_ROCM_ARCH=gfx1250 +ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} + +# Install build dependencies and utilities +RUN apt-get update -q -y && apt-get install -q -y \ + sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \ + apt-transport-https ca-certificates wget curl \ + build-essential libnuma-dev ccache mold +RUN --mount=type=cache,target=/root/.cache/pip \ + python3 -m pip install --upgrade pip +# Note: mold is installed but not set as the system default linker because +# some packages use JIT compilation at runtime with flags mold does not support. +# Build stages opt in via LDFLAGS="-fuse-ld=mold". +# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base) +ARG USE_SCCACHE +RUN if [ "$USE_SCCACHE" != "1" ]; then \ + apt-get purge -y sccache || true; \ + python3 -m pip uninstall -y sccache || true; \ + rm -f "$(which sccache)" || true; \ + fi + +# Install UV — download first, then run, so a curl failure is not masked by the pipe +RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \ + && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \ + && rm -f /tmp/uv-install.sh \ + && uv --version + +# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out +# Reference: https://github.com/astral-sh/uv/pull/1694 +ENV UV_HTTP_TIMEOUT=500 +ENV UV_INDEX_STRATEGY="unsafe-best-match" +# Use copy mode to avoid hardlink failures with Docker cache mounts +ENV UV_LINK_MODE=copy +# python binary fall back for non venv builds +ENV UV_PYTHON=${VIRTUAL_ENV:-/usr}/bin/python3 +# Expose paths from wheel installation +ENV PKG_CONFIG_PATH=${ROCM_PATH}/lib/rocm_sysdeps/lib/pkgconfig:${PKG_CONFIG_PATH} +# ccache directory - persisted across layer rebuilds via cache mounts. +ENV CCACHE_DIR=/root/.cache/ccache +ENV CCACHE_COMPILERCHECK=content +# Empty by default so build steps fall back to $(nproc); CI can override. +ARG max_jobs +ENV MAX_JOBS=${max_jobs} + +# Install sccache if USE_SCCACHE is enabled (for release builds) +ARG USE_SCCACHE +ARG SCCACHE_DOWNLOAD_URL +ARG SCCACHE_ENDPOINT +ARG SCCACHE_BUCKET_NAME +ARG SCCACHE_REGION_NAME +ARG SCCACHE_S3_NO_CREDENTIALS +RUN if [ "$USE_SCCACHE" = "1" ]; then \ + if command -v sccache >/dev/null 2>&1; then \ + echo "sccache already installed, skipping installation"; \ + sccache --version; \ + else \ + echo "Installing sccache..." \ + && SCCACHE_ARCH="x86_64" \ + && SCCACHE_VERSION="v0.8.1" \ + && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \ + && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \ + && tar -xzf /tmp/sccache.tar.gz -C /tmp \ + && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \ + && chmod +x /usr/bin/sccache \ + && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \ + && sccache --version; \ + fi; \ + fi + +# Set sccache environment variables only when USE_SCCACHE=1 +# This prevents S3 config from leaking into images when sccache is not used +ARG USE_SCCACHE +ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}} +ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}} +ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}} +ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0} + +ARG COMMON_WORKDIR +WORKDIR ${COMMON_WORKDIR} + + +# ----------------------- +# vLLM fetch stages +FROM base AS fetch_vllm_0 +ONBUILD COPY ./ vllm/ +FROM base AS fetch_vllm_1 +ARG VLLM_REPO="https://github.com/ROCm/vllm.git" +ARG VLLM_BRANCH="455_wip" +ENV VLLM_REPO=${VLLM_REPO} +ENV VLLM_BRANCH=${VLLM_BRANCH} +ONBUILD RUN git clone ${VLLM_REPO} \ + && cd vllm \ + && git fetch -v --prune -- origin ${VLLM_BRANCH} \ + && git checkout FETCH_HEAD \ + && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \ + git remote add upstream "https://github.com/vllm-project/vllm.git" \ + && git fetch upstream ; fi +FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm + +# ----------------------- +# Rust build stage +# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages +# don't need the rust toolchain or protoc. +FROM fetch_vllm AS rust-build +ARG COMMON_WORKDIR + +# protoc is used by tonic-build/prost-build. +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ + ca-certificates curl unzip \ + && rm -rf /var/lib/apt/lists/* + +COPY tools/install_protoc.sh /tmp/install_protoc.sh +RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh + +# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit +# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). +ENV CARGO_BUILD_JOBS=4 +ENV CARGO_NET_RETRY=10 +ENV RUSTUP_MAX_RETRIES=10 + +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd ${COMMON_WORKDIR}/vllm \ + && uv pip install --system -r requirements/build/rust.txt + +# Build the release binary. Cargo's registry/git caches can be written by +# concurrent BuildKit jobs on shared workers, so lock those cache mounts while +# keeping the cache benefit. Do not cache target/, because stale target metadata +# can outlive source updates across BuildKit cache reuse. +RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ + --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ + cd ${COMMON_WORKDIR}/vllm \ + && bash build_rust.sh \ + && test -x vllm/vllm-rs + +# ----------------------- +# vLLM native build stages +# +# csrc-build intentionally copies only files that affect ROCm native extension +# compilation. That keeps unrelated CI/test/docs edits from invalidating the +# expensive HIP/C++ build layer. +FROM base AS csrc-build +ARG COMMON_WORKDIR +WORKDIR ${COMMON_WORKDIR}/vllm + +COPY requirements/rocm.txt requirements/rocm.txt +COPY requirements/common.txt requirements/common.txt +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + uv pip install --system -r requirements/rocm.txt + +# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do +# not invalidate the expensive native build layer. +COPY setup.py CMakeLists.txt ./ +COPY tools/build_rust.py tools/build_rust.py +COPY cmake cmake/ +COPY csrc csrc/ +COPY vllm/envs.py vllm/envs.py +COPY vllm/__init__.py vllm/__init__.py + +ENV VLLM_TARGET_DEVICE=rocm +ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" + +RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ + export CCACHE_BASEDIR="$PWD" \ + && echo "=== ccache stats before ROCm native build ===" \ + && (ccache --show-stats || true) \ + && (ccache --zero-stats || true) \ + && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \ + && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \ + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \ + && test -d dist \ + && ls dist/*.whl >/dev/null \ + && echo "=== ccache stats after ROCm native build ===" \ + && (ccache --show-stats || true) + +# Build the full vLLM ROCm wheel by reusing the native extension wheel from +# csrc-build. This stage still rebuilds for Python/package changes, but skips +# the expensive HIP/C++ compile when native inputs are unchanged. +FROM fetch_vllm AS build_vllm +ARG COMMON_WORKDIR +ENV VLLM_TARGET_DEVICE=rocm + +COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels + +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ + +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd vllm \ + && uv pip install --system -r requirements/rocm.txt \ + && export VLLM_USE_PRECOMPILED=1 \ + && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ + && export VLLM_DOCKER_BUILD_CONTEXT=1 \ + && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \ + && python3 setup.py bdist_wheel --dist-dir=dist \ + && test -d dist \ + && ls dist/*.whl >/dev/null +FROM scratch AS export_vllm +ARG COMMON_WORKDIR +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl / +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 + +# RIXL/UCX build stages +FROM base AS build_rixl +ARG RIXL_BRANCH="39be1de8" +ARG RIXL_REPO="https://github.com/ROCm/RIXL.git" +ARG UCX_BRANCH="bfb51733" +ARG UCX_REPO="https://github.com/openucx/ucx.git" +# ENV ROCM_PATH=/opt/rocm -> correct ROCM_PATH is set in base image +ENV UCX_HOME=/usr/local/ucx +ENV RIXL_HOME=/usr/local/rixl +ENV RIXL_BENCH_HOME=/usr/local/rixl_bench + +# RIXL build system dependences and RDMA support +RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \ + libgrpc-dev \ + libgrpc++-dev \ + libprotobuf-dev \ + protobuf-compiler-grpc \ + libcpprest-dev \ + libaio-dev \ + librdmacm1 \ + librdmacm-dev \ + libibverbs1 \ + libibverbs-dev \ + ibverbs-utils \ + rdmacm-utils \ + ibverbs-providers \ + && rm -rf /var/lib/apt/lists/* + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system meson auditwheel patchelf tomlkit + +RUN --mount=type=cache,target=/root/.cache/ccache \ + cd /usr/local/src && \ + git clone ${UCX_REPO} && \ + cd ucx && \ + git checkout ${UCX_BRANCH} && \ + ./autogen.sh && \ + mkdir build && cd build && \ + CC="ccache gcc" CXX="ccache g++" \ + ../configure \ + --prefix=/usr/local/ucx \ + --enable-shared \ + --disable-static \ + --disable-doxygen-doc \ + --enable-optimizations \ + --enable-devel-headers \ + --with-rocm=${ROCM_PATH} \ + --with-verbs \ + --with-dm \ + --enable-mt && \ + make -j$(nproc) && \ + make install + +ENV PATH=/usr/local/ucx/bin:$PATH +ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH} + +RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone ${RIXL_REPO} /opt/rixl && \ + cd /opt/rixl && \ + git checkout ${RIXL_BRANCH} && \ + CC="ccache gcc" CXX="ccache g++" \ + meson setup build --prefix=${RIXL_HOME} \ + --force-fallback-for=abseil-cpp \ + -Ducx_path=${UCX_HOME} \ + -Drocm_path=${ROCM_PATH} && \ + cd build && \ + ninja -j$(nproc) && \ + ninja install + +# Generate RIXL wheel +# Exclude libcore and libpull from auditwheel: transitive dependencies +# that are not shipped in the wheel and vary across base images. +RUN cd /opt/rixl && \ + sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \ + contrib/build-wheel.sh && \ + # The wheel build re-runs meson via meson-python; force the bundled abseil + sed -i 's|setup = \["-Dinstall_headers=false"\]|setup = ["-Dinstall_headers=false", "--force-fallback-for=abseil-cpp"]|' \ + pyproject.toml && \ + grep -q 'force-fallback-for' pyproject.toml && \ + mkdir -p /app/install && \ + _ucx_install_dir=${UCX_HOME} \ + ./contrib/build-wheel.sh \ + --output-dir /app/install \ + --rocm-dir ${ROCM_PATH} \ + --ucx-plugins-dir ${UCX_HOME}/lib/ucx \ + --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins + +# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not +# invalidate the slow ROCShmem build. +FROM base AS build_rocshmem +ARG ROCSHMEM_BRANCH="f0acb0c6" +ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git" +# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so +# it can be linked against DeepEP without arch mismatches. +ARG DEEPEP_ROCM_ARCH="gfx942;gfx950" +# ENV ROCM_PATH=/opt/rocm -> Correct rocm_path is set in base image +ENV ROCSHMEM_DIR=/opt/rocshmem + +RUN --mount=type=cache,target=/root/.cache/ccache \ + git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \ + && cd rocm-systems \ + && git sparse-checkout set --cone projects/rocshmem \ + && git checkout ${ROCSHMEM_BRANCH} \ + && mkdir -p projects/rocshmem/build \ + && cd projects/rocshmem/build \ + && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \ + bash ../scripts/build_configs/all_backends \ + -DROCM_PATH=${ROCM_PATH} \ + -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \ + -DUSE_EXTERNAL_MPI=OFF + +# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel. +FROM build_rocshmem AS build_deepep +ARG DEEPEP_BRANCH="a9ea9774" +ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git" +ARG DEEPEP_NIC="cx7" + +# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR. +# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list. +RUN --mount=type=cache,target=/root/.cache/ccache \ + export PYTORCH_ROCM_ARCH="gfx942;gfx950" \ + && git clone ${DEEPEP_REPO} \ + && cd DeepEP \ + && git checkout ${DEEPEP_BRANCH} \ + && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install + +# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do +# not force users to rebuild the long-lived Dockerfile.rocm_base image. +FROM base AS mori_base +ARG NIC_BACKEND +ARG AINIC_VERSION +ARG UBUNTU_CODENAME +RUN /bin/bash -lc 'set -euo pipefail; \ + \ + install_ainic() { \ + apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \ + rm -rf /var/lib/apt/lists/*; \ + mkdir -p /etc/apt/keyrings; \ + curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ + > /etc/apt/sources.list.d/amdainic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + libionic-dev \ + ionic-common \ + ; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \ + # bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \ + install_bnxt() { \ + install -m 0755 -d /etc/apt/keyrings; \ + curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \ + -o /etc/apt/keyrings/broadcom-nic.asc; \ + chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \ + > /etc/apt/sources.list.d/broadcom-nic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + bnxt-rocelib=235.2.86.0 \ + ; \ + cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \ + ldconfig; \ + rm -rf /var/lib/apt/lists/*; \ + }; \ + \ + echo "[MORI] Install MoRI proxy deps"; \ + pip install --quiet --ignore-installed blinker && \ + pip install --quiet quart msgpack aiohttp pyzmq; \ + echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \ + \ + # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \ + # Only vendor packages are installed here for dlopen; no compile-time flags needed. \ + case "${NIC_BACKEND}" in \ + none) ;; \ + all) install_ainic; install_bnxt ;; \ + ainic) install_ainic ;; \ + bnxt) install_bnxt ;; \ + *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \ + esac' + +# ----------------------- +# vLLM wheel release build stage (for building distributable wheels) +# This stage pins dependencies to custom ROCm wheel versions and handles version detection +FROM fetch_vllm AS build_vllm_wheel_release + +ARG COMMON_WORKDIR + +# Drop the pre-built Rust artifacts into the source tree. setup.py detects +# them and ships them as-is, skipping the local Rust build. +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs +COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ + +# Create /install directory for custom wheels +RUN mkdir -p /install + +# Copy custom ROCm wheels from docker/context if they exist +# COPY ensures Docker cache is invalidated when wheels change +# .keep file ensures directory always exists for COPY to work +COPY docker/context/base-wheels/ /tmp/base-wheels/ +# This is how we know if we are building for a wheel release or not. +# If there are not wheels found there, we are not building for a wheel release. +# So we exit with an error. To skip this stage. +RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \ + echo "Found custom wheels - copying to /install"; \ + cp /tmp/base-wheels/*.whl /install/ && \ + echo "Copied custom wheels:"; \ + ls -lh /install/; \ + else \ + echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \ + echo "Wheel releases require pre-built ROCm wheels."; \ + exit 1; \ + fi + +# GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds) +# This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm +ARG GIT_REPO_CHECK=0 +RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ + echo "Running repository checks..."; \ + cd vllm && bash tools/check_repo.sh; \ + fi + +# Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) +# This ensures setuptools_scm sees clean repo state for version detection +RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,target=/root/.cache/uv \ + cd vllm \ + && uv pip install --system setuptools_scm regex \ + && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \ + && echo "Detected vLLM version: ${VLLM_VERSION}" \ + && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt + +# Fail if git-based package dependencies are found in requirements files +# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI) +# Extra notes: pip install is able to handle git+ URLs, but uv doesn't. +RUN echo "Checking for git-based packages in requirements files..." \ + && echo "Checking common.txt for git-based packages:" \ + && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \ + echo "ERROR: Git-based packages found in common.txt:"; \ + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \ + echo "Please publish these packages to PyPI instead of using git dependencies."; \ + exit 1; \ + else \ + echo " ✓ No git-based packages found in common.txt"; \ + fi \ + && echo "Checking rocm.txt for git-based packages:" \ + && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \ + echo "ERROR: Git-based packages found in rocm.txt:"; \ + grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \ + echo "Please publish these packages to PyPI instead of using git dependencies."; \ + exit 1; \ + else \ + echo " ✓ No git-based packages found in rocm.txt"; \ + fi \ + && echo "All requirements files are clean - no git-based packages found" + +# Pin vLLM dependencies to exact versions of custom ROCm wheels +# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi +COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py +RUN echo "Pinning vLLM dependencies to custom wheel versions..." \ + && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt + +# Install dependencies using custom wheels from /install +RUN --mount=type=cache,target=/root/.cache/uv \ + cd vllm \ + && echo "Building vLLM with custom wheels from /install" \ + && uv pip install --system --find-links /install -r requirements/rocm.txt + +# Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt +# (setup.py auto-detects ccache/sccache in PATH) +RUN --mount=type=bind,source=.git,target=vllm/.git \ + --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ + cd vllm \ + && export CCACHE_BASEDIR="$PWD" \ + && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \ + && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ + && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist + +FROM scratch AS export_vllm_wheel_release +ARG COMMON_WORKDIR +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/dist/*.whl / +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requirements +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/ +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml +COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1 + +# ----------------------- +# CI base image (Tier 1) - stable, rarely changing CI dependencies. +# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers +# in the vLLM artifacts for the current commit. +FROM mori_base AS ci_base +ARG COMMON_WORKDIR + +# Update rdma-core to support latest rocshmem. +ARG DEEPEP_NIC +RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \ + git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \ + cd /tmp/rdma-core && \ + mkdir -p build && cd build && \ + cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \ + ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \ +fi + +# Install RIXL + DeepEP wheels. +RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ + --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \ + uv pip install --system /rixl_install/*.whl /deep_install/*.whl + +# Copy ROCShmem runtime libraries. +COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem + +# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec. +RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \ + librdmacm1 \ + libibverbs1 \ + ibverbs-providers \ + ibverbs-utils \ + pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \ + libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \ + && rm -rf /var/lib/apt/lists/* + +# Install torchcodec from source for ROCm/torch ABI compatibility. +COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh +RUN --mount=type=cache,target=/root/.cache/uv \ + --mount=type=cache,target=/root/.cache/pip \ + --mount=type=cache,target=/root/.cache/torchcodec-wheels \ + bash /tmp/install_torchcodec.sh \ + && rm /tmp/install_torchcodec.sh \ + && apt-get clean && rm -rf /var/lib/apt/lists/* + +# Pre-install shared ROCm runtime dependencies. +COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/ +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \ + && rm -rf /tmp/ci-base-requirements + +# Enable fast and less brittle model downloads in tests. +ENV HF_XET_HIGH_PERFORMANCE=1 +ENV HF_HUB_DOWNLOAD_TIMEOUT=60 + +# Pre-install vLLM test dependencies. +COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system -r /tmp/rocm-test-reqs.txt + +# Rebuild fastsafetensors from source so its C++ extension is compiled with +# USE_ROCM and can detect libamdhip64.so at runtime. +RUN --mount=type=cache,target=/root/.cache/pip \ + FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \ + && test -n "${FASTSAFETENSORS_REQ}" \ + && python3 -m pip install --force-reinstall --no-deps \ + --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \ + && rm /tmp/rocm-test-reqs.txt + +# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel. +# See: https://github.com/pytorch/pytorch/issues/169857 +ENV MIOPEN_DEBUG_CONV_DIRECT=0 +ENV MIOPEN_DEBUG_CONV_GEMM=0 + +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc. +# See: https://github.com/ROCm/rocm-libraries/issues/6266 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1 + +# ROCm profiler limits workaround. +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" + +# Install vllm_test_utils in ci_base for ci_base + wheel parity. +COPY tests/vllm_test_utils /tmp/vllm_test_utils +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system /tmp/vllm_test_utils \ + && rm -rf /tmp/vllm_test_utils + +# ----------------------- +# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base. +FROM ${CI_BASE_IMAGE} AS test +ARG COMMON_WORKDIR + +# Install the vLLM wheel (--no-deps: all deps already in ci_base). +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ + --mount=type=cache,target=/root/.cache/uv \ + cd /install \ + && uv pip install --system --no-deps *.whl + +# Store the vLLM wheel in the image for python-only install tests. +COPY --from=export_vllm /*.whl /opt/vllm-wheels/ + +WORKDIR /vllm-workspace +COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace + +# Copy in the v1 package (for python-only install test group). +COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1 + +# Hide source under src/ so it won't shadow the installed package in tests. +RUN mkdir src && mv vllm src/vllm + +# ----------------------- +# Final vLLM image +FROM mori_base AS final + +RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/* + +# Clean up sccache from release image (not needed at runtime) +# This removes the binary and wrappers that may have been installed during build +RUN rm -f /usr/bin/sccache || true \ + && rm -rf /opt/sccache-wrappers || true + +# Unset sccache environment variables for the release image +# This prevents S3 bucket config from leaking into production images +ENV SCCACHE_BUCKET= +ENV SCCACHE_REGION= +ENV SCCACHE_ENDPOINT= +ENV SCCACHE_S3_NO_CREDENTIALS= +ENV SCCACHE_IDLE_TIMEOUT= + +# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt. +# Manually remove it so that later steps of numpy upgrade can continue +RUN case "$(which python3)" in \ + *"/opt/conda/envs/py_3.9"*) \ + rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \ + *) ;; esac + +RUN --mount=type=cache,target=/root/.cache/uv \ + uv pip install --system --upgrade huggingface-hub[cli] + +# Install vLLM using uv (inherited from base stage) +# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version +RUN --mount=type=bind,from=export_vllm,src=/,target=/install \ + --mount=type=cache,target=/root/.cache/uv \ + cd /install \ + && uv pip install --system -r requirements/rocm.txt \ + && pip uninstall -y vllm \ + && uv pip install --system *.whl + +# Install RIXL wheel +RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \ + uv pip install --system /rixl_install/*.whl + +ARG COMMON_WORKDIR +ARG BASE_IMAGE +ARG NIC_BACKEND +ARG AINIC_VERSION + +# Copy over the benchmark scripts as well +COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks +COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples +COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker + +# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc +# See: https://github.com/ROCm/rocm-libraries/issues/6266 +ENV HSA_ENABLE_IPC_MODE_LEGACY=1 + +ENV TOKENIZERS_PARALLELISM=false + +# ENV that can improve safe tensor loading, and end-to-end time +ENV SAFETENSORS_FAST_GPU=1 + +# Performance environment variable. +ENV HIP_FORCE_DEV_KERNARG=1 + +# Workaround for ROCm profiler limits +RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf +ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf" +RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \ + && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \ + && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt + +# Download bench scripts for gfx1250 +RUN curl -k -O https://raw.githubusercontent.com/ROCm/vllm/refs/heads/gfx1250_bench/vllm_smoketest.sh + + +### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better +RUN pip3 uninstall -y triton && \ + git clone https://github.com/triton-lang/triton.git && \ + cd triton && \ + git checkout c517f38c && \ + TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install . + +CMD ["/bin/bash"] + +#Set entrypoint for vllm-openai official images +FROM final AS vllm-openai +ENTRYPOINT ["vllm", "serve"] diff --git a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py index bbfde877b9..63357195d0 100644 --- a/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py +++ b/tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py @@ -31,7 +31,7 @@ def _get_attention_backend_params() -> list[str | None]: falls back to ROCM_AITER_UNIFIED_ATTN or TRITON_ATTN for cross-attention since ROCM_ATTN doesn't support ENCODER_DECODER) - TRITON_ATTN: always available on ROCm - - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950 + - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950/gfx1250 On non-ROCm platforms, we just run with the default backend. """ @@ -40,9 +40,9 @@ def _get_attention_backend_params() -> list[str | None]: if current_platform.is_rocm(): backends: list[str | None] = [None, "TRITON_ATTN"] - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import get_cdna_version - if _ON_MI3XX: + if get_cdna_version() > 2: backends.append("ROCM_AITER_UNIFIED_ATTN") return backends except Exception: diff --git a/tests/entrypoints/speech_to_text/translation/test_translation_validation.py b/tests/entrypoints/speech_to_text/translation/test_translation_validation.py index ed3cff5f1c..25f26404d0 100644 --- a/tests/entrypoints/speech_to_text/translation/test_translation_validation.py +++ b/tests/entrypoints/speech_to_text/translation/test_translation_validation.py @@ -37,13 +37,13 @@ def _get_rocm_attention_config(model_name): if "whisper" in model_name.lower(): try: - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import get_cdna_version - if _ON_MI3XX: + if get_cdna_version() > 2: return {"backend": "ROCM_AITER_UNIFIED_ATTN"} except ImportError: logger.warning( - "Could not import _ON_MI3XX from rocm platform, " + "Could not check cdna version from rocm platform, " "falling back to TRITON_ATTN for Whisper." ) return {"backend": "TRITON_ATTN"} diff --git a/tests/kernels/attention/test_rocm_aiter_unified_attn.py b/tests/kernels/attention/test_rocm_aiter_unified_attn.py index 9e33f24ea2..a03c00e552 100644 --- a/tests/kernels/attention/test_rocm_aiter_unified_attn.py +++ b/tests/kernels/attention/test_rocm_aiter_unified_attn.py @@ -15,15 +15,15 @@ from tests.kernels.attention.test_triton_unified_attention import ref_paged_attn from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed -_SKIP_NON_MI3XX = True +_SKIP_NON_CDNA_2_PLUS = True if current_platform.is_rocm(): - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - _SKIP_NON_MI3XX = not on_mi3xx() + _SKIP_NON_CDNA_2_PLUS = get_cdna_version() < 2 pytestmark = [ pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"), - pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"), + pytest.mark.skipif(_SKIP_NON_CDNA_2_PLUS, reason="CDNA 2+ ROCm only"), ] NUM_Q_HEADS = 8 diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 0315d8d89e..b3a42912da 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -4,6 +4,8 @@ Test modular OAI Triton MoE """ +from __future__ import annotations + import pytest import torch @@ -48,6 +50,35 @@ MNK = [ ] +def deepseek_v4_flash_moe_topology(): + """MoE sizes for the DeepSeek-V4-Flash`. + Default weights from: ``deepseek-ai/DeepSeek-V4-Flash`. + ``moe_intermediate_size``, ``n_routed_experts``, and ``num_experts_per_tok``. + """ + defaults = { + "hidden_size": 4096, + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "num_experts_per_tok": 6, + } + + return defaults + + +def scaled_deepseek_v4_flash_problem( + *, + dim_scale: int = 8, + expert_scale: int = 8, +): + """Smaller K/N/E for kernel tests; keeps production top_k and K:N ratio (~2:1).""" + t = deepseek_v4_flash_moe_topology() + k = max(128, t["hidden_size"] // dim_scale) + n = max(64, t["moe_intermediate_size"] // dim_scale) + num_experts = max(t["num_experts_per_tok"], t["n_routed_experts"] // expert_scale) + topk = t["num_experts_per_tok"] + return k, n, num_experts, topk + + def unshuffle_weight(w: torch.Tensor): first = w[..., ::2] second = w[..., 1::2] @@ -160,8 +191,8 @@ def oai_triton_moe_impl( x: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, - w1_scale: "PrecisionConfig", - w2_scale: "PrecisionConfig", + w1_scale: PrecisionConfig, + w2_scale: PrecisionConfig, w1_bias: torch.Tensor | None, w2_bias: torch.Tensor | None, num_experts: int, @@ -261,3 +292,93 @@ def test_oai_triton_moe( ) assert_close(ref=out_ref, tri=out, maxtol=0.025, rmstol=0.005) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." +) +def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_init): + """Exercise ``UnfusedOAITritonExperts.apply`` with explicit workspaces. + + Same MoE topology as ``launch_dsv4.sh`` / DeepSeek-V4-Flash ``config.json``, + with linear dimensions and expert count scaled down for test GPU memory. + """ + wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1) + set_random_seed(0) + + k, n, num_experts, topk = scaled_deepseek_v4_flash_problem() + m = 7 + dtype = torch.bfloat16 + + ( + w1, + w2, + w1_bias, + w2_bias, + w1_tri, + w2_tri, + w1_bias_tri, + w2_bias_tri, + w1_precision_config, + w2_precision_config, + ) = make_weights(dtype, k, n, num_experts) + + x = torch.randn((m, k), dtype=dtype, device="cuda") + router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype) + topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True) + topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1) + + quant_config = mxfp4_w4a16_moe_quant_config( + w1_bias=w1_bias_tri, + w2_bias=w2_bias_tri, + w1_scale=w1_precision_config, + w2_scale=w2_precision_config, + ) + moe_config = make_dummy_moe_config( + num_experts=num_experts, + experts_per_token=topk, + hidden_dim=k, + intermediate_size=n, + ) + experts = UnfusedOAITritonExperts(moe_config, quant_config) + + if not UnfusedOAITritonExperts._supports_current_device(): + pytest.skip("UnfusedOAITritonExperts does not support this device") + + _, _, N, K, top_k = experts.moe_problem_size(x, w1_tri, w2_tri, topk_ids) + assert top_k == topk + ws13_shape, ws2_shape, out_shape = experts.workspace_shapes( + m, + N, + K, + topk, + num_experts, + num_experts, + None, + MoEActivation.SWIGLUOAI, + ) + workspace13 = torch.empty(ws13_shape, dtype=dtype, device="cuda") + workspace2 = torch.empty(ws2_shape, dtype=dtype, device="cuda") + output = torch.empty(out_shape, dtype=dtype, device="cuda") + + with set_current_vllm_config(VllmConfig()): + out_ref = torch_moe_impl(x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids) + experts.apply( + output=output, + hidden_states=x, + w1=w1_tri, + w2=w2_tri, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SWIGLUOAI, + global_num_experts=num_experts, + expert_map=None, + a1q_scale=None, + a2_scale=None, + workspace13=workspace13, + workspace2=workspace2, + expert_tokens_meta=None, + apply_router_weight_on_input=False, + ) + + assert_close(ref=out_ref, tri=output, maxtol=0.025, rmstol=0.005) diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index 9b12fef045..0eded09b1a 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -29,13 +29,13 @@ from vllm.platforms import current_platform def _is_aiter_capable() -> bool: - """Check if the platform supports AITER (gfx942/gfx950).""" + """Check if the platform supports AITER (gfx942/gfx950/gfx1250).""" if not current_platform.is_rocm(): return False try: - from vllm.platforms.rocm import _ON_MI3XX + from vllm.platforms.rocm import get_cdna_version - return _ON_MI3XX + return get_cdna_version() > 2 except ImportError: return False diff --git a/tests/model_executor/layers/test_rocm_unquantized_gemm.py b/tests/model_executor/layers/test_rocm_unquantized_gemm.py index f4de9bc903..f20be7cfb6 100644 --- a/tests/model_executor/layers/test_rocm_unquantized_gemm.py +++ b/tests/model_executor/layers/test_rocm_unquantized_gemm.py @@ -26,6 +26,7 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t()) @@ -52,6 +53,7 @@ def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t()) @@ -76,6 +78,7 @@ def test_rocm_unquantized_gemm_gfx950_wvsplitkrc_path(monkeypatch): monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False) monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True) + monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: True) monkeypatch.setattr(utils, "num_compute_units", lambda: 120) wvsplitkrc_mock = MagicMock(side_effect=lambda x_view, w, _, __: x_view @ w.t()) diff --git a/tests/models/multimodal/generation/test_granite_speech.py b/tests/models/multimodal/generation/test_granite_speech.py index 3019f5f22d..e1fee635ff 100644 --- a/tests/models/multimodal/generation/test_granite_speech.py +++ b/tests/models/multimodal/generation/test_granite_speech.py @@ -45,9 +45,10 @@ models: dict[str, str | None] = { def granite_speech_attention_config(): """Return attention config for Granite Speech tests on ROCm.""" if current_platform.is_rocm(): - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - if on_mi3xx(): + # -1 (unknown arch) is truthy; gate on CDNA3+ like other call sites. + if get_cdna_version() > 2: return {"backend": "ROCM_AITER_FA"} return {"backend": "TRITON_ATTN"} return None diff --git a/tests/models/quantization/test_bitsandbytes.py b/tests/models/quantization/test_bitsandbytes.py index 3c2fd0dbd4..c2f1a89df2 100644 --- a/tests/models/quantization/test_bitsandbytes.py +++ b/tests/models/quantization/test_bitsandbytes.py @@ -21,11 +21,11 @@ from ...utils import compare_two_settings, multi_gpu_test from ..utils import check_embeddings_close, check_logprobs_close if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx9 + from vllm.platforms.rocm import on_cdna pytestmark = pytest.mark.skipif( - on_gfx9(), - reason="bitsandbytes not supported on gfx9 (warp size 64 limitation)", + on_cdna(), + reason="bitsandbytes not supported on CDNA (warp size 64 limitation)", ) models_4bit_to_test = [ diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 8f9e8acac6..36d5cca9a4 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -28,16 +28,9 @@ def mock_vllm_config(): @pytest.fixture -def mock_on_gfx9(): - """Mock gfx9 arch detection to return True.""" - with patch("vllm.platforms.rocm.on_gfx9", return_value=True): - yield - - -@pytest.fixture -def mock_on_mi3xx(): - """Mock mi3xx arch detection to return True.""" - with patch("vllm.platforms.rocm.on_mi3xx", return_value=True): +def mock_get_cdna_version(): + """Mock cdna version arch detection to return True.""" + with patch("vllm.platforms.rocm.get_cdna_version", return_value=3): yield @@ -111,8 +104,7 @@ def test_standard_attention_backend_selection( selected_backend, expected_backend_path, mock_vllm_config, - mock_on_gfx9, - mock_on_mi3xx, + mock_get_cdna_version, monkeypatch, ): """Test standard attention backend selection with various configurations.""" @@ -305,12 +297,12 @@ def test_mla_backend_selection( def test_aiter_fa_requires_mi3xx(mock_vllm_config): - """Test that ROCM_AITER_FA requires mi3xx architecture.""" + """Test that ROCM_AITER_FA requires CDNA3+ architecture.""" from vllm.platforms.rocm import RocmPlatform - # Mock on_mi3xx to return False (used by supports_compute_capability) + # Mock cdna version to return 1 (used by supports_compute_capability) with ( - patch("vllm.platforms.rocm.on_mi3xx", return_value=False), + patch("vllm.platforms.rocm.get_cdna_version", return_value=1), pytest.raises( ValueError, match="compute capability not supported", diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 5d70d97b23..30bfcb80b0 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -52,7 +52,7 @@ IS_AITER_FOUND = is_aiter_found() def is_aiter_found_and_supported() -> bool: """Check if AITER library is available and platform supports it. - Checks: platform (ROCm), device arch (gfx9), and library existence. + Checks: platform (ROCm), device arch is CDNA 3 or better, and library existence. Does NOT check environment variables - that's handled by rocm_aiter_ops.is_enabled(). This function determines if aiter CAN be used, not if it SHOULD be used. @@ -66,9 +66,9 @@ def is_aiter_found_and_supported() -> bool: VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery. """ if current_platform.is_rocm() and IS_AITER_FOUND: - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - return on_mi3xx() + return get_cdna_version() > 2 return False @@ -1798,16 +1798,21 @@ class rocm_aiter_ops: @classmethod @if_aiter_supported def is_fp4bmm_enabled(cls) -> bool: - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import get_cdna_version - return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and on_gfx950() + # TODO GFX1250: Enable for cdna 4+ when aiter supports batched_gemm_a16wfp4 on gfx1250 + return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and get_cdna_version() == 4 @classmethod @if_aiter_supported def is_linear_hipbmm_enabled(cls) -> bool: - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - return cls.is_linear_enabled() and on_mi3xx() and cls._LINEAR_HIPBMM_ENABLED + return ( + cls.is_linear_enabled() + and (get_cdna_version() > 2) + and cls._LINEAR_HIPBMM_ENABLED + ) @classmethod @if_aiter_supported diff --git a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py index da8f69fa97..ab396d026a 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/aiter.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/aiter.py @@ -374,9 +374,14 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): super().__init__(config) n, k = config.weight_shape - self.use_triton = ( - not current_platform.is_fp8_fnuz() - and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) + _on_gfx1250 = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 + + _on_gfx1250 = on_gfx1250() + + self.use_triton = not current_platform.is_fp8_fnuz() and ( + rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) or _on_gfx1250 ) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: diff --git a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py index 2b6d3ed736..a04be2e097 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/pytorch.py @@ -104,10 +104,10 @@ class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel): if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version - if not on_mi3xx(): - return False, "requires MI3xx." + if get_cdna_version() <= 2: + return False, "requires CDNA3+" if compute_capability is not None and compute_capability < 94: return False, "requires compute capability 94 and above." diff --git a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py index 64bc5b6c8b..4a35f49a22 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/rocm.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/rocm.py @@ -79,10 +79,15 @@ class ROCmFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): if not current_platform.is_rocm(): return False, "requires ROCm." - from vllm.platforms.rocm import on_gfx12x, on_mi3xx + from vllm.platforms.rocm import get_cdna_version, on_gfx12x, on_gfx1250 - if not (on_mi3xx() or on_gfx12x()): - return False, "requires MI3xx or gfx12x" + # wvSplitKQ (skinny GEMM) is excluded from the gfx1250 build. + if on_gfx1250(): + return False, "wvSplitKQ (skinny GEMM) is not built on gfx1250" + + # Restore RDNA4 (gfx12x) dropped by the get_cdna_version()>2 refactor. + if get_cdna_version() <= 2 and not on_gfx12x(): + return False, "requires CDNA3+ (gfx942/gfx950) or RDNA4 (gfx12x)" if not envs.VLLM_ROCM_USE_SKINNY_GEMM: return False, "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled." diff --git a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py index 7c3fe5831f..53880cb916 100644 --- a/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py @@ -20,7 +20,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( __all__ = [ "AiterW4A8ExpertsMonolithic", + "AiterW4A16ExpertsMonolithic", "aiter_triton_kernel_w4a8_moe_forward", + "aiter_triton_kernel_w4a16_moe_forward", ] @@ -46,11 +48,27 @@ def aiter_triton_kernel_w4a8_moe_forward( and quant_config.use_mxfp4_w4a8 and rocm_aiter_ops.is_enabled() ) - from aiter.ops.triton.moe_routing.routing import routing as aiter_routing + from vllm.platforms.rocm import on_gfx1250 + + try: + from aiter.ops.triton.moe.moe_routing import routing as _routing_mod + except ImportError: + from aiter.ops.triton.moe_routing import routing as _routing_mod + + if on_gfx1250(): + _routing_mod.is_tdm_avail = lambda: False + aiter_routing = _routing_mod.routing routing_data, gather_idx, scatter_idx = aiter_routing( gating_output, topk, sm_first=not renormalize ) + + # gfx1250: aiter's in-kernel gather is numerically broken + if on_gfx1250(): + gather_src = gather_idx.to(torch.long) // topk + hidden_states = hidden_states[gather_src] + gather_idx = None + return triton_kernel_fused_mxfp4_w4a8_experts( None, hidden_states, @@ -199,12 +217,11 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): @staticmethod def _supports_current_device() -> bool: - # Requires AITER and GFX950 if not rocm_aiter_ops.is_enabled(): return False - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx950, on_gfx1250 - return on_gfx950() + return on_gfx950() or on_gfx1250() @staticmethod def _supports_no_act_and_mul() -> bool: @@ -293,3 +310,358 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): unpadded_N_w2=self.moe_config.hidden_dim_unpadded, unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded, ) + + +def _aiter_raw(t): + if t is None or isinstance(t, torch.Tensor): + return t + return t.storage.data if hasattr(t, "storage") else t + + +def _aiter_w4a16_silu_via_a8w4( + hidden_states: torch.Tensor, + w1_data, + w2_data, + w1_wscale, + w2_wscale, + w1_bias, + w2_bias, + routing_data, + gather_idx, + scatter_idx, + gammas, + apply_router_weight_on_input: bool, + swiglu_limit: float, + unpadded_N_w1, + unpadded_K_w1, + unpadded_N_w2, + unpadded_K_w2, +) -> torch.Tensor: + """ + MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation. + """ + from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul + from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4 + from aiter.ops.triton.quant import dynamic_mxfp8_quant + + from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( + should_use_cdna4_mx_scale_swizzle, + ) + + swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None + quant_dtype = torch.float8_e4m3fn + + g1_gammas = gammas if apply_router_weight_on_input else None + g2_gammas = None if apply_router_weight_on_input else gammas + + hidden_q, a1_scale = dynamic_mxfp8_quant(hidden_states, quant_dtype=quant_dtype) + raw_gate_up = moe_gemm_a8w4( + hidden_q, + w1_data, + a1_scale, + w1_wscale, + None, + None, + w1_bias, + routing_data, + gather_indx=gather_idx, + gammas=g1_gammas, + swizzle_mx_scale=swz, + out_dtype=torch.bfloat16, + apply_swiglu=False, + unpadded_N=unpadded_N_w1, + unpadded_K=unpadded_K_w1, + ) + if unpadded_N_w1 is not None: + raw_gate_up = raw_gate_up[:, :unpadded_N_w1] + + interim_fp8, a2_scale = fused_clamp_act_mul( + raw_gate_up, + swiglu_limit=swiglu_limit, + activation="silu", + dtype_quant=quant_dtype, + scale_dtype_fmt="ue8m0", + quant_block_size=32, + ) + + out = moe_gemm_a8w4( + interim_fp8, + w2_data, + a2_scale, + w2_wscale, + None, + None, + w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=g2_gammas, + swizzle_mx_scale=swz, + unpadded_N=unpadded_N_w2, + unpadded_K=unpadded_K_w2, + ) + return out + + +def aiter_triton_kernel_w4a16_moe_forward( + hidden_states: torch.Tensor, + w1, + w2, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + activation: MoEActivation = MoEActivation.SWIGLUOAI, + quant_config: FusedMoEQuantConfig | None = None, + apply_router_weight_on_input: bool = False, + global_num_experts: int = -1, + expert_map: torch.Tensor | None = None, + unpadded_N_w1=None, + unpadded_K_w1=None, + unpadded_N_w2=None, + unpadded_K_w2=None, + num_expert_group: int | None = None, + topk_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + score_mode: str | None = None, +): + assert quant_config is not None and rocm_aiter_ops.is_enabled() + from vllm.platforms.rocm import on_gfx1250 + + try: + from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 + from aiter.ops.triton.moe.moe_routing import routing as _routing_mod + except ImportError: + from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 + from aiter.ops.triton.moe_routing import routing as _routing_mod + + if on_gfx1250(): + _routing_mod.is_tdm_avail = lambda: False + aiter_routing = _routing_mod.routing + + if score_mode is not None: + use_grouped_topk = num_expert_group is not None and num_expert_group > 1 + routing_data, gather_idx, scatter_idx = aiter_routing( + gating_output, + topk, + score_mode=score_mode, + bias=e_score_correction_bias, + renorm=renormalize, + routed_scaling_factor=( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ), + use_grouped_topk=use_grouped_topk, + num_expert_group=num_expert_group, + topk_group=topk_group, + ) + else: + routing_data, gather_idx, scatter_idx = aiter_routing( + gating_output, topk, sm_first=not renormalize + ) + + if on_gfx1250(): + gather_src = gather_idx.to(torch.long) // topk + hidden_states = hidden_states[gather_src] + gather_idx = None + + assert quant_config.w1_precision is not None + assert quant_config.w2_precision is not None + + w1_data = _aiter_raw(w1) + w2_data = _aiter_raw(w2) + w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale) + w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale) + + gammas = routing_data.gate_scal if routing_data else None + + swiglu_alpha = ( + quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0 + ) + swiglu_limit = ( + quant_config.gemm1_clamp_limit + if quant_config.gemm1_clamp_limit is not None + else 7.0 + ) + + # SILU on gfx1250: use the verified a8w4 kernel (dynamic MXFP8); a16w4 faults. + if activation == MoEActivation.SILU and on_gfx1250(): + return _aiter_w4a16_silu_via_a8w4( + hidden_states, + w1_data, + w2_data, + w1_wscale, + w2_wscale, + quant_config.w1_bias, + quant_config.w2_bias, + routing_data, + gather_idx, + scatter_idx, + gammas, + apply_router_weight_on_input, + swiglu_limit, + unpadded_N_w1, + unpadded_K_w1, + unpadded_N_w2, + unpadded_K_w2, + ) + + # SILU: silu(gate) * up — same kernel, just no "+1" residual in swiglu. + swiglu_add_residual = activation != MoEActivation.SILU + + intermediate = moe_gemm_a16w4( + hidden_states, + w1_data, + None, + w1_wscale, + None, + None, + quant_config.w1_bias, + routing_data, + gather_indx=gather_idx, + gammas=gammas if apply_router_weight_on_input else None, + swizzle_mx_scale=None, + apply_swiglu=True, + alpha=swiglu_alpha, + limit=swiglu_limit, + swiglu_add_residual=swiglu_add_residual, + unpadded_N=unpadded_N_w1, + unpadded_K=unpadded_K_w1, + ) + + out = moe_gemm_a16w4( + intermediate, + w2_data, + None, + w2_wscale, + None, + None, + quant_config.w2_bias, + routing_data, + scatter_indx=scatter_idx, + gammas=None if apply_router_weight_on_input else gammas, + swizzle_mx_scale=None, + unpadded_N=unpadded_N_w2, + unpadded_K=unpadded_K_w2, + ) + + return out + + +class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): + def __init__( + self, + moe_config: FusedMoEConfig, + quant_config: FusedMoEQuantConfig, + ): + super().__init__(moe_config, quant_config) + self.topk = moe_config.experts_per_token + self.renormalize = moe_config.routing_method in ( + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepseekV4, + ) + + @staticmethod + def activation_format() -> mk.FusedMoEActivationFormat: + return mk.FusedMoEActivationFormat.Standard + + @staticmethod + def _supports_current_device() -> bool: + if not rocm_aiter_ops.is_enabled(): + return False + from vllm.platforms.rocm import on_gfx950, on_gfx1250 + + return on_gfx950() or on_gfx1250() + + @staticmethod + def _supports_no_act_and_mul() -> bool: + return False + + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return (weight_key, activation_key) == (kMxfp4Static, None) + + @staticmethod + def _supports_activation(activation: MoEActivation) -> bool: + return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU) + + @staticmethod + def _supports_parallel_config( + moe_parallel_config: FusedMoEParallelConfig, + ) -> bool: + return ( + not moe_parallel_config.use_all2all_kernels + and not moe_parallel_config.enable_eplb + and moe_parallel_config.dp_size <= 1 + ) + + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + return routing_method in [ + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + RoutingMethodType.DeepseekV4, + ] + + @staticmethod + def _supports_router_logits_dtype( + router_logits_dtype: torch.dtype | None, + routing_method: RoutingMethodType, + ) -> bool: + return True + + @property + def expects_unquantized_inputs(self) -> bool: + return True + + def apply( + self, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + router_logits: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + apply_router_weight_on_input: bool, + num_expert_group: int | None = None, + e_score_correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + topk_group: int | None = None, + ) -> torch.Tensor: + assert self.moe_config.intermediate_size_per_partition_unpadded is not None + assert self.moe_config.hidden_dim_unpadded is not None + score_mode = ( + "sqrtsoftplus" + if self.moe_config.routing_method == RoutingMethodType.DeepseekV4 + else None + ) + return aiter_triton_kernel_w4a16_moe_forward( + hidden_states=hidden_states, + w1=w1, + w2=w2, + gating_output=router_logits, + topk=self.topk, + renormalize=self.renormalize, + activation=activation, + global_num_experts=global_num_experts, + expert_map=expert_map, + quant_config=self.quant_config, + apply_router_weight_on_input=apply_router_weight_on_input, + unpadded_N_w1=self.moe_config.intermediate_size_per_partition_unpadded * 2, + unpadded_K_w1=self.moe_config.hidden_dim_unpadded, + unpadded_N_w2=self.moe_config.hidden_dim_unpadded, + unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded, + num_expert_group=num_expert_group, + topk_group=topk_group, + e_score_correction_bias=e_score_correction_bias, + routed_scaling_factor=routed_scaling_factor, + score_mode=score_mode, + ) diff --git a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py index 609b28a959..5cc15866d5 100644 --- a/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py @@ -814,13 +814,13 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular): ) -> bool: p = current_platform if p.is_rocm(): - from vllm.platforms.rocm import on_gfx9 + from vllm.platforms.rocm import get_cdna_version - is_rocm_on_gfx9 = on_gfx9() + _rocm_support_fp8 = get_cdna_version() > 2 else: - is_rocm_on_gfx9 = False + _rocm_support_fp8 = False - device_supports_fp8 = is_rocm_on_gfx9 or ( + device_supports_fp8 = _rocm_support_fp8 or ( p.is_cuda() and p.has_device_capability((8, 9)) ) diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index bc770d0746..096e729480 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -472,11 +472,10 @@ class AiterExperts(mk.FusedMoEExpertsModular): ] if (weight_key, activation_key) not in SUPPORTED_W_A: return False - # CK MXFP4 MoE kernels are only supported on gfx950. if weight_key == kMxfp4Static: - from vllm.platforms.rocm import on_gfx950 + from vllm.platforms.rocm import on_gfx950, on_gfx1250 - if not on_gfx950(): + if not on_gfx950() or on_gfx1250(): return False return True diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 0aed6cbf40..64eb2afb1e 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -222,11 +222,14 @@ def backend_to_kernel_cls( return [BatchedMarlinExperts] elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: + from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import ( + AiterW4A16ExpertsMonolithic, + ) from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( AiterExperts, ) - return [AiterExperts] + return [AiterExperts, AiterW4A16ExpertsMonolithic] elif backend == Mxfp4MoeBackend.AITER_MXFP4_FP8: from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import ( @@ -1263,6 +1266,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( Supports DeepGEMM, TRTLLM MXFP8, Triton and Marlin backends. """ + from vllm.platforms.rocm import on_gfx1250 if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4: w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales( @@ -1435,7 +1439,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias, ) - elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not on_gfx1250(): # Initially introduced for DeepSeekV4 if w13_bias is not None: @@ -1492,7 +1496,9 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias, ) - elif mxfp4_backend in TRITON_BACKENDS: + elif mxfp4_backend in TRITON_BACKENDS or ( + mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250() + ): from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig if mxfp4_backend == Mxfp4MoeBackend.TRITON: diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 76faec0e01..0e1bd4086d 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -480,6 +480,10 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase): global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, ) @@ -735,7 +739,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): # For TRITON backends, weights are wrapped tensors from triton_kernels # that don't support .detach(). Manually assign parameters. - if self.mxfp4_backend not in TRITON_BACKENDS: + from vllm.platforms.rocm import on_gfx1250 + + uses_triton_weight_format = self.mxfp4_backend in TRITON_BACKENDS or ( + self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250() + ) + if not uses_triton_weight_format: replace_parameter(layer, "w13_weight", w13) replace_parameter(layer, "w2_weight", w2) replace_parameter(layer, "w13_weight_scale", w13_scale) @@ -747,7 +756,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self.w2_precision_config = w2_scale # AITER backend requires weights to be marked as shuffled. - if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16: + if ( + self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 + and not uses_triton_weight_format + ): layer.w13_weight.is_shuffled = True layer.w2_weight.is_shuffled = True @@ -841,7 +853,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): w2_bias = getattr(layer, "w2_bias", None) swiglu_limit = getattr(layer, "swiglu_limit", None) - if self.mxfp4_backend in TRITON_BACKENDS: + from vllm.platforms.rocm import on_gfx1250 + + if self.mxfp4_backend in TRITON_BACKENDS or ( + self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250() + ): # TRITON backends free w13/w2_weight_scale after swizzling; the # swizzled scales live inside the precision configs instead. assert self.w13_precision_config is not None diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py index 29283c7bbd..e92ebf2508 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py @@ -64,9 +64,7 @@ class QuarkW4A8_MXFP4_FP8(QuarkScheme): kernel_supported_gpu = False if current_platform.is_rocm(): - from vllm.platforms.rocm import on_gfx950 - - kernel_supported_gpu = on_gfx950() + kernel_supported_gpu = current_platform.supports_mx() self.use_aiter_kernel = ( is_aiter_found_and_supported() diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 2e4fbdf4c6..1fce368247 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -875,10 +875,40 @@ def w8a8_triton_block_scaled_mm( Bs: The per-block quantization scale for `B`. block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128]. - output_dtype: The dtype of the returned tensor. + output_dytpe: The dtype of the returned tensor. Returns: torch.Tensor: The result of matmul. """ + + from vllm.platforms.rocm import on_gfx1250 + + if on_gfx1250(): + # Torch upcast reference: dequantize A,B to fp32 and matmul in fp32. + # Avoids the gfx1250 native-fp8 block GEMM NaN bug. Correct but slow. + _bn, _bk = block_size[0], block_size[1] + _As = ( + _upcast_e8m0_to_fp32(As) + if As.dtype == torch.float8_e8m0fnu + else As.to(torch.float32) + ) + _Bs = ( + _upcast_e8m0_to_fp32(Bs) + if Bs.dtype == torch.float8_e8m0fnu + else Bs.to(torch.float32) + ) + _K = A.shape[-1] + _N = B.shape[0] + _Af = A.to(torch.float32).reshape(-1, _K) + _Asf = ( + _As.to(torch.float32) + .reshape(-1, _As.shape[-1]) + .repeat_interleave(_bk, dim=1)[:, :_K] + ) + _Bf = B.to(torch.float32) + _Bsf = _Bs.repeat_interleave(_bn, dim=0).repeat_interleave(_bk, dim=1)[:_N, :_K] + _out = (_Af * _Asf) @ (_Bf * _Bsf).t() + return _out.to(output_dtype).reshape(*A.shape[:-1], _N) + assert len(block_size) == 2 block_n, block_k = block_size[0], block_size[1] diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 2b19b3d3c4..f7e0b65105 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -122,7 +122,7 @@ def use_aiter_triton_gemm(n, m, k, dtype): def rocm_unquantized_gemm_impl( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None ) -> torch.Tensor: - from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950 + from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250 n = x.numel() // x.size(-1) m = weight.shape[0] @@ -164,7 +164,11 @@ def rocm_unquantized_gemm_impl( if use_skinny_reduce_counting: return ops.wvSplitKrc(x, weight, cu_count, bias) - if use_aiter_triton_gemm(n, m, k, x.dtype): + # gfx1250's aiter gemm_a16w16 uses the gluon backend, which requires + # K % 256 == 0 (it walks K with fixed-size descriptors and won't pad a + # partial last tile). Some whitelisted shapes have K=2880 (e.g. gpt-oss-120b + # hidden), so skip aiter there and fall back to the torch GEMM path below. + if use_aiter_triton_gemm(n, m, k, x.dtype) and not (on_gfx1250() and k % 256 != 0): from aiter.ops.triton.gemm_a16w16 import gemm_a16w16 return gemm_a16w16(x, weight, bias) @@ -172,6 +176,8 @@ def rocm_unquantized_gemm_impl( use_skinny = ( envs.VLLM_ROCM_USE_SKINNY_GEMM and (on_gfx9() or on_gfx1x()) + # build (gfx9/gfx11 ISA); fall back to torch GEMM there. + # TODO GFX1250: Include once skinny GEMM is supported on gfx1250 and x.dtype in [torch.float16, torch.bfloat16] and k % 8 == 0 ) diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index a5ca055f74..e39b279124 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -194,11 +194,6 @@ def _get_gcn_arch() -> str: return _query_gcn_arch_from_amdsmi() except Exception as e: logger.debug("Failed to get GCN arch via amdsmi: %s", e) - logger.warning_once( - "Failed to get GCN arch via amdsmi, falling back to torch.cuda. " - "This will initialize CUDA and may cause " - "issues if CUDA_VISIBLE_DEVICES is not set yet." - ) # Ultimate fallback: use torch.cuda (will initialize CUDA) return torch.cuda.get_device_properties("cuda").gcnArchName @@ -218,6 +213,11 @@ _ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"]) _ON_GFX90A = "gfx90a" in _GCN_ARCH _ON_GFX942 = "gfx942" in _GCN_ARCH _ON_GFX950 = "gfx950" in _GCN_ARCH +_ON_GFX1250 = "gfx1250" in _GCN_ARCH + +_ON_CDNA = any(arch in _GCN_ARCH for arch in ["gfx9", "gfx1250"]) +# RDNA = gfx11/gfx12 minus the CDNA-classified gfx1250. +_ON_RDNA = _ON_GFX1X and not _ON_CDNA def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: @@ -292,7 +292,7 @@ def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None: def on_gfx1x() -> bool: - return _ON_GFX1X + return _ON_GFX1X and not _ON_CDNA def on_gfx11() -> bool: @@ -308,7 +308,11 @@ def on_gfx1151() -> bool: def on_gfx12x() -> bool: - return _ON_GFX12X + return _ON_GFX12X and not _ON_CDNA + + +def on_gfx1250() -> bool: + return _ON_GFX1250 def on_mi3xx() -> bool: @@ -331,13 +335,33 @@ def on_gfx950() -> bool: return _ON_GFX950 +def on_cdna() -> bool: + return _ON_CDNA + + +def on_rdna() -> bool: + return _ON_RDNA + + +def get_cdna_version() -> int: + if on_gfx90a(): + return 2 + if on_gfx942(): + return 3 + if on_gfx950(): + return 4 + if on_gfx1250(): + return 5 + return 0 + + # Enable HIP online tuning early, before hipBLASLt initializes. # Turn on hipBLASLt online tuning if use AITER hipBLASLt GEMM. if ( envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_LINEAR and envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM - and on_mi3xx() + and get_cdna_version() > 2 ): os.environ["HIP_ONLINE_TUNING"] = "1" @@ -356,7 +380,7 @@ def use_rocm_custom_paged_attention( ) -> bool: # custom paged attn always supported on V0. On V1, requires sliding window # disabled due to observed numerical discrepancy. - if _ON_GFX9: + if on_cdna(): return ( (sliding_window == 0 or sliding_window == (-1, -1)) and (qtype == torch.half or qtype == torch.bfloat16) @@ -657,12 +681,12 @@ class RocmPlatform(Platform): from vllm._aiter_ops import rocm_aiter_ops - if rocm_aiter_ops.is_mha_enabled() and on_gfx9(): + if rocm_aiter_ops.is_mha_enabled() and on_cdna(): logger.info_once("Using AITER Flash Attention backend for ViT model.") return AttentionBackendEnum.ROCM_AITER_FA if ( - on_gfx9() + on_cdna() and find_spec("flash_attn") is not None and (dtype == torch.float16 or dtype == torch.bfloat16) ): @@ -878,11 +902,11 @@ class RocmPlatform(Platform): @classmethod def supports_mx(cls) -> bool: - return any(gfx in _GCN_ARCH for gfx in ["gfx95"]) + return any(gfx in _GCN_ARCH for gfx in ["gfx95", "gfx1250"]) @classmethod def supports_fp8(cls) -> bool: - return on_gfx9() or on_gfx12x() + return on_cdna() or on_gfx12x() @classmethod def is_fp8_fnuz(cls) -> bool: diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 812adb59e6..3df83148ce 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -505,7 +505,18 @@ def has_triton_kernels() -> bool: @cache def has_tilelang() -> bool: """Whether the optional `tilelang` package is available.""" - return _has_module("tilelang") + if not _has_module("tilelang"): + return False + # ROCm-only guard, imported lazily to avoid loading rocm on CUDA. + from vllm.platforms import current_platform + + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 + + # TODO: Re-enable when tilelang supports gfx1250 + if on_gfx1250(): + return False + return True def has_arctic_inference() -> bool: diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index 4c3327283a..cc823f704b 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -231,10 +231,26 @@ def _log_triton_jit_compile(fn_name: str, kwargs) -> None: def _setup_triton_jit_hook() -> None: - """Register a ``jit_post_compile_hook`` that warns on compilation.""" + """Register a jit_post_compile_hook that warns on compilation.""" if not HAS_TRITON: return - from triton import knobs # type: ignore[import-untyped] + from triton import knobs + from triton.runtime import jit as _triton_jit + + # kernels pass non-JSON-serializable constexprs + # make that serialization non-fatal + _orig = _triton_jit.serialize_specialization_data + if not getattr(_orig, "_vllm_guarded", False): + + @functools.wraps(_orig) + def _guarded(*args, **kwargs): + try: + return _orig(*args, **kwargs) + except (TypeError, ValueError): + return None # best-effort metadata; monitor ignores it + + cast(Any, _guarded)._vllm_guarded = True + _triton_jit.serialize_specialization_data = _guarded existing_hook = knobs.runtime.jit_post_compile_hook diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 9299af09a0..c01aa0181a 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/vllm/v1/attention/backends/fa_utils.py @@ -35,23 +35,31 @@ elif current_platform.is_xpu(): compile_flash_attn_varlen_func_from_specs = None # type: ignore[assignment] get_scheduler_metadata = xpu_ops.get_scheduler_metadata # type: ignore[assignment] elif current_platform.is_rocm(): + # On ROCm we use AITER's Triton flash-attention; the upstream flash-attn + # package is not installed/available. (Same source as aiter_triton_mla.py.) + # The FA4 compile-from-specs API is CUDA-only, so it is unavailable on ROCm + # regardless of whether AITER is present. + from vllm.platforms.rocm import on_gfx1250 + + compile_flash_attn_varlen_func_from_specs = None # type: ignore[assignment] try: - from flash_attn import flash_attn_varlen_func # type: ignore[no-redef] + if on_gfx1250(): + from aiter.ops.triton.mha import ( # type: ignore[no-redef] + flash_attn_varlen_func, + ) + else: + from flash_attn import flash_attn_varlen_func # type: ignore[no-redef] - compile_flash_attn_varlen_func_from_specs = None # type: ignore[assignment] - - # Mark that upstream flash-attn is available on ROCm _ROCM_FLASH_ATTN_AVAILABLE = True except ImportError: def flash_attn_varlen_func(*args: Any, **kwargs: Any) -> Any: # type: ignore[no-redef,misc] + package = "aiter" if on_gfx1250() else "flash-attn" raise ImportError( - "ROCm platform requires upstream flash-attn " - "to be installed. Please install flash-attn first." + f"ROCm platform requires upstream {package} " + f"to be installed. Please install {package} first." ) - compile_flash_attn_varlen_func_from_specs = None # type: ignore[assignment] - # ROCm doesn't use scheduler metadata (FA3 feature), provide stub def get_scheduler_metadata(*args: Any, **kwargs: Any) -> None: # type: ignore[misc] return None @@ -348,7 +356,8 @@ def is_flash_attn_varlen_func_available() -> bool: Platform-specific sources: - CUDA: vllm.vllm_flash_attn.flash_attn_varlen_func - XPU: xpu_ops.flash_attn_varlen_func - - ROCm: upstream flash_attn.flash_attn_varlen_func (if available) + - ROCm: aiter.ops.triton.mha.flash_attn_varlen_func (if AITER available) or + upstream flash_attn.flash_attn_varlen_func Note: This is separate from the AITER flash attention backend (rocm_aiter_fa.py) which uses rocm_aiter_ops.flash_attn_varlen_func. The condition to use AITER is diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 1f46dddabf..5a204c7ed6 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -776,12 +776,12 @@ class AiterFlashAttentionBackend(AttentionBackend): @classmethod def supports_compute_capability(cls, capability: DeviceCapability) -> bool: - from vllm.platforms.rocm import on_mi3xx + from vllm.platforms.rocm import get_cdna_version # DeviceCapability is currently created using torch.cuda.get_device_capability() - # which is known to be buggy on rocm systems. on_mi3xx uses amd-smi which is + # which is known to be buggy on rocm systems. on CDNA uses amd-smi which is # more reliable. - return on_mi3xx() + return get_cdna_version() > 2 @classmethod def supports_non_causal(cls) -> bool: diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index f608a0583d..c83ac7c4b9 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -18,6 +18,13 @@ if HAS_TRITON: logger = init_logger(__name__) +def _skip_aiter_sampler_on_gfx1250() -> bool: + # Lazy ROCm-only import; keeps arch detection out of import time on CUDA/CPU. + from vllm.platforms.rocm import on_gfx1250 + + return on_gfx1250() + + def flashinfer_sampler_supported() -> bool: """Decide whether FlashInfer's top-p/top-k sampler can be used. @@ -110,6 +117,7 @@ class TopKTopPSampler(nn.Module): elif ( logprobs_mode not in PROCESSED_LOGPROBS_MODES and rocm_aiter_ops.is_enabled() + and not _skip_aiter_sampler_on_gfx1250() # TODO (JPVILLAM): Enable ): self.aiter_ops = None self._aiter_ops_import_failed = False