Enable gfx1250 ROCm architecture (#46516)

Signed-off-by: Douglas Lehr <[email protected]> Signed-off-by: jpvillam <[email protected]> Signed-off-by: Daniel <[email protected]> Signed-off-by: Jaden Mathias <[email protected]> Signed-off-by: jpvillam <[email protected]> Signed-off-by: Juan Villamizar <[email protected]> Signed-off-by: Douglas Lehr <[email protected]> Signed-off-by: Jaden Mathias <[email protected]> Co-authored-by: Douglas Lehr <[email protected]> Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]> Co-authored-by: Daniel <[email protected]> Co-authored-by: Jaden Mathias <[email protected]> Co-authored-by: jpvillam <[email protected]> Co-authored-by: root <[email protected]> Co-authored-by: Douglas Lehr <[email protected]> Co-authored-by: dllehr <[email protected]> Co-authored-by: Daniel Chan <[email protected]> Co-authored-by: Andreas Karatzas <[email protected]>
2026-08-15 02:08:10 +00:00 · 2026-07-31 01:41:12 +00:00
parent 553fcb82d5
commit 4f1da84eb5
33 changed files with 1850 additions and 104 deletions
@@ -49,7 +49,7 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
 set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14")

 # Supported AMD GPU architectures.
-set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
+set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1250;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")

 # ROCm installation prefix. Default to /opt/rocm but allow override via
 # -DROCM_PATH=/your/rocm/path when invoking cmake.
@@ -1422,6 +1422,38 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/rocm/skinny_gemms.cu"
    "csrc/rocm/skinny_gemms_int4.cu"
    "csrc/rocm/attention.cu")
+  set(VLLM_ROCM_EXT_FLAGS ${VLLM_GPU_FLAGS})
+
+  # skinny_gemms*.cu are built on gfx9/gfx11 ISA (MFMA, dot2/dot4, legacy
+  # s_waitcnt asm) that gfx1250 (gfx12) does not provide. Exclude them from the
+  # gfx1250 build and disable their op registrations (VLLM_SKIP_SKINNY_GEMMS);
+  # vLLM falls back to default/Triton GEMM for those ops on gfx1250.
+  list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms.cu")
+  list(REMOVE_ITEM VLLM_ROCM_EXT_SRC "csrc/rocm/skinny_gemms_int4.cu")
+  set(VLLM_SKINNY_ARCHES ${VLLM_GPU_ARCHES})
+  list(FILTER VLLM_SKINNY_ARCHES EXCLUDE REGEX "gfx1250")
+  if(VLLM_SKINNY_ARCHES)
+    message(STATUS "Building skinny_gemms for archs: ${VLLM_SKINNY_ARCHES}")
+    set(VLLM_SKINNY_SRC
+      "csrc/rocm/skinny_gemms.cu"
+      "csrc/rocm/skinny_gemms_int4.cu")
+    hipify_sources_target(VLLM_SKINNY_HIP_SRCS _rocm_C_skinny "${VLLM_SKINNY_SRC}")
+    unset(_VLLM_LAST_HIPIFY_TARGET)
+    add_library(_rocm_C_skinny OBJECT ${VLLM_SKINNY_HIP_SRCS})
+    add_dependencies(_rocm_C_skinny hipify_all)
+    set_source_files_properties(${VLLM_SKINNY_HIP_SRCS} PROPERTIES LANGUAGE ${VLLM_GPU_LANG})
+    set_target_properties(_rocm_C_skinny PROPERTIES
+      ${VLLM_GPU_LANG}_ARCHITECTURES "${VLLM_SKINNY_ARCHES}"
+      POSITION_INDEPENDENT_CODE ON)
+    target_include_directories(_rocm_C_skinny PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/csrc)
+    target_compile_options(_rocm_C_skinny PRIVATE
+      $<$<COMPILE_LANGUAGE:${VLLM_GPU_LANG}>:${VLLM_ROCM_EXT_FLAGS}>)
+    target_compile_definitions(_rocm_C_skinny PRIVATE "-DTORCH_EXTENSION_NAME=_rocm_C")
+    target_link_libraries(_rocm_C_skinny PRIVATE torch)
+  else()
+    message(STATUS "Only gfx1250, skipping skinny_gemms")
+    list(APPEND VLLM_ROCM_EXT_FLAGS "-DVLLM_SKIP_SKINNY_GEMMS")
+  endif()

  set(VLLM_ROCM_HAS_GFX1100 OFF)
  if(VLLM_GPU_ARCHES MATCHES "gfx1100")
@@ -1437,10 +1469,16 @@ if(VLLM_GPU_LANG STREQUAL "HIP")
    DESTINATION vllm
    LANGUAGE ${VLLM_GPU_LANG}
    SOURCES ${VLLM_ROCM_EXT_SRC}
-    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    COMPILE_FLAGS ${VLLM_ROCM_EXT_FLAGS}
    ARCHITECTURES ${VLLM_GPU_ARCHES}
    USE_SABI 3
    WITH_SOABI)
+  
+  if(TARGET _rocm_C_skinny)
+    target_link_libraries(_rocm_C PRIVATE _rocm_C_skinny)
+  else()
+    target_compile_definitions(_rocm_C PRIVATE VLLM_SKIP_SKINNY_GEMMS)
+  endif()

  if(VLLM_ROCM_HAS_GFX1100)
    target_compile_definitions(_rocm_C PRIVATE VLLM_ROCM_GFX1100)
@@ -79,6 +79,7 @@ union BufferResource {
  };
 };

+#if !defined(__gfx1250__)
 __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
    int32x4_t srsrc, int32_t voffset, int32_t soffset,
    int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32");
@@ -86,6 +87,16 @@ __quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
 __quickreduce_device_inline__ static void buffer_store_dwordx4(
    int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset,
    int32_t aux) __asm("llvm.amdgcn.raw.buffer.store.v4i32");
+#else
+__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
+    int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) {}
+
+__quickreduce_device_inline__ static void buffer_store_dwordx4(int32x4_t data,
+                                                               int32x4_t srsrc,
+                                                               int32_t voffset,
+                                                               int32_t soffset,
+                                                               int32_t aux) {}
+#endif

 __quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) {
 #if defined(__gfx942__)
@@ -2405,6 +2405,16 @@ template <typename T, int absz, int cbid, int blgp>
 __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA,
                                                          const bit16x8& inpB,
                                                          const floatx8& inpC) {
+  #if defined(__gfx1250__)
+  // gfx1250 (gfx12 family) does not provide the gfx12 WMMA variant used by
+  // gfx1200/1201 (needs wmma-128b-insts). This custom-attention WMMA path is
+  // unsupported on gfx1250; trap if ever launched (fail loud, not
+  // silent-wrong).
+  (void)inpA;
+  (void)inpB;
+  __builtin_trap();
+  return inpC;
+  #else
  if constexpr (std::is_same<T, _Float16>::value) {
    return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(inpA, inpB, inpC);
  } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
@@ -2412,6 +2422,7 @@ __device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA,
  } else {
    static_assert(false, "unsupported 16b dtype");
  }
+  #endif
 }

 template <typename T>
@@ -14,6 +14,10 @@
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
  // vLLM custom ops for rocm

+// skinny_gemms.cu (LLMM1/wvSplitK/wvSplitKrc/wvSplitKQ) is excluded on gfx1250
+// (gfx9/gfx11 ISA, unsupported there); skip these registrations to avoid
+// undefined symbols. vLLM uses default/Triton GEMM for these ops on gfx1250.
+#ifndef VLLM_SKIP_SKINNY_GEMMS
  // Custom gemm op for matrix-vector multiplication
  rocm_ops.def(
      "LLMM1(Tensor in_a, Tensor in_b, int rows_per_block) -> "
@@ -46,6 +50,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
      "Tensor scale_a, "
      "          Tensor scale_b, int CuCount) -> ()");
  rocm_ops.impl("wvSplitKQ", torch::kCUDA, &wvSplitKQ);
+#endif  // VLLM_SKIP_SKINNY_GEMMS

 #ifdef VLLM_ROCM_GFX1100
  // W4A16 GPTQ kernels for AMD RDNA3 (gfx1100).
