Enable gfx1250 ROCm architecture (#46516)
Signed-off-by: Douglas Lehr <Doug.Lehr@amd.com> Signed-off-by: jpvillam <juan.villamizar@amd.com> Signed-off-by: Daniel <danichan@amd.com> Signed-off-by: Jaden Mathias <jaden.mathias@amd.com> Signed-off-by: jpvillam <jpvillam@amd.com> Signed-off-by: Juan Villamizar <100237675+jpvillam-amd@users.noreply.github.com> Signed-off-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Signed-off-by: Jaden Mathias <Jaden.Mathias@amd.com> Co-authored-by: Douglas Lehr <Doug.Lehr@amd.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Daniel <danichan@amd.com> Co-authored-by: Jaden Mathias <jaden.mathias@amd.com> Co-authored-by: jpvillam <jpvillam@amd.com> Co-authored-by: root <root@ctheliosr-1b114-f01-2.mnb.dcgpu> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: dllehr <douglehr@amd.com> Co-authored-by: Daniel Chan <Daniel.Chan@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Juan Villamizar committed
Jul 30, 2026 at 20:41 UTC
4f1da84eb55db17a20227e22db59dd0082ac18a0
33 files changed
+1850
-104
CMakeLists.txt
+40
-2
@@ -49,7 +49,7 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
49
set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14")
50
51
# Supported AMD GPU architectures.
52
-set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
52
+set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1250;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
53
54
# ROCm installation prefix. Default to /opt/rocm but allow override via
55
# -DROCM_PATH=/your/rocm/path when invoking cmake.
@@ -1422,6 +1422,38 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
1422
"csrc/rocm/skinny_gemms.cu"
1423
"csrc/rocm/skinny_gemms_int4.cu"
1424
"csrc/rocm/attention.cu")
1425
+ set(VLLM_ROCM_EXT_FLAGS ${VLLM_GPU_FLAGS})
1426
+
1427
+ # skinny_gemms*.cu are built on gfx9/gfx11 ISA (MFMA, dot2/dot4, legacy
1428
+ # s_waitcnt asm) that gfx1250 (gfx12) does not provide. Exclude them from the
1429
+ # gfx1250 build and disable their op registrations (VLLM_SKIP_SKINNY_GEMMS);
1430
+ # vLLM falls back to default/Triton GEMM for those ops on gfx1250.
1431
+ list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms.cu")
1432
+ list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms_int4.cu")
1433
+ set(VLLM_SKINNY_ARCHES ${VLLM_GPU_ARCHES})
1434
+ list(FILTER VLLM_SKINNY_ARCHES EXCLUDE REGEX "gfx1250")
1435
+ if(VLLM_SKINNY_ARCHES)
1436
+ message(STATUS "Building skinny_gemms for archs: ${VLLM_SKINNY_ARCHES}")
1437
+ set(VLLM_SKINNY_SRC
1438
+ "csrc/rocm/skinny_gemms.cu"
1439
+ "csrc/rocm/skinny_gemms_int4.cu")
1440
+ hipify_sources_target(VLLM_SKINNY_HIP_SRCS _rocm_C_skinny "${VLLM_SKINNY_SRC}")
1441
+ unset(_VLLM_LAST_HIPIFY_TARGET)
1442
+ add_library(_rocm_C_skinny OBJECT ${VLLM_SKINNY_HIP_SRCS})
1443
+ add_dependencies(_rocm_C_skinny hipify_all)
1444
+ set_source_files_properties(${VLLM_SKINNY_HIP_SRCS} PROPERTIES LANGUAGE ${VLLM_GPU_LANG})
1445
+ set_target_properties(_rocm_C_skinny PROPERTIES
1446
+ ${VLLM_GPU_LANG}_ARCHITECTURES "${VLLM_SKINNY_ARCHES}"
1447
+ POSITION_INDEPENDENT_CODE ON)
1448
+ target_include_directories(_rocm_C_skinny PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/csrc)
1449
+ target_compile_options(_rocm_C_skinny PRIVATE
1450
+ $<$<COMPILE_LANGUAGE:${VLLM_GPU_LANG}>:${VLLM_ROCM_EXT_FLAGS}>)
1451
+ target_compile_definitions(_rocm_C_skinny PRIVATE "-DTORCH_EXTENSION_NAME=_rocm_C")
1452
+ target_link_libraries(_rocm_C_skinny PRIVATE torch)
1453
+ else()
1454
+ message(STATUS "Only gfx1250, skipping skinny_gemms")
1455
+ list(APPEND VLLM_ROCM_EXT_FLAGS "-DVLLM_SKIP_SKINNY_GEMMS")
1456
+ endif()
1457
1458
set(VLLM_ROCM_HAS_GFX1100 OFF)
1459
if(VLLM_GPU_ARCHES MATCHES "gfx1100")
@@ -1437,10 +1469,16 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
1469
DESTINATION vllm
1470
LANGUAGE ${VLLM_GPU_LANG}
1471
SOURCES ${VLLM_ROCM_EXT_SRC}
1440
- COMPILE_FLAGS ${VLLM_GPU_FLAGS}
1472
+ COMPILE_FLAGS ${VLLM_ROCM_EXT_FLAGS}
1473
ARCHITECTURES ${VLLM_GPU_ARCHES}
1474
USE_SABI 3
1475
WITH_SOABI)
1476
+
1477
+ if(TARGET _rocm_C_skinny)
1478
+ target_link_libraries(_rocm_C PRIVATE _rocm_C_skinny)
1479
+ else()
1480
+ target_compile_definitions(_rocm_C PRIVATE VLLM_SKIP_SKINNY_GEMMS)
1481
+ endif()
1482
1483
if(VLLM_ROCM_HAS_GFX1100)
1484
target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100)
csrc/quickreduce/base.h
+11
@@ -79,6 +79,7 @@ union BufferResource {
79
};
80
};
81
82
+#if !defined(__gfx1250__)
83
__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
84
int32x4_t srsrc, int32_t voffset, int32_t soffset,
85
int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32");
@@ -86,6 +87,16 @@ __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
87
__quickreduce_device_inline__ static void buffer_store_dwordx4(
88
int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset,
89
int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32");
90
+#else
91
+__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
92
+ int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) {}
93
+
94
+__quickreduce_device_inline__ static void buffer_store_dwordx4(int32x4_t data,
95
+ int32x4_t srsrc,
96
+ int32_t voffset,
97
+ int32_t soffset,
98
+ int32_t aux) {}
99
+#endif
100
101
__quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) {
102
#if defined(__gfx942__)
csrc/rocm/attention.cu
+11
@@ -2405,6 +2405,16 @@ template <typename T, int absz, int cbid, int blgp>
2405
__device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA,
2406
const bit16x8& inpB,
2407
const floatx8& inpC) {
2408
+ #if defined(__gfx1250__)
2409
+ // gfx1250 (gfx12 family) does not provide the gfx12 WMMA variant used by
2410
+ // gfx1200/1201 (needs wmma-128b-insts). This custom-attention WMMA path is
2411
+ // unsupported on gfx1250; trap if ever launched (fail loud, not
2412
+ // silent-wrong).
2413
+ (void)inpA;
2414
+ (void)inpB;
2415
+ __builtin_trap();
2416
+ return inpC;
2417
+ #else
2418
if constexpr (std::is_same<T, _Float16>::value) {
2419
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(inpA, inpB, inpC);
2420
} else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
@@ -2412,6 +2422,7 @@ __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA,
2422
} else {
2423
static_assert(false, "unsupported 16b dtype");
2424
}
2425
+ #endif
2426
}
2427
2428
template <typename T>
csrc/rocm/torch_bindings.cpp
+5
@@ -14,6 +14,10 @@
14
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
15
// vLLM custom ops for rocm
16
17
+// skinny_gemms.cu (LLMM1/wvSplitK/wvSplitKrc/wvSplitKQ) is excluded on gfx1250
18
+// (gfx9/gfx11 ISA, unsupported there); skip these registrations to avoid
19
+// undefined symbols. vLLM uses default/Triton GEMM for these ops on gfx1250.
20
+#ifndef VLLM_SKIP_SKINNY_GEMMS
21
// Custom gemm op for matrix-vector multiplication
22
rocm_ops.def(
23
"LLMM1(Tensor in_a, Tensor in_b, int rows_per_block) -> "
@@ -46,6 +50,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
50
"Tensor scale_a, "
51
" Tensor scale_b, int CuCount) -> ()");
52
rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ);
53
+#endif // VLLM_SKIP_SKINNY_GEMMS
54
55
#ifdef VLLM_ROCM_GFX1100
56
// W4A16 GPTQ kernels for AMD RDNA3 (gfx1100).
docker/Dockerfile.rocm_base_gfx1250
new
+325
@@ -0,0 +1,325 @@
1
+ARG BASE_IMAGE=ubuntu:24.04
2
+ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/
3
+ARG ROCM_SDK_VERSION=7.14.0a20260623
4
+ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623
5
+ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623
6
+ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623
7
+ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623
8
+ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623
9
+
10
+
11
+ARG FA_BRANCH="jpvillam/gfx1250_wip"
12
+ARG FA_REPO="https://github.com/jpvillam-amd/flash-attention.git"
13
+ARG AITER_BRANCH="main"
14
+ARG AITER_REPO="https://github.com/ROCm/aiter.git"
15
+ARG MORI_BRANCH="v1.1.0"
16
+ARG MORI_REPO="https://github.com/ROCm/mori.git"
17
+
18
+# Sccache configuration (only used in release pipeline)
19
+ARG USE_SCCACHE
20
+ARG SCCACHE_DOWNLOAD_URL
21
+ARG SCCACHE_ENDPOINT
22
+ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
23
+ARG SCCACHE_REGION_NAME=us-west-2
24
+ARG SCCACHE_S3_NO_CREDENTIALS=0
25
+
26
+FROM ${BASE_IMAGE} AS base
27
+
28
+ARG PYTORCH_ROCM_ARCH=gfx1250
29
+ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
30
+ENV AITER_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
31
+ENV MORI_GPU_ARCHS=gfx942;gfx950
32
+ENV FA_GPU_ARCHS=gfx942;gfx950;gfx1250
33
+
34
+# TODO: Unset these when support is available for gfx1250
35
+ENV ENABLE_CK=0
36
+ENV SKIP_CK_BUILD="TRUE"
37
+ENV FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"
38
+ARG PREBUILD_KERNELS=0
39
+ARG TRITON_KERNEL_BRANCH="padroute"
40
+ARG TRITON_KERNEL_REPO="https://github.com/jpvillam-amd/triton.git"
41
+
42
+# Required for RCCL in ROCm7.1
43
+ENV HSA_NO_SCRATCH_RECLAIM=1
44
+
45
+ARG PYTHON_VERSION=3.12
46
+ENV PYTHON_VERSION=${PYTHON_VERSION}
47
+
48
+RUN mkdir -p /app
49
+WORKDIR /app
50
+ENV DEBIAN_FRONTEND=noninteractive
51
+
52
+# Install Python and other dependencies
53
+RUN apt-get update -y \
54
+ && 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++ \
55
+ && for i in 1 2 3; do \
56
+ add-apt-repository -y ppa:deadsnakes/ppa && break || \
57
+ { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
58
+ done \
59
+ && apt-get update -y \
60
+ && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \
61
+ python${PYTHON_VERSION}-lib2to3 python-is-python3 \
62
+ && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
63
+ && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
64
+ && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
65
+ && python3 --version
66
+
67
+ENV VIRTUAL_ENV=/opt/venv
68
+RUN python${PYTHON_VERSION} -m venv "${VIRTUAL_ENV}" && \
69
+ "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML
70
+ENV PATH=${VIRTUAL_ENV}/bin:$PATH
71
+
72
+RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython
73
+RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/*
74
+
75
+# Install sccache if USE_SCCACHE is enabled (for release builds)
76
+ARG USE_SCCACHE
77
+ARG SCCACHE_DOWNLOAD_URL
78
+ARG SCCACHE_ENDPOINT
79
+ARG SCCACHE_BUCKET_NAME
80
+ARG SCCACHE_REGION_NAME
81
+ARG SCCACHE_S3_NO_CREDENTIALS
82
+RUN if [ "$USE_SCCACHE" = "1" ]; then \
83
+ echo "Installing sccache..." \
84
+ && SCCACHE_ARCH="x86_64" \
85
+ && SCCACHE_VERSION="v0.8.1" \
86
+ && 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}" \
87
+ && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
88
+ && tar -xzf /tmp/sccache.tar.gz -C /tmp \
89
+ && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
90
+ && chmod +x /usr/bin/sccache \
91
+ && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
92
+ && sccache --version; \
93
+ fi
94
+
95
+##
96
+## Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index
97
+##
98
+ARG ROCM_WHEEL_INDEX
99
+ARG TORCH_VERSION
100
+ARG TORCHVISION_VERSION
101
+ARG TORCHAUDIO_VERSION
102
+ARG ROCM_SDK_VERSION
103
+ARG APEX_VERSION
104
+ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages
105
+ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel
106
+ENV ROCM_HOME=${ROCM_PATH}
107
+ENV ROCM_SOURCE_DIR=${ROCM_PATH}
108
+ENV ROCM_BIN=${ROCM_PATH}/bin
109
+ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake
110
+ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode
111
+ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH
112
+ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib
113
+ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake
114
+ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi
115
+
116
+# torch/torchvision/torchaudio must be pinned to mutually-consistent builds
117
+# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm
118
+# sdk version is derived from torch's own dependency pin unless overridden,
119
+# which keeps the set consistent and avoids pip backtracking.
120
+RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \
121
+ --extra-index-url https://pypi.org/simple \
122
+ "torch[device-all]==${TORCH_VERSION}" \
123
+ "torchvision==${TORCHVISION_VERSION}" \
124
+ "torchaudio==${TORCHAUDIO_VERSION}" \
125
+ "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \
126
+ rocm-sdk init
127
+
128
+# Torch runtime deps that may not be published on the ROCm wheel index;
129
+# install them from PyPI afterwards.
130
+RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \
131
+ "networkx>=2.5.1" jinja2 "fsspec>=0.8.5"
132
+
133
+
134
+
135
+# Expose the rocm-sdk wheel as a conventional /opt/rocm install so downstream
136
+# builds (Dockerfile.rocm: vLLM csrc, RIXL/UCX, ROCShmem/DeepEP) keep working.
137
+RUN ln -sfn "${ROCM_PATH}" /opt/rocm;
138
+
139
+RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \
140
+ sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \
141
+ "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \
142
+ fi
143
+
144
+# The ROCm SDK wheel ships a broken CMake export for hsakmt
145
+# This patch includes the right paths for the numa build target
146
+RUN <<'EOF'
147
+set -eu
148
+TARGETS="/opt/rocm/lib/cmake/hsakmt/hsakmtTargets.cmake"
149
+[ -f "$TARGETS" ] || exit 0 # nothing to patch
150
+grep -q NUMA_LIBRARY "$TARGETS" && exit 0 # already patched
151
+
152
+# 1. Point libdrm's -L at the copy bundled in the wheel, not the builder path.
153
+sed -i 's|-L/__w/[^;"]*|-L${_IMPORT_PREFIX}/lib/rocm_sysdeps/lib|g' "$TARGETS"
154
+
155
+# 2. Drop the nonexistent RHEL libc path (libc is linked implicitly anyway).
156
+sed -i 's|/usr/lib64/libc.so;||g' "$TARGETS"
157
+
158
+# 3. Define the numa::numa target the export references but forgot to create.
159
+cat >> "$TARGETS" <<'CMAKE'
160
+
161
+if(NOT TARGET numa::numa)
162
+ find_library(NUMA_LIBRARY NAMES numa REQUIRED)
163
+ add_library(numa::numa UNKNOWN IMPORTED)
164
+ set_target_properties(numa::numa PROPERTIES IMPORTED_LOCATION "${NUMA_LIBRARY}")
165
+endif()
166
+CMAKE
167
+EOF
168
+
169
+# Clone custom triton_kernels for install later
170
+RUN mkdir -p /app/patched_triton_kernels; \
171
+ cd /app/patched_triton_kernels \
172
+ && git clone ${TRITON_KERNEL_REPO} \
173
+ && cd triton \
174
+ && git checkout ${TRITON_KERNEL_BRANCH} \
175
+ && git submodule update --init --recursive
176
+ENV TRITON_KERNELS_SRC_DIR="/app/patched_triton_kernels/triton/python/triton_kernels/triton_kernels/"
177
+
178
+
179
+# Setup sccache for HIP compilation via HIP_CLANG_PATH
180
+# This creates wrapper scripts in a separate directory and points HIP to use them
181
+# This avoids modifying the original ROCm binaries which can break detection
182
+# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm)
183
+# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1
184
+RUN if [ "$USE_SCCACHE" = "1" ]; then \
185
+ echo "Setting up sccache wrappers for HIP compilation..." \
186
+ && mkdir -p /opt/sccache-wrappers \
187
+ && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \
188
+ && chmod +x /opt/sccache-wrappers/clang++ \
189
+ && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \
190
+ && chmod +x /opt/sccache-wrappers/clang \
191
+ && echo "sccache wrappers created in /opt/sccache-wrappers"; \
192
+ fi
193
+
194
+# Set sccache environment variables only when USE_SCCACHE=1
195
+# This prevents S3 config from leaking into images when sccache is not used
196
+ARG USE_SCCACHE
197
+ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
198
+ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
199
+ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
200
+ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
201
+
202
+
203
+###
204
+### AMD SMI Build
205
+###
206
+FROM base AS build_amdsmi
207
+RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \
208
+ && pip wheel . --wheel-dir=dist
209
+RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install
210
+
211
+
212
+###
213
+### MORI Build TODO(Build needs fixing)
214
+###
215
+FROM base AS build_mori
216
+ARG MORI_BRANCH
217
+ARG MORI_REPO
218
+ARG MORI_GPU_ARCHS
219
+RUN mkdir -p /app/install; \
220
+ git clone ${MORI_REPO} \
221
+ && cd mori \
222
+ && git checkout ${MORI_BRANCH} \
223
+ && git submodule update --init --recursive \
224
+ && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \
225
+ && cp /app/mori/dist/*.whl /app/install;
226
+
227
+
228
+###
229
+### FlashAttention Build
230
+###
231
+# Remove && git submodule update --init \ for CK
232
+FROM base AS build_fa
233
+ARG FA_BRANCH
234
+ARG FA_REPO
235
+ARG USE_SCCACHE
236
+RUN mkdir -p /app/install; \
237
+ git clone ${FA_REPO} \
238
+ && cd flash-attention \
239
+ && git checkout ${FA_BRANCH} \
240
+ && if [ "$USE_SCCACHE" = "1" ]; then \
241
+ export HIP_CLANG_PATH=/opt/sccache-wrappers \
242
+ && sccache --show-stats; \
243
+ fi \
244
+ && GPU_ARCHS=$(echo ${FA_GPU_ARCHS} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \
245
+ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
246
+ && cp dist/*.whl /app/install;
247
+
248
+
249
+###
250
+### AITER Build
251
+###
252
+FROM base AS build_aiter
253
+ARG AITER_BRANCH
254
+ARG AITER_REPO
255
+ARG USE_SCCACHE
256
+RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO}
257
+RUN cd aiter \
258
+ && git submodule update --init --recursive \
259
+ && pip install -r requirements.txt
260
+RUN pip install pyyaml && cd aiter \
261
+ && if [ "$USE_SCCACHE" = "1" ]; then \
262
+ export HIP_CLANG_PATH=/opt/sccache-wrappers \
263
+ && sccache --show-stats; \
264
+ fi \
265
+ && AITER_USE_SYSTEM_TRITON=1 PREBUILD_KERNELS=${PREBUILD_KERNELS} GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
266
+ && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
267
+ && ls /app/aiter/dist/*.whl
268
+RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
269
+
270
+
271
+###
272
+### Final Build
273
+###
274
+
275
+# Wheel release stage -
276
+# only includes dependencies used by wheel release pipeline
277
+FROM base AS debs_wheel_release
278
+RUN mkdir /app/debs
279
+RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
280
+ if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi
281
+RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
282
+ cp /install/*.whl /app/debs
283
+RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
284
+ cp /install/*.whl /app/debs
285
+
286
+# Full debs stage - includes Mori (used by Docker releases)
287
+FROM base AS debs
288
+RUN mkdir /app/debs
289
+RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
290
+ if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi
291
+RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
292
+ cp /install/*.whl /app/debs
293
+RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
294
+ cp /install/*.whl /app/debs
295
+RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
296
+ if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi
297
+
298
+FROM base AS final
299
+RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
300
+ pip install /install/*.whl
301
+
302
+ARG BASE_IMAGE
303
+ARG ROCM_WHEEL_INDEX
304
+ARG ROCM_SDK_VERSION
305
+ARG TORCH_VERSION
306
+ARG TORCHVISION_VERSION
307
+ARG TORCHAUDIO_VERSION
308
+ARG FA_BRANCH
309
+ARG FA_REPO
310
+ARG AITER_BRANCH
311
+ARG AITER_REPO
312
+ARG MORI_BRANCH
313
+ARG MORI_REPO
314
+RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
315
+ && echo "ROCM_WHEEL_INDEX: ${ROCM_WHEEL_INDEX}" >> /app/versions.txt \
316
+ && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \
317
+ && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \
318
+ && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \
319
+ && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \
320
+ && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
321
+ && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \
322
+ && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
323
+ && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \
324
+ && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \
325
+ && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt
docker/Dockerfile.rocm_gfx1250
new
+729
@@ -0,0 +1,729 @@
1
+# default base image
2
+ARG REMOTE_VLLM="0"
3
+ARG COMMON_WORKDIR=/app
4
+ARG BASE_IMAGE=rocm/vllm-dev:base
5
+ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base
6
+# NIC backend for MoRI RDMA support.
7
+# By default (all), drivers and userspace libraries for all supported NIC types
8
+# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
9
+# To install drivers for a single NIC type only, set NIC_BACKEND explicitly:
10
+# --build-arg NIC_BACKEND=ainic # AMD AINIC (Pensando) only
11
+# --build-arg NIC_BACKEND=bnxt # Broadcom Thor-2 only
12
+# --build-arg NIC_BACKEND=none # Install nothing.
13
+ARG NIC_BACKEND=all
14
+# AMD AINIC apt repo settings
15
+# Users can specify a custom version compatible with their host drivers.
16
+# The default version has been tested with ioinic-dkms=25.11.1.001
17
+ARG AINIC_VERSION=1.117.3-hydra
18
+ARG UBUNTU_CODENAME=jammy
19
+
20
+# Sccache configuration. Release builds use this today; CI can opt in when a
21
+# shared S3-compatible cache backend is available.
22
+ARG USE_SCCACHE
23
+ARG SCCACHE_DOWNLOAD_URL
24
+ARG SCCACHE_ENDPOINT
25
+ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
26
+ARG SCCACHE_REGION_NAME=us-west-2
27
+ARG SCCACHE_S3_NO_CREDENTIALS=0
28
+
29
+FROM ${BASE_IMAGE} AS base
30
+
31
+ARG ARG_PYTORCH_ROCM_ARCH=gfx1250
32
+ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
33
+
34
+# Install build dependencies and utilities
35
+RUN apt-get update -q -y && apt-get install -q -y \
36
+ sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \
37
+ apt-transport-https ca-certificates wget curl \
38
+ build-essential libnuma-dev ccache mold
39
+RUN --mount=type=cache,target=/root/.cache/pip \
40
+ python3 -m pip install --upgrade pip
41
+# Note: mold is installed but not set as the system default linker because
42
+# some packages use JIT compilation at runtime with flags mold does not support.
43
+# Build stages opt in via LDFLAGS="-fuse-ld=mold".
44
+# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base)
45
+ARG USE_SCCACHE
46
+RUN if [ "$USE_SCCACHE" != "1" ]; then \
47
+ apt-get purge -y sccache || true; \
48
+ python3 -m pip uninstall -y sccache || true; \
49
+ rm -f "$(which sccache)" || true; \
50
+ fi
51
+
52
+# Install UV — download first, then run, so a curl failure is not masked by the pipe
53
+RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \
54
+ && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \
55
+ && rm -f /tmp/uv-install.sh \
56
+ && uv --version
57
+
58
+# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
59
+# Reference: https://github.com/astral-sh/uv/pull/1694
60
+ENV UV_HTTP_TIMEOUT=500
61
+ENV UV_INDEX_STRATEGY="unsafe-best-match"
62
+# Use copy mode to avoid hardlink failures with Docker cache mounts
63
+ENV UV_LINK_MODE=copy
64
+# python binary fall back for non venv builds
65
+ENV UV_PYTHON=${VIRTUAL_ENV:-/usr}/bin/python3
66
+# Expose paths from wheel installation
67
+ENV PKG_CONFIG_PATH=${ROCM_PATH}/lib/rocm_sysdeps/lib/pkgconfig:${PKG_CONFIG_PATH}
68
+# ccache directory - persisted across layer rebuilds via cache mounts.
69
+ENV CCACHE_DIR=/root/.cache/ccache
70
+ENV CCACHE_COMPILERCHECK=content
71
+# Empty by default so build steps fall back to $(nproc); CI can override.
72
+ARG max_jobs
73
+ENV MAX_JOBS=${max_jobs}
74
+
75
+# Install sccache if USE_SCCACHE is enabled (for release builds)
76
+ARG USE_SCCACHE
77
+ARG SCCACHE_DOWNLOAD_URL
78
+ARG SCCACHE_ENDPOINT
79
+ARG SCCACHE_BUCKET_NAME
80
+ARG SCCACHE_REGION_NAME
81
+ARG SCCACHE_S3_NO_CREDENTIALS
82
+RUN if [ "$USE_SCCACHE" = "1" ]; then \
83
+ if command -v sccache >/dev/null 2>&1; then \
84
+ echo "sccache already installed, skipping installation"; \
85
+ sccache --version; \
86
+ else \
87
+ echo "Installing sccache..." \
88
+ && SCCACHE_ARCH="x86_64" \
89
+ && SCCACHE_VERSION="v0.8.1" \
90
+ && 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}" \
91
+ && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
92
+ && tar -xzf /tmp/sccache.tar.gz -C /tmp \
93
+ && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
94
+ && chmod +x /usr/bin/sccache \
95
+ && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
96
+ && sccache --version; \
97
+ fi; \
98
+ fi
99
+
100
+# Set sccache environment variables only when USE_SCCACHE=1
101
+# This prevents S3 config from leaking into images when sccache is not used
102
+ARG USE_SCCACHE
103
+ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
104
+ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
105
+ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
106
+ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
107
+
108
+ARG COMMON_WORKDIR
109
+WORKDIR ${COMMON_WORKDIR}
110
+
111
+
112
+# -----------------------
113
+# vLLM fetch stages
114
+FROM base AS fetch_vllm_0
115
+ONBUILD COPY ./ vllm/
116
+FROM base AS fetch_vllm_1
117
+ARG VLLM_REPO="https://github.com/ROCm/vllm.git"
118
+ARG VLLM_BRANCH="455_wip"
119
+ENV VLLM_REPO=${VLLM_REPO}
120
+ENV VLLM_BRANCH=${VLLM_BRANCH}
121
+ONBUILD RUN git clone ${VLLM_REPO} \
122
+ && cd vllm \
123
+ && git fetch -v --prune -- origin ${VLLM_BRANCH} \
124
+ && git checkout FETCH_HEAD \
125
+ && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \
126
+ git remote add upstream "https://github.com/vllm-project/vllm.git" \
127
+ && git fetch upstream ; fi
128
+FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
129
+
130
+# -----------------------
131
+# Rust build stage
132
+# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages
133
+# don't need the rust toolchain or protoc.
134
+FROM fetch_vllm AS rust-build
135
+ARG COMMON_WORKDIR
136
+
137
+# protoc is used by tonic-build/prost-build.
138
+RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
139
+ ca-certificates curl unzip \
140
+ && rm -rf /var/lib/apt/lists/*
141
+
142
+COPY tools/install_protoc.sh /tmp/install_protoc.sh
143
+RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
144
+
145
+# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
146
+# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
147
+ENV CARGO_BUILD_JOBS=4
148
+ENV CARGO_NET_RETRY=10
149
+ENV RUSTUP_MAX_RETRIES=10
150
+
151
+RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
152
+ cd ${COMMON_WORKDIR}/vllm \
153
+ && uv pip install --system -r requirements/build/rust.txt
154
+
155
+# Build the release binary. Cargo's registry/git caches can be written by
156
+# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
157
+# keeping the cache benefit. Do not cache target/, because stale target metadata
158
+# can outlive source updates across BuildKit cache reuse.
159
+RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
160
+ --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
161
+ cd ${COMMON_WORKDIR}/vllm \
162
+ && bash build_rust.sh \
163
+ && test -x vllm/vllm-rs
164
+
165
+# -----------------------
166
+# vLLM native build stages
167
+#
168
+# csrc-build intentionally copies only files that affect ROCm native extension
169
+# compilation. That keeps unrelated CI/test/docs edits from invalidating the
170
+# expensive HIP/C++ build layer.
171
+FROM base AS csrc-build
172
+ARG COMMON_WORKDIR
173
+WORKDIR ${COMMON_WORKDIR}/vllm
174
+
175
+COPY requirements/rocm.txt requirements/rocm.txt
176
+COPY requirements/common.txt requirements/common.txt
177
+RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
178
+ uv pip install --system -r requirements/rocm.txt
179
+
180
+# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do
181
+# not invalidate the expensive native build layer.
182
+COPY setup.py CMakeLists.txt ./
183
+COPY tools/build_rust.py tools/build_rust.py
184
+COPY cmake cmake/
185
+COPY csrc csrc/
186
+COPY vllm/envs.py vllm/envs.py
187
+COPY vllm/__init__.py vllm/__init__.py
188
+
189
+ENV VLLM_TARGET_DEVICE=rocm
190
+ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build"
191
+
192
+RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \
193
+ --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
194
+ export CCACHE_BASEDIR="$PWD" \
195
+ && echo "=== ccache stats before ROCm native build ===" \
196
+ && (ccache --show-stats || true) \
197
+ && (ccache --zero-stats || true) \
198
+ && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \
199
+ && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \
200
+ && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \
201
+ && test -d dist \
202
+ && ls dist/*.whl >/dev/null \
203
+ && echo "=== ccache stats after ROCm native build ===" \
204
+ && (ccache --show-stats || true)
205
+
206
+# Build the full vLLM ROCm wheel by reusing the native extension wheel from
207
+# csrc-build. This stage still rebuilds for Python/package changes, but skips
208
+# the expensive HIP/C++ compile when native inputs are unchanged.
209
+FROM fetch_vllm AS build_vllm
210
+ARG COMMON_WORKDIR
211
+ENV VLLM_TARGET_DEVICE=rocm
212
+
213
+COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels
214
+
215
+# Drop the pre-built Rust artifacts into the source tree. setup.py detects
216
+# them and ships them as-is, skipping the local Rust build.
217
+COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
218
+COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
219
+
220
+RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
221
+ cd vllm \
222
+ && uv pip install --system -r requirements/rocm.txt \
223
+ && export VLLM_USE_PRECOMPILED=1 \
224
+ && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \
225
+ && export VLLM_DOCKER_BUILD_CONTEXT=1 \
226
+ && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \
227
+ && python3 setup.py bdist_wheel --dist-dir=dist \
228
+ && test -d dist \
229
+ && ls dist/*.whl >/dev/null
230
+FROM scratch AS export_vllm
231
+ARG COMMON_WORKDIR
232
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl /
233
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements
234
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
235
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
236
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
237
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
238
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
239
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
240
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
241
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
242
+
243
+# RIXL/UCX build stages
244
+FROM base AS build_rixl
245
+ARG RIXL_BRANCH="39be1de8"
246
+ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
247
+ARG UCX_BRANCH="bfb51733"
248
+ARG UCX_REPO="https://github.com/openucx/ucx.git"
249
+# ENV ROCM_PATH=/opt/rocm -> correct ROCM_PATH is set in base image
250
+ENV UCX_HOME=/usr/local/ucx
251
+ENV RIXL_HOME=/usr/local/rixl
252
+ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
253
+
254
+# RIXL build system dependences and RDMA support
255
+RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
256
+ libgrpc-dev \
257
+ libgrpc++-dev \
258
+ libprotobuf-dev \
259
+ protobuf-compiler-grpc \
260
+ libcpprest-dev \
261
+ libaio-dev \
262
+ librdmacm1 \
263
+ librdmacm-dev \
264
+ libibverbs1 \
265
+ libibverbs-dev \
266
+ ibverbs-utils \
267
+ rdmacm-utils \
268
+ ibverbs-providers \
269
+ && rm -rf /var/lib/apt/lists/*
270
+
271
+RUN --mount=type=cache,target=/root/.cache/uv \
272
+ uv pip install --system meson auditwheel patchelf tomlkit
273
+
274
+RUN --mount=type=cache,target=/root/.cache/ccache \
275
+ cd /usr/local/src && \
276
+ git clone ${UCX_REPO} && \
277
+ cd ucx && \
278
+ git checkout ${UCX_BRANCH} && \
279
+ ./autogen.sh && \
280
+ mkdir build && cd build && \
281
+ CC="ccache gcc" CXX="ccache g++" \
282
+ ../configure \
283
+ --prefix=/usr/local/ucx \
284
+ --enable-shared \
285
+ --disable-static \
286
+ --disable-doxygen-doc \
287
+ --enable-optimizations \
288
+ --enable-devel-headers \
289
+ --with-rocm=${ROCM_PATH} \
290
+ --with-verbs \
291
+ --with-dm \
292
+ --enable-mt && \
293
+ make -j$(nproc) && \
294
+ make install
295
+
296
+ENV PATH=/usr/local/ucx/bin:$PATH
297
+ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
298
+
299
+RUN --mount=type=cache,target=/root/.cache/ccache \
300
+ git clone ${RIXL_REPO} /opt/rixl && \
301
+ cd /opt/rixl && \
302
+ git checkout ${RIXL_BRANCH} && \
303
+ CC="ccache gcc" CXX="ccache g++" \
304
+ meson setup build --prefix=${RIXL_HOME} \
305
+ --force-fallback-for=abseil-cpp \
306
+ -Ducx_path=${UCX_HOME} \
307
+ -Drocm_path=${ROCM_PATH} && \
308
+ cd build && \
309
+ ninja -j$(nproc) && \
310
+ ninja install
311
+
312
+# Generate RIXL wheel
313
+# Exclude libcore and libpull from auditwheel: transitive dependencies
314
+# that are not shipped in the wheel and vary across base images.
315
+RUN cd /opt/rixl && \
316
+ sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
317
+ contrib/build-wheel.sh && \
318
+ # The wheel build re-runs meson via meson-python; force the bundled abseil
319
+ sed -i 's|setup = \["-Dinstall_headers=false"\]|setup = ["-Dinstall_headers=false", "--force-fallback-for=abseil-cpp"]|' \
320
+ pyproject.toml && \
321
+ grep -q 'force-fallback-for' pyproject.toml && \
322
+ mkdir -p /app/install && \
323
+ _ucx_install_dir=${UCX_HOME} \
324
+ ./contrib/build-wheel.sh \
325
+ --output-dir /app/install \
326
+ --rocm-dir ${ROCM_PATH} \
327
+ --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
328
+ --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
329
+
330
+# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
331
+# invalidate the slow ROCShmem build.
332
+FROM base AS build_rocshmem
333
+ARG ROCSHMEM_BRANCH="f0acb0c6"
334
+ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
335
+# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so
336
+# it can be linked against DeepEP without arch mismatches.
337
+ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
338
+# ENV ROCM_PATH=/opt/rocm -> Correct rocm_path is set in base image
339
+ENV ROCSHMEM_DIR=/opt/rocshmem
340
+
341
+RUN --mount=type=cache,target=/root/.cache/ccache \
342
+ git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \
343
+ && cd rocm-systems \
344
+ && git sparse-checkout set --cone projects/rocshmem \
345
+ && git checkout ${ROCSHMEM_BRANCH} \
346
+ && mkdir -p projects/rocshmem/build \
347
+ && cd projects/rocshmem/build \
348
+ && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \
349
+ bash ../scripts/build_configs/all_backends \
350
+ -DROCM_PATH=${ROCM_PATH} \
351
+ -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
352
+ -DUSE_EXTERNAL_MPI=OFF
353
+
354
+# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel.
355
+FROM build_rocshmem AS build_deepep
356
+ARG DEEPEP_BRANCH="a9ea9774"
357
+ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
358
+ARG DEEPEP_NIC="cx7"
359
+
360
+# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR.
361
+# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list.
362
+RUN --mount=type=cache,target=/root/.cache/ccache \
363
+ export PYTORCH_ROCM_ARCH="gfx942;gfx950" \
364
+ && git clone ${DEEPEP_REPO} \
365
+ && cd DeepEP \
366
+ && git checkout ${DEEPEP_BRANCH} \
367
+ && 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
368
+
369
+# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
370
+# not force users to rebuild the long-lived Dockerfile.rocm_base image.
371
+FROM base AS mori_base
372
+ARG NIC_BACKEND
373
+ARG AINIC_VERSION
374
+ARG UBUNTU_CODENAME
375
+RUN /bin/bash -lc 'set -euo pipefail; \
376
+ \
377
+ install_ainic() { \
378
+ apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \
379
+ rm -rf /var/lib/apt/lists/*; \
380
+ mkdir -p /etc/apt/keyrings; \
381
+ curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
382
+ echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
383
+ > /etc/apt/sources.list.d/amdainic.list; \
384
+ apt-get update && apt-get install -y --no-install-recommends \
385
+ libionic-dev \
386
+ ionic-common \
387
+ ; \
388
+ rm -rf /var/lib/apt/lists/*; \
389
+ }; \
390
+ \
391
+ # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \
392
+ # bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \
393
+ install_bnxt() { \
394
+ install -m 0755 -d /etc/apt/keyrings; \
395
+ curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \
396
+ -o /etc/apt/keyrings/broadcom-nic.asc; \
397
+ chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \
398
+ echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \
399
+ > /etc/apt/sources.list.d/broadcom-nic.list; \
400
+ apt-get update && apt-get install -y --no-install-recommends \
401
+ bnxt-rocelib=235.2.86.0 \
402
+ ; \
403
+ cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \
404
+ ldconfig; \
405
+ rm -rf /var/lib/apt/lists/*; \
406
+ }; \
407
+ \
408
+ echo "[MORI] Install MoRI proxy deps"; \
409
+ pip install --quiet --ignore-installed blinker && \
410
+ pip install --quiet quart msgpack aiohttp pyzmq; \
411
+ echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
412
+ \
413
+ # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \
414
+ # Only vendor packages are installed here for dlopen; no compile-time flags needed. \
415
+ case "${NIC_BACKEND}" in \
416
+ none) ;; \
417
+ all) install_ainic; install_bnxt ;; \
418
+ ainic) install_ainic ;; \
419
+ bnxt) install_bnxt ;; \
420
+ *) echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \
421
+ esac'
422
+
423
+# -----------------------
424
+# vLLM wheel release build stage (for building distributable wheels)
425
+# This stage pins dependencies to custom ROCm wheel versions and handles version detection
426
+FROM fetch_vllm AS build_vllm_wheel_release
427
+
428
+ARG COMMON_WORKDIR
429
+
430
+# Drop the pre-built Rust artifacts into the source tree. setup.py detects
431
+# them and ships them as-is, skipping the local Rust build.
432
+COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
433
+COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
434
+
435
+# Create /install directory for custom wheels
436
+RUN mkdir -p /install
437
+
438
+# Copy custom ROCm wheels from docker/context if they exist
439
+# COPY ensures Docker cache is invalidated when wheels change
440
+# .keep file ensures directory always exists for COPY to work
441
+COPY docker/context/base-wheels/ /tmp/base-wheels/
442
+# This is how we know if we are building for a wheel release or not.
443
+# If there are not wheels found there, we are not building for a wheel release.
444
+# So we exit with an error. To skip this stage.
445
+RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \
446
+ echo "Found custom wheels - copying to /install"; \
447
+ cp /tmp/base-wheels/*.whl /install/ && \
448
+ echo "Copied custom wheels:"; \
449
+ ls -lh /install/; \
450
+ else \
451
+ echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \
452
+ echo "Wheel releases require pre-built ROCm wheels."; \
453
+ exit 1; \
454
+ fi
455
+
456
+# GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds)
457
+# This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm
458
+ARG GIT_REPO_CHECK=0
459
+RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
460
+ echo "Running repository checks..."; \
461
+ cd vllm && bash tools/check_repo.sh; \
462
+ fi
463
+
464
+# Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt)
465
+# This ensures setuptools_scm sees clean repo state for version detection
466
+RUN --mount=type=bind,source=.git,target=vllm/.git \
467
+ --mount=type=cache,target=/root/.cache/uv \
468
+ cd vllm \
469
+ && uv pip install --system setuptools_scm regex \
470
+ && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \
471
+ && echo "Detected vLLM version: ${VLLM_VERSION}" \
472
+ && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
473
+
474
+# Fail if git-based package dependencies are found in requirements files
475
+# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
476
+# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
477
+RUN echo "Checking for git-based packages in requirements files..." \
478
+ && echo "Checking common.txt for git-based packages:" \
479
+ && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
480
+ echo "ERROR: Git-based packages found in common.txt:"; \
481
+ grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
482
+ echo "Please publish these packages to PyPI instead of using git dependencies."; \
483
+ exit 1; \
484
+ else \
485
+ echo " ✓ No git-based packages found in common.txt"; \
486
+ fi \
487
+ && echo "Checking rocm.txt for git-based packages:" \
488
+ && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
489
+ echo "ERROR: Git-based packages found in rocm.txt:"; \
490
+ grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
491
+ echo "Please publish these packages to PyPI instead of using git dependencies."; \
492
+ exit 1; \
493
+ else \
494
+ echo " ✓ No git-based packages found in rocm.txt"; \
495
+ fi \
496
+ && echo "All requirements files are clean - no git-based packages found"
497
+
498
+# Pin vLLM dependencies to exact versions of custom ROCm wheels
499
+# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
500
+COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
501
+RUN echo "Pinning vLLM dependencies to custom wheel versions..." \
502
+ && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt
503
+
504
+# Install dependencies using custom wheels from /install
505
+RUN --mount=type=cache,target=/root/.cache/uv \
506
+ cd vllm \
507
+ && echo "Building vLLM with custom wheels from /install" \
508
+ && uv pip install --system --find-links /install -r requirements/rocm.txt
509
+
510
+# Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt
511
+# (setup.py auto-detects ccache/sccache in PATH)
512
+RUN --mount=type=bind,source=.git,target=vllm/.git \
513
+ --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
514
+ cd vllm \
515
+ && export CCACHE_BASEDIR="$PWD" \
516
+ && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \
517
+ && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \
518
+ && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist
519
+
520
+FROM scratch AS export_vllm_wheel_release
521
+ARG COMMON_WORKDIR
522
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/dist/*.whl /
523
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requirements
524
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
525
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests
526
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples
527
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
528
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
529
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
530
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
531
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
532
+
533
+# -----------------------
534
+# CI base image (Tier 1) - stable, rarely changing CI dependencies.
535
+# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers
536
+# in the vLLM artifacts for the current commit.
537
+FROM mori_base AS ci_base
538
+ARG COMMON_WORKDIR
539
+
540
+# Update rdma-core to support latest rocshmem.
541
+ARG DEEPEP_NIC
542
+RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
543
+ git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
544
+ cd /tmp/rdma-core && \
545
+ mkdir -p build && cd build && \
546
+ cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
547
+ ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
548
+fi
549
+
550
+# Install RIXL + DeepEP wheels.
551
+RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
552
+ --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
553
+ uv pip install --system /rixl_install/*.whl /deep_install/*.whl
554
+
555
+# Copy ROCShmem runtime libraries.
556
+COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
557
+
558
+# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec.
559
+RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
560
+ librdmacm1 \
561
+ libibverbs1 \
562
+ ibverbs-providers \
563
+ ibverbs-utils \
564
+ pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
565
+ libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \
566
+ && rm -rf /var/lib/apt/lists/*
567
+
568
+# Install torchcodec from source for ROCm/torch ABI compatibility.
569
+COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh
570
+RUN --mount=type=cache,target=/root/.cache/uv \
571
+ --mount=type=cache,target=/root/.cache/pip \
572
+ --mount=type=cache,target=/root/.cache/torchcodec-wheels \
573
+ bash /tmp/install_torchcodec.sh \
574
+ && rm /tmp/install_torchcodec.sh \
575
+ && apt-get clean && rm -rf /var/lib/apt/lists/*
576
+
577
+# Pre-install shared ROCm runtime dependencies.
578
+COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/
579
+RUN --mount=type=cache,target=/root/.cache/uv \
580
+ uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \
581
+ && rm -rf /tmp/ci-base-requirements
582
+
583
+# Enable fast and less brittle model downloads in tests.
584
+ENV HF_XET_HIGH_PERFORMANCE=1
585
+ENV HF_HUB_DOWNLOAD_TIMEOUT=60
586
+
587
+# Pre-install vLLM test dependencies.
588
+COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt
589
+RUN --mount=type=cache,target=/root/.cache/uv \
590
+ uv pip install --system -r /tmp/rocm-test-reqs.txt
591
+
592
+# Rebuild fastsafetensors from source so its C++ extension is compiled with
593
+# USE_ROCM and can detect libamdhip64.so at runtime.
594
+RUN --mount=type=cache,target=/root/.cache/pip \
595
+ FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \
596
+ && test -n "${FASTSAFETENSORS_REQ}" \
597
+ && python3 -m pip install --force-reinstall --no-deps \
598
+ --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \
599
+ && rm /tmp/rocm-test-reqs.txt
600
+
601
+# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel.
602
+# See: https://github.com/pytorch/pytorch/issues/169857
603
+ENV MIOPEN_DEBUG_CONV_DIRECT=0
604
+ENV MIOPEN_DEBUG_CONV_GEMM=0
605
+
606
+# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
607
+# See: https://github.com/ROCm/rocm-libraries/issues/6266
608
+ENV HSA_ENABLE_IPC_MODE_LEGACY=1
609
+
610
+# ROCm profiler limits workaround.
611
+RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
612
+ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
613
+
614
+# Install vllm_test_utils in ci_base for ci_base + wheel parity.
615
+COPY tests/vllm_test_utils /tmp/vllm_test_utils
616
+RUN --mount=type=cache,target=/root/.cache/uv \
617
+ uv pip install --system /tmp/vllm_test_utils \
618
+ && rm -rf /tmp/vllm_test_utils
619
+
620
+# -----------------------
621
+# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base.
622
+FROM ${CI_BASE_IMAGE} AS test
623
+ARG COMMON_WORKDIR
624
+
625
+# Install the vLLM wheel (--no-deps: all deps already in ci_base).
626
+RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
627
+ --mount=type=cache,target=/root/.cache/uv \
628
+ cd /install \
629
+ && uv pip install --system --no-deps *.whl
630
+
631
+# Store the vLLM wheel in the image for python-only install tests.
632
+COPY --from=export_vllm /*.whl /opt/vllm-wheels/
633
+
634
+WORKDIR /vllm-workspace
635
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
636
+
637
+# Copy in the v1 package (for python-only install test group).
638
+COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
639
+
640
+# Hide source under src/ so it won't shadow the installed package in tests.
641
+RUN mkdir src && mv vllm src/vllm
642
+
643
+# -----------------------
644
+# Final vLLM image
645
+FROM mori_base AS final
646
+
647
+RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
648
+
649
+# Clean up sccache from release image (not needed at runtime)
650
+# This removes the binary and wrappers that may have been installed during build
651
+RUN rm -f /usr/bin/sccache || true \
652
+ && rm -rf /opt/sccache-wrappers || true
653
+
654
+# Unset sccache environment variables for the release image
655
+# This prevents S3 bucket config from leaking into production images
656
+ENV SCCACHE_BUCKET=
657
+ENV SCCACHE_REGION=
658
+ENV SCCACHE_ENDPOINT=
659
+ENV SCCACHE_S3_NO_CREDENTIALS=
660
+ENV SCCACHE_IDLE_TIMEOUT=
661
+
662
+# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt.
663
+# Manually remove it so that later steps of numpy upgrade can continue
664
+RUN case "$(which python3)" in \
665
+ *"/opt/conda/envs/py_3.9"*) \
666
+ rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \
667
+ *) ;; esac
668
+
669
+RUN --mount=type=cache,target=/root/.cache/uv \
670
+ uv pip install --system --upgrade huggingface-hub[cli]
671
+
672
+# Install vLLM using uv (inherited from base stage)
673
+# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version
674
+RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
675
+ --mount=type=cache,target=/root/.cache/uv \
676
+ cd /install \
677
+ && uv pip install --system -r requirements/rocm.txt \
678
+ && pip uninstall -y vllm \
679
+ && uv pip install --system *.whl
680
+
681
+# Install RIXL wheel
682
+RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
683
+ uv pip install --system /rixl_install/*.whl
684
+
685
+ARG COMMON_WORKDIR
686
+ARG BASE_IMAGE
687
+ARG NIC_BACKEND
688
+ARG AINIC_VERSION
689
+
690
+# Copy over the benchmark scripts as well
691
+COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
692
+COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples
693
+COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
694
+
695
+# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
696
+# See: https://github.com/ROCm/rocm-libraries/issues/6266
697
+ENV HSA_ENABLE_IPC_MODE_LEGACY=1
698
+
699
+ENV TOKENIZERS_PARALLELISM=false
700
+
701
+# ENV that can improve safe tensor loading, and end-to-end time
702
+ENV SAFETENSORS_FAST_GPU=1
703
+
704
+# Performance environment variable.
705
+ENV HIP_FORCE_DEV_KERNARG=1
706
+
707
+# Workaround for ROCm profiler limits
708
+RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
709
+ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
710
+RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
711
+ && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
712
+ && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
713
+
714
+# Download bench scripts for gfx1250
715
+RUN curl -k -O https://raw.githubusercontent.com/ROCm/vllm/refs/heads/gfx1250_bench/vllm_smoketest.sh
716
+
717
+
718
+### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better
719
+RUN pip3 uninstall -y triton && \
720
+ git clone https://github.com/triton-lang/triton.git && \
721
+ cd triton && \
722
+ git checkout c517f38c && \
723
+ TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install .
724
+
725
+CMD ["/bin/bash"]
726
+
727
+#Set entrypoint for vllm-openai official images
728
+FROM final AS vllm-openai
729
+ENTRYPOINT ["vllm", "serve"]
tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py
+3
-3
@@ -31,7 +31,7 @@ def _get_attention_backend_params() -> list[str | None]:
31
falls back to ROCM_AITER_UNIFIED_ATTN or TRITON_ATTN for
32
cross-attention since ROCM_ATTN doesn't support ENCODER_DECODER)
33
- TRITON_ATTN: always available on ROCm
34
- - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950
34
+ - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950/gfx1250
35
36
On non-ROCm platforms, we just run with the default backend.
37
"""
@@ -40,9 +40,9 @@ def _get_attention_backend_params() -> list[str | None]:
40
41
if current_platform.is_rocm():
42
backends: list[str | None] = [None, "TRITON_ATTN"]
43
- from vllm.platforms.rocm import _ON_MI3XX
43
+ from vllm.platforms.rocm import get_cdna_version
44
45
- if _ON_MI3XX:
45
+ if get_cdna_version() > 2:
46
backends.append("ROCM_AITER_UNIFIED_ATTN")
47
return backends
48
except Exception:
tests/entrypoints/speech_to_text/translation/test_translation_validation.py
+3
-3
@@ -37,13 +37,13 @@ def _get_rocm_attention_config(model_name):
37
38
if "whisper" in model_name.lower():
39
try:
40
- from vllm.platforms.rocm import _ON_MI3XX
40
+ from vllm.platforms.rocm import get_cdna_version
41
42
- if _ON_MI3XX:
42
+ if get_cdna_version() > 2:
43
return {"backend": "ROCM_AITER_UNIFIED_ATTN"}
44
except ImportError:
45
logger.warning(
46
- "Could not import _ON_MI3XX from rocm platform, "
46
+ "Could not check cdna version from rocm platform, "
47
"falling back to TRITON_ATTN for Whisper."
48
)
49
return {"backend": "TRITON_ATTN"}
tests/kernels/attention/test_rocm_aiter_unified_attn.py
+4
-4
@@ -15,15 +15,15 @@ from tests.kernels.attention.test_triton_unified_attention import ref_paged_attn
15
from vllm.platforms import current_platform
16
from vllm.utils.torch_utils import set_random_seed
17
18
-_SKIP_NON_MI3XX = True
18
+_SKIP_NON_CDNA_2_PLUS = True
19
if current_platform.is_rocm():
20
- from vllm.platforms.rocm import on_mi3xx
20
+ from vllm.platforms.rocm import get_cdna_version
21
22
- _SKIP_NON_MI3XX = not on_mi3xx()
22
+ _SKIP_NON_CDNA_2_PLUS = get_cdna_version() < 2
23
24
pytestmark = [
25
pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"),
26
- pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"),
26
+ pytest.mark.skipif(_SKIP_NON_CDNA_2_PLUS, reason="CDNA 2+ ROCm only"),
27
]
28
29
NUM_Q_HEADS = 8
tests/kernels/moe/test_modular_oai_triton_moe.py
+123
-2
@@ -4,6 +4,8 @@
4
Test modular OAI Triton MoE
5
"""
6
7
+from __future__ import annotations
8
+
9
import pytest
10
import torch
11
@@ -48,6 +50,35 @@ MNK = [
50
]
51
52
53
+def deepseek_v4_flash_moe_topology():
54
+ """MoE sizes for the DeepSeek-V4-Flash`.
55
+ Default weights from: ``deepseek-ai/DeepSeek-V4-Flash`.
56
+ ``moe_intermediate_size``, ``n_routed_experts``, and ``num_experts_per_tok``.
57
+ """
58
+ defaults = {
59
+ "hidden_size": 4096,
60
+ "moe_intermediate_size": 2048,
61
+ "n_routed_experts": 256,
62
+ "num_experts_per_tok": 6,
63
+ }
64
+
65
+ return defaults
66
+
67
+
68
+def scaled_deepseek_v4_flash_problem(
69
+ *,
70
+ dim_scale: int = 8,
71
+ expert_scale: int = 8,
72
+):
73
+ """Smaller K/N/E for kernel tests; keeps production top_k and K:N ratio (~2:1)."""
74
+ t = deepseek_v4_flash_moe_topology()
75
+ k = max(128, t["hidden_size"] // dim_scale)
76
+ n = max(64, t["moe_intermediate_size"] // dim_scale)
77
+ num_experts = max(t["num_experts_per_tok"], t["n_routed_experts"] // expert_scale)
78
+ topk = t["num_experts_per_tok"]
79
+ return k, n, num_experts, topk
80
+
81
+
82
def unshuffle_weight(w: torch.Tensor):
83
first = w[..., ::2]
84
second = w[..., 1::2]
@@ -160,8 +191,8 @@ def oai_triton_moe_impl(
191
x: torch.Tensor,
192
w1: torch.Tensor,
193
w2: torch.Tensor,
163
- w1_scale: "PrecisionConfig",
164
- w2_scale: "PrecisionConfig",
194
+ w1_scale: PrecisionConfig,
195
+ w2_scale: PrecisionConfig,
196
w1_bias: torch.Tensor | None,
197
w2_bias: torch.Tensor | None,
198
num_experts: int,
@@ -261,3 +292,93 @@ def test_oai_triton_moe(
292
)
293
294
assert_close(ref=out_ref, tri=out, maxtol=0.025, rmstol=0.005)
295
+
296
+
297
+@pytest.mark.skipif(
298
+ not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
299
+)
300
+def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_init):
301
+ """Exercise ``UnfusedOAITritonExperts.apply`` with explicit workspaces.
302
+
303
+ Same MoE topology as ``launch_dsv4.sh`` / DeepSeek-V4-Flash ``config.json``,
304
+ with linear dimensions and expert count scaled down for test GPU memory.
305
+ """
306
+ wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1)
307
+ set_random_seed(0)
308
+
309
+ k, n, num_experts, topk = scaled_deepseek_v4_flash_problem()
310
+ m = 7
311
+ dtype = torch.bfloat16
312
+
313
+ (
314
+ w1,
315
+ w2,
316
+ w1_bias,
317
+ w2_bias,
318
+ w1_tri,
319
+ w2_tri,
320
+ w1_bias_tri,
321
+ w2_bias_tri,
322
+ w1_precision_config,
323
+ w2_precision_config,
324
+ ) = make_weights(dtype, k, n, num_experts)
325
+
326
+ x = torch.randn((m, k), dtype=dtype, device="cuda")
327
+ router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype)
328
+ topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True)
329
+ topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
330
+
331
+ quant_config = mxfp4_w4a16_moe_quant_config(
332
+ w1_bias=w1_bias_tri,
333
+ w2_bias=w2_bias_tri,
334
+ w1_scale=w1_precision_config,
335
+ w2_scale=w2_precision_config,
336
+ )
337
+ moe_config = make_dummy_moe_config(
338
+ num_experts=num_experts,
339
+ experts_per_token=topk,
340
+ hidden_dim=k,
341
+ intermediate_size=n,
342
+ )
343
+ experts = UnfusedOAITritonExperts(moe_config, quant_config)
344
+
345
+ if not UnfusedOAITritonExperts._supports_current_device():
346
+ pytest.skip("UnfusedOAITritonExperts does not support this device")
347
+
348
+ _, _, N, K, top_k = experts.moe_problem_size(x, w1_tri, w2_tri, topk_ids)
349
+ assert top_k == topk
350
+ ws13_shape, ws2_shape, out_shape = experts.workspace_shapes(
351
+ m,
352
+ N,
353
+ K,
354
+ topk,
355
+ num_experts,
356
+ num_experts,
357
+ None,
358
+ MoEActivation.SWIGLUOAI,
359
+ )
360
+ workspace13 = torch.empty(ws13_shape, dtype=dtype, device="cuda")
361
+ workspace2 = torch.empty(ws2_shape, dtype=dtype, device="cuda")
362
+ output = torch.empty(out_shape, dtype=dtype, device="cuda")
363
+
364
+ with set_current_vllm_config(VllmConfig()):
365
+ out_ref = torch_moe_impl(x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids)
366
+ experts.apply(
367
+ output=output,
368
+ hidden_states=x,
369
+ w1=w1_tri,
370
+ w2=w2_tri,
371
+ topk_weights=topk_weights,
372
+ topk_ids=topk_ids,
373
+ activation=MoEActivation.SWIGLUOAI,
374
+ global_num_experts=num_experts,
375
+ expert_map=None,
376
+ a1q_scale=None,
377
+ a2_scale=None,
378
+ workspace13=workspace13,
379
+ workspace2=workspace2,
380
+ expert_tokens_meta=None,
381
+ apply_router_weight_on_input=False,
382
+ )
383
+
384
+ assert_close(ref=out_ref, tri=output, maxtol=0.025, rmstol=0.005)
tests/kernels/moe/test_routing.py
+3
-3
@@ -29,13 +29,13 @@ from vllm.platforms import current_platform
29
30
31
def _is_aiter_capable() -> bool:
32
- """Check if the platform supports AITER (gfx942/gfx950)."""
32
+ """Check if the platform supports AITER (gfx942/gfx950/gfx1250)."""
33
if not current_platform.is_rocm():
34
return False
35
try:
36
- from vllm.platforms.rocm import _ON_MI3XX
36
+ from vllm.platforms.rocm import get_cdna_version
37
38
- return _ON_MI3XX
38
+ return get_cdna_version() > 2
39
except ImportError:
40
return False
41
tests/model_executor/layers/test_rocm_unquantized_gemm.py
+3
@@ -26,6 +26,7 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
26
monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
27
monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
28
monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
29
+ monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
30
monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
31
32
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):
53
monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
54
monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
55
monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
56
+ monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
57
monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
58
59
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):
78
monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False)
79
monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
80
monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True)
81
+ monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: True)
82
monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
83
84
wvsplitkrc_mock = MagicMock(side_effect=lambda x_view, w, _, __: x_view @ w.t())
tests/models/multimodal/generation/test_granite_speech.py
+3
-2
@@ -45,9 +45,10 @@ models: dict[str, str | None] = {
45
def granite_speech_attention_config():
46
"""Return attention config for Granite Speech tests on ROCm."""
47
if current_platform.is_rocm():
48
- from vllm.platforms.rocm import on_mi3xx
48
+ from vllm.platforms.rocm import get_cdna_version
49
50
- if on_mi3xx():
50
+ # -1 (unknown arch) is truthy; gate on CDNA3+ like other call sites.
51
+ if get_cdna_version() > 2:
52
return {"backend": "ROCM_AITER_FA"}
53
return {"backend": "TRITON_ATTN"}
54
return None
tests/models/quantization/test_bitsandbytes.py
+3
-3
@@ -21,11 +21,11 @@ from ...utils import compare_two_settings, multi_gpu_test
21
from ..utils import check_embeddings_close, check_logprobs_close
22
23
if current_platform.is_rocm():
24
- from vllm.platforms.rocm import on_gfx9
24
+ from vllm.platforms.rocm import on_cdna
25
26
pytestmark = pytest.mark.skipif(
27
- on_gfx9(),
28
- reason="bitsandbytes not supported on gfx9 (warp size 64 limitation)",
27
+ on_cdna(),
28
+ reason="bitsandbytes not supported on CDNA (warp size 64 limitation)",
29
)
30
31
models_4bit_to_test = [
tests/v1/attention/test_rocm_attention_backends_selection.py
+7
-15
@@ -28,16 +28,9 @@ def mock_vllm_config():
28
29
30
@pytest.fixture
31
-def mock_on_gfx9():
32
- """Mock gfx9 arch detection to return True."""
33
- with patch("vllm.platforms.rocm.on_gfx9", return_value=True):
34
- yield
35
-
36
-
37
-@pytest.fixture
38
-def mock_on_mi3xx():
39
- """Mock mi3xx arch detection to return True."""
40
- with patch("vllm.platforms.rocm.on_mi3xx", return_value=True):
31
+def mock_get_cdna_version():
32
+ """Mock cdna version arch detection to return True."""
33
+ with patch("vllm.platforms.rocm.get_cdna_version", return_value=3):
34
yield
35
36
@@ -111,8 +104,7 @@ def test_standard_attention_backend_selection(
104
selected_backend,
105
expected_backend_path,
106
mock_vllm_config,
114
- mock_on_gfx9,
115
- mock_on_mi3xx,
107
+ mock_get_cdna_version,
108
monkeypatch,
109
):
110
"""Test standard attention backend selection with various configurations."""
@@ -305,12 +297,12 @@ def test_mla_backend_selection(
297
298
299
def test_aiter_fa_requires_mi3xx(mock_vllm_config):
308
- """Test that ROCM_AITER_FA requires mi3xx architecture."""
300
+ """Test that ROCM_AITER_FA requires CDNA3+ architecture."""
301
from vllm.platforms.rocm import RocmPlatform
302
311
- # Mock on_mi3xx to return False (used by supports_compute_capability)
303
+ # Mock cdna version to return 1 (used by supports_compute_capability)
304
with (
313
- patch("vllm.platforms.rocm.on_mi3xx", return_value=False),
305
+ patch("vllm.platforms.rocm.get_cdna_version", return_value=1),
306
pytest.raises(
307
ValueError,
308
match="compute capability not supported",
vllm/_aiter_ops.py
+12
-7
@@ -52,7 +52,7 @@ IS_AITER_FOUND = is_aiter_found()
52
def is_aiter_found_and_supported() -> bool:
53
"""Check if AITER library is available and platform supports it.
54
55
- Checks: platform (ROCm), device arch (gfx9), and library existence.
55
+ Checks: platform (ROCm), device arch is CDNA 3 or better, and library existence.
56
Does NOT check environment variables - that's handled by rocm_aiter_ops.is_enabled().
57
58
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:
66
VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery.
67
"""
68
if current_platform.is_rocm() and IS_AITER_FOUND:
69
- from vllm.platforms.rocm import on_mi3xx
69
+ from vllm.platforms.rocm import get_cdna_version
70
71
- return on_mi3xx()
71
+ return get_cdna_version() > 2
72
return False
73
74
@@ -1798,16 +1798,21 @@ class rocm_aiter_ops:
1798
@classmethod
1799
@if_aiter_supported
1800
def is_fp4bmm_enabled(cls) -> bool:
1801
- from vllm.platforms.rocm import on_gfx950
1801
+ from vllm.platforms.rocm import get_cdna_version
1802
1803
- return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and on_gfx950()
1803
+ # TODO GFX1250: Enable for cdna 4+ when aiter supports batched_gemm_a16wfp4 on gfx1250
1804
+ return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and get_cdna_version() == 4
1805
1806
@classmethod
1807
@if_aiter_supported
1808
def is_linear_hipbmm_enabled(cls) -> bool:
1808
- from vllm.platforms.rocm import on_mi3xx
1809
+ from vllm.platforms.rocm import get_cdna_version
1810
1810
- return cls.is_linear_enabled() and on_mi3xx() and cls._LINEAR_HIPBMM_ENABLED
1811
+ return (
1812
+ cls.is_linear_enabled()
1813
+ and (get_cdna_version() > 2)
1814
+ and cls._LINEAR_HIPBMM_ENABLED
1815
+ )
1816
1817
@classmethod
1818
@if_aiter_supported
vllm/model_executor/kernels/linear/scaled_mm/aiter.py
+8
-3
@@ -374,9 +374,14 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
374
super().__init__(config)
375
n, k = config.weight_shape
376
377
- self.use_triton = (
378
- not current_platform.is_fp8_fnuz()
379
- and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k)
377
+ _on_gfx1250 = False
378
+ if current_platform.is_rocm():
379
+ from vllm.platforms.rocm import on_gfx1250
380
+
381
+ _on_gfx1250 = on_gfx1250()
382
+
383
+ self.use_triton = not current_platform.is_fp8_fnuz() and (
384
+ rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) or _on_gfx1250
385
)
386
387
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
vllm/model_executor/kernels/linear/scaled_mm/pytorch.py
+3
-3
@@ -104,10 +104,10 @@ class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
104
if not current_platform.is_rocm():
105
return False, "requires ROCm."
106
107
- from vllm.platforms.rocm import on_mi3xx
107
+ from vllm.platforms.rocm import get_cdna_version
108
109
- if not on_mi3xx():
110
- return False, "requires MI3xx."
109
+ if get_cdna_version() <= 2:
110
+ return False, "requires CDNA3+"
111
112
if compute_capability is not None and compute_capability < 94:
113
return False, "requires compute capability 94 and above."
vllm/model_executor/kernels/linear/scaled_mm/rocm.py
+8
-3
@@ -79,10 +79,15 @@ class ROCmFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
79
if not current_platform.is_rocm():
80
return False, "requires ROCm."
81
82
- from vllm.platforms.rocm import on_gfx12x, on_mi3xx
82
+ from vllm.platforms.rocm import get_cdna_version, on_gfx12x, on_gfx1250
83
84
- if not (on_mi3xx() or on_gfx12x()):
85
- return False, "requires MI3xx or gfx12x"
84
+ # wvSplitKQ (skinny GEMM) is excluded from the gfx1250 build.
85
+ if on_gfx1250():
86
+ return False, "wvSplitKQ (skinny GEMM) is not built on gfx1250"
87
+
88
+ # Restore RDNA4 (gfx12x) dropped by the get_cdna_version()>2 refactor.
89
+ if get_cdna_version() <= 2 and not on_gfx12x():
90
+ return False, "requires CDNA3+ (gfx942/gfx950) or RDNA4 (gfx12x)"
91
92
if not envs.VLLM_ROCM_USE_SKINNY_GEMM:
93
return False, "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled."
vllm/model_executor/layers/fused_moe/experts/aiter_mxfp4_w4a8_moe.py
+376
-4
@@ -20,7 +20,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
21
__all__ = [
22
"AiterW4A8ExpertsMonolithic",
23
+ "AiterW4A16ExpertsMonolithic",
24
"aiter_triton_kernel_w4a8_moe_forward",
25
+ "aiter_triton_kernel_w4a16_moe_forward",
26
]
27
28
@@ -46,11 +48,27 @@ def aiter_triton_kernel_w4a8_moe_forward(
48
and quant_config.use_mxfp4_w4a8
49
and rocm_aiter_ops.is_enabled()
50
)
49
- from aiter.ops.triton.moe_routing.routing import routing as aiter_routing
51
+ from vllm.platforms.rocm import on_gfx1250
52
+
53
+ try:
54
+ from aiter.ops.triton.moe.moe_routing import routing as _routing_mod
55
+ except ImportError:
56
+ from aiter.ops.triton.moe_routing import routing as _routing_mod
57
+
58
+ if on_gfx1250():
59
+ _routing_mod.is_tdm_avail = lambda: False
60
+ aiter_routing = _routing_mod.routing
61
62
routing_data, gather_idx, scatter_idx = aiter_routing(
63
gating_output, topk, sm_first=not renormalize
64
)
65
+
66
+ # gfx1250: aiter's in-kernel gather is numerically broken
67
+ if on_gfx1250():
68
+ gather_src = gather_idx.to(torch.long) // topk
69
+ hidden_states = hidden_states[gather_src]
70
+ gather_idx = None
71
+
72
return triton_kernel_fused_mxfp4_w4a8_experts(
73
None,
74
hidden_states,
@@ -199,12 +217,11 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
217
218
@staticmethod
219
def _supports_current_device() -> bool:
202
- # Requires AITER and GFX950
220
if not rocm_aiter_ops.is_enabled():
221
return False
205
- from vllm.platforms.rocm import on_gfx950
222
+ from vllm.platforms.rocm import on_gfx950, on_gfx1250
223
207
- return on_gfx950()
224
+ return on_gfx950() or on_gfx1250()
225
226
@staticmethod
227
def _supports_no_act_and_mul() -> bool:
@@ -293,3 +310,358 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
310
unpadded_N_w2=self.moe_config.hidden_dim_unpadded,
311
unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded,
312
)
313
+
314
+
315
+def _aiter_raw(t):
316
+ if t is None or isinstance(t, torch.Tensor):
317
+ return t
318
+ return t.storage.data if hasattr(t, "storage") else t
319
+
320
+
321
+def _aiter_w4a16_silu_via_a8w4(
322
+ hidden_states: torch.Tensor,
323
+ w1_data,
324
+ w2_data,
325
+ w1_wscale,
326
+ w2_wscale,
327
+ w1_bias,
328
+ w2_bias,
329
+ routing_data,
330
+ gather_idx,
331
+ scatter_idx,
332
+ gammas,
333
+ apply_router_weight_on_input: bool,
334
+ swiglu_limit: float,
335
+ unpadded_N_w1,
336
+ unpadded_K_w1,
337
+ unpadded_N_w2,
338
+ unpadded_K_w2,
339
+) -> torch.Tensor:
340
+ """
341
+ MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation.
342
+ """
343
+ from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul
344
+ from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
345
+ from aiter.ops.triton.quant import dynamic_mxfp8_quant
346
+
347
+ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
348
+ should_use_cdna4_mx_scale_swizzle,
349
+ )
350
+
351
+ swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None
352
+ quant_dtype = torch.float8_e4m3fn
353
+
354
+ g1_gammas = gammas if apply_router_weight_on_input else None
355
+ g2_gammas = None if apply_router_weight_on_input else gammas
356
+
357
+ hidden_q, a1_scale = dynamic_mxfp8_quant(hidden_states, quant_dtype=quant_dtype)
358
+ raw_gate_up = moe_gemm_a8w4(
359
+ hidden_q,
360
+ w1_data,
361
+ a1_scale,
362
+ w1_wscale,
363
+ None,
364
+ None,
365
+ w1_bias,
366
+ routing_data,
367
+ gather_indx=gather_idx,
368
+ gammas=g1_gammas,
369
+ swizzle_mx_scale=swz,
370
+ out_dtype=torch.bfloat16,
371
+ apply_swiglu=False,
372
+ unpadded_N=unpadded_N_w1,
373
+ unpadded_K=unpadded_K_w1,
374
+ )
375
+ if unpadded_N_w1 is not None:
376
+ raw_gate_up = raw_gate_up[:, :unpadded_N_w1]
377
+
378
+ interim_fp8, a2_scale = fused_clamp_act_mul(
379
+ raw_gate_up,
380
+ swiglu_limit=swiglu_limit,
381
+ activation="silu",
382
+ dtype_quant=quant_dtype,
383
+ scale_dtype_fmt="ue8m0",
384
+ quant_block_size=32,
385
+ )
386
+
387
+ out = moe_gemm_a8w4(
388
+ interim_fp8,
389
+ w2_data,
390
+ a2_scale,
391
+ w2_wscale,
392
+ None,
393
+ None,
394
+ w2_bias,
395
+ routing_data,
396
+ scatter_indx=scatter_idx,
397
+ gammas=g2_gammas,
398
+ swizzle_mx_scale=swz,
399
+ unpadded_N=unpadded_N_w2,
400
+ unpadded_K=unpadded_K_w2,
401
+ )
402
+ return out
403
+
404
+
405
+def aiter_triton_kernel_w4a16_moe_forward(
406
+ hidden_states: torch.Tensor,
407
+ w1,
408
+ w2,
409
+ gating_output: torch.Tensor,
410
+ topk: int,
411
+ renormalize: bool,
412
+ activation: MoEActivation = MoEActivation.SWIGLUOAI,
413
+ quant_config: FusedMoEQuantConfig | None = None,
414
+ apply_router_weight_on_input: bool = False,
415
+ global_num_experts: int = -1,
416
+ expert_map: torch.Tensor | None = None,
417
+ unpadded_N_w1=None,
418
+ unpadded_K_w1=None,
419
+ unpadded_N_w2=None,
420
+ unpadded_K_w2=None,
421
+ num_expert_group: int | None = None,
422
+ topk_group: int | None = None,
423
+ e_score_correction_bias: torch.Tensor | None = None,
424
+ routed_scaling_factor: float | None = None,
425
+ score_mode: str | None = None,
426
+):
427
+ assert quant_config is not None and rocm_aiter_ops.is_enabled()
428
+ from vllm.platforms.rocm import on_gfx1250
429
+
430
+ try:
431
+ from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4
432
+ from aiter.ops.triton.moe.moe_routing import routing as _routing_mod
433
+ except ImportError:
434
+ from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4
435
+ from aiter.ops.triton.moe_routing import routing as _routing_mod
436
+
437
+ if on_gfx1250():
438
+ _routing_mod.is_tdm_avail = lambda: False
439
+ aiter_routing = _routing_mod.routing
440
+
441
+ if score_mode is not None:
442
+ use_grouped_topk = num_expert_group is not None and num_expert_group > 1
443
+ routing_data, gather_idx, scatter_idx = aiter_routing(
444
+ gating_output,
445
+ topk,
446
+ score_mode=score_mode,
447
+ bias=e_score_correction_bias,
448
+ renorm=renormalize,
449
+ routed_scaling_factor=(
450
+ routed_scaling_factor if routed_scaling_factor is not None else 1.0
451
+ ),
452
+ use_grouped_topk=use_grouped_topk,
453
+ num_expert_group=num_expert_group,
454
+ topk_group=topk_group,
455
+ )
456
+ else:
457
+ routing_data, gather_idx, scatter_idx = aiter_routing(
458
+ gating_output, topk, sm_first=not renormalize
459
+ )
460
+
461
+ if on_gfx1250():
462
+ gather_src = gather_idx.to(torch.long) // topk
463
+ hidden_states = hidden_states[gather_src]
464
+ gather_idx = None
465
+
466
+ assert quant_config.w1_precision is not None
467
+ assert quant_config.w2_precision is not None
468
+
469
+ w1_data = _aiter_raw(w1)
470
+ w2_data = _aiter_raw(w2)
471
+ w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale)
472
+ w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale)
473
+
474
+ gammas = routing_data.gate_scal if routing_data else None
475
+
476
+ swiglu_alpha = (
477
+ quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
478
+ )
479
+ swiglu_limit = (
480
+ quant_config.gemm1_clamp_limit
481
+ if quant_config.gemm1_clamp_limit is not None
482
+ else 7.0
483
+ )
484
+
485
+ # SILU on gfx1250: use the verified a8w4 kernel (dynamic MXFP8); a16w4 faults.
486
+ if activation == MoEActivation.SILU and on_gfx1250():
487
+ return _aiter_w4a16_silu_via_a8w4(
488
+ hidden_states,
489
+ w1_data,
490
+ w2_data,
491
+ w1_wscale,
492
+ w2_wscale,
493
+ quant_config.w1_bias,
494
+ quant_config.w2_bias,
495
+ routing_data,
496
+ gather_idx,
497
+ scatter_idx,
498
+ gammas,
499
+ apply_router_weight_on_input,
500
+ swiglu_limit,
501
+ unpadded_N_w1,
502
+ unpadded_K_w1,
503
+ unpadded_N_w2,
504
+ unpadded_K_w2,
505
+ )
506
+
507
+ # SILU: silu(gate) * up — same kernel, just no "+1" residual in swiglu.
508
+ swiglu_add_residual = activation != MoEActivation.SILU
509
+
510
+ intermediate = moe_gemm_a16w4(
511
+ hidden_states,
512
+ w1_data,
513
+ None,
514
+ w1_wscale,
515
+ None,
516
+ None,
517
+ quant_config.w1_bias,
518
+ routing_data,
519
+ gather_indx=gather_idx,
520
+ gammas=gammas if apply_router_weight_on_input else None,
521
+ swizzle_mx_scale=None,
522
+ apply_swiglu=True,
523
+ alpha=swiglu_alpha,
524
+ limit=swiglu_limit,
525
+ swiglu_add_residual=swiglu_add_residual,
526
+ unpadded_N=unpadded_N_w1,
527
+ unpadded_K=unpadded_K_w1,
528
+ )
529
+
530
+ out = moe_gemm_a16w4(
531
+ intermediate,
532
+ w2_data,
533
+ None,
534
+ w2_wscale,
535
+ None,
536
+ None,
537
+ quant_config.w2_bias,
538
+ routing_data,
539
+ scatter_indx=scatter_idx,
540
+ gammas=None if apply_router_weight_on_input else gammas,
541
+ swizzle_mx_scale=None,
542
+ unpadded_N=unpadded_N_w2,
543
+ unpadded_K=unpadded_K_w2,
544
+ )
545
+
546
+ return out
547
+
548
+
549
+class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
550
+ def __init__(
551
+ self,
552
+ moe_config: FusedMoEConfig,
553
+ quant_config: FusedMoEQuantConfig,
554
+ ):
555
+ super().__init__(moe_config, quant_config)
556
+ self.topk = moe_config.experts_per_token
557
+ self.renormalize = moe_config.routing_method in (
558
+ RoutingMethodType.Renormalize,
559
+ RoutingMethodType.RenormalizeNaive,
560
+ RoutingMethodType.DeepseekV4,
561
+ )
562
+
563
+ @staticmethod
564
+ def activation_format() -> mk.FusedMoEActivationFormat:
565
+ return mk.FusedMoEActivationFormat.Standard
566
+
567
+ @staticmethod
568
+ def _supports_current_device() -> bool:
569
+ if not rocm_aiter_ops.is_enabled():
570
+ return False
571
+ from vllm.platforms.rocm import on_gfx950, on_gfx1250
572
+
573
+ return on_gfx950() or on_gfx1250()
574
+
575
+ @staticmethod
576
+ def _supports_no_act_and_mul() -> bool:
577
+ return False
578
+
579
+ @staticmethod
580
+ def _supports_quant_scheme(
581
+ weight_key: QuantKey | None,
582
+ activation_key: QuantKey | None,
583
+ ) -> bool:
584
+ return (weight_key, activation_key) == (kMxfp4Static, None)
585
+
586
+ @staticmethod
587
+ def _supports_activation(activation: MoEActivation) -> bool:
588
+ return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU)
589
+
590
+ @staticmethod
591
+ def _supports_parallel_config(
592
+ moe_parallel_config: FusedMoEParallelConfig,
593
+ ) -> bool:
594
+ return (
595
+ not moe_parallel_config.use_all2all_kernels
596
+ and not moe_parallel_config.enable_eplb
597
+ and moe_parallel_config.dp_size <= 1
598
+ )
599
+
600
+ @staticmethod
601
+ def _supports_routing_method(
602
+ routing_method: RoutingMethodType,
603
+ weight_key: QuantKey | None,
604
+ activation_key: QuantKey | None,
605
+ ) -> bool:
606
+ return routing_method in [
607
+ RoutingMethodType.Renormalize,
608
+ RoutingMethodType.RenormalizeNaive,
609
+ RoutingMethodType.DeepseekV4,
610
+ ]
611
+
612
+ @staticmethod
613
+ def _supports_router_logits_dtype(
614
+ router_logits_dtype: torch.dtype | None,
615
+ routing_method: RoutingMethodType,
616
+ ) -> bool:
617
+ return True
618
+
619
+ @property
620
+ def expects_unquantized_inputs(self) -> bool:
621
+ return True
622
+
623
+ def apply(
624
+ self,
625
+ hidden_states: torch.Tensor,
626
+ w1: torch.Tensor,
627
+ w2: torch.Tensor,
628
+ router_logits: torch.Tensor,
629
+ activation: MoEActivation,
630
+ global_num_experts: int,
631
+ expert_map: torch.Tensor | None,
632
+ a1q_scale: torch.Tensor | None,
633
+ apply_router_weight_on_input: bool,
634
+ num_expert_group: int | None = None,
635
+ e_score_correction_bias: torch.Tensor | None = None,
636
+ routed_scaling_factor: float | None = None,
637
+ topk_group: int | None = None,
638
+ ) -> torch.Tensor:
639
+ assert self.moe_config.intermediate_size_per_partition_unpadded is not None
640
+ assert self.moe_config.hidden_dim_unpadded is not None
641
+ score_mode = (
642
+ "sqrtsoftplus"
643
+ if self.moe_config.routing_method == RoutingMethodType.DeepseekV4
644
+ else None
645
+ )
646
+ return aiter_triton_kernel_w4a16_moe_forward(
647
+ hidden_states=hidden_states,
648
+ w1=w1,
649
+ w2=w2,
650
+ gating_output=router_logits,
651
+ topk=self.topk,
652
+ renormalize=self.renormalize,
653
+ activation=activation,
654
+ global_num_experts=global_num_experts,
655
+ expert_map=expert_map,
656
+ quant_config=self.quant_config,
657
+ apply_router_weight_on_input=apply_router_weight_on_input,
658
+ unpadded_N_w1=self.moe_config.intermediate_size_per_partition_unpadded * 2,
659
+ unpadded_K_w1=self.moe_config.hidden_dim_unpadded,
660
+ unpadded_N_w2=self.moe_config.hidden_dim_unpadded,
661
+ unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded,
662
+ num_expert_group=num_expert_group,
663
+ topk_group=topk_group,
664
+ e_score_correction_bias=e_score_correction_bias,
665
+ routed_scaling_factor=routed_scaling_factor,
666
+ score_mode=score_mode,
667
+ )
vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
+4
-4
@@ -814,13 +814,13 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
814
) -> bool:
815
p = current_platform
816
if p.is_rocm():
817
- from vllm.platforms.rocm import on_gfx9
817
+ from vllm.platforms.rocm import get_cdna_version
818
819
- is_rocm_on_gfx9 = on_gfx9()
819
+ _rocm_support_fp8 = get_cdna_version() > 2
820
else:
821
- is_rocm_on_gfx9 = False
821
+ _rocm_support_fp8 = False
822
823
- device_supports_fp8 = is_rocm_on_gfx9 or (
823
+ device_supports_fp8 = _rocm_support_fp8 or (
824
p.is_cuda() and p.has_device_capability((8, 9))
825
)
826
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+2
-3
@@ -472,11 +472,10 @@ class AiterExperts(mk.FusedMoEExpertsModular):
472
]
473
if (weight_key, activation_key) not in SUPPORTED_W_A:
474
return False
475
- # CK MXFP4 MoE kernels are only supported on gfx950.
475
if weight_key == kMxfp4Static:
477
- from vllm.platforms.rocm import on_gfx950
476
+ from vllm.platforms.rocm import on_gfx950, on_gfx1250
477
479
- if not on_gfx950():
478
+ if not on_gfx950() or on_gfx1250():
479
return False
480
return True
481
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+9
-3
@@ -222,11 +222,14 @@ def backend_to_kernel_cls(
222
return [BatchedMarlinExperts]
223
224
elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
225
+ from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import (
226
+ AiterW4A16ExpertsMonolithic,
227
+ )
228
from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
229
AiterExperts,
230
)
231
229
- return [AiterExperts]
232
+ return [AiterExperts, AiterW4A16ExpertsMonolithic]
233
234
elif backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
235
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(
1266
1267
Supports DeepGEMM, TRTLLM MXFP8, Triton and Marlin backends.
1268
"""
1269
+ from vllm.platforms.rocm import on_gfx1250
1270
1271
if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
1272
w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
@@ -1435,7 +1439,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1439
w2_bias,
1440
)
1441
1438
- elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
1442
+ elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not on_gfx1250():
1443
# Initially introduced for DeepSeekV4
1444
1445
if w13_bias is not None:
@@ -1492,7 +1496,9 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1496
w2_bias,
1497
)
1498
1495
- elif mxfp4_backend in TRITON_BACKENDS:
1499
+ elif mxfp4_backend in TRITON_BACKENDS or (
1500
+ mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250()
1501
+ ):
1502
from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
1503
1504
if mxfp4_backend == Mxfp4MoeBackend.TRITON:
vllm/model_executor/layers/quantization/mxfp4.py
+19
-3
@@ -480,6 +480,10 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
480
global_num_experts=layer.global_num_experts,
481
expert_map=layer.expert_map,
482
apply_router_weight_on_input=layer.apply_router_weight_on_input,
483
+ num_expert_group=layer.num_expert_group,
484
+ topk_group=layer.topk_group,
485
+ e_score_correction_bias=layer.e_score_correction_bias,
486
+ routed_scaling_factor=layer.routed_scaling_factor,
487
)
488
489
@@ -735,7 +739,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
739
740
# For TRITON backends, weights are wrapped tensors from triton_kernels
741
# that don't support .detach(). Manually assign parameters.
738
- if self.mxfp4_backend not in TRITON_BACKENDS:
742
+ from vllm.platforms.rocm import on_gfx1250
743
+
744
+ uses_triton_weight_format = self.mxfp4_backend in TRITON_BACKENDS or (
745
+ self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250()
746
+ )
747
+ if not uses_triton_weight_format:
748
replace_parameter(layer, "w13_weight", w13)
749
replace_parameter(layer, "w2_weight", w2)
750
replace_parameter(layer, "w13_weight_scale", w13_scale)
@@ -747,7 +756,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
756
self.w2_precision_config = w2_scale
757
758
# AITER backend requires weights to be marked as shuffled.
750
- if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
759
+ if (
760
+ self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16
761
+ and not uses_triton_weight_format
762
+ ):
763
layer.w13_weight.is_shuffled = True
764
layer.w2_weight.is_shuffled = True
765
@@ -841,7 +853,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
853
w2_bias = getattr(layer, "w2_bias", None)
854
swiglu_limit = getattr(layer, "swiglu_limit", None)
855
844
- if self.mxfp4_backend in TRITON_BACKENDS:
856
+ from vllm.platforms.rocm import on_gfx1250
857
+
858
+ if self.mxfp4_backend in TRITON_BACKENDS or (
859
+ self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250()
860
+ ):
861
# TRITON backends free w13/w2_weight_scale after swizzling; the
862
# swizzled scales live inside the precision configs instead.
863
assert self.w13_precision_config is not None
vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py
+1
-3
@@ -64,9 +64,7 @@ class QuarkW4A8_MXFP4_FP8(QuarkScheme):
64
65
kernel_supported_gpu = False
66
if current_platform.is_rocm():
67
- from vllm.platforms.rocm import on_gfx950
68
-
69
- kernel_supported_gpu = on_gfx950()
67
+ kernel_supported_gpu = current_platform.supports_mx()
68
69
self.use_aiter_kernel = (
70
is_aiter_found_and_supported()
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+31
-1
@@ -875,10 +875,40 @@ def w8a8_triton_block_scaled_mm(
875
Bs: The per-block quantization scale for `B`.
876
block_size: The block size for per-block quantization. It should
877
be 2-dim, e.g., [128, 128].
878
- output_dtype: The dtype of the returned tensor.
878
+ output_dytpe: The dtype of the returned tensor.
879
Returns:
880
torch.Tensor: The result of matmul.
881
"""
882
+
883
+ from vllm.platforms.rocm import on_gfx1250
884
+
885
+ if on_gfx1250():
886
+ # Torch upcast reference: dequantize A,B to fp32 and matmul in fp32.
887
+ # Avoids the gfx1250 native-fp8 block GEMM NaN bug. Correct but slow.
888
+ _bn, _bk = block_size[0], block_size[1]
889
+ _As = (
890
+ _upcast_e8m0_to_fp32(As)
891
+ if As.dtype == torch.float8_e8m0fnu
892
+ else As.to(torch.float32)
893
+ )
894
+ _Bs = (
895
+ _upcast_e8m0_to_fp32(Bs)
896
+ if Bs.dtype == torch.float8_e8m0fnu
897
+ else Bs.to(torch.float32)
898
+ )
899
+ _K = A.shape[-1]
900
+ _N = B.shape[0]
901
+ _Af = A.to(torch.float32).reshape(-1, _K)
902
+ _Asf = (
903
+ _As.to(torch.float32)
904
+ .reshape(-1, _As.shape[-1])
905
+ .repeat_interleave(_bk, dim=1)[:, :_K]
906
+ )
907
+ _Bf = B.to(torch.float32)
908
+ _Bsf = _Bs.repeat_interleave(_bn, dim=0).repeat_interleave(_bk, dim=1)[:_N, :_K]
909
+ _out = (_Af * _Asf) @ (_Bf * _Bsf).t()
910
+ return _out.to(output_dtype).reshape(*A.shape[:-1], _N)
911
+
912
assert len(block_size) == 2
913
block_n, block_k = block_size[0], block_size[1]
914
vllm/model_executor/layers/utils.py
+8
-2
@@ -122,7 +122,7 @@ def use_aiter_triton_gemm(n, m, k, dtype):
122
def rocm_unquantized_gemm_impl(
123
x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
124
) -> torch.Tensor:
125
- from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950
125
+ from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250
126
127
n = x.numel() // x.size(-1)
128
m = weight.shape[0]
@@ -164,7 +164,11 @@ def rocm_unquantized_gemm_impl(
164
if use_skinny_reduce_counting:
165
return ops.wvSplitKrc(x, weight, cu_count, bias)
166
167
- if use_aiter_triton_gemm(n, m, k, x.dtype):
167
+ # gfx1250's aiter gemm_a16w16 uses the gluon backend, which requires
168
+ # K % 256 == 0 (it walks K with fixed-size descriptors and won't pad a
169
+ # partial last tile). Some whitelisted shapes have K=2880 (e.g. gpt-oss-120b
170
+ # hidden), so skip aiter there and fall back to the torch GEMM path below.
171
+ if use_aiter_triton_gemm(n, m, k, x.dtype) and not (on_gfx1250() and k % 256 != 0):
172
from aiter.ops.triton.gemm_a16w16 import gemm_a16w16
173
174
return gemm_a16w16(x, weight, bias)
@@ -172,6 +176,8 @@ def rocm_unquantized_gemm_impl(
176
use_skinny = (
177
envs.VLLM_ROCM_USE_SKINNY_GEMM
178
and (on_gfx9() or on_gfx1x())
179
+ # build (gfx9/gfx11 ISA); fall back to torch GEMM there.
180
+ # TODO GFX1250: Include once skinny GEMM is supported on gfx1250
181
and x.dtype in [torch.float16, torch.bfloat16]
182
and k % 8 == 0
183
)
vllm/platforms/rocm.py
+37
-13
@@ -194,11 +194,6 @@ def _get_gcn_arch() -> str:
194
return _query_gcn_arch_from_amdsmi()
195
except Exception as e:
196
logger.debug("Failed to get GCN arch via amdsmi: %s", e)
197
- logger.warning_once(
198
- "Failed to get GCN arch via amdsmi, falling back to torch.cuda. "
199
- "This will initialize CUDA and may cause "
200
- "issues if CUDA_VISIBLE_DEVICES is not set yet."
201
- )
197
# Ultimate fallback: use torch.cuda (will initialize CUDA)
198
return torch.cuda.get_device_properties("cuda").gcnArchName
199
@@ -218,6 +213,11 @@ _ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
213
_ON_GFX90A = "gfx90a" in _GCN_ARCH
214
_ON_GFX942 = "gfx942" in _GCN_ARCH
215
_ON_GFX950 = "gfx950" in _GCN_ARCH
216
+_ON_GFX1250 = "gfx1250" in _GCN_ARCH
217
+
218
+_ON_CDNA = any(arch in _GCN_ARCH for arch in ["gfx9", "gfx1250"])
219
+# RDNA = gfx11/gfx12 minus the CDNA-classified gfx1250.
220
+_ON_RDNA = _ON_GFX1X and not _ON_CDNA
221
222
223
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:
292
293
294
def on_gfx1x() -> bool:
295
- return _ON_GFX1X
295
+ return _ON_GFX1X and not _ON_CDNA
296
297
298
def on_gfx11() -> bool:
@@ -308,7 +308,11 @@ def on_gfx1151() -> bool:
308
309
310
def on_gfx12x() -> bool:
311
- return _ON_GFX12X
311
+ return _ON_GFX12X and not _ON_CDNA
312
+
313
+
314
+def on_gfx1250() -> bool:
315
+ return _ON_GFX1250
316
317
318
def on_mi3xx() -> bool:
@@ -331,13 +335,33 @@ def on_gfx950() -> bool:
335
return _ON_GFX950
336
337
338
+def on_cdna() -> bool:
339
+ return _ON_CDNA
340
+
341
+
342
+def on_rdna() -> bool:
343
+ return _ON_RDNA
344
+
345
+
346
+def get_cdna_version() -> int:
347
+ if on_gfx90a():
348
+ return 2
349
+ if on_gfx942():
350
+ return 3
351
+ if on_gfx950():
352
+ return 4
353
+ if on_gfx1250():
354
+ return 5
355
+ return 0
356
+
357
+
358
# Enable HIP online tuning early, before hipBLASLt initializes.
359
# Turn on hipBLASLt online tuning if use AITER hipBLASLt GEMM.
360
if (
361
envs.VLLM_ROCM_USE_AITER
362
and envs.VLLM_ROCM_USE_AITER_LINEAR
363
and envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
340
- and on_mi3xx()
364
+ and get_cdna_version() > 2
365
):
366
os.environ["HIP_ONLINE_TUNING"] = "1"
367
@@ -356,7 +380,7 @@ def use_rocm_custom_paged_attention(
380
) -> bool:
381
# custom paged attn always supported on V0. On V1, requires sliding window
382
# disabled due to observed numerical discrepancy.
359
- if _ON_GFX9:
383
+ if on_cdna():
384
return (
385
(sliding_window == 0 or sliding_window == (-1, -1))
386
and (qtype == torch.half or qtype == torch.bfloat16)
@@ -657,12 +681,12 @@ class RocmPlatform(Platform):
681
682
from vllm._aiter_ops import rocm_aiter_ops
683
660
- if rocm_aiter_ops.is_mha_enabled() and on_gfx9():
684
+ if rocm_aiter_ops.is_mha_enabled() and on_cdna():
685
logger.info_once("Using AITER Flash Attention backend for ViT model.")
686
return AttentionBackendEnum.ROCM_AITER_FA
687
688
if (
665
- on_gfx9()
689
+ on_cdna()
690
and find_spec("flash_attn") is not None
691
and (dtype == torch.float16 or dtype == torch.bfloat16)
692
):
@@ -878,11 +902,11 @@ class RocmPlatform(Platform):
902
903
@classmethod
904
def supports_mx(cls) -> bool:
881
- return any(gfx in _GCN_ARCH for gfx in ["gfx95"])
905
+ return any(gfx in _GCN_ARCH for gfx in ["gfx95", "gfx1250"])
906
907
@classmethod
908
def supports_fp8(cls) -> bool:
885
- return on_gfx9() or on_gfx12x()
909
+ return on_cdna() or on_gfx12x()
910
911
@classmethod
912
def is_fp8_fnuz(cls) -> bool:
vllm/utils/import_utils.py
+12
-1
@@ -505,7 +505,18 @@ def has_triton_kernels() -> bool:
505
@cache
506
def has_tilelang() -> bool:
507
"""Whether the optional `tilelang` package is available."""
508
- return _has_module("tilelang")
508
+ if not _has_module("tilelang"):
509
+ return False
510
+ # ROCm-only guard, imported lazily to avoid loading rocm on CUDA.
511
+ from vllm.platforms import current_platform
512
+
513
+ if current_platform.is_rocm():
514
+ from vllm.platforms.rocm import on_gfx1250
515
+
516
+ # TODO: Re-enable when tilelang supports gfx1250
517
+ if on_gfx1250():
518
+ return False
519
+ return True
520
521
522
def has_arctic_inference() -> bool:
vllm/utils/jit_monitor.py
+18
-2
@@ -231,10 +231,26 @@ def _log_triton_jit_compile(fn_name: str, kwargs) -> None:
231
232
233
def _setup_triton_jit_hook() -> None:
234
- """Register a ``jit_post_compile_hook`` that warns on compilation."""
234
+ """Register a jit_post_compile_hook that warns on compilation."""
235
if not HAS_TRITON:
236
return
237
- from triton import knobs # type: ignore[import-untyped]
237
+ from triton import knobs
238
+ from triton.runtime import jit as _triton_jit
239
+
240
+ # kernels pass non-JSON-serializable constexprs
241
+ # make that serialization non-fatal
242
+ _orig = _triton_jit.serialize_specialization_data
243
+ if not getattr(_orig, "_vllm_guarded", False):
244
+
245
+ @functools.wraps(_orig)
246
+ def _guarded(*args, **kwargs):
247
+ try:
248
+ return _orig(*args, **kwargs)
249
+ except (TypeError, ValueError):
250
+ return None # best-effort metadata; monitor ignores it
251
+
252
+ cast(Any, _guarded)._vllm_guarded = True
253
+ _triton_jit.serialize_specialization_data = _guarded
254
255
existing_hook = knobs.runtime.jit_post_compile_hook
256
vllm/v1/attention/backends/fa_utils.py
+18
-9
@@ -35,23 +35,31 @@ elif current_platform.is_xpu():
35
compile_flash_attn_varlen_func_from_specs = None # type: ignore[assignment]
36
get_scheduler_metadata = xpu_ops.get_scheduler_metadata # type: ignore[assignment]
37
elif current_platform.is_rocm():
38
- try:
39
- from flash_attn import flash_attn_varlen_func # type: ignore[no-redef]
38
+ # On ROCm we use AITER's Triton flash-attention; the upstream flash-attn
39
+ # package is not installed/available. (Same source as aiter_triton_mla.py.)
40
+ # The FA4 compile-from-specs API is CUDA-only, so it is unavailable on ROCm
41
+ # regardless of whether AITER is present.
42
+ from vllm.platforms.rocm import on_gfx1250
43
41
- compile_flash_attn_varlen_func_from_specs = None # type: ignore[assignment]
44
+ compile_flash_attn_varlen_func_from_specs = None # type: ignore[assignment]
45
+ try:
46
+ if on_gfx1250():
47
+ from aiter.ops.triton.mha import ( # type: ignore[no-redef]
48
+ flash_attn_varlen_func,
49
+ )
50
+ else:
51
+ from flash_attn import flash_attn_varlen_func # type: ignore[no-redef]
52
43
- # Mark that upstream flash-attn is available on ROCm
53
_ROCM_FLASH_ATTN_AVAILABLE = True
54
except ImportError:
55
56
def flash_attn_varlen_func(*args: Any, **kwargs: Any) -> Any: # type: ignore[no-redef,misc]
57
+ package = "aiter" if on_gfx1250() else "flash-attn"
58
raise ImportError(
49
- "ROCm platform requires upstream flash-attn "
50
- "to be installed. Please install flash-attn first."
59
+ f"ROCm platform requires upstream {package} "
60
+ f"to be installed. Please install {package} first."
61
)
62
53
- compile_flash_attn_varlen_func_from_specs = None # type: ignore[assignment]
54
-
63
# ROCm doesn't use scheduler metadata (FA3 feature), provide stub
64
def get_scheduler_metadata(*args: Any, **kwargs: Any) -> None: # type: ignore[misc]
65
return None
@@ -348,7 +356,8 @@ def is_flash_attn_varlen_func_available() -> bool:
356
Platform-specific sources:
357
- CUDA: vllm.vllm_flash_attn.flash_attn_varlen_func
358
- XPU: xpu_ops.flash_attn_varlen_func
351
- - ROCm: upstream flash_attn.flash_attn_varlen_func (if available)
359
+ - ROCm: aiter.ops.triton.mha.flash_attn_varlen_func (if AITER available) or
360
+ upstream flash_attn.flash_attn_varlen_func
361
362
Note: This is separate from the AITER flash attention backend (rocm_aiter_fa.py)
363
which uses rocm_aiter_ops.flash_attn_varlen_func. The condition to use AITER is
vllm/v1/attention/backends/rocm_aiter_fa.py
+3
-3
@@ -776,12 +776,12 @@ class AiterFlashAttentionBackend(AttentionBackend):
776
777
@classmethod
778
def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
779
- from vllm.platforms.rocm import on_mi3xx
779
+ from vllm.platforms.rocm import get_cdna_version
780
781
# DeviceCapability is currently created using torch.cuda.get_device_capability()
782
- # which is known to be buggy on rocm systems. on_mi3xx uses amd-smi which is
782
+ # which is known to be buggy on rocm systems. on CDNA uses amd-smi which is
783
# more reliable.
784
- return on_mi3xx()
784
+ return get_cdna_version() > 2
785
786
@classmethod
787
def supports_non_causal(cls) -> bool:
vllm/v1/sample/ops/topk_topp_sampler.py
+8
@@ -18,6 +18,13 @@ if HAS_TRITON:
18
logger = init_logger(__name__)
19
20
21
+def _skip_aiter_sampler_on_gfx1250() -> bool:
22
+ # Lazy ROCm-only import; keeps arch detection out of import time on CUDA/CPU.
23
+ from vllm.platforms.rocm import on_gfx1250
24
+
25
+ return on_gfx1250()
26
+
27
+
28
def flashinfer_sampler_supported() -> bool:
29
"""Decide whether FlashInfer's top-p/top-k sampler can be used.
30
@@ -110,6 +117,7 @@ class TopKTopPSampler(nn.Module):
117
elif (
118
logprobs_mode not in PROCESSED_LOGPROBS_MODES
119
and rocm_aiter_ops.is_enabled()
120
+ and not _skip_aiter_sampler_on_gfx1250() # TODO (JPVILLAM): Enable
121
):
122
self.aiter_ops = None
123
self._aiter_ops_import_failed = False