Skip to content

Commit b32f626

Browse files
author
Aleksandr Malyshev
committed
torch.compile passes, but correctness is wrong
1 parent 9c6587b commit b32f626

2 files changed

Lines changed: 53 additions & 85 deletions

File tree

‎docker/Dockerfile.rocm_base‎

Lines changed: 28 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,13 @@ ARG BASE_IMAGE=compute-artifactory.amd.com:5000/rocm-plus-docker/framework/compu
44
# ARG LEGACY_HIPBLASLT_OPTION=
55
ARG TRITON_BRANCH="916969a"
66
ARG TRITON_REPO="https://github.com/ROCm/triton.git"
7-
# ARG PYTORCH_BRANCH="37f92bb"
8-
# ARG PYTORCH_VISION_BRANCH="v0.21.0"
9-
# ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
10-
# ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
7+
ARG PYTORCH_BRANCH="3a7dec2"
8+
ARG PYTORCH_VISION_BRANCH="0.22.0a0"
9+
ARG PYTORCH_REPO="https://github.com/ROCm/pytorch.git"
10+
ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
1111
ARG FA_BRANCH="8ede036"
1212
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
13-
ARG AITER_BRANCH="d765e80"
13+
ARG AITER_BRANCH="498ff21"
1414
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
1515

1616
FROM ${BASE_IMAGE} AS base
@@ -20,7 +20,7 @@ ENV ROCM_PATH=/opt/rocm
2020
ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
2121
ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201
2222
ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
23-
ENV AITER_ROCM_ARCH=gfx950
23+
ENV AITER_ROCM_ARCH=gfx942;gfx950
2424

2525
ARG PYTHON_VERSION=3.12
2626

@@ -85,22 +85,22 @@ RUN cd /opt/rocm/share/amd_smi \
8585
RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install
8686

