@tduncan / Nvidia-Optimization-Stack-VLLM / commits / cbb5f045be

[ROCm][CI] Refresh ROCm base images when docker rocm_base changes (#46904)

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com> Signed-off-by: Codex <codex@example.invalid> Signed-off-by: Micah Williamson <micah.williamson@amd.com> Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Rohan138 <rohanpotdar138@gmail.com> Co-authored-by: Codex <codex@example.invalid> Co-authored-by: Micah Williamson <micah.williamson@amd.com> Co-authored-by: Matthew Wong <Matthew.Wong2@amd.com>

Andreas Karatzas committed Jul 7, 2026 at 05:10 UTC cbb5f045beb1f82bc139e14a4d3b2a904a8b358e
14 files changed +724 -52
.buildkite/ci_config_rocm.yaml
+1
@@ -8,6 +8,7 @@ run_all_patterns:
8 - "docker/docker-bake-rocm.hcl"
9 - ".buildkite/hardware_tests/amd.yaml"
10 - ".buildkite/scripts/ci-bake-rocm.sh"
11 + - ".buildkite/scripts/rocm/"
12 - ".buildkite/scripts/hardware_ci/run-amd-test.py"
13 - ".buildkite/scripts/hardware_ci/run-amd-test.sh"
14 - "CMakeLists.txt"
.buildkite/hardware_tests/amd.yaml
+32 -29
@@ -1,5 +1,28 @@
1 group: Hardware - AMD Build
2 +
3 +# ROCm image flow:
4 +# 1. Refresh the long-lived ROCm base image only when Dockerfile.rocm_base changes.
5 +# 2. Build ci_base from either the stable base or the freshly refreshed base.
6 +# 3. Build the per-commit ROCm CI image and smoke-test it before GPU jobs run.
7 steps:
8 + - label: "AMD: :docker: refresh ROCm base"
9 + key: refresh-rocm-base-amd
10 + depends_on: []
11 + device: amd_cpu
12 + no_plugin: true
13 + commands:
14 + - bash .buildkite/scripts/rocm/refresh-base-image.sh
15 + env:
16 + DOCKER_BUILDKIT: "1"
17 + BUILDKIT_PROGRESS: "tty"
18 + TERM: "xterm-256color"
19 + retry:
20 + automatic:
21 + - exit_status: -1 # Agent was lost
22 + limit: 1
23 + - exit_status: -10 # Agent was lost
24 + limit: 1
25 +
26 # Ensure ci_base is up-to-date before building the test image.
27 # Compares a content hash of ci_base-affecting files against the remote
28 # image label. If hashes match the build is skipped (< 30 s); if they
@@ -7,13 +30,16 @@ steps:
30 - label: "AMD: :docker: ensure ci_base"
31 key: ensure-ci-base-amd
32 soft_fail: false
10 - depends_on: []
33 + depends_on:
34 + - refresh-rocm-base-amd
35 device: amd_cpu
36 no_plugin: true
37 commands:
14 - - bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps
38 + - bash .buildkite/scripts/rocm/build-ci-base.sh
39 env:
40 DOCKER_BUILDKIT: "1"
41 + BUILDKIT_PROGRESS: "tty"
42 + TERM: "xterm-256color"
43 VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
44 PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
45 REMOTE_VLLM: "1"
@@ -33,35 +59,12 @@ steps:
59 device: amd_cpu
60 no_plugin: true
61 commands:
36 - - |
37 - if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" ]]; then
38 - echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only"
39 - IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts
40 - else
41 - bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel
42 - fi
43 - - |
44 - docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
45 - if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
46 - if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
47 - if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
48 - if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
49 - command -v python3
50 - command -v uv
51 - command -v pytest
52 - if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
53 - echo No ROCm CLI found in image >&2
54 - exit 1
55 - fi
56 - python3 - <<PY
57 - import torch, vllm
58 - print(torch.__version__)
59 - print(vllm.__version__)
60 - PY
61 - echo AMD image smoke OK
62 - '
62 + - bash .buildkite/scripts/rocm/build-test-image.sh
63 + - bash .buildkite/scripts/rocm/smoke-test-image.sh
64 env:
65 DOCKER_BUILDKIT: "1"
66 + BUILDKIT_PROGRESS: "tty"
67 + TERM: "xterm-256color"
68 VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
69 PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
70 IMAGE_TAG: "rocm/vllm-ci:$BUILDKITE_COMMIT"
.buildkite/scripts/ci-bake-rocm.sh
+25 -3
@@ -15,7 +15,7 @@ set -euo pipefail
15
16 DEFAULT_REPO_SLUG="vllm-project/vllm"
17 DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
18 -DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh"
18 +DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
19 DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
20 DEFAULT_CI_BASE_DOCKERFILE_STAGES="base build_rixl build_rocshmem build_deepep mori_base ci_base"
21 DEFAULT_CI_BASE_METADATA_VERSION="1"
@@ -393,6 +393,16 @@ should_upload_wheel_artifacts() {
393 || "${TARGET}" == *"artifact"* ]]
394 }
395
396 +set_buildkite_metadata() {
397 + local key="$1"
398 + local value="$2"
399 +
400 + [[ -n "${value}" ]] || return 0
401 + if command -v buildkite-agent >/dev/null 2>&1; then
402 + buildkite-agent meta-data set "${key}" "${value}" || true
403 + fi
404 +}
405 +
406 get_remote_image_label() {
407 local image_ref="$1"
408 local label_key="$2"
@@ -733,6 +743,8 @@ configure_ci_base_image_refs() {
743
744 if is_ci_base_target; then
745 IMAGE_TAG="${primary_tag}"
746 + CI_BASE_IMAGE="${primary_tag}"
747 + export CI_BASE_IMAGE
748 export IMAGE_TAG
749
750 echo "ci_base primary image tag: ${CI_BASE_IMAGE_TAG}"
@@ -750,6 +762,10 @@ configure_ci_base_image_refs() {
762 echo "ci_base stable alias will not be pushed for this build"
763 echo "Set NIGHTLY=1 on ${CI_BASE_STABLE_BRANCH:-main} to refresh ${stable_tag}"
764 fi
765 + set_buildkite_metadata "rocm-ci-base-image" "${CI_BASE_IMAGE_TAG}"
766 + set_buildkite_metadata "rocm-ci-base-image-content" "${content_tag}"
767 + set_buildkite_metadata "rocm-ci-base-image-commit" "${CI_BASE_IMAGE_TAG_COMMIT:-}"
768 + set_buildkite_metadata "rocm-ci-base-image-stable" "${CI_BASE_IMAGE_TAG_STABLE:-}"
769 return 0
770 fi
771
@@ -1779,7 +1795,10 @@ seed_dependency_caches_if_needed() {
1795
1796 echo "--- :docker: Seeding ${target}"
1797 echo "Expected cache ref: ${cache_ref}"
1782 - docker buildx bake "${BAKE_FILES[@]}" --progress plain "${target}"
1798 + docker buildx bake \
1799 + "${BAKE_FILES[@]}" \
1800 + --progress "${BUILDKIT_PROGRESS:-plain}" \
1801 + "${target}"
1802 verify_dependency_cache_ref "${cache_ref}"
1803 done
1804 }
@@ -1807,7 +1826,10 @@ run_bake() {
1826 local build_rc=0
1827
1828 echo "--- :docker: Building ${TARGET}"
1810 - docker buildx bake "${BAKE_FILES[@]}" --progress plain "${BAKE_TARGETS[@]}" || build_rc=$?
1829 + docker buildx bake \
1830 + "${BAKE_FILES[@]}" \
1831 + --progress "${BUILDKIT_PROGRESS:-plain}" \
1832 + "${BAKE_TARGETS[@]}" || build_rc=$?
1833
1834 if [[ ${build_rc} -eq 0 ]]; then
1835 echo "--- :white_check_mark: Build complete"
.buildkite/scripts/hardware_ci/run-amd-test.sh
+25 -1
@@ -28,6 +28,17 @@
28 ###############################################################################
29 set -o pipefail
30
31 +: "${BUILDKIT_PROGRESS:=plain}"
32 +: "${TERM:=xterm-256color}"
33 +: "${FORCE_COLOR:=1}"
34 +: "${CLICOLOR_FORCE:=1}"
35 +: "${PY_COLORS:=1}"
36 +: "${ROCM_DOCKER_TTY:=1}"
37 +if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then
38 + PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes"
39 +fi
40 +export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS ROCM_DOCKER_TTY
41 +
42 # Export Python path for commands that run directly on the host. Containerized
43 # tests set this to /vllm-workspace below so spawned Python processes do not
44 # depend on their current working directory.
@@ -149,6 +160,7 @@ EOF
160 echo "--- Building local ROCm test image"
161 docker build \
162 --pull=false \
163 + --progress "${BUILDKIT_PROGRESS}" \
164 --build-arg "BASE_IMAGE=${base_image}" \
165 -t "${artifact_image}" \
166 "${context_dir}" || return 1
@@ -535,6 +547,13 @@ if is_multi_node "$commands"; then
547 else
548 echo "--- Single-node job"
549 echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
550 + docker_run_terminal_args=(-i)
551 + if [[ "${ROCM_DOCKER_TTY}" == "1" ]]; then
552 + docker_run_terminal_args+=(-t)
553 + echo "Docker interactive stdin: enabled; TTY allocation: enabled"
554 + else
555 + echo "Docker interactive stdin: enabled; TTY allocation: disabled"
556 + fi
557
558 ulimit_core_hard=$(ulimit -H -c)
559 if [[ "$ulimit_core_hard" == "unlimited" ]]; then
@@ -551,7 +570,7 @@ else
570 fi
571
572 docker run \
554 - -t -i \
573 + "${docker_run_terminal_args[@]}" \
574 --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
575 $RDMA_FLAGS \
576 --network=host \
@@ -566,6 +585,11 @@ else
585 -e AWS_SECRET_ACCESS_KEY \
586 -e BUILDKITE_PARALLEL_JOB \
587 -e BUILDKITE_PARALLEL_JOB_COUNT \
588 + -e TERM \
589 + -e FORCE_COLOR \
590 + -e CLICOLOR_FORCE \
591 + -e PY_COLORS \
592 + -e PYTEST_ADDOPTS \
593 -v "${HF_CACHE}:${HF_MOUNT}" \
594 -e "HF_HOME=${HF_MOUNT}" \
595 -e "PYTHONPATH=${MYPYTHONPATH}" \
.buildkite/scripts/rocm/build-ci-base.sh new
+32
@@ -0,0 +1,32 @@
1 +#!/usr/bin/env bash
2 +# Build the ROCm ci_base image, optionally from a freshly rebuilt ROCm base.
3 +
4 +set -euo pipefail
5 +
6 +metadata_get() {
7 + local key="$1"
8 + if command -v buildkite-agent >/dev/null 2>&1; then
9 + buildkite-agent meta-data get "${key}" 2>/dev/null || true
10 + fi
11 +}
12 +
13 +main() {
14 + local base_refreshed=""
15 +
16 + base_refreshed="$(metadata_get rocm-base-refresh)"
17 + if [[ "${base_refreshed}" == "1" ]]; then
18 + export BASE_IMAGE
19 + export CI_BASE_PUSH_STABLE_TAG
20 +
21 + BASE_IMAGE="$(metadata_get rocm-base-image)"
22 + CI_BASE_PUSH_STABLE_TAG="$(metadata_get rocm-base-push-stable-tag)"
23 + CI_BASE_PUSH_STABLE_TAG="${CI_BASE_PUSH_STABLE_TAG:-0}"
24 +
25 + echo "Using refreshed ROCm base image for ci_base: ${BASE_IMAGE}"
26 + echo "Push stable ci_base tag: ${CI_BASE_PUSH_STABLE_TAG}"
27 + fi
28 +
29 + bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps
30 +}
31 +
32 +main "$@"
.buildkite/scripts/rocm/build-test-image.sh new
+57
@@ -0,0 +1,57 @@
1 +#!/usr/bin/env bash
2 +# Build the ROCm CI test image or wheel artifact.
3 +#
4 +# When Dockerfile.rocm_base changes, always build the full image so downstream
5 +# ROCm tests can validate the freshly rebuilt base -> ci_base -> ci image chain.
6 +
7 +set -euo pipefail
8 +
9 +metadata_get() {
10 + local key="$1"
11 + if command -v buildkite-agent >/dev/null 2>&1; then
12 + buildkite-agent meta-data get "${key}" 2>/dev/null || true
13 + fi
14 +}
15 +
16 +use_refreshed_base_if_present() {
17 + local base_refreshed=""
18 +
19 + base_refreshed="$(metadata_get rocm-base-refresh)"
20 + if [[ "${base_refreshed}" != "1" ]]; then
21 + return 1
22 + fi
23 +
24 + export BASE_IMAGE
25 + export CI_BASE_IMAGE
26 + export IMAGE_TAG_LATEST
27 +
28 + BASE_IMAGE="$(metadata_get rocm-base-image)"
29 + CI_BASE_IMAGE="$(metadata_get rocm-ci-base-image)"
30 + IMAGE_TAG_LATEST="$(metadata_get rocm-ci-image-descriptive)"
31 +
32 + echo "Using refreshed ROCm base image for test image: ${BASE_IMAGE}"
33 + echo "Using refreshed ROCm ci_base image for test image: ${CI_BASE_IMAGE}"
34 + if [[ -n "${IMAGE_TAG_LATEST}" ]]; then
35 + echo "Also tagging full ROCm CI image as: ${IMAGE_TAG_LATEST}"
36 + fi
37 +
38 + return 0
39 +}
40 +
41 +main() {
42 + local base_refreshed=0
43 +
44 + if use_refreshed_base_if_present; then
45 + base_refreshed=1
46 + fi
47 +
48 + if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" && "${base_refreshed}" != "1" ]]; then
49 + echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only"
50 + IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts
51 + return
52 + fi
53 +
54 + bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel
55 +}
56 +
57 +main "$@"
.buildkite/scripts/rocm/refresh-base-image.sh new
+513
@@ -0,0 +1,513 @@
1 +#!/usr/bin/env bash
2 +# Build and publish a fresh ROCm base image when Dockerfile.rocm_base changes.
3 +#
4 +# Normal AMD CI builds should not pay for this path. The script no-ops unless
5 +# docker/Dockerfile.rocm_base changed relative to the branch base, the previous
6 +# main commit, or ROCM_BASE_REFRESH_FORCE=1 is set.
7 +
8 +set -euo pipefail
9 +
10 +DOCKERFILE="${ROCM_BASE_DOCKERFILE:-docker/Dockerfile.rocm_base}"
11 +BASE_REPO="${ROCM_BASE_IMAGE_REPO:-rocm/vllm-dev}"
12 +CI_IMAGE_REPO="${ROCM_CI_IMAGE_REPO:-rocm/vllm-ci}"
13 +BUILDER_NAME="${ROCM_BASE_BUILDER_NAME:-vllm-rocm-base-builder}"
14 +DEFAULT_ROCM_BASE_METADATA_VERSION="1"
15 +DEFAULT_ROCM_BASE_CONTENT_FILES="${DOCKERFILE}"
16 +DEFAULT_ROCM_BASE_CONTENT_ARGS="BASE_IMAGE TRITON_BRANCH TRITON_REPO PYTORCH_BRANCH PYTORCH_REPO PYTORCH_VISION_BRANCH PYTORCH_VISION_REPO PYTORCH_AUDIO_BRANCH PYTORCH_AUDIO_REPO FA_BRANCH FA_REPO AITER_BRANCH AITER_REPO MORI_BRANCH MORI_REPO PYTORCH_ROCM_ARCH PYTHON_VERSION USE_SCCACHE"
17 +
18 +metadata_set() {
19 + local key="$1"
20 + local value="$2"
21 +
22 + [[ -n "${value}" ]] || return 0
23 + if command -v buildkite-agent >/dev/null 2>&1; then
24 + buildkite-agent meta-data set "${key}" "${value}" || true
25 + fi
26 +}
27 +
28 +compute_content_hash() {
29 + local path=""
30 + local file=""
31 +
32 + for path in "$@"; do
33 + if [[ -d "${path}" ]]; then
34 + while IFS= read -r -d '' file; do
35 + printf 'file:%s\n' "${file}"
36 + sha256sum "${file}"
37 + done < <(find "${path}" -type f -print0 | sort -z)
38 + elif [[ -f "${path}" ]]; then
39 + printf 'file:%s\n' "${path}"
40 + sha256sum "${path}"
41 + else
42 + printf 'missing:%s\n' "${path}"
43 + fi
44 + done | sha256sum | cut -d' ' -f1
45 +}
46 +
47 +clean_docker_tag() {
48 + local input="$1"
49 + echo "${input}" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128
50 +}
51 +
52 +tag_component() {
53 + local input="$1"
54 + local max_chars="${2:-24}"
55 +
56 + clean_docker_tag "${input:-unknown}" | cut -c1-"${max_chars}"
57 +}
58 +
59 +extract_arg_default() {
60 + local arg_name="$1"
61 +
62 + sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
63 + "${DOCKERFILE}" | head -1
64 +}
65 +
66 +resolve_image_digest() {
67 + local image_ref="$1"
68 +
69 + docker buildx imagetools inspect "${image_ref}" 2>/dev/null \
70 + | sed -n -E 's/^Digest:[[:space:]]+//p' \
71 + | head -1 || true
72 +}
73 +
74 +resolve_rocm_base_arg_value() {
75 + local arg_name="$1"
76 + local use_sccache="$2"
77 +
78 + case "${arg_name}" in
79 + USE_SCCACHE)
80 + printf '%s\n' "${use_sccache}"
81 + ;;
82 + *)
83 + extract_arg_default "${arg_name}"
84 + ;;
85 + esac
86 +}
87 +
88 +hash_rocm_base_arg_values() {
89 + local use_sccache="$1"
90 + local base_image_digest="$2"
91 + local arg_name=""
92 + local arg_value=""
93 + shift 2 || true
94 +
95 + for arg_name in "$@"; do
96 + [[ -n "${arg_name}" ]] || continue
97 + arg_value=$(resolve_rocm_base_arg_value "${arg_name}" "${use_sccache}")
98 + printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-<empty>}"
99 + if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then
100 + printf 'arg:%s.digest=%s\n' "${arg_name}" "${base_image_digest:-unknown}"
101 + fi
102 + done
103 +}
104 +
105 +rocm_version_from_base_image() {
106 + local base_image="$1"
107 + local version=""
108 +
109 + version="$(sed -n -E 's/.*:([0-9]+\.[0-9]+(\.[0-9]+)?)-.*/\1/p' <<<"${base_image}")"
110 + tag_component "${version:-${base_image}}" 16
111 +}
112 +
113 +git_diff_changed_base() {
114 + local range="$1"
115 + [[ -n "$(git diff --name-only "${range}" -- "${DOCKERFILE}" 2>/dev/null)" ]]
116 +}
117 +
118 +short_git_ref() {
119 + local ref="$1"
120 +
121 + git rev-parse --short "${ref}" 2>/dev/null || printf '%s\n' "${ref}"
122 +}
123 +
124 +extract_arg_default_from_ref() {
125 + local ref="$1"
126 + local arg_name="$2"
127 + local content=""
128 +
129 + content="$(git show "${ref}:${DOCKERFILE}" 2>/dev/null || true)"
130 + sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
131 + <<<"${content}" | head -1
132 +}
133 +
134 +log_arg_default_changes() {
135 + local old_ref="$1"
136 + local new_ref="$2"
137 + local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}"
138 + local arg_name=""
139 + local old_value=""
140 + local new_value=""
141 + local changed=0
142 +
143 + echo "Changed ROCm base ARG defaults:"
144 + for arg_name in ${content_args}; do
145 + old_value="$(extract_arg_default_from_ref "${old_ref}" "${arg_name}")"
146 + new_value="$(extract_arg_default_from_ref "${new_ref}" "${arg_name}")"
147 + if [[ "${old_value}" != "${new_value}" ]]; then
148 + echo " - ${arg_name}: ${old_value:-<unset>} -> ${new_value:-<unset>}"
149 + changed=1
150 + fi
151 + done
152 +
153 + if [[ "${changed}" == "0" ]]; then
154 + echo " - none detected; Dockerfile instructions changed outside tracked ARG defaults"
155 + fi
156 +}
157 +
158 +log_arg_line_diff() {
159 + local range="$1"
160 + local arg_diff=""
161 +
162 + arg_diff="$(
163 + git diff --unified=0 "${range}" -- "${DOCKERFILE}" 2>/dev/null \
164 + | awk '/^[+-][[:space:]]*ARG[[:space:]]/ && $0 !~ /^(---|\+\+\+)/ { print " " $0 }' \
165 + || true
166 + )"
167 +
168 + if [[ -n "${arg_diff}" ]]; then
169 + echo "Changed Dockerfile ARG lines:"
170 + printf '%s\n' "${arg_diff}"
171 + fi
172 +}
173 +
174 +log_rocm_base_change_check() {
175 + local context="$1"
176 + local range="$2"
177 + local old_ref="$3"
178 + local old_short=""
179 + local head_short=""
180 +
181 + old_short="$(short_git_ref "${old_ref}")"
182 + head_short="$(short_git_ref HEAD)"
183 +
184 + echo "--- :mag: ROCm base refresh check"
185 + echo "Context: ${context}"
186 + echo "Dockerfile: ${DOCKERFILE}"
187 + echo "Base revision: ${old_short}"
188 + echo "Head revision: ${head_short}"
189 + echo "Git diff range: ${range}"
190 +}
191 +
192 +log_rocm_base_rebuild_reason() {
193 + local context="$1"
194 + local range="$2"
195 + local old_ref="$3"
196 + local changed_files=""
197 +
198 + log_rocm_base_change_check "${context}" "${range}" "${old_ref}"
199 +
200 + changed_files="$(git diff --name-only "${range}" -- "${DOCKERFILE}" 2>/dev/null || true)"
201 + echo "Changed files:"
202 + if [[ -n "${changed_files}" ]]; then
203 + sed 's/^/ - /' <<<"${changed_files}"
204 + else
205 + echo " - ${DOCKERFILE}"
206 + fi
207 + log_arg_default_changes "${old_ref}" HEAD
208 + log_arg_line_diff "${range}"
209 + echo "Decision: rebuilding ROCm base image because ${DOCKERFILE} changed."
210 +}
211 +
212 +rocm_base_changed_in_range() {
213 + local context="$1"
214 + local range="$2"
215 + local old_ref="$3"
216 +
217 + if git_diff_changed_base "${range}"; then
218 + log_rocm_base_rebuild_reason "${context}" "${range}" "${old_ref}"
219 + return 0
220 + fi
221 +
222 + log_rocm_base_change_check "${context}" "${range}" "${old_ref}"
223 + echo "Decision: ROCm base refresh not required; ${DOCKERFILE} is unchanged."
224 + return 1
225 +}
226 +
227 +rocm_base_changed() {
228 + local base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
229 + local base_ref="refs/remotes/origin/${base_branch}"
230 + local merge_base=""
231 +
232 + if [[ "${ROCM_BASE_REFRESH_SKIP:-0}" == "1" ]]; then
233 + echo "ROCM_BASE_REFRESH_SKIP=1 set; skipping ROCm base refresh"
234 + return 1
235 + fi
236 +
237 + if [[ "${ROCM_BASE_REFRESH_FORCE:-0}" == "1" ]]; then
238 + echo "ROCM_BASE_REFRESH_FORCE=1 set; refreshing ROCm base image"
239 + return 0
240 + fi
241 +
242 + if ! git rev-parse --is-inside-work-tree >/dev/null 2>&1; then
243 + echo "Not in a git checkout; skipping ROCm base refresh unless forced"
244 + return 1
245 + fi
246 +
247 + if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
248 + git fetch --no-tags --depth=200 origin \
249 + "+refs/heads/${base_branch}:${base_ref}" >/dev/null 2>&1 || true
250 + merge_base=$(git merge-base HEAD "${base_ref}" 2>/dev/null || true)
251 + if [[ -z "${merge_base}" ]]; then
252 + echo "Unable to determine merge base with PR base ${base_ref}; skipping ROCm base refresh unless forced"
253 + return 1
254 + fi
255 + if rocm_base_changed_in_range \
256 + "pull request build against ${base_ref}" \
257 + "${merge_base}...HEAD" \
258 + "${merge_base}"; then
259 + return 0
260 + fi
261 + elif [[ "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]] \
262 + && git rev-parse --verify HEAD~1 >/dev/null 2>&1; then
263 + if rocm_base_changed_in_range \
264 + "stable branch build; comparing against previous ${ROCM_BASE_STABLE_BRANCH:-main} commit" \
265 + "HEAD~1..HEAD" \
266 + "HEAD~1"; then
267 + return 0
268 + fi
269 + else
270 + git fetch --no-tags --depth=200 origin \
271 + "+refs/heads/${base_branch}:${base_ref}" >/dev/null 2>&1 || true
272 + merge_base=$(git merge-base HEAD "${base_ref}" 2>/dev/null || true)
273 + if [[ -z "${merge_base}" ]]; then
274 + echo "Unable to determine merge base with branch base ${base_ref}; skipping ROCm base refresh unless forced"
275 + return 1
276 + fi
277 + if rocm_base_changed_in_range \
278 + "branch build against ${base_ref}" \
279 + "${merge_base}...HEAD" \
280 + "${merge_base}"; then
281 + return 0
282 + fi
283 + fi
284 +
285 + return 1
286 +}
287 +
288 +should_push_stable_tag() {
289 + if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
290 + return 1
291 + fi
292 +
293 + if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "1" ]]; then
294 + return 0
295 + fi
296 + if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "0" ]]; then
297 + return 1
298 + fi
299 +
300 + [[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" \
301 + && "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]]
302 +}
303 +
304 +setup_builder() {
305 + echo "--- :buildkite: Setting up buildx builder for ROCm base"
306 + if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
307 + docker buildx use "${BUILDER_NAME}"
308 + else
309 + docker buildx create --name "${BUILDER_NAME}" --driver docker-container --use
310 + fi
311 + docker buildx inspect --bootstrap
312 +}
313 +
314 +compute_base_content_hash() {
315 + local use_sccache="$1"
316 + local base_image_digest="$2"
317 + local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}"
318 + local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}"
319 + local -a content_paths=()
320 + local -a content_arg_names=()
321 +
322 + read -r -a content_paths <<< "${content_files}"
323 + read -r -a content_arg_names <<< "${content_args}"
324 +
325 + {
326 + printf 'content-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")"
327 + printf 'dockerfile:%s\n' "${DOCKERFILE}"
328 + printf 'resolved-build-args:\n'
329 + hash_rocm_base_arg_values \
330 + "${use_sccache}" "${base_image_digest}" "${content_arg_names[@]}"
331 + } | sha256sum | cut -d' ' -f1
332 +}
333 +
334 +build_base_image() {
335 + local use_sccache="${ROCM_BASE_USE_SCCACHE:-${USE_SCCACHE:-0}}"
336 + local base_hash=""
337 + local build_date=""
338 + local build_suffix=""
339 + local base_image_arg=""
340 + local base_image_digest=""
341 + local rocm_version=""
342 + local triton_arg=""
343 + local pytorch_arg=""
344 + local pytorch_vision_arg=""
345 + local pytorch_audio_arg=""
346 + local fa_arg=""
347 + local aiter_arg=""
348 + local mori_arg=""
349 + local python_version_arg=""
350 + local pytorch_rocm_arch_arg=""
351 + local pytorch_branch=""
352 + local aiter_branch=""
353 + local dependency_summary=""
354 + local descriptor=""
355 + local ci_descriptor=""
356 + local descriptive_tag=""
357 + local stable_tag="${BASE_REPO}:base"
358 + local ci_descriptive_tag=""
359 + local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}"
360 + local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}"
361 + local content_files_hash=""
362 + local metadata_version="${ROCM_BASE_METADATA_VERSION:-${DEFAULT_ROCM_BASE_METADATA_VERSION}}"
363 + local -a tags=()
364 + local -a no_cache_args=()
365 + local -a sccache_args=()
366 + local -a content_paths=()
367 +
368 + if [[ ! -f "${DOCKERFILE}" ]]; then
369 + echo "Error: ROCm base Dockerfile not found: ${DOCKERFILE}" >&2
370 + exit 1
371 + fi
372 +
373 + build_date="${ROCM_BASE_TAG_DATE:-$(date -u +%Y%m%d)}"
374 + if [[ -n "${BUILDKITE_BUILD_NUMBER:-}" ]]; then
375 + build_suffix="_bk_${BUILDKITE_BUILD_NUMBER}"
376 + fi
377 + base_image_arg="$(extract_arg_default BASE_IMAGE)"
378 + base_image_digest="$(resolve_image_digest "${base_image_arg}")"
379 + read -r -a content_paths <<< "${content_files}"
380 + content_files_hash="$(compute_content_hash "${content_paths[@]}")"
381 + base_hash=$(compute_base_content_hash "${use_sccache}" "${base_image_digest}")
382 + rocm_version="$(rocm_version_from_base_image "${base_image_arg}")"
383 + triton_arg="$(extract_arg_default TRITON_BRANCH)"
384 + pytorch_arg="$(extract_arg_default PYTORCH_BRANCH)"
385 + pytorch_vision_arg="$(extract_arg_default PYTORCH_VISION_BRANCH)"
386 + pytorch_audio_arg="$(extract_arg_default PYTORCH_AUDIO_BRANCH)"
387 + fa_arg="$(extract_arg_default FA_BRANCH)"
388 + aiter_arg="$(extract_arg_default AITER_BRANCH)"
389 + mori_arg="$(extract_arg_default MORI_BRANCH)"
390 + python_version_arg="$(extract_arg_default PYTHON_VERSION)"
391 + pytorch_rocm_arch_arg="$(extract_arg_default PYTORCH_ROCM_ARCH)"
392 + pytorch_branch="$(tag_component "${pytorch_arg}" 16)"
393 + aiter_branch="$(tag_component "${aiter_arg}" 24)"
394 + dependency_summary="base=${base_image_arg},rocm=${rocm_version},python=${python_version_arg},pytorch=${pytorch_arg},torchvision=${pytorch_vision_arg},torchaudio=${pytorch_audio_arg},triton=${triton_arg},flash-attn=${fa_arg},aiter=${aiter_arg},mori=${mori_arg},pytorch-rocm-arch=${pytorch_rocm_arch_arg}"
395 + descriptor="$(clean_docker_tag "base_custom_aiter_${aiter_branch}_torch_${pytorch_branch}_${build_date}${build_suffix}")"
396 + ci_descriptor="$(clean_docker_tag "ci_custom_aiter_${aiter_branch}_torch_${pytorch_branch}_${build_date}${build_suffix}")"
397 +
398 + descriptive_tag="${BASE_REPO}:${descriptor}"
399 + ci_descriptive_tag="${CI_IMAGE_REPO}:${ci_descriptor}"
400 +
401 + tags=(-t "${descriptive_tag}")
402 + if should_push_stable_tag; then
403 + tags+=(-t "${stable_tag}")
404 + metadata_set "rocm-base-push-stable-tag" "1"
405 + else
406 + metadata_set "rocm-base-push-stable-tag" "0"
407 + fi
408 +
409 + if [[ "${ROCM_BASE_NO_CACHE:-1}" == "1" ]]; then
410 + no_cache_args=(--no-cache)
411 + fi
412 +
413 + for env_name in \
414 + SCCACHE_DOWNLOAD_URL \
415 + SCCACHE_ENDPOINT \
416 + SCCACHE_BUCKET_NAME \
417 + SCCACHE_REGION_NAME \
418 + SCCACHE_S3_NO_CREDENTIALS; do
419 + if [[ -n "${!env_name:-}" ]]; then
420 + sccache_args+=(--build-arg "${env_name}=${!env_name}")
421 + fi
422 + done
423 +
424 + echo "--- :docker: Building ROCm base image"
425 + echo "Dockerfile: ${DOCKERFILE}"
426 + echo "Descriptive tag: ${descriptive_tag}"
427 + echo "Stable tag: ${stable_tag} ($(should_push_stable_tag && echo enabled || echo disabled))"
428 + echo "Content hash: ${base_hash}"
429 + echo "Dependency summary: ${dependency_summary}"
430 + echo "USE_SCCACHE: ${use_sccache}"
431 +
432 + docker buildx build \
433 + "${no_cache_args[@]}" \
434 + --pull \
435 + --progress "${BUILDKIT_PROGRESS:-plain}" \
436 + --file "${DOCKERFILE}" \
437 + --build-arg "USE_SCCACHE=${use_sccache}" \
438 + "${sccache_args[@]}" \
439 + --label "org.opencontainers.image.source=https://github.com/vllm-project/vllm" \
440 + --label "org.opencontainers.image.vendor=vLLM" \
441 + --label "org.opencontainers.image.title=vLLM ROCm base" \
442 + --label "org.opencontainers.image.revision=${BUILDKITE_COMMIT:-}" \
443 + --label "vllm.rocm_base.metadata_version=${metadata_version}" \
444 + --label "vllm.rocm_base.content_hash=${base_hash}" \
445 + --label "vllm.rocm_base.content_files_hash=${content_files_hash}" \
446 + --label "vllm.rocm_base.dockerfile=${DOCKERFILE}" \
447 + --label "vllm.rocm_base.image.descriptive=${descriptive_tag}" \
448 + --label "vllm.rocm_base.image.stable=${stable_tag}" \
449 + --label "vllm.rocm_base.git_commit=${BUILDKITE_COMMIT:-}" \
450 + --label "vllm.rocm_base.stable_branch=${ROCM_BASE_STABLE_BRANCH:-main}" \
451 + --label "vllm.rocm_base.descriptor=${descriptor}" \
452 + --label "vllm.rocm_base.dependency_summary=${dependency_summary}" \
453 + --label "vllm.rocm_base.base_image=${base_image_arg}" \
454 + --label "vllm.rocm_base.base_image_digest=${base_image_digest}" \
455 + --label "vllm.rocm_base.dependency.rocm=${rocm_version}" \
456 + --label "vllm.rocm_base.dependency.python=${python_version_arg}" \
457 + --label "vllm.rocm_base.dependency.pytorch=${pytorch_arg}" \
458 + --label "vllm.rocm_base.dependency.torchvision=${pytorch_vision_arg}" \
459 + --label "vllm.rocm_base.dependency.torchaudio=${pytorch_audio_arg}" \
460 + --label "vllm.rocm_base.dependency.triton=${triton_arg}" \
461 + --label "vllm.rocm_base.dependency.flash_attention=${fa_arg}" \
462 + --label "vllm.rocm_base.dependency.aiter=${aiter_arg}" \
463 + --label "vllm.rocm_base.dependency.mori=${mori_arg}" \
464 + --label "vllm.rocm_base.pytorch_rocm_arch=${pytorch_rocm_arch_arg}" \
465 + "${tags[@]}" \
466 + --push \
467 + .
468 +
469 + docker buildx imagetools inspect "${descriptive_tag}" >/dev/null
470 +
471 + metadata_set "rocm-base-refresh" "1"
472 + metadata_set "rocm-base-image" "${descriptive_tag}"
473 + metadata_set "rocm-base-image-descriptive" "${descriptive_tag}"
474 + metadata_set "rocm-base-image-stable" "${stable_tag}"
475 + metadata_set "rocm-base-image-ci-descriptive" "${ci_descriptive_tag}"
476 + metadata_set "rocm-base-metadata-version" "${metadata_version}"
477 + metadata_set "rocm-base-content-hash" "${base_hash}"
478 + metadata_set "rocm-base-content-files-hash" "${content_files_hash}"
479 + metadata_set "rocm-base-content-files" "${content_files}"
480 + metadata_set "rocm-base-content-args" "${content_args}"
481 + metadata_set "rocm-base-base-image-digest" "${base_image_digest}"
482 + metadata_set "rocm-base-dockerfile" "${DOCKERFILE}"
483 + metadata_set "rocm-base-descriptor" "${descriptor}"
484 + metadata_set "rocm-base-dependency-summary" "${dependency_summary}"
485 + metadata_set "rocm-base-dependency-rocm" "${rocm_version}"
486 + metadata_set "rocm-base-dependency-python" "${python_version_arg}"
487 + metadata_set "rocm-base-dependency-pytorch" "${pytorch_arg}"
488 + metadata_set "rocm-base-dependency-torchvision" "${pytorch_vision_arg}"
489 + metadata_set "rocm-base-dependency-torchaudio" "${pytorch_audio_arg}"
490 + metadata_set "rocm-base-dependency-triton" "${triton_arg}"
491 + metadata_set "rocm-base-dependency-flash-attention" "${fa_arg}"
492 + metadata_set "rocm-base-dependency-aiter" "${aiter_arg}"
493 + metadata_set "rocm-base-dependency-mori" "${mori_arg}"
494 + metadata_set "rocm-base-pytorch-rocm-arch" "${pytorch_rocm_arch_arg}"
495 + metadata_set "rocm-ci-image-descriptive" "${ci_descriptive_tag}"
496 +
497 + echo "--- :white_check_mark: ROCm base image published"
498 + echo "Use BASE_IMAGE=${descriptive_tag} for downstream ROCm CI builds"
499 +}
500 +
501 +main() {
502 + metadata_set "rocm-base-refresh" "0"
503 +
504 + if ! rocm_base_changed; then
505 + echo "ROCm base Dockerfile did not change; skipping base image refresh"
506 + return 0
507 + fi
508 +
509 + setup_builder
510 + build_base_image
511 +}
512 +
513 +main "$@"
.buildkite/scripts/rocm/smoke-test-image.sh new
+32
@@ -0,0 +1,32 @@
1 +#!/usr/bin/env bash
2 +# Fast structural smoke test for the full ROCm CI image.
3 +
4 +set -euo pipefail
5 +
6 +image_ref="${VLLM_CI_SMOKE_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}}"
7 +
8 +docker run --rm --network=none --entrypoint /bin/bash "${image_ref}" -ec '
9 + if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
10 + if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
11 + if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
12 + if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
13 +
14 + command -v python3
15 + command -v uv
16 + command -v pytest
17 +
18 + if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
19 + echo No ROCm CLI found in image >&2
20 + exit 1
21 + fi
22 +
23 + python3 - <<PY
24 +import torch
25 +import vllm
26 +
27 +print(torch.__version__)
28 +print(vllm.__version__)
29 +PY
30 +
31 + echo AMD image smoke OK
32 +'
docker/Dockerfile.rocm_base
+1 -1
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
9 ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
10 ARG FA_BRANCH="0e60e394"
11 ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
12 -ARG AITER_BRANCH="v0.1.16.post2"
12 +ARG AITER_BRANCH="v0.1.16.post3"
13 ARG AITER_REPO="https://github.com/ROCm/aiter.git"
14 ARG MORI_BRANCH="v1.1.0"
15 ARG MORI_REPO="https://github.com/ROCm/mori.git"
docs/design/attention_backends.md
+1 -1
@@ -170,7 +170,7 @@ Priority is **1 = highest** (tried first).
170 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
171 | `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
172 | `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
173 -| `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
173 +| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
174 | `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
175 | `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |
176 | `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
tests/compile/passes/test_fusion_attn.py
-5
@@ -306,11 +306,6 @@ def test_attention_quant_pattern(
306 torch.manual_seed(42)
307
308 backend_cls = backend.get_class()
309 -
310 - # TODO: drop once AITER reenables fp16 unified attention.
311 - if dtype not in backend_cls.supported_dtypes:
312 - pytest.skip(f"{backend.name} does not support dtype {dtype}")
313 -
309 block_size = backend_cls.get_preferred_block_size(16)
310
311 model_config = ModelConfig(
tests/kernels/attention/test_rocm_aiter_unified_attn.py
+1 -2
@@ -30,8 +30,7 @@ NUM_Q_HEADS = 8
30 NUM_KV_HEADS = 8
31 HEAD_SIZES = [128, 256]
32 BLOCK_SIZES = [16, 64]
33 -# TODO: re-add torch.float16 once AITER reenables fp16 unified attention.
34 -DTYPES = [torch.bfloat16]
33 +DTYPES = [torch.bfloat16, torch.float16]
34 FP8_DTYPE = current_platform.fp8_dtype()
35
36 # (query_len, kv_len) per sequence
tests/v1/attention/test_rocm_attention_backends_selection.py
+1 -9
@@ -136,17 +136,9 @@ def test_standard_attention_backend_selection(
136 # Get the backend class path
137 from vllm.platforms.rocm import RocmPlatform
138
139 - # The AITER unified attention kernel only supports BF16/FP8 KV caches
140 - # (its 3D kernel asserts on fp16), so it must be selected with bf16.
141 - dtype = (
142 - torch.bfloat16
143 - if selected_backend == "ROCM_AITER_UNIFIED_ATTN"
144 - else torch.float16
145 - )
146 -
139 attn_selector_config = AttentionSelectorConfig(
140 head_size=128,
149 - dtype=dtype,
141 + dtype=torch.float16,
142 kv_cache_dtype="auto",
143 block_size=16,
144 use_mla=False,
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+3 -1
@@ -27,12 +27,14 @@ logger = init_logger(__name__)
27
28
29 class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
30 - supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
30 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
31 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
32 "auto",
33 + "float16",
34 "bfloat16",
35 "fp8",
36 "fp8_e4m3",
37 + "fp8_e5m2",
38 ]
39
40 @staticmethod