@@ -0,0 +1,325 @@
+ARG BASE_IMAGE=ubuntu:24.04
+ARG ROCM_WHEEL_INDEX=https://rocm.devreleases.amd.com/whl-multi-arch/
+ARG ROCM_SDK_VERSION=7.14.0a20260623
+ARG TORCH_VERSION=2.11.0+rocm7.14.0a20260623
+ARG TORCHVISION_VERSION=0.26.0+rocm7.14.0a20260623
+ARG TORCHAUDIO_VERSION=2.11.0+rocm7.14.0a20260623
+ARG TRITON_VERSION=3.7.1+git110cd8e2.rocm7.14.0a20260623
+ARG APEX_VERSION=1.11.0+rocm7.14.0a20260623
+
+
+ARG FA_BRANCH="jpvillam/gfx1250_wip"
+ARG FA_REPO="https://github.com/jpvillam-amd/flash-attention.git"
+ARG AITER_BRANCH="main"
+ARG AITER_REPO="https://github.com/ROCm/aiter.git"
+ARG MORI_BRANCH="v1.1.0"
+ARG MORI_REPO="https://github.com/ROCm/mori.git"
+
+# Sccache configuration (only used in release pipeline)
+ARG USE_SCCACHE
+ARG SCCACHE_DOWNLOAD_URL
+ARG SCCACHE_ENDPOINT
+ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
+ARG SCCACHE_REGION_NAME=us-west-2
+ARG SCCACHE_S3_NO_CREDENTIALS=0
+
+FROM ${BASE_IMAGE} AS base
+
+ARG PYTORCH_ROCM_ARCH=gfx1250
+ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
+ENV AITER_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
+ENV MORI_GPU_ARCHS=gfx942;gfx950
+ENV FA_GPU_ARCHS=gfx942;gfx950;gfx1250
+
+# TODO: Unset these when support is available for gfx1250
+ENV ENABLE_CK=0
+ENV SKIP_CK_BUILD="TRUE"
+ENV FLASH_ATTENTION_TRITON_AMD_ENABLE="TRUE"
+ARG PREBUILD_KERNELS=0
+ARG TRITON_KERNEL_BRANCH="padroute"
+ARG TRITON_KERNEL_REPO="https://github.com/jpvillam-amd/triton.git"
+
+# Required for RCCL in ROCm7.1
+ENV HSA_NO_SCRATCH_RECLAIM=1
+
+ARG PYTHON_VERSION=3.12
+ENV PYTHON_VERSION=${PYTHON_VERSION}
+
+RUN mkdir -p /app
+WORKDIR /app
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install Python and other dependencies
+RUN apt-get update -y \
+    && apt-get install -y software-properties-common git curl sudo vim less libgfortran5 libopenmpi-dev libpci-dev liblzma-dev libnuma-dev libdrm-dev pkg-config g++ \
+    && for i in 1 2 3; do \
+    add-apt-repository -y ppa:deadsnakes/ppa && break || \
+    { echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
+    done \
+    && apt-get update -y \
+    && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \
+    python${PYTHON_VERSION}-lib2to3 python-is-python3  \
+    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
+    && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
+    && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
+    && python3 --version
+
+ENV VIRTUAL_ENV=/opt/venv
+RUN python${PYTHON_VERSION} -m venv "${VIRTUAL_ENV}" && \
+    "${VIRTUAL_ENV}/bin/python" -m pip install --upgrade pip setuptools PyYAML
+ENV PATH=${VIRTUAL_ENV}/bin:$PATH
+
+RUN pip install -U packaging 'cmake<4' ninja wheel 'setuptools<80' pybind11 Cython
+RUN apt-get update && apt-get install -y libjpeg-dev libsox-dev libsox-fmt-all sox && rm -rf /var/lib/apt/lists/*
+
+# Install sccache if USE_SCCACHE is enabled (for release builds)
+ARG USE_SCCACHE
+ARG SCCACHE_DOWNLOAD_URL
+ARG SCCACHE_ENDPOINT
+ARG SCCACHE_BUCKET_NAME
+ARG SCCACHE_REGION_NAME
+ARG SCCACHE_S3_NO_CREDENTIALS
+RUN if [ "$USE_SCCACHE" = "1" ]; then \
+    echo "Installing sccache..." \
+    && SCCACHE_ARCH="x86_64" \
+    && SCCACHE_VERSION="v0.8.1" \
+    && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
+    && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
+    && tar -xzf /tmp/sccache.tar.gz -C /tmp \
+    && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
+    && chmod +x /usr/bin/sccache \
+    && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
+    && sccache --version; \
+    fi
+
+##
+## Install PyTorch w/ Triton + ROCM_SDK from ROCM wheel index
+##
+ARG ROCM_WHEEL_INDEX
+ARG TORCH_VERSION
+ARG TORCHVISION_VERSION
+ARG TORCHAUDIO_VERSION
+ARG ROCM_SDK_VERSION
+ARG APEX_VERSION
+ENV SITE_PACKAGES=${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages
+ENV ROCM_PATH=${SITE_PACKAGES}/_rocm_sdk_devel
+ENV ROCM_HOME=${ROCM_PATH}
+ENV ROCM_SOURCE_DIR=${ROCM_PATH}
+ENV ROCM_BIN=${ROCM_PATH}/bin
+ENV ROCM_CMAKE_PREFIX=${ROCM_PATH}/lib/cmake
+ENV HIP_DEVICE_LIB_PATH=${SITE_PACKAGES}/_rocm_sdk_core/lib/llvm/amdgcn/bitcode
+ENV PATH=${ROCM_PATH}/bin:${ROCM_PATH}/llvm/bin:$PATH
+ENV LD_LIBRARY_PATH=${ROCM_PATH}/lib:${SITE_PACKAGES}/_rocm_sdk_core/lib
+ENV CMAKE_PREFIX_PATH=${ROCM_PATH}/lib/cmake:${SITE_PACKAGES}/torch/share/cmake
+ENV PYTHONPATH=${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi
+
+# torch/torchvision/torchaudio must be pinned to mutually-consistent builds
+# (same +rocm... suffix) or the C++ ops break at import (ABI skew). The rocm
+# sdk version is derived from torch's own dependency pin unless overridden,
+# which keeps the set consistent and avoids pip backtracking.
+RUN pip install --pre --index-url "${ROCM_WHEEL_INDEX}" \
+    --extra-index-url https://pypi.org/simple \
+    "torch[device-all]==${TORCH_VERSION}" \
+    "torchvision==${TORCHVISION_VERSION}" \
+    "torchaudio==${TORCHAUDIO_VERSION}" \
+    "rocm[libraries,devel,device-all]==${ROCM_SDK_VERSION}" && \
+    rocm-sdk init
+ 
+# Torch runtime deps that may not be published on the ROCm wheel index;
+# install them from PyPI afterwards.
+RUN pip install filelock "typing-extensions>=4.10.0" "sympy>=1.13.3" \
+    "networkx>=2.5.1" jinja2 "fsspec>=0.8.5"
+
+
+
+# Expose the rocm-sdk wheel as a conventional /opt/rocm install so downstream
+# builds (Dockerfile.rocm: vLLM csrc, RIXL/UCX, ROCShmem/DeepEP) keep working.
+RUN ln -sfn "${ROCM_PATH}" /opt/rocm;  
+
+RUN if [ -f "${SITE_PACKAGES}/rocm_sdk/__init__.py" ]; then \
+    sed -i 's/rtld_global: bool = True/rtld_global: bool = False/g' \
+    "${SITE_PACKAGES}/rocm_sdk/__init__.py"; \
+    fi
+
+# The ROCm SDK wheel ships a broken CMake export for hsakmt
+# This patch includes the right paths for the numa build target
+RUN <<'EOF'
+set -eu
+TARGETS="/opt/rocm/lib/cmake/hsakmt/hsakmtTargets.cmake"
+[ -f "$TARGETS" ] || exit 0            # nothing to patch
+grep -q NUMA_LIBRARY "$TARGETS" && exit 0   # already patched
+
+# 1. Point libdrm's -L at the copy bundled in the wheel, not the builder path.
+sed -i 's|-L/__w/[^;"]*|-L${_IMPORT_PREFIX}/lib/rocm_sysdeps/lib|g' "$TARGETS"
+
+# 2. Drop the nonexistent RHEL libc path (libc is linked implicitly anyway).
+sed -i 's|/usr/lib64/libc.so;||g' "$TARGETS"
+
+# 3. Define the numa::numa target the export references but forgot to create.
+cat >> "$TARGETS" <<'CMAKE'
+
+if(NOT TARGET numa::numa)
+  find_library(NUMA_LIBRARY NAMES numa REQUIRED)
+  add_library(numa::numa UNKNOWN IMPORTED)
+  set_target_properties(numa::numa PROPERTIES IMPORTED_LOCATION "${NUMA_LIBRARY}")
+endif()
+CMAKE
+EOF
+
+# Clone custom triton_kernels for install later
+RUN mkdir -p /app/patched_triton_kernels; \
+    cd /app/patched_triton_kernels \
+    && git clone ${TRITON_KERNEL_REPO} \
+    && cd triton \
+    && git checkout ${TRITON_KERNEL_BRANCH} \
+    && git submodule update --init --recursive
+ENV TRITON_KERNELS_SRC_DIR="/app/patched_triton_kernels/triton/python/triton_kernels/triton_kernels/"
+
+
+# Setup sccache for HIP compilation via HIP_CLANG_PATH
+# This creates wrapper scripts in a separate directory and points HIP to use them
+# This avoids modifying the original ROCm binaries which can break detection
+# NOTE: HIP_CLANG_PATH is NOT set as ENV to avoid affecting downstream images (Dockerfile.rocm)
+# Instead, each build stage should export HIP_CLANG_PATH=/opt/sccache-wrappers if USE_SCCACHE=1
+RUN if [ "$USE_SCCACHE" = "1" ]; then \
+    echo "Setting up sccache wrappers for HIP compilation..." \
+    && mkdir -p /opt/sccache-wrappers \
+    && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang++ "$@"\n' > /opt/sccache-wrappers/clang++ \
+    && chmod +x /opt/sccache-wrappers/clang++ \
+    && printf '#!/bin/bash\nexec sccache ${ROCM_PATH}/lib/llvm/bin/clang "$@"\n' > /opt/sccache-wrappers/clang \
+    && chmod +x /opt/sccache-wrappers/clang \
+    && echo "sccache wrappers created in /opt/sccache-wrappers"; \
+    fi
+
+# Set sccache environment variables only when USE_SCCACHE=1
+# This prevents S3 config from leaking into images when sccache is not used
+ARG USE_SCCACHE
+ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
+ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
+ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
+ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
+
+
+###
+### AMD SMI Build
+###
+FROM base AS build_amdsmi
+RUN cd ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi \
+    && pip wheel . --wheel-dir=dist
+RUN mkdir -p /app/install && cp ${SITE_PACKAGES}/_rocm_sdk_core/share/amd_smi/dist/*.whl /app/install
+
+
+###
+### MORI Build TODO(Build needs fixing)
+###
+FROM base AS build_mori
+ARG MORI_BRANCH
+ARG MORI_REPO
+ARG MORI_GPU_ARCHS
+RUN mkdir -p /app/install; \
+    git clone ${MORI_REPO} \
+    && cd mori \
+    && git checkout ${MORI_BRANCH} \
+    && git submodule update --init --recursive \
+    && python3 setup.py bdist_wheel --dist-dir=dist && ls /app/mori/dist/*.whl \
+    && cp /app/mori/dist/*.whl /app/install;
+
+
+###
+### FlashAttention Build
+###
+# Remove  && git submodule update --init \ for CK
+FROM base AS build_fa
+ARG FA_BRANCH
+ARG FA_REPO
+ARG USE_SCCACHE
+RUN mkdir -p /app/install; \
+    git clone ${FA_REPO} \
+    && cd flash-attention \
+    && git checkout ${FA_BRANCH} \
+    && if [ "$USE_SCCACHE" = "1" ]; then \
+    export HIP_CLANG_PATH=/opt/sccache-wrappers \
+    && sccache --show-stats; \
+    fi \
+    && GPU_ARCHS=$(echo ${FA_GPU_ARCHS} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist \
+    && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
+    && cp dist/*.whl /app/install;
+
+
+###
+### AITER Build
+###
+FROM base AS build_aiter
+ARG AITER_BRANCH
+ARG AITER_REPO
+ARG USE_SCCACHE
+RUN git clone --recursive --branch ${AITER_BRANCH} ${AITER_REPO}
+RUN cd aiter \
+    && git submodule update --init --recursive \
+    && pip install -r requirements.txt
+RUN pip install pyyaml && cd aiter \
+    && if [ "$USE_SCCACHE" = "1" ]; then \
+    export HIP_CLANG_PATH=/opt/sccache-wrappers \
+    && sccache --show-stats; \
+    fi \
+    && AITER_USE_SYSTEM_TRITON=1 PREBUILD_KERNELS=${PREBUILD_KERNELS} GPU_ARCHS=${AITER_ROCM_ARCH} python3 setup.py bdist_wheel --dist-dir=dist \
+    && if [ "$USE_SCCACHE" = "1" ]; then sccache --show-stats; fi \
+    && ls /app/aiter/dist/*.whl
+RUN mkdir -p /app/install && cp /app/aiter/dist/*.whl /app/install
+
+
+###
+### Final Build
+###
+
+# Wheel release stage - 
+# only includes dependencies used by wheel release pipeline
+FROM base AS debs_wheel_release
+RUN mkdir /app/debs
+RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
+    if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi
+RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
+    cp /install/*.whl /app/debs
+RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
+    cp /install/*.whl /app/debs
+
+# Full debs stage - includes Mori (used by Docker releases)
+FROM base AS debs
+RUN mkdir /app/debs
+RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
+    if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi
+RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
+    cp /install/*.whl /app/debs
+RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
+    cp /install/*.whl /app/debs
+RUN --mount=type=bind,from=build_mori,src=/app/install/,target=/install \
+    if ls /install/*.whl >/dev/null 2>&1; then cp /install/*.whl /app/debs; fi
+
+FROM base AS final
+RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
+    pip install /install/*.whl
+
+ARG BASE_IMAGE
+ARG ROCM_WHEEL_INDEX
+ARG ROCM_SDK_VERSION
+ARG TORCH_VERSION
+ARG TORCHVISION_VERSION
+ARG TORCHAUDIO_VERSION
+ARG FA_BRANCH
+ARG FA_REPO
+ARG AITER_BRANCH
+ARG AITER_REPO
+ARG MORI_BRANCH
+ARG MORI_REPO
+RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
+    && echo "ROCM_WHEEL_INDEX: ${ROCM_WHEEL_INDEX}" >> /app/versions.txt \
+    && echo "ROCM_SDK_VERSION: ${ROCM_SDK_VERSION}" >> /app/versions.txt \
+    && echo "TORCH_VERSION: ${TORCH_VERSION}" >> /app/versions.txt \
+    && echo "TORCHVISION_VERSION: ${TORCHVISION_VERSION}" >> /app/versions.txt \
+    && echo "TORCHAUDIO_VERSION: ${TORCHAUDIO_VERSION}" >> /app/versions.txt \
+    && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
+    && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt \
+    && echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
+    && echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt \
+    && echo "MORI_BRANCH: ${MORI_BRANCH}" >> /app/versions.txt \
+    && echo "MORI_REPO: ${MORI_REPO}" >> /app/versions.txt
@@ -0,0 +1,729 @@
+# default base image
+ARG REMOTE_VLLM="0"
+ARG COMMON_WORKDIR=/app
+ARG BASE_IMAGE=rocm/vllm-dev:base
+ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base
+# NIC backend for MoRI RDMA support.
+# By default (all), drivers and userspace libraries for all supported NIC types
+# (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime.
+# To install drivers for a single NIC type only, set NIC_BACKEND explicitly:
+#   --build-arg NIC_BACKEND=ainic   # AMD AINIC (Pensando) only
+#   --build-arg NIC_BACKEND=bnxt    # Broadcom Thor-2 only
+#   --build-arg NIC_BACKEND=none    # Install nothing.
+ARG NIC_BACKEND=all
+# AMD AINIC apt repo settings
+# Users can specify a custom version compatible with their host drivers.
+# The default version has been tested with ioinic-dkms=25.11.1.001
+ARG AINIC_VERSION=1.117.3-hydra
+ARG UBUNTU_CODENAME=jammy
+
+# Sccache configuration. Release builds use this today; CI can opt in when a
+# shared S3-compatible cache backend is available.
+ARG USE_SCCACHE
+ARG SCCACHE_DOWNLOAD_URL
+ARG SCCACHE_ENDPOINT
+ARG SCCACHE_BUCKET_NAME=vllm-build-sccache
+ARG SCCACHE_REGION_NAME=us-west-2
+ARG SCCACHE_S3_NO_CREDENTIALS=0
+
+FROM ${BASE_IMAGE} AS base
+
+ARG ARG_PYTORCH_ROCM_ARCH=gfx1250
+ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
+
+# Install build dependencies and utilities
+RUN apt-get update -q -y && apt-get install -q -y \
+    sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \
+    apt-transport-https ca-certificates wget curl \
+    build-essential libnuma-dev ccache mold
+RUN --mount=type=cache,target=/root/.cache/pip \
+    python3 -m pip install --upgrade pip
+# Note: mold is installed but not set as the system default linker because
+# some packages use JIT compilation at runtime with flags mold does not support.
+# Build stages opt in via LDFLAGS="-fuse-ld=mold".
+# Remove sccache only if not using sccache (it exists in base image from Dockerfile.rocm_base)
+ARG USE_SCCACHE
+RUN if [ "$USE_SCCACHE" != "1" ]; then \
+        apt-get purge -y sccache || true; \
+        python3 -m pip uninstall -y sccache || true; \
+        rm -f "$(which sccache)" || true; \
+    fi
+
+# Install UV — download first, then run, so a curl failure is not masked by the pipe
+RUN curl -LsSf --retry 3 --retry-delay 5 https://astral.sh/uv/install.sh -o /tmp/uv-install.sh \
+    && env UV_INSTALL_DIR="/usr/local/bin" sh /tmp/uv-install.sh \
+    && rm -f /tmp/uv-install.sh \
+    && uv --version
+
+# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
+# Reference: https://github.com/astral-sh/uv/pull/1694
+ENV UV_HTTP_TIMEOUT=500
+ENV UV_INDEX_STRATEGY="unsafe-best-match"
+# Use copy mode to avoid hardlink failures with Docker cache mounts
+ENV UV_LINK_MODE=copy
+# python binary fall back for non venv builds
+ENV UV_PYTHON=${VIRTUAL_ENV:-/usr}/bin/python3
+# Expose paths from wheel installation
+ENV PKG_CONFIG_PATH=${ROCM_PATH}/lib/rocm_sysdeps/lib/pkgconfig:${PKG_CONFIG_PATH}
+# ccache directory - persisted across layer rebuilds via cache mounts.
+ENV CCACHE_DIR=/root/.cache/ccache
+ENV CCACHE_COMPILERCHECK=content
+# Empty by default so build steps fall back to $(nproc); CI can override.
+ARG max_jobs
+ENV MAX_JOBS=${max_jobs}
+
+# Install sccache if USE_SCCACHE is enabled (for release builds)
+ARG USE_SCCACHE
+ARG SCCACHE_DOWNLOAD_URL
+ARG SCCACHE_ENDPOINT
+ARG SCCACHE_BUCKET_NAME
+ARG SCCACHE_REGION_NAME
+ARG SCCACHE_S3_NO_CREDENTIALS
+RUN if [ "$USE_SCCACHE" = "1" ]; then \
+        if command -v sccache >/dev/null 2>&1; then \
+            echo "sccache already installed, skipping installation"; \
+            sccache --version; \
+        else \
+            echo "Installing sccache..." \
+            && SCCACHE_ARCH="x86_64" \
+            && SCCACHE_VERSION="v0.8.1" \
+            && SCCACHE_DL_URL="${SCCACHE_DOWNLOAD_URL:-https://github.com/mozilla/sccache/releases/download/${SCCACHE_VERSION}/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl.tar.gz}" \
+            && curl -L -o /tmp/sccache.tar.gz ${SCCACHE_DL_URL} \
+            && tar -xzf /tmp/sccache.tar.gz -C /tmp \
+            && mv /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl/sccache /usr/bin/sccache \
+            && chmod +x /usr/bin/sccache \
+            && rm -rf /tmp/sccache.tar.gz /tmp/sccache-${SCCACHE_VERSION}-${SCCACHE_ARCH}-unknown-linux-musl \
+            && sccache --version; \
+        fi; \
+    fi
+
+# Set sccache environment variables only when USE_SCCACHE=1
+# This prevents S3 config from leaking into images when sccache is not used
+ARG USE_SCCACHE
+ENV SCCACHE_BUCKET=${USE_SCCACHE:+${SCCACHE_BUCKET_NAME}}
+ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
+ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
+ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
+
+ARG COMMON_WORKDIR
+WORKDIR ${COMMON_WORKDIR}
+
+
+# -----------------------
+# vLLM fetch stages
+FROM base AS fetch_vllm_0
+ONBUILD COPY ./ vllm/
+FROM base AS fetch_vllm_1
+ARG VLLM_REPO="https://github.com/ROCm/vllm.git"
+ARG VLLM_BRANCH="455_wip"
+ENV VLLM_REPO=${VLLM_REPO}
+ENV VLLM_BRANCH=${VLLM_BRANCH}
+ONBUILD RUN git clone ${VLLM_REPO} \
+	    && cd vllm \
+	    && git fetch -v --prune -- origin ${VLLM_BRANCH} \
+	    && git checkout FETCH_HEAD \
+        && if [ ${VLLM_REPO} != "https://github.com/vllm-project/vllm.git" ] ; then \
+               git remote add upstream "https://github.com/vllm-project/vllm.git" \
+               && git fetch upstream ; fi
+FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
+
+# -----------------------
+# Rust build stage
+# Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages
+# don't need the rust toolchain or protoc.
+FROM fetch_vllm AS rust-build
+ARG COMMON_WORKDIR
+
+# protoc is used by tonic-build/prost-build.
+RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
+        ca-certificates curl unzip \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY tools/install_protoc.sh /tmp/install_protoc.sh
+RUN /tmp/install_protoc.sh && rm /tmp/install_protoc.sh
+
+# Cap cargo parallelism to avoid exhausting the AMD CI host's open-file limit
+# (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise).
+ENV CARGO_BUILD_JOBS=4
+ENV CARGO_NET_RETRY=10
+ENV RUSTUP_MAX_RETRIES=10
+
+RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
+    cd ${COMMON_WORKDIR}/vllm \
+    && uv pip install --system -r requirements/build/rust.txt
+
+# Build the release binary. Cargo's registry/git caches can be written by
+# concurrent BuildKit jobs on shared workers, so lock those cache mounts while
+# keeping the cache benefit. Do not cache target/, because stale target metadata
+# can outlive source updates across BuildKit cache reuse.
+RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \
+    --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \
+    cd ${COMMON_WORKDIR}/vllm \
+    && bash build_rust.sh \
+    && test -x vllm/vllm-rs
+
+# -----------------------
+# vLLM native build stages
+#
+# csrc-build intentionally copies only files that affect ROCm native extension
+# compilation. That keeps unrelated CI/test/docs edits from invalidating the
+# expensive HIP/C++ build layer.
+FROM base AS csrc-build
+ARG COMMON_WORKDIR
+WORKDIR ${COMMON_WORKDIR}/vllm
+
+COPY requirements/rocm.txt requirements/rocm.txt
+COPY requirements/common.txt requirements/common.txt
+RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
+    uv pip install --system -r requirements/rocm.txt
+
+# pyproject.toml is bind-mounted in the RUN step so metadata-only changes do
+# not invalidate the expensive native build layer.
+COPY setup.py CMakeLists.txt ./
+COPY tools/build_rust.py tools/build_rust.py
+COPY cmake cmake/
+COPY csrc csrc/
+COPY vllm/envs.py vllm/envs.py
+COPY vllm/__init__.py vllm/__init__.py
+
+ENV VLLM_TARGET_DEVICE=rocm
+ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build"
+
+RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \
+    --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
+    export CCACHE_BASEDIR="$PWD" \
+    && echo "=== ccache stats before ROCm native build ===" \
+    && (ccache --show-stats || true) \
+    && (ccache --zero-stats || true) \
+    && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \
+    && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \
+    && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \
+    && test -d dist \
+    && ls dist/*.whl >/dev/null \
+    && echo "=== ccache stats after ROCm native build ===" \
+    && (ccache --show-stats || true)
+
+# Build the full vLLM ROCm wheel by reusing the native extension wheel from
+# csrc-build. This stage still rebuilds for Python/package changes, but skips
+# the expensive HIP/C++ compile when native inputs are unchanged.
+FROM fetch_vllm AS build_vllm
+ARG COMMON_WORKDIR
+ENV VLLM_TARGET_DEVICE=rocm
+
+COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels
+
+# Drop the pre-built Rust artifacts into the source tree. setup.py detects
+# them and ships them as-is, skipping the local Rust build.
+COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
+COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
+
+RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \
+    cd vllm \
+    && uv pip install --system -r requirements/rocm.txt \
+    && export VLLM_USE_PRECOMPILED=1 \
+    && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \
+    && export VLLM_DOCKER_BUILD_CONTEXT=1 \
+    && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \
+    && python3 setup.py bdist_wheel --dist-dir=dist \
+    && test -d dist \
+    && ls dist/*.whl >/dev/null
+FROM scratch AS export_vllm
+ARG COMMON_WORKDIR
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl /
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
+
+# RIXL/UCX build stages
+FROM base AS build_rixl
+ARG RIXL_BRANCH="39be1de8"
+ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
+ARG UCX_BRANCH="bfb51733"
+ARG UCX_REPO="https://github.com/openucx/ucx.git"
+# ENV ROCM_PATH=/opt/rocm -> correct ROCM_PATH is set in base image
+ENV UCX_HOME=/usr/local/ucx
+ENV RIXL_HOME=/usr/local/rixl
+ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
+
+# RIXL build system dependences and RDMA support
+RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
+    libgrpc-dev \
+    libgrpc++-dev \
+    libprotobuf-dev \
+    protobuf-compiler-grpc \
+    libcpprest-dev \
+    libaio-dev \
+    librdmacm1 \
+    librdmacm-dev \
+    libibverbs1 \
+    libibverbs-dev \
+    ibverbs-utils \
+    rdmacm-utils \
+    ibverbs-providers \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system meson auditwheel patchelf tomlkit
+
+RUN --mount=type=cache,target=/root/.cache/ccache \
+    cd /usr/local/src && \
+    git clone ${UCX_REPO} &&  \
+    cd ucx  && \
+    git checkout ${UCX_BRANCH} && \
+    ./autogen.sh && \
+    mkdir build && cd build && \
+    CC="ccache gcc" CXX="ccache g++" \
+    ../configure \
+        --prefix=/usr/local/ucx \
+        --enable-shared \
+        --disable-static \
+        --disable-doxygen-doc \
+        --enable-optimizations \
+        --enable-devel-headers \
+        --with-rocm=${ROCM_PATH} \
+        --with-verbs \
+        --with-dm \
+        --enable-mt && \
+    make -j$(nproc) && \
+    make install
+
+ENV PATH=/usr/local/ucx/bin:$PATH
+ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
+
+RUN --mount=type=cache,target=/root/.cache/ccache \
+    git clone ${RIXL_REPO} /opt/rixl && \
+    cd /opt/rixl && \
+    git checkout ${RIXL_BRANCH} && \
+    CC="ccache gcc" CXX="ccache g++" \
+    meson setup build --prefix=${RIXL_HOME} \
+                     --force-fallback-for=abseil-cpp \
+                     -Ducx_path=${UCX_HOME} \
+                     -Drocm_path=${ROCM_PATH} && \
+    cd build && \
+    ninja -j$(nproc) && \
+    ninja install
+
+# Generate RIXL wheel
+# Exclude libcore and libpull from auditwheel: transitive dependencies
+# that are not shipped in the wheel and vary across base images.
+RUN cd /opt/rixl && \
+    sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
+        contrib/build-wheel.sh && \
+    # The wheel build re-runs meson via meson-python; force the bundled abseil
+    sed -i 's|setup = \["-Dinstall_headers=false"\]|setup = ["-Dinstall_headers=false", "--force-fallback-for=abseil-cpp"]|' \
+        pyproject.toml && \
+    grep -q 'force-fallback-for' pyproject.toml && \
+    mkdir -p /app/install && \
+    _ucx_install_dir=${UCX_HOME} \
+    ./contrib/build-wheel.sh \
+        --output-dir /app/install \
+        --rocm-dir ${ROCM_PATH} \
+        --ucx-plugins-dir ${UCX_HOME}/lib/ucx \
+        --nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
+
+# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
+# invalidate the slow ROCShmem build.
+FROM base AS build_rocshmem
+ARG ROCSHMEM_BRANCH="f0acb0c6"
+ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
+# DeepEP only supports gfx942 and gfx950; build ROCShmem for the same set so
+# it can be linked against DeepEP without arch mismatches.
+ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
+# ENV ROCM_PATH=/opt/rocm -> Correct rocm_path is set in base image
+ENV ROCSHMEM_DIR=/opt/rocshmem
+
+RUN --mount=type=cache,target=/root/.cache/ccache \
+    git clone --no-checkout --filter=blob:none ${ROCSHMEM_REPO} \
+ && cd rocm-systems \
+ && git sparse-checkout set --cone projects/rocshmem \
+ && git checkout ${ROCSHMEM_BRANCH} \
+ && mkdir -p projects/rocshmem/build \
+ && cd projects/rocshmem/build \
+ && CC="ccache gcc" CXX="ccache g++" INSTALL_PREFIX=${ROCSHMEM_DIR} \
+    bash ../scripts/build_configs/all_backends \
+      -DROCM_PATH=${ROCM_PATH} \
+      -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
+      -DUSE_EXTERNAL_MPI=OFF
+
+# DeepEP build stage - depends on ROCShmem, builds the HIP kernel wheel.
+FROM build_rocshmem AS build_deepep
+ARG DEEPEP_BRANCH="a9ea9774"
+ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
+ARG DEEPEP_NIC="cx7"
+
+# Build DeepEP wheel. DeepEP looks for rocshmem at ROCSHMEM_DIR.
+# DeepEP only supports gfx942 and gfx950, so avoid gfx90a in the default list.
+RUN --mount=type=cache,target=/root/.cache/ccache \
+    export PYTORCH_ROCM_ARCH="gfx942;gfx950" \
+ && git clone ${DEEPEP_REPO} \
+ && cd DeepEP \
+ && git checkout ${DEEPEP_BRANCH} \
+ && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
+
+# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
+# not force users to rebuild the long-lived Dockerfile.rocm_base image.
+FROM base AS mori_base
+ARG NIC_BACKEND
+ARG AINIC_VERSION
+ARG UBUNTU_CODENAME
+RUN /bin/bash -lc 'set -euo pipefail; \
+ \
+ install_ainic() { \
+   apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https; \
+   rm -rf /var/lib/apt/lists/*; \
+   mkdir -p /etc/apt/keyrings; \
+   curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \
+   echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \
+     > /etc/apt/sources.list.d/amdainic.list; \
+   apt-get update && apt-get install -y --no-install-recommends \
+     libionic-dev \
+     ionic-common \
+   ; \
+   rm -rf /var/lib/apt/lists/*; \
+ }; \
+ \
+ # NOTE: requires FW 235.2.86.0 and kernel drivers on the host: \
+ #   bnxt-en-dkms=1.10.3.235.2.86.0 bnxt-re-dkms=235.2.86.0 (from packages.broadcom.com PPA) \
+ install_bnxt() { \
+   install -m 0755 -d /etc/apt/keyrings; \
+   curl -fsSL https://packages.broadcom.com/artifactory/api/security/keypair/PackagesKey/public \
+     -o /etc/apt/keyrings/broadcom-nic.asc; \
+   chmod a+r /etc/apt/keyrings/broadcom-nic.asc; \
+   echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/broadcom-nic.asc] https://packages.broadcom.com/artifactory/ethernet-nic-debian-public jammy main" \
+     > /etc/apt/sources.list.d/broadcom-nic.list; \
+   apt-get update && apt-get install -y --no-install-recommends \
+     bnxt-rocelib=235.2.86.0 \
+   ; \
+   cp -a /usr/local/lib/x86_64-linux-gnu/libbnxt_re* /usr/local/lib/; \
+   ldconfig; \
+   rm -rf /var/lib/apt/lists/*; \
+ }; \
+ \
+ echo "[MORI] Install MoRI proxy deps"; \
+ pip install --quiet --ignore-installed blinker && \
+ pip install --quiet quart msgpack aiohttp pyzmq; \
+ echo "[MORI] NIC_BACKEND=${NIC_BACKEND}"; \
+ \
+ # NIC backend deps — mori auto-detects NIC at runtime (MORI_DEVICE_NIC env var override). \
+ # Only vendor packages are installed here for dlopen; no compile-time flags needed. \
+ case "${NIC_BACKEND}" in \
+   none)  ;; \
+   all)   install_ainic; install_bnxt ;; \
+   ainic) install_ainic ;; \
+   bnxt)  install_bnxt ;; \
+   *)     echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic, bnxt, all"; exit 2 ;; \
+ esac'
+
+# -----------------------
+# vLLM wheel release build stage (for building distributable wheels)
+# This stage pins dependencies to custom ROCm wheel versions and handles version detection
+FROM fetch_vllm AS build_vllm_wheel_release
+
+ARG COMMON_WORKDIR
+
+# Drop the pre-built Rust artifacts into the source tree. setup.py detects
+# them and ships them as-is, skipping the local Rust build.
+COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vllm/vllm/vllm-rs
+COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/
+
+# Create /install directory for custom wheels
+RUN mkdir -p /install
+
+# Copy custom ROCm wheels from docker/context if they exist
+# COPY ensures Docker cache is invalidated when wheels change
+# .keep file ensures directory always exists for COPY to work
+COPY docker/context/base-wheels/ /tmp/base-wheels/
+# This is how we know if we are building for a wheel release or not.
+# If there are not wheels found there, we are not building for a wheel release. 
+# So we exit with an error. To skip this stage.
+RUN if [ -n "$(ls /tmp/base-wheels/*.whl 2>/dev/null)" ]; then \
+        echo "Found custom wheels - copying to /install"; \
+        cp /tmp/base-wheels/*.whl /install/ && \
+        echo "Copied custom wheels:"; \
+        ls -lh /install/; \
+    else \
+        echo "ERROR: No custom wheels found in docker/context/base-wheels/"; \
+        echo "Wheel releases require pre-built ROCm wheels."; \
+        exit 1; \
+    fi
+
+# GIT_REPO_CHECK: Verify repo is clean and tags are available (for release builds)
+# This matches CUDA's Dockerfile behavior for proper version detection via setuptools_scm
+ARG GIT_REPO_CHECK=0
+RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
+        echo "Running repository checks..."; \
+        cd vllm && bash tools/check_repo.sh; \
+    fi
+
+# Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt)
+# This ensures setuptools_scm sees clean repo state for version detection
+RUN --mount=type=bind,source=.git,target=vllm/.git \
+    --mount=type=cache,target=/root/.cache/uv \
+    cd vllm \
+    && uv pip install --system setuptools_scm regex \
+    && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \
+    && echo "Detected vLLM version: ${VLLM_VERSION}" \
+    && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
+
+# Fail if git-based package dependencies are found in requirements files
+# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
+# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
+RUN echo "Checking for git-based packages in requirements files..." \
+    && echo "Checking common.txt for git-based packages:" \
+    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
+         echo "ERROR: Git-based packages found in common.txt:"; \
+         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
+         echo "Please publish these packages to PyPI instead of using git dependencies."; \
+         exit 1; \
+       else \
+         echo "  ✓ No git-based packages found in common.txt"; \
+       fi \
+    && echo "Checking rocm.txt for git-based packages:" \
+    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
+         echo "ERROR: Git-based packages found in rocm.txt:"; \
+         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
+         echo "Please publish these packages to PyPI instead of using git dependencies."; \
+         exit 1; \
+       else \
+         echo "  ✓ No git-based packages found in rocm.txt"; \
+       fi \
+    && echo "All requirements files are clean - no git-based packages found"
+
+# Pin vLLM dependencies to exact versions of custom ROCm wheels
+# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
+COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
+RUN echo "Pinning vLLM dependencies to custom wheel versions..." \
+    && python3 /tmp/pin_rocm_dependencies.py /install ${COMMON_WORKDIR}/vllm/requirements/rocm.txt
+
+# Install dependencies using custom wheels from /install
+RUN --mount=type=cache,target=/root/.cache/uv \
+    cd vllm \
+    && echo "Building vLLM with custom wheels from /install" \
+    && uv pip install --system --find-links /install -r requirements/rocm.txt
+
+# Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt
+# (setup.py auto-detects ccache/sccache in PATH)
+RUN --mount=type=bind,source=.git,target=vllm/.git \
+    --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \
+    cd vllm \
+    && export CCACHE_BASEDIR="$PWD" \
+    && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \
+    && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \
+    && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist
+
+FROM scratch AS export_vllm_wheel_release
+ARG COMMON_WORKDIR
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/dist/*.whl /
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/requirements /requirements
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tests /tests
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/examples /examples
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/tools/install_torchcodec_rocm.sh /tools/install_torchcodec_rocm.sh
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/Dockerfile.rocm /docker/
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
+COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
+
+# -----------------------
+# CI base image (Tier 1) - stable, rarely changing CI dependencies.
+# Per-PR test builds pull this as CI_BASE_IMAGE so the test stage only layers
+# in the vLLM artifacts for the current commit.
+FROM mori_base AS ci_base
+ARG COMMON_WORKDIR
+
+# Update rdma-core to support latest rocshmem.
+ARG DEEPEP_NIC
+RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
+    git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
+    cd /tmp/rdma-core && \
+    mkdir -p build && cd build && \
+    cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
+    ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
+fi
+
+# Install RIXL + DeepEP wheels.
+RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
+    --mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
+    uv pip install --system /rixl_install/*.whl /deep_install/*.whl
+
+# Copy ROCShmem runtime libraries.
+COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
+
+# RDMA userspace libraries plus FFmpeg dev libs needed by torchcodec.
+RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
+    librdmacm1 \
+    libibverbs1 \
+    ibverbs-providers \
+    ibverbs-utils \
+    pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
+    libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install torchcodec from source for ROCm/torch ABI compatibility.
+COPY tools/install_torchcodec_rocm.sh /tmp/install_torchcodec.sh
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=cache,target=/root/.cache/pip \
+    --mount=type=cache,target=/root/.cache/torchcodec-wheels \
+    bash /tmp/install_torchcodec.sh \
+    && rm /tmp/install_torchcodec.sh \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+# Pre-install shared ROCm runtime dependencies.
+COPY requirements/common.txt requirements/rocm.txt /tmp/ci-base-requirements/
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system -r /tmp/ci-base-requirements/rocm.txt \
+    && rm -rf /tmp/ci-base-requirements
+
+# Enable fast and less brittle model downloads in tests.
+ENV HF_XET_HIGH_PERFORMANCE=1
+ENV HF_HUB_DOWNLOAD_TIMEOUT=60
+
+# Pre-install vLLM test dependencies.
+COPY requirements/test/rocm.txt /tmp/rocm-test-reqs.txt
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system -r /tmp/rocm-test-reqs.txt
+
+# Rebuild fastsafetensors from source so its C++ extension is compiled with
+# USE_ROCM and can detect libamdhip64.so at runtime.
+RUN --mount=type=cache,target=/root/.cache/pip \
+    FASTSAFETENSORS_REQ="$(grep -E '^fastsafetensors(==| @ )' /tmp/rocm-test-reqs.txt | head -1)" \
+    && test -n "${FASTSAFETENSORS_REQ}" \
+    && python3 -m pip install --force-reinstall --no-deps \
+        --no-binary fastsafetensors "${FASTSAFETENSORS_REQ}" \
+    && rm /tmp/rocm-test-reqs.txt
+
+# Set MIOPEN ENVS to resolve performance regressions in MIOpen 3D convolution kernel.
+# See: https://github.com/pytorch/pytorch/issues/169857
+ENV MIOPEN_DEBUG_CONV_DIRECT=0
+ENV MIOPEN_DEBUG_CONV_GEMM=0
+
+# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
+# See: https://github.com/ROCm/rocm-libraries/issues/6266
+ENV HSA_ENABLE_IPC_MODE_LEGACY=1
+
+# ROCm profiler limits workaround.
+RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
+ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
+
+# Install vllm_test_utils in ci_base for ci_base + wheel parity.
+COPY tests/vllm_test_utils /tmp/vllm_test_utils
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system /tmp/vllm_test_utils \
+    && rm -rf /tmp/vllm_test_utils
+
+# -----------------------
+# Test vLLM image (Tier 2) - vLLM-only layer on top of ci_base.
+FROM ${CI_BASE_IMAGE} AS test
+ARG COMMON_WORKDIR
+
+# Install the vLLM wheel (--no-deps: all deps already in ci_base).
+RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
+    --mount=type=cache,target=/root/.cache/uv \
+    cd /install \
+    && uv pip install --system --no-deps *.whl
+
+# Store the vLLM wheel in the image for python-only install tests.
+COPY --from=export_vllm /*.whl /opt/vllm-wheels/
+
+WORKDIR /vllm-workspace
+COPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace
+
+# Copy in the v1 package (for python-only install test group).
+COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-packages/vllm/v1
+
+# Hide source under src/ so it won't shadow the installed package in tests.
+RUN mkdir src && mv vllm src/vllm
+
+# -----------------------
+# Final vLLM image
+FROM mori_base AS final
+
+RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
+
+# Clean up sccache from release image (not needed at runtime)
+# This removes the binary and wrappers that may have been installed during build
+RUN rm -f /usr/bin/sccache || true \
+    && rm -rf /opt/sccache-wrappers || true
+
+# Unset sccache environment variables for the release image
+# This prevents S3 bucket config from leaking into production images
+ENV SCCACHE_BUCKET=
+ENV SCCACHE_REGION=
+ENV SCCACHE_ENDPOINT=
+ENV SCCACHE_S3_NO_CREDENTIALS=
+ENV SCCACHE_IDLE_TIMEOUT=
+
+# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt.
+# Manually remove it so that later steps of numpy upgrade can continue
+RUN case "$(which python3)" in \
+        *"/opt/conda/envs/py_3.9"*) \
+            rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \
+        *) ;; esac
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --system --upgrade huggingface-hub[cli]
+
+# Install vLLM using uv (inherited from base stage)
+# Note: No -U flag to avoid upgrading PyTorch ROCm to CUDA version
+RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
+    --mount=type=cache,target=/root/.cache/uv \
+    cd /install \
+    && uv pip install --system -r requirements/rocm.txt \
+    && pip uninstall -y vllm \
+    && uv pip install --system *.whl
+
+# Install RIXL wheel
+RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
+    uv pip install --system /rixl_install/*.whl
+
+ARG COMMON_WORKDIR
+ARG BASE_IMAGE
+ARG NIC_BACKEND
+ARG AINIC_VERSION
+
+# Copy over the benchmark scripts as well
+COPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks
+COPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples
+COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
+
+# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
+# See: https://github.com/ROCm/rocm-libraries/issues/6266
+ENV HSA_ENABLE_IPC_MODE_LEGACY=1
+
+ENV TOKENIZERS_PARALLELISM=false
+
+# ENV that can improve safe tensor loading, and end-to-end time
+ENV SAFETENSORS_FAST_GPU=1
+
+# Performance environment variable.
+ENV HIP_FORCE_DEV_KERNARG=1
+
+# Workaround for ROCm profiler limits
+RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
+ENV KINETO_CONFIG="${COMMON_WORKDIR}/libkineto.conf"
+RUN echo "VLLM_BASE_IMAGE=${BASE_IMAGE}" >> ${COMMON_WORKDIR}/versions.txt \
+    && echo "MORI_NIC_BACKEND=${NIC_BACKEND}" >> ${COMMON_WORKDIR}/versions.txt \
+    && echo "AINIC_VERSION=${AINIC_VERSION}" >> ${COMMON_WORKDIR}/versions.txt
+
+# Download bench scripts for gfx1250
+RUN curl -k -O https://raw.githubusercontent.com/ROCm/vllm/refs/heads/gfx1250_bench/vllm_smoketest.sh
+
+
+### Install triton from upstream for AITER Deps TODO: (JPVILLAM) If possible to get this on whls it would be better
+RUN pip3 uninstall -y triton && \
+    git clone https://github.com/triton-lang/triton.git && \
+    cd triton && \
+    git checkout c517f38c && \
+    TRITON_APPEND_CMAKE_ARGS="-DCMAKE_FIND_USE_CMAKE_ENVIRONMENT_PATH=FALSE" pip3 install .
+
+CMD ["/bin/bash"]
+
+#Set entrypoint for vllm-openai official images
+FROM final AS vllm-openai
+ENTRYPOINT ["vllm", "serve"]
@@ -31,7 +31,7 @@ def _get_attention_backend_params() -> list[str | None]:
              falls back to ROCM_AITER_UNIFIED_ATTN or TRITON_ATTN for
              cross-attention since ROCM_ATTN doesn't support ENCODER_DECODER)
      - TRITON_ATTN: always available on ROCm
-      - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950
+      - ROCM_AITER_UNIFIED_ATTN: only on gfx942/gfx950/gfx1250

    On non-ROCm platforms, we just run with the default backend.
    """
@@ -40,9 +40,9 @@ def _get_attention_backend_params() -> list[str | None]:

        if current_platform.is_rocm():
            backends: list[str | None] = [None, "TRITON_ATTN"]
-            from vllm.platforms.rocm import _ON_MI3XX
+            from vllm.platforms.rocm import get_cdna_version

-            if _ON_MI3XX:
+            if get_cdna_version() > 2:
                backends.append("ROCM_AITER_UNIFIED_ATTN")
            return backends
    except Exception:
@@ -37,13 +37,13 @@ def _get_rocm_attention_config(model_name):

    if "whisper" in model_name.lower():
        try:
-            from vllm.platforms.rocm import _ON_MI3XX
+            from vllm.platforms.rocm import get_cdna_version

-            if _ON_MI3XX:
+            if get_cdna_version() > 2:
                return {"backend": "ROCM_AITER_UNIFIED_ATTN"}
        except ImportError:
            logger.warning(
-                "Could not import _ON_MI3XX from rocm platform, "
+                "Could not check cdna version from rocm platform, "
                "falling back to TRITON_ATTN for Whisper."
            )
        return {"backend": "TRITON_ATTN"}
@@ -15,15 +15,15 @@ from tests.kernels.attention.test_triton_unified_attention import ref_paged_attn
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed

-_SKIP_NON_MI3XX = True
+_SKIP_NON_CDNA_2_PLUS = True
 if current_platform.is_rocm():
-    from vllm.platforms.rocm import on_mi3xx
+    from vllm.platforms.rocm import get_cdna_version

-    _SKIP_NON_MI3XX = not on_mi3xx()
+    _SKIP_NON_CDNA_2_PLUS = get_cdna_version() < 2

 pytestmark = [
    pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"),
-    pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"),
+    pytest.mark.skipif(_SKIP_NON_CDNA_2_PLUS, reason="CDNA 2+ ROCm only"),
 ]

 NUM_Q_HEADS = 8
@@ -4,6 +4,8 @@
 Test modular OAI Triton MoE
 """

+from __future__ import annotations
+
 import pytest
 import torch

@@ -48,6 +50,35 @@ MNK = [
 ]


+def deepseek_v4_flash_moe_topology():
+    """MoE sizes for the DeepSeek-V4-Flash`.
+    Default weights from: ``deepseek-ai/DeepSeek-V4-Flash`.
+    ``moe_intermediate_size``, ``n_routed_experts``, and ``num_experts_per_tok``.
+    """
+    defaults = {
+        "hidden_size": 4096,
+        "moe_intermediate_size": 2048,
+        "n_routed_experts": 256,
+        "num_experts_per_tok": 6,
+    }
+
+    return defaults
+
+
+def scaled_deepseek_v4_flash_problem(
+    *,
+    dim_scale: int = 8,
+    expert_scale: int = 8,
+):
+    """Smaller K/N/E for kernel tests; keeps production top_k and K:N ratio (~2:1)."""
+    t = deepseek_v4_flash_moe_topology()
+    k = max(128, t["hidden_size"] // dim_scale)
+    n = max(64, t["moe_intermediate_size"] // dim_scale)
+    num_experts = max(t["num_experts_per_tok"], t["n_routed_experts"] // expert_scale)
+    topk = t["num_experts_per_tok"]
+    return k, n, num_experts, topk
+
+
 def unshuffle_weight(w: torch.Tensor):
    first = w[..., ::2]
    second = w[..., 1::2]
@@ -160,8 +191,8 @@ def oai_triton_moe_impl(
    x: torch.Tensor,
    w1: torch.Tensor,
    w2: torch.Tensor,
-    w1_scale: "PrecisionConfig",
-    w2_scale: "PrecisionConfig",
+    w1_scale: PrecisionConfig,
+    w2_scale: PrecisionConfig,
    w1_bias: torch.Tensor | None,
    w2_bias: torch.Tensor | None,
    num_experts: int,
@@ -261,3 +292,93 @@ def test_oai_triton_moe(
        )

    assert_close(ref=out_ref, tri=out, maxtol=0.025, rmstol=0.005)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
+)
+def test_unfused_oai_triton_experts_apply_direct_deepseek_v4_topology(workspace_init):
+    """Exercise ``UnfusedOAITritonExperts.apply`` with explicit workspaces.
+
+    Same MoE topology as ``launch_dsv4.sh`` / DeepSeek-V4-Flash ``config.json``,
+    with linear dimensions and expert count scaled down for test GPU memory.
+    """
+    wait_for_gpu_memory_to_clear(devices=[0], threshold_ratio=0.1)
+    set_random_seed(0)
+
+    k, n, num_experts, topk = scaled_deepseek_v4_flash_problem()
+    m = 7
+    dtype = torch.bfloat16
+
+    (
+        w1,
+        w2,
+        w1_bias,
+        w2_bias,
+        w1_tri,
+        w2_tri,
+        w1_bias_tri,
+        w2_bias_tri,
+        w1_precision_config,
+        w2_precision_config,
+    ) = make_weights(dtype, k, n, num_experts)
+
+    x = torch.randn((m, k), dtype=dtype, device="cuda")
+    router_logits = torch.randn(m, num_experts, device="cuda", dtype=dtype)
+    topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1, sorted=True)
+    topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
+
+    quant_config = mxfp4_w4a16_moe_quant_config(
+        w1_bias=w1_bias_tri,
+        w2_bias=w2_bias_tri,
+        w1_scale=w1_precision_config,
+        w2_scale=w2_precision_config,
+    )
+    moe_config = make_dummy_moe_config(
+        num_experts=num_experts,
+        experts_per_token=topk,
+        hidden_dim=k,
+        intermediate_size=n,
+    )
+    experts = UnfusedOAITritonExperts(moe_config, quant_config)
+
+    if not UnfusedOAITritonExperts._supports_current_device():
+        pytest.skip("UnfusedOAITritonExperts does not support this device")
+
+    _, _, N, K, top_k = experts.moe_problem_size(x, w1_tri, w2_tri, topk_ids)
+    assert top_k == topk
+    ws13_shape, ws2_shape, out_shape = experts.workspace_shapes(
+        m,
+        N,
+        K,
+        topk,
+        num_experts,
+        num_experts,
+        None,
+        MoEActivation.SWIGLUOAI,
+    )
+    workspace13 = torch.empty(ws13_shape, dtype=dtype, device="cuda")
+    workspace2 = torch.empty(ws2_shape, dtype=dtype, device="cuda")
+    output = torch.empty(out_shape, dtype=dtype, device="cuda")
+
+    with set_current_vllm_config(VllmConfig()):
+        out_ref = torch_moe_impl(x, w1, w2, w1_bias, w2_bias, topk_weights, topk_ids)
+        experts.apply(
+            output=output,
+            hidden_states=x,
+            w1=w1_tri,
+            w2=w2_tri,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=MoEActivation.SWIGLUOAI,
+            global_num_experts=num_experts,
+            expert_map=None,
+            a1q_scale=None,
+            a2_scale=None,
+            workspace13=workspace13,
+            workspace2=workspace2,
+            expert_tokens_meta=None,
+            apply_router_weight_on_input=False,
+        )
+
+    assert_close(ref=out_ref, tri=output, maxtol=0.025, rmstol=0.005)
@@ -29,13 +29,13 @@ from vllm.platforms import current_platform


 def _is_aiter_capable() -> bool:
-    """Check if the platform supports AITER (gfx942/gfx950)."""
+    """Check if the platform supports AITER (gfx942/gfx950/gfx1250)."""
    if not current_platform.is_rocm():
        return False
    try:
-        from vllm.platforms.rocm import _ON_MI3XX
+        from vllm.platforms.rocm import get_cdna_version

-        return _ON_MI3XX
+        return get_cdna_version() > 2
    except ImportError:
        return False

@@ -26,6 +26,7 @@ def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
    monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
    monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
    monkeypatch.setattr(utils, "num_compute_units", lambda: 120)

    wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
@@ -52,6 +53,7 @@ def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch):
    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
    monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
    monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: False)
    monkeypatch.setattr(utils, "num_compute_units", lambda: 120)

    wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
@@ -76,6 +78,7 @@ def test_rocm_unquantized_gemm_gfx950_wvsplitkrc_path(monkeypatch):
    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False)
    monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
    monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1250", lambda: True)
    monkeypatch.setattr(utils, "num_compute_units", lambda: 120)

    wvsplitkrc_mock = MagicMock(side_effect=lambda x_view, w, _, __: x_view @ w.t())
@@ -45,9 +45,10 @@ models: dict[str, str | None] = {
 def granite_speech_attention_config():
    """Return attention config for Granite Speech tests on ROCm."""
    if current_platform.is_rocm():
-        from vllm.platforms.rocm import on_mi3xx
+        from vllm.platforms.rocm import get_cdna_version

-        if on_mi3xx():
+        # -1 (unknown arch) is truthy; gate on CDNA3+ like other call sites.
+        if get_cdna_version() > 2:
            return {"backend": "ROCM_AITER_FA"}
        return {"backend": "TRITON_ATTN"}
    return None
@@ -21,11 +21,11 @@ from ...utils import compare_two_settings, multi_gpu_test
 from ..utils import check_embeddings_close, check_logprobs_close

 if current_platform.is_rocm():
-    from vllm.platforms.rocm import on_gfx9
+    from vllm.platforms.rocm import on_cdna

    pytestmark = pytest.mark.skipif(
-        on_gfx9(),
-        reason="bitsandbytes not supported on gfx9 (warp size 64 limitation)",
+        on_cdna(),
+        reason="bitsandbytes not supported on CDNA (warp size 64 limitation)",
    )

 models_4bit_to_test = [
@@ -28,16 +28,9 @@ def mock_vllm_config():


@pytest.fixture
-def mock_on_gfx9():
-    """Mock gfx9 arch detection to return True."""
-    with patch("vllm.platforms.rocm.on_gfx9", return_value=True):
-        yield
-
-
-@pytest.fixture
-def mock_on_mi3xx():
-    """Mock mi3xx arch detection to return True."""
-    with patch("vllm.platforms.rocm.on_mi3xx", return_value=True):
+def mock_get_cdna_version():
+    """Mock cdna version arch detection to return True."""
+    with patch("vllm.platforms.rocm.get_cdna_version", return_value=3):
        yield


@@ -111,8 +104,7 @@ def test_standard_attention_backend_selection(
    selected_backend,
    expected_backend_path,
    mock_vllm_config,
-    mock_on_gfx9,
-    mock_on_mi3xx,
+    mock_get_cdna_version,
    monkeypatch,
 ):
    """Test standard attention backend selection with various configurations."""
@@ -305,12 +297,12 @@ def test_mla_backend_selection(


 def test_aiter_fa_requires_mi3xx(mock_vllm_config):
-    """Test that ROCM_AITER_FA requires mi3xx architecture."""
+    """Test that ROCM_AITER_FA requires CDNA3+ architecture."""
    from vllm.platforms.rocm import RocmPlatform

-    # Mock on_mi3xx to return False (used by supports_compute_capability)
+    # Mock cdna version to return 1 (used by supports_compute_capability)
    with (
-        patch("vllm.platforms.rocm.on_mi3xx", return_value=False),
+        patch("vllm.platforms.rocm.get_cdna_version", return_value=1),
        pytest.raises(
            ValueError,
            match="compute capability not supported",
@@ -52,7 +52,7 @@ IS_AITER_FOUND = is_aiter_found()
 def is_aiter_found_and_supported() -> bool:
    """Check if AITER library is available and platform supports it.

-    Checks: platform (ROCm), device arch (gfx9), and library existence.
+    Checks: platform (ROCm), device arch is CDNA 3 or better, and library existence.
    Does NOT check environment variables - that's handled by rocm_aiter_ops.is_enabled().

    This function determines if aiter CAN be used, not if it SHOULD be used.
@@ -66,9 +66,9 @@ def is_aiter_found_and_supported() -> bool:
    VLLM_ROCM_USE_AITER=0, while preventing unwanted JIT warnings for auto-discovery.
    """
    if current_platform.is_rocm() and IS_AITER_FOUND:
-        from vllm.platforms.rocm import on_mi3xx
+        from vllm.platforms.rocm import get_cdna_version

-        return on_mi3xx()
+        return get_cdna_version() > 2
    return False


@@ -1798,16 +1798,21 @@ class rocm_aiter_ops:
    @classmethod
    @if_aiter_supported
    def is_fp4bmm_enabled(cls) -> bool:
-        from vllm.platforms.rocm import on_gfx950
+        from vllm.platforms.rocm import get_cdna_version

-        return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and on_gfx950()
+        # TODO GFX1250: Enable for cdna 4+ when aiter supports batched_gemm_a16wfp4 on gfx1250
+        return cls._AITER_ENABLED and cls._FP4BMM_ENABLED and get_cdna_version() == 4

    @classmethod
    @if_aiter_supported
    def is_linear_hipbmm_enabled(cls) -> bool:
-        from vllm.platforms.rocm import on_mi3xx
+        from vllm.platforms.rocm import get_cdna_version

-        return cls.is_linear_enabled() and on_mi3xx() and cls._LINEAR_HIPBMM_ENABLED
+        return (
+            cls.is_linear_enabled()
+            and (get_cdna_version() > 2)
+            and cls._LINEAR_HIPBMM_ENABLED
+        )

    @classmethod
    @if_aiter_supported
@@ -374,9 +374,14 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
        super().__init__(config)
        n, k = config.weight_shape

-        self.use_triton = (
-            not current_platform.is_fp8_fnuz()
-            and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k)
+        _on_gfx1250 = False
+        if current_platform.is_rocm():
+            from vllm.platforms.rocm import on_gfx1250
+
+            _on_gfx1250 = on_gfx1250()
+
+        self.use_triton = not current_platform.is_fp8_fnuz() and (
+            rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k) or _on_gfx1250
        )

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -104,10 +104,10 @@ class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
        if not current_platform.is_rocm():
            return False, "requires ROCm."

-        from vllm.platforms.rocm import on_mi3xx
+        from vllm.platforms.rocm import get_cdna_version

-        if not on_mi3xx():
-            return False, "requires MI3xx."
+        if get_cdna_version() <= 2:
+            return False, "requires CDNA3+"

        if compute_capability is not None and compute_capability < 94:
            return False, "requires compute capability 94 and above."
@@ -79,10 +79,15 @@ class ROCmFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
        if not current_platform.is_rocm():
            return False, "requires ROCm."

-        from vllm.platforms.rocm import on_gfx12x, on_mi3xx
+        from vllm.platforms.rocm import get_cdna_version, on_gfx12x, on_gfx1250

-        if not (on_mi3xx() or on_gfx12x()):
-            return False, "requires MI3xx or gfx12x"
+        # wvSplitKQ (skinny GEMM) is excluded from the gfx1250 build.
+        if on_gfx1250():
+            return False, "wvSplitKQ (skinny GEMM) is not built on gfx1250"
+
+        # Restore RDNA4 (gfx12x) dropped by the get_cdna_version()>2 refactor.
+        if get_cdna_version() <= 2 and not on_gfx12x():
+            return False, "requires CDNA3+ (gfx942/gfx950) or RDNA4 (gfx12x)"

        if not envs.VLLM_ROCM_USE_SKINNY_GEMM:
            return False, "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled."
@@ -20,7 +20,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

 __all__ = [
    "AiterW4A8ExpertsMonolithic",
+    "AiterW4A16ExpertsMonolithic",
    "aiter_triton_kernel_w4a8_moe_forward",
+    "aiter_triton_kernel_w4a16_moe_forward",
 ]


@@ -46,11 +48,27 @@ def aiter_triton_kernel_w4a8_moe_forward(
        and quant_config.use_mxfp4_w4a8
        and rocm_aiter_ops.is_enabled()
    )
-    from aiter.ops.triton.moe_routing.routing import routing as aiter_routing
+    from vllm.platforms.rocm import on_gfx1250
+
+    try:
+        from aiter.ops.triton.moe.moe_routing import routing as _routing_mod
+    except ImportError:
+        from aiter.ops.triton.moe_routing import routing as _routing_mod
+
+    if on_gfx1250():
+        _routing_mod.is_tdm_avail = lambda: False
+    aiter_routing = _routing_mod.routing

    routing_data, gather_idx, scatter_idx = aiter_routing(
        gating_output, topk, sm_first=not renormalize
    )
+
+    # gfx1250: aiter's in-kernel gather is numerically broken
+    if on_gfx1250():
+        gather_src = gather_idx.to(torch.long) // topk
+        hidden_states = hidden_states[gather_src]
+        gather_idx = None
+
    return triton_kernel_fused_mxfp4_w4a8_experts(
        None,
        hidden_states,
@@ -199,12 +217,11 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):

    @staticmethod
    def _supports_current_device() -> bool:
-        # Requires AITER and GFX950
        if not rocm_aiter_ops.is_enabled():
            return False
-        from vllm.platforms.rocm import on_gfx950
+        from vllm.platforms.rocm import on_gfx950, on_gfx1250

-        return on_gfx950()
+        return on_gfx950() or on_gfx1250()

    @staticmethod
    def _supports_no_act_and_mul() -> bool:
@@ -293,3 +310,358 @@ class AiterW4A8ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
            unpadded_N_w2=self.moe_config.hidden_dim_unpadded,
            unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded,
        )
+
+
+def _aiter_raw(t):
+    if t is None or isinstance(t, torch.Tensor):
+        return t
+    return t.storage.data if hasattr(t, "storage") else t
+
+
+def _aiter_w4a16_silu_via_a8w4(
+    hidden_states: torch.Tensor,
+    w1_data,
+    w2_data,
+    w1_wscale,
+    w2_wscale,
+    w1_bias,
+    w2_bias,
+    routing_data,
+    gather_idx,
+    scatter_idx,
+    gammas,
+    apply_router_weight_on_input: bool,
+    swiglu_limit: float,
+    unpadded_N_w1,
+    unpadded_K_w1,
+    unpadded_N_w2,
+    unpadded_K_w2,
+) -> torch.Tensor:
+    """
+    MXFP4 w4a16 MoE with a SILU (concatenated ``[gate | up]``) activation.
+    """
+    from aiter.ops.triton.fusions.fused_clamp_act_mul import fused_clamp_act_mul
+    from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
+    from aiter.ops.triton.quant import dynamic_mxfp8_quant
+
+    from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
+        should_use_cdna4_mx_scale_swizzle,
+    )
+
+    swz = "CDNA4_SCALE" if should_use_cdna4_mx_scale_swizzle() else None
+    quant_dtype = torch.float8_e4m3fn
+
+    g1_gammas = gammas if apply_router_weight_on_input else None
+    g2_gammas = None if apply_router_weight_on_input else gammas
+
+    hidden_q, a1_scale = dynamic_mxfp8_quant(hidden_states, quant_dtype=quant_dtype)
+    raw_gate_up = moe_gemm_a8w4(
+        hidden_q,
+        w1_data,
+        a1_scale,
+        w1_wscale,
+        None,
+        None,
+        w1_bias,
+        routing_data,
+        gather_indx=gather_idx,
+        gammas=g1_gammas,
+        swizzle_mx_scale=swz,
+        out_dtype=torch.bfloat16,
+        apply_swiglu=False,
+        unpadded_N=unpadded_N_w1,
+        unpadded_K=unpadded_K_w1,
+    )
+    if unpadded_N_w1 is not None:
+        raw_gate_up = raw_gate_up[:, :unpadded_N_w1]
+
+    interim_fp8, a2_scale = fused_clamp_act_mul(
+        raw_gate_up,
+        swiglu_limit=swiglu_limit,
+        activation="silu",
+        dtype_quant=quant_dtype,
+        scale_dtype_fmt="ue8m0",
+        quant_block_size=32,
+    )
+
+    out = moe_gemm_a8w4(
+        interim_fp8,
+        w2_data,
+        a2_scale,
+        w2_wscale,
+        None,
+        None,
+        w2_bias,
+        routing_data,
+        scatter_indx=scatter_idx,
+        gammas=g2_gammas,
+        swizzle_mx_scale=swz,
+        unpadded_N=unpadded_N_w2,
+        unpadded_K=unpadded_K_w2,
+    )
+    return out
+
+
+def aiter_triton_kernel_w4a16_moe_forward(
+    hidden_states: torch.Tensor,
+    w1,
+    w2,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    activation: MoEActivation = MoEActivation.SWIGLUOAI,
+    quant_config: FusedMoEQuantConfig | None = None,
+    apply_router_weight_on_input: bool = False,
+    global_num_experts: int = -1,
+    expert_map: torch.Tensor | None = None,
+    unpadded_N_w1=None,
+    unpadded_K_w1=None,
+    unpadded_N_w2=None,
+    unpadded_K_w2=None,
+    num_expert_group: int | None = None,
+    topk_group: int | None = None,
+    e_score_correction_bias: torch.Tensor | None = None,
+    routed_scaling_factor: float | None = None,
+    score_mode: str | None = None,
+):
+    assert quant_config is not None and rocm_aiter_ops.is_enabled()
+    from vllm.platforms.rocm import on_gfx1250
+
+    try:
+        from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4
+        from aiter.ops.triton.moe.moe_routing import routing as _routing_mod
+    except ImportError:
+        from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4
+        from aiter.ops.triton.moe_routing import routing as _routing_mod
+
+    if on_gfx1250():
+        _routing_mod.is_tdm_avail = lambda: False
+    aiter_routing = _routing_mod.routing
+
+    if score_mode is not None:
+        use_grouped_topk = num_expert_group is not None and num_expert_group > 1
+        routing_data, gather_idx, scatter_idx = aiter_routing(
+            gating_output,
+            topk,
+            score_mode=score_mode,
+            bias=e_score_correction_bias,
+            renorm=renormalize,
+            routed_scaling_factor=(
+                routed_scaling_factor if routed_scaling_factor is not None else 1.0
+            ),
+            use_grouped_topk=use_grouped_topk,
+            num_expert_group=num_expert_group,
+            topk_group=topk_group,
+        )
+    else:
+        routing_data, gather_idx, scatter_idx = aiter_routing(
+            gating_output, topk, sm_first=not renormalize
+        )
+
+    if on_gfx1250():
+        gather_src = gather_idx.to(torch.long) // topk
+        hidden_states = hidden_states[gather_src]
+        gather_idx = None
+
+    assert quant_config.w1_precision is not None
+    assert quant_config.w2_precision is not None
+
+    w1_data = _aiter_raw(w1)
+    w2_data = _aiter_raw(w2)
+    w1_wscale = _aiter_raw(quant_config.w1_precision.weight_scale)
+    w2_wscale = _aiter_raw(quant_config.w2_precision.weight_scale)
+
+    gammas = routing_data.gate_scal if routing_data else None
+
+    swiglu_alpha = (
+        quant_config.gemm1_alpha if quant_config.gemm1_alpha is not None else 1.0
+    )
+    swiglu_limit = (
+        quant_config.gemm1_clamp_limit
+        if quant_config.gemm1_clamp_limit is not None
+        else 7.0
+    )
+
+    # SILU on gfx1250: use the verified a8w4 kernel (dynamic MXFP8); a16w4 faults.
+    if activation == MoEActivation.SILU and on_gfx1250():
+        return _aiter_w4a16_silu_via_a8w4(
+            hidden_states,
+            w1_data,
+            w2_data,
+            w1_wscale,
+            w2_wscale,
+            quant_config.w1_bias,
+            quant_config.w2_bias,
+            routing_data,
+            gather_idx,
+            scatter_idx,
+            gammas,
+            apply_router_weight_on_input,
+            swiglu_limit,
+            unpadded_N_w1,
+            unpadded_K_w1,
+            unpadded_N_w2,
+            unpadded_K_w2,
+        )
+
+    # SILU: silu(gate) * up — same kernel, just no "+1" residual in swiglu.
+    swiglu_add_residual = activation != MoEActivation.SILU
+
+    intermediate = moe_gemm_a16w4(
+        hidden_states,
+        w1_data,
+        None,
+        w1_wscale,
+        None,
+        None,
+        quant_config.w1_bias,
+        routing_data,
+        gather_indx=gather_idx,
+        gammas=gammas if apply_router_weight_on_input else None,
+        swizzle_mx_scale=None,
+        apply_swiglu=True,
+        alpha=swiglu_alpha,
+        limit=swiglu_limit,
+        swiglu_add_residual=swiglu_add_residual,
+        unpadded_N=unpadded_N_w1,
+        unpadded_K=unpadded_K_w1,
+    )
+
+    out = moe_gemm_a16w4(
+        intermediate,
+        w2_data,
+        None,
+        w2_wscale,
+        None,
+        None,
+        quant_config.w2_bias,
+        routing_data,
+        scatter_indx=scatter_idx,
+        gammas=None if apply_router_weight_on_input else gammas,
+        swizzle_mx_scale=None,
+        unpadded_N=unpadded_N_w2,
+        unpadded_K=unpadded_K_w2,
+    )
+
+    return out
+
+
+class AiterW4A16ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
+    def __init__(
+        self,
+        moe_config: FusedMoEConfig,
+        quant_config: FusedMoEQuantConfig,
+    ):
+        super().__init__(moe_config, quant_config)
+        self.topk = moe_config.experts_per_token
+        self.renormalize = moe_config.routing_method in (
+            RoutingMethodType.Renormalize,
+            RoutingMethodType.RenormalizeNaive,
+            RoutingMethodType.DeepseekV4,
+        )
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        if not rocm_aiter_ops.is_enabled():
+            return False
+        from vllm.platforms.rocm import on_gfx950, on_gfx1250
+
+        return on_gfx950() or on_gfx1250()
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        return False
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return (weight_key, activation_key) == (kMxfp4Static, None)
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return activation in (MoEActivation.SWIGLUOAI, MoEActivation.SILU)
+
+    @staticmethod
+    def _supports_parallel_config(
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> bool:
+        return (
+            not moe_parallel_config.use_all2all_kernels
+            and not moe_parallel_config.enable_eplb
+            and moe_parallel_config.dp_size <= 1
+        )
+
+    @staticmethod
+    def _supports_routing_method(
+        routing_method: RoutingMethodType,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return routing_method in [
+            RoutingMethodType.Renormalize,
+            RoutingMethodType.RenormalizeNaive,
+            RoutingMethodType.DeepseekV4,
+        ]
+
+    @staticmethod
+    def _supports_router_logits_dtype(
+        router_logits_dtype: torch.dtype | None,
+        routing_method: RoutingMethodType,
+    ) -> bool:
+        return True
+
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        return True
+
+    def apply(
+        self,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        router_logits: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        num_expert_group: int | None = None,
+        e_score_correction_bias: torch.Tensor | None = None,
+        routed_scaling_factor: float | None = None,
+        topk_group: int | None = None,
+    ) -> torch.Tensor:
+        assert self.moe_config.intermediate_size_per_partition_unpadded is not None
+        assert self.moe_config.hidden_dim_unpadded is not None
+        score_mode = (
+            "sqrtsoftplus"
+            if self.moe_config.routing_method == RoutingMethodType.DeepseekV4
+            else None
+        )
+        return aiter_triton_kernel_w4a16_moe_forward(
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            gating_output=router_logits,
+            topk=self.topk,
+            renormalize=self.renormalize,
+            activation=activation,
+            global_num_experts=global_num_experts,
+            expert_map=expert_map,
+            quant_config=self.quant_config,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+            unpadded_N_w1=self.moe_config.intermediate_size_per_partition_unpadded * 2,
+            unpadded_K_w1=self.moe_config.hidden_dim_unpadded,
+            unpadded_N_w2=self.moe_config.hidden_dim_unpadded,
+            unpadded_K_w2=self.moe_config.intermediate_size_per_partition_unpadded,
+            num_expert_group=num_expert_group,
+            topk_group=topk_group,
+            e_score_correction_bias=e_score_correction_bias,
+            routed_scaling_factor=routed_scaling_factor,
+            score_mode=score_mode,
+        )
@@ -814,13 +814,13 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
    ) -> bool:
        p = current_platform
        if p.is_rocm():
-            from vllm.platforms.rocm import on_gfx9
+            from vllm.platforms.rocm import get_cdna_version

-            is_rocm_on_gfx9 = on_gfx9()
+            _rocm_support_fp8 = get_cdna_version() > 2
        else:
-            is_rocm_on_gfx9 = False
+            _rocm_support_fp8 = False

-        device_supports_fp8 = is_rocm_on_gfx9 or (
+        device_supports_fp8 = _rocm_support_fp8 or (
            p.is_cuda() and p.has_device_capability((8, 9))
        )

@@ -472,11 +472,10 @@ class AiterExperts(mk.FusedMoEExpertsModular):
        ]
        if (weight_key, activation_key) not in SUPPORTED_W_A:
            return False
-        # CK MXFP4 MoE kernels are only supported on gfx950.
        if weight_key == kMxfp4Static:
-            from vllm.platforms.rocm import on_gfx950
+            from vllm.platforms.rocm import on_gfx950, on_gfx1250

-            if not on_gfx950():
+            if not on_gfx950() or on_gfx1250():
                return False
        return True

@@ -222,11 +222,14 @@ def backend_to_kernel_cls(
        return [BatchedMarlinExperts]

    elif backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
+        from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import (
+            AiterW4A16ExpertsMonolithic,
+        )
        from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
            AiterExperts,
        )

-        return [AiterExperts]
+        return [AiterExperts, AiterW4A16ExpertsMonolithic]

    elif backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:
        from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import (
@@ -1263,6 +1266,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(

    Supports DeepGEMM, TRTLLM MXFP8, Triton and Marlin backends.
    """
+    from vllm.platforms.rocm import on_gfx1250

    if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
        w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
@@ -1435,7 +1439,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
            w2_bias,
        )

-    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
+    elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not on_gfx1250():
        # Initially introduced for DeepSeekV4

        if w13_bias is not None:
@@ -1492,7 +1496,9 @@ def convert_weight_to_mxfp4_moe_kernel_format(
            w2_bias,
        )

-    elif mxfp4_backend in TRITON_BACKENDS:
+    elif mxfp4_backend in TRITON_BACKENDS or (
+        mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250()
+    ):
        from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig

        if mxfp4_backend == Mxfp4MoeBackend.TRITON:
@@ -480,6 +480,10 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
            global_num_experts=layer.global_num_experts,
            expert_map=layer.expert_map,
            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            num_expert_group=layer.num_expert_group,
+            topk_group=layer.topk_group,
+            e_score_correction_bias=layer.e_score_correction_bias,
+            routed_scaling_factor=layer.routed_scaling_factor,
        )


@@ -735,7 +739,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

        # For TRITON backends, weights are wrapped tensors from triton_kernels
        # that don't support .detach(). Manually assign parameters.
-        if self.mxfp4_backend not in TRITON_BACKENDS:
+        from vllm.platforms.rocm import on_gfx1250
+
+        uses_triton_weight_format = self.mxfp4_backend in TRITON_BACKENDS or (
+            self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250()
+        )
+        if not uses_triton_weight_format:
            replace_parameter(layer, "w13_weight", w13)
            replace_parameter(layer, "w2_weight", w2)
            replace_parameter(layer, "w13_weight_scale", w13_scale)
@@ -747,7 +756,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
            self.w2_precision_config = w2_scale

        # AITER backend requires weights to be marked as shuffled.
-        if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:
+        if (
+            self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16
+            and not uses_triton_weight_format
+        ):
            layer.w13_weight.is_shuffled = True
            layer.w2_weight.is_shuffled = True

@@ -841,7 +853,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
        w2_bias = getattr(layer, "w2_bias", None)
        swiglu_limit = getattr(layer, "swiglu_limit", None)

-        if self.mxfp4_backend in TRITON_BACKENDS:
+        from vllm.platforms.rocm import on_gfx1250
+
+        if self.mxfp4_backend in TRITON_BACKENDS or (
+            self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250()
+        ):
            # TRITON backends free w13/w2_weight_scale after swizzling; the
            # swizzled scales live inside the precision configs instead.
            assert self.w13_precision_config is not None
@@ -64,9 +64,7 @@ class QuarkW4A8_MXFP4_FP8(QuarkScheme):

        kernel_supported_gpu = False
        if current_platform.is_rocm():
-            from vllm.platforms.rocm import on_gfx950
-
-            kernel_supported_gpu = on_gfx950()
+            kernel_supported_gpu = current_platform.supports_mx()

        self.use_aiter_kernel = (
            is_aiter_found_and_supported()
@@ -875,10 +875,40 @@ def w8a8_triton_block_scaled_mm(
        Bs: The per-block quantization scale for `B`.
        block_size: The block size for per-block quantization. It should
        be 2-dim, e.g., [128, 128].
-        output_dtype: The dtype of the returned tensor.
+        output_dytpe: The dtype of the returned tensor.
    Returns:
        torch.Tensor: The result of matmul.
    """
+
+    from vllm.platforms.rocm import on_gfx1250
+
+    if on_gfx1250():
+        # Torch upcast reference: dequantize A,B to fp32 and matmul in fp32.
+        # Avoids the gfx1250 native-fp8 block GEMM NaN bug. Correct but slow.
+        _bn, _bk = block_size[0], block_size[1]
+        _As = (
+            _upcast_e8m0_to_fp32(As)
+            if As.dtype == torch.float8_e8m0fnu
+            else As.to(torch.float32)
+        )
+        _Bs = (
+            _upcast_e8m0_to_fp32(Bs)
+            if Bs.dtype == torch.float8_e8m0fnu
+            else Bs.to(torch.float32)
+        )
+        _K = A.shape[-1]
+        _N = B.shape[0]
+        _Af = A.to(torch.float32).reshape(-1, _K)
+        _Asf = (
+            _As.to(torch.float32)
+            .reshape(-1, _As.shape[-1])
+            .repeat_interleave(_bk, dim=1)[:, :_K]
+        )
+        _Bf = B.to(torch.float32)
+        _Bsf = _Bs.repeat_interleave(_bn, dim=0).repeat_interleave(_bk, dim=1)[:_N, :_K]
+        _out = (_Af * _Asf) @ (_Bf * _Bsf).t()
+        return _out.to(output_dtype).reshape(*A.shape[:-1], _N)
+
    assert len(block_size) == 2
    block_n, block_k = block_size[0], block_size[1]

@@ -122,7 +122,7 @@ def use_aiter_triton_gemm(n, m, k, dtype):
 def rocm_unquantized_gemm_impl(
    x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
 ) -> torch.Tensor:
-    from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950
+    from vllm.platforms.rocm import on_gfx1x, on_gfx9, on_gfx950, on_gfx1250

    n = x.numel() // x.size(-1)
    m = weight.shape[0]
@@ -164,7 +164,11 @@ def rocm_unquantized_gemm_impl(
    if use_skinny_reduce_counting:
        return ops.wvSplitKrc(x, weight, cu_count, bias)

-    if use_aiter_triton_gemm(n, m, k, x.dtype):
+    # gfx1250's aiter gemm_a16w16 uses the gluon backend, which requires
+    # K % 256 == 0 (it walks K with fixed-size descriptors and won't pad a
+    # partial last tile). Some whitelisted shapes have K=2880 (e.g. gpt-oss-120b
+    # hidden), so skip aiter there and fall back to the torch GEMM path below.
+    if use_aiter_triton_gemm(n, m, k, x.dtype) and not (on_gfx1250() and k % 256 != 0):
        from aiter.ops.triton.gemm_a16w16 import gemm_a16w16

        return gemm_a16w16(x, weight, bias)
@@ -172,6 +176,8 @@ def rocm_unquantized_gemm_impl(
    use_skinny = (
        envs.VLLM_ROCM_USE_SKINNY_GEMM
        and (on_gfx9() or on_gfx1x())
+        # build (gfx9/gfx11 ISA); fall back to torch GEMM there.
+        # TODO GFX1250: Include once skinny GEMM is supported on gfx1250
        and x.dtype in [torch.float16, torch.bfloat16]
        and k % 8 == 0
    )
@@ -194,11 +194,6 @@ def _get_gcn_arch() -> str:
        return _query_gcn_arch_from_amdsmi()
    except Exception as e:
        logger.debug("Failed to get GCN arch via amdsmi: %s", e)
-        logger.warning_once(
-            "Failed to get GCN arch via amdsmi, falling back to torch.cuda. "
-            "This will initialize CUDA and may cause "
-            "issues if CUDA_VISIBLE_DEVICES is not set yet."
-        )
    # Ultimate fallback: use torch.cuda (will initialize CUDA)
    return torch.cuda.get_device_properties("cuda").gcnArchName

@@ -218,6 +213,11 @@ _ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
 _ON_GFX90A = "gfx90a" in _GCN_ARCH
 _ON_GFX942 = "gfx942" in _GCN_ARCH
 _ON_GFX950 = "gfx950" in _GCN_ARCH
+_ON_GFX1250 = "gfx1250" in _GCN_ARCH
+
+_ON_CDNA = any(arch in _GCN_ARCH for arch in ["gfx9", "gfx1250"])
+# RDNA = gfx11/gfx12 minus the CDNA-classified gfx1250.
+_ON_RDNA = _ON_GFX1X and not _ON_CDNA


 def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None:
@@ -292,7 +292,7 @@ def _capability_from_gcn_arch(gcn_arch: str) -> tuple[int, int] | None:


 def on_gfx1x() -> bool:
-    return _ON_GFX1X
+    return _ON_GFX1X and not _ON_CDNA


 def on_gfx11() -> bool:
@@ -308,7 +308,11 @@ def on_gfx1151() -> bool:


 def on_gfx12x() -> bool:
-    return _ON_GFX12X
+    return _ON_GFX12X and not _ON_CDNA
+
+
+def on_gfx1250() -> bool:
+    return _ON_GFX1250


 def on_mi3xx() -> bool:
@@ -331,13 +335,33 @@ def on_gfx950() -> bool:
    return _ON_GFX950


+def on_cdna() -> bool:
+    return _ON_CDNA
+
+
+def on_rdna() -> bool:
+    return _ON_RDNA
+
+
+def get_cdna_version() -> int:
+    if on_gfx90a():
+        return 2
+    if on_gfx942():
+        return 3
+    if on_gfx950():
+        return 4
+    if on_gfx1250():
+        return 5
+    return 0
+
+
 # Enable HIP online tuning early, before hipBLASLt initializes.
 # Turn on hipBLASLt online tuning if use AITER hipBLASLt GEMM.
 if (
    envs.VLLM_ROCM_USE_AITER
    and envs.VLLM_ROCM_USE_AITER_LINEAR
    and envs.VLLM_ROCM_USE_AITER_LINEAR_HIPBMM
-    and on_mi3xx()
+    and get_cdna_version() > 2
 ):
    os.environ["HIP_ONLINE_TUNING"] = "1"

@@ -356,7 +380,7 @@ def use_rocm_custom_paged_attention(
 ) -> bool:
    # custom paged attn always supported on V0. On V1, requires sliding window
    # disabled due to observed numerical discrepancy.
-    if _ON_GFX9:
+    if on_cdna():
        return (
            (sliding_window == 0 or sliding_window == (-1, -1))
            and (qtype == torch.half or qtype == torch.bfloat16)
@@ -657,12 +681,12 @@ class RocmPlatform(Platform):

        from vllm._aiter_ops import rocm_aiter_ops

-        if rocm_aiter_ops.is_mha_enabled() and on_gfx9():
+        if rocm_aiter_ops.is_mha_enabled() and on_cdna():
            logger.info_once("Using AITER Flash Attention backend for ViT model.")
            return AttentionBackendEnum.ROCM_AITER_FA

        if (
-            on_gfx9()
+            on_cdna()
            and find_spec("flash_attn") is not None
            and (dtype == torch.float16 or dtype == torch.bfloat16)
        ):
@@ -878,11 +902,11 @@ class RocmPlatform(Platform):

    @classmethod
    def supports_mx(cls) -> bool:
-        return any(gfx in _GCN_ARCH for gfx in ["gfx95"])
+        return any(gfx in _GCN_ARCH for gfx in ["gfx95", "gfx1250"])

    @classmethod
    def supports_fp8(cls) -> bool:
-        return on_gfx9() or on_gfx12x()
+        return on_cdna() or on_gfx12x()

    @classmethod
    def is_fp8_fnuz(cls) -> bool:
@@ -505,7 +505,18 @@ def has_triton_kernels() -> bool:
@cache
 def has_tilelang() -> bool:
    """Whether the optional `tilelang` package is available."""
-    return _has_module("tilelang")
+    if not _has_module("tilelang"):
+        return False
+    # ROCm-only guard, imported lazily to avoid loading rocm on CUDA.
+    from vllm.platforms import current_platform
+
+    if current_platform.is_rocm():
+        from vllm.platforms.rocm import on_gfx1250
+
+        # TODO: Re-enable when tilelang supports gfx1250
+        if on_gfx1250():
+            return False
+    return True


 def has_arctic_inference() -> bool:
@@ -231,10 +231,26 @@ def _log_triton_jit_compile(fn_name: str, kwargs) -> None:


 def _setup_triton_jit_hook() -> None:
-    """Register a ``jit_post_compile_hook`` that warns on compilation."""
+    """Register a jit_post_compile_hook that warns on compilation."""
    if not HAS_TRITON:
        return
-    from triton import knobs  # type: ignore[import-untyped]
+    from triton import knobs
+    from triton.runtime import jit as _triton_jit
+
+    # kernels pass non-JSON-serializable constexprs
+    # make that serialization non-fatal
+    _orig = _triton_jit.serialize_specialization_data
+    if not getattr(_orig, "_vllm_guarded", False):
+
+        @functools.wraps(_orig)
+        def _guarded(*args, **kwargs):
+            try:
+                return _orig(*args, **kwargs)
+            except (TypeError, ValueError):
+                return None  # best-effort metadata; monitor ignores it
+
+        cast(Any, _guarded)._vllm_guarded = True
+        _triton_jit.serialize_specialization_data = _guarded

    existing_hook = knobs.runtime.jit_post_compile_hook

@@ -35,23 +35,31 @@ elif current_platform.is_xpu():
    compile_flash_attn_varlen_func_from_specs = None  # type: ignore[assignment]
    get_scheduler_metadata = xpu_ops.get_scheduler_metadata  # type: ignore[assignment]
 elif current_platform.is_rocm():
+    # On ROCm we use AITER's Triton flash-attention; the upstream flash-attn
+    # package is not installed/available. (Same source as aiter_triton_mla.py.)
+    # The FA4 compile-from-specs API is CUDA-only, so it is unavailable on ROCm
+    # regardless of whether AITER is present.
+    from vllm.platforms.rocm import on_gfx1250
+
+    compile_flash_attn_varlen_func_from_specs = None  # type: ignore[assignment]
    try:
-        from flash_attn import flash_attn_varlen_func  # type: ignore[no-redef]
+        if on_gfx1250():
+            from aiter.ops.triton.mha import (  # type: ignore[no-redef]
+                flash_attn_varlen_func,
+            )
+        else:
+            from flash_attn import flash_attn_varlen_func  # type: ignore[no-redef]

-        compile_flash_attn_varlen_func_from_specs = None  # type: ignore[assignment]
-
-        # Mark that upstream flash-attn is available on ROCm
        _ROCM_FLASH_ATTN_AVAILABLE = True
    except ImportError:

        def flash_attn_varlen_func(*args: Any, **kwargs: Any) -> Any:  # type: ignore[no-redef,misc]
+            package = "aiter" if on_gfx1250() else "flash-attn"
            raise ImportError(
-                "ROCm platform requires upstream flash-attn "
-                "to be installed. Please install flash-attn first."
+                f"ROCm platform requires upstream {package} "
+                f"to be installed. Please install {package} first."
            )

-        compile_flash_attn_varlen_func_from_specs = None  # type: ignore[assignment]
-
    # ROCm doesn't use scheduler metadata (FA3 feature), provide stub
    def get_scheduler_metadata(*args: Any, **kwargs: Any) -> None:  # type: ignore[misc]
        return None
@@ -348,7 +356,8 @@ def is_flash_attn_varlen_func_available() -> bool:
    Platform-specific sources:
    - CUDA: vllm.vllm_flash_attn.flash_attn_varlen_func
    - XPU: xpu_ops.flash_attn_varlen_func
-    - ROCm: upstream flash_attn.flash_attn_varlen_func (if available)
+    - ROCm: aiter.ops.triton.mha.flash_attn_varlen_func (if AITER available) or
+    upstream flash_attn.flash_attn_varlen_func

    Note: This is separate from the AITER flash attention backend (rocm_aiter_fa.py)
    which uses rocm_aiter_ops.flash_attn_varlen_func. The condition to use AITER is
@@ -776,12 +776,12 @@ class AiterFlashAttentionBackend(AttentionBackend):

    @classmethod
    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
-        from vllm.platforms.rocm import on_mi3xx
+        from vllm.platforms.rocm import get_cdna_version

        # DeviceCapability is currently created using torch.cuda.get_device_capability()
-        # which is known to be buggy on rocm systems. on_mi3xx uses amd-smi which is
+        # which is known to be buggy on rocm systems. on CDNA uses amd-smi which is
        # more reliable.
-        return on_mi3xx()
+        return get_cdna_version() > 2

    @classmethod
    def supports_non_causal(cls) -> bool:
@@ -18,6 +18,13 @@ if HAS_TRITON:
 logger = init_logger(__name__)


+def _skip_aiter_sampler_on_gfx1250() -> bool:
+    # Lazy ROCm-only import; keeps arch detection out of import time on CUDA/CPU.
+    from vllm.platforms.rocm import on_gfx1250
+
+    return on_gfx1250()
+
+
 def flashinfer_sampler_supported() -> bool:
    """Decide whether FlashInfer's top-p/top-k sampler can be used.

@@ -110,6 +117,7 @@ class TopKTopPSampler(nn.Module):
        elif (
            logprobs_mode not in PROCESSED_LOGPROBS_MODES
            and rocm_aiter_ops.is_enabled()
+            and not _skip_aiter_sampler_on_gfx1250()  # TODO (JPVILLAM): Enable
        ):
            self.aiter_ops = None
            self._aiter_ops_import_failed = False