8787
FROM base AS build_pytorch
88-
# ARG PYTORCH_BRANCH
89-
# ARG PYTORCH_VISION_BRANCH
90-
# ARG PYTORCH_REPO
91-
# ARG PYTORCH_VISION_REPO
88+
ARG PYTORCH_BRANCH
89+
ARG PYTORCH_VISION_BRANCH
90+
ARG PYTORCH_REPO
91+
ARG PYTORCH_VISION_REPO
9292
ARG FA_BRANCH
9393
ARG FA_REPO
94-
# RUN git clone ${PYTORCH_REPO} pytorch
95-
# RUN cd pytorch && git checkout ${PYTORCH_BRANCH} && \
96-
# pip install -r requirements.txt && git submodule update --init --recursive \
97-
# && python3 tools/amd_build/build_amd.py \
98-
# && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \
99-
# && pip install dist/*.whl
100-
# RUN git clone ${PYTORCH_VISION_REPO} vision
101-
# RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \
102-
# && python3 setup.py bdist_wheel --dist-dir=dist \
103-
# && pip install dist/*.whl
94+
RUN git clone ${PYTORCH_REPO} pytorch
95+
RUN cd pytorch && git checkout ${PYTORCH_BRANCH} && \
96+
pip install -r requirements.txt && git submodule update --init --recursive \
97+
&& python3 tools/amd_build/build_amd.py \
98+
&& CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \
99+
&& pip install dist/*.whl
100+
RUN git clone ${PYTORCH_VISION_REPO} vision
101+
RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \
102+
&& python3 setup.py bdist_wheel --dist-dir=dist \
103+
&& pip install dist/*.whl
104104
RUN git clone ${FA_REPO}
105105
RUN cd flash-attention \
106106
&& git checkout ${FA_BRANCH} \
@@ -132,8 +132,8 @@ RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
132132
cp /install/*.whl /app/debs
133133
RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
134134
cp /install/*.whl /app/debs
135-
# RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
136-
# cp /install/*.whl /app/debs
135+
RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
136+
cp /install/*.whl /app/debs
137137
RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
138138
cp /install/*.whl /app/debs
139139

@@ -147,8 +147,8 @@ RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
147147
pip install /install/*.whl
148148
RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
149149
pip install /install/*.whl
150-
# RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
151-
# pip install /install/*.whl
150+
RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
151+
pip install /install/*.whl
152152
RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
153153
pip install /install/*.whl
154154

@@ -172,10 +172,10 @@ RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
172172
# && echo "LEGACY_HIPBLASLT_OPTION: ${LEGACY_HIPBLASLT_OPTION}" >> /app/versions.txt \
173173
&& echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
174174
&& echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
175-
# && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \
176-
# && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \
177-
# && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \
178-
# && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \
175+
&& echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \
176+
&& echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \
177+
&& echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \
178+
&& echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \
179179
&& echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
180180
&& echo "AITER_BRANCH: ${AITER_BRANCH}" >> /app/versions.txt \
181181
&& echo "AITER_REPO: ${AITER_REPO}" >> /app/versions.txt

‎vllm/attention/backends/rocm_flash_attn.py‎

Lines changed: 25 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -561,7 +561,8 @@ def __init__(
561561

562562
self.use_naive_attn = False
563563
# NOTE: Allow for switching between Triton and CK. Defaulting to triton.
564-
self.use_triton_flash_attn = envs.VLLM_USE_TRITON_FLASH_ATTN
564+
self.use_triton_flash_attn = envs.VLLM_USE_TRITON_FLASH_ATTN and \
565+
envs.VLLM_ROCM_USE_AITER
565566
if self.use_triton_flash_attn:
566567
if logits_soft_cap is not None:
567568
raise ValueError(
@@ -571,16 +572,11 @@ def __init__(
571572
"FA backend instead by setting the env var "
572573
"`VLLM_USE_TRITON_FLASH_ATTN=0`")
573574

574-
if self.kv_cache_dtype in ["int8", "fp8_e4m3"]:
575-
from vllm.attention.ops.triton_flash_attention import ( # noqa: F401
576-
triton_attention)
577-
self.triton_attn_func = triton_attention
578-
else:
579-
from aiter.ops.triton.mha import flash_attn_varlen_func
580-
from aiter.ops.triton.mha import (
581-
mha_set_use_int64_strides as set_triton_fa_strides)
582-
set_triton_fa_strides(True)
583-
self.triton_attn_func = flash_attn_varlen_func
575+
from aiter.ops.triton.mha import flash_attn_varlen_func
576+
from aiter.ops.triton.mha import (mha_set_use_int64_strides as
577+
set_triton_fa_strides)
578+
set_triton_fa_strides(True)
579+
self.triton_attn_func = flash_attn_varlen_func
584580
logger.debug("Using Triton FA in ROCmBackend")
585581
if self.sliding_window != (-1, -1):
586582
logger.warning("ROCm Triton FA does not currently support "
@@ -803,52 +799,24 @@ def forward(
803799
query.dtype,
804800
seq_lens,
805801
make_attn_mask=causal_mask) # type: ignore
806-
if self.kv_cache_dtype in ["int8", "fp8_e4m3"]:
807-
use_fp8_scales = (layer._q_scale is not None
808-
and layer._k_scale is not None
809-
and layer._v_scale is not None
810-
and layer._prob_scale is not None and
811-
envs.VLLM_USE_ROCM_FP8_FLASH_ATTN)
812-
full_scales = (layer._q_scale.item(),
813-
layer._k_scale.item(),
814-
layer._v_scale.item(),
815-
layer._prob_scale.item()
816-
) if use_fp8_scales else None
817-
self.triton_attn_func(
818-
query,
819-
key,
820-
value,
821-
output[:num_prefill_tokens],
822-
query_seq_start_loc,
823-
key_seq_start_loc,
824-
query_max_seq_len,
825-
key_max_seq_len,
826-
causal_mask,
827-
self.scale,
828-
attn_masks[0][None]
829-
if attn_masks is not None else None,
830-
full_scales,
831-
layer._out_scale,
832-
)
833-
else:
834-
output[:num_prefill_tokens] = self.triton_attn_func(
835-
q=query,
836-
k=key,
837-
v=value,
838-
cu_seqlens_q=query_seq_start_loc,
839-
cu_seqlens_k=key_seq_start_loc,
840-
max_seqlen_q=query_max_seq_len,
841-
max_seqlen_k=key_max_seq_len,
842-
dropout_p=0.0,
843-
softmax_scale=self.scale,
844-
causal=causal_mask,
845-
window_size=self.sliding_window,
846-
alibi_slopes=self.alibi_slopes,
847-
deterministic=False,
848-
return_lse=False,
849-
return_attn_probs=False,
850-
block_table=None,
851-
)
802+
output[:num_prefill_tokens] = self.triton_attn_func(
803+
q=query,
804+
k=key,
805+
v=value,
806+
cu_seqlens_q=query_seq_start_loc,
807+
cu_seqlens_k=key_seq_start_loc,
808+
max_seqlen_q=query_max_seq_len,
809+
max_seqlen_k=key_max_seq_len,
810+
dropout_p=0.0,
811+
softmax_scale=self.scale,
812+
causal=causal_mask,
813+
window_size=self.sliding_window,
814+
alibi_slopes=self.alibi_slopes,
815+
deterministic=False,
816+
return_lse=False,
817+
return_attn_probs=False,
818+
block_table=None,
819+
)
852820
elif self.use_naive_attn:
853821
if self.num_kv_heads != self.num_heads:
854822
# Interleave for MQA workaround.

0 commit comments

Comments
 (0)