Skip to content

[BUGFIX][ROCM] ViT FlashAttention on ROCm (no GFX9) and contiguous on qwen3vl ROCm TORCH_SDPA #27190

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
DarkLight1337 merged 10 commits into from
Oct 26, 2025
Merged

[BUGFIX][ROCM] ViT FlashAttention on ROCm (no GFX9) and contiguous on qwen3vl ROCm TORCH_SDPA #27190

Changes from 2 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
e87f29c
fixbug vit_flash_attn_on_rocm_no_gfx9 use rocm selection for it
JartX Oct 20, 2025
c8735e7
refactor precommit
JartX Oct 20, 2025
889cde8
Merge branch 'main' into fixbug/torch_sdpa_rdna3_for_vl_models
JartX Oct 22, 2025
633792c
refactor maybe_get_vit_flash_attn_backend
JartX Oct 22, 2025
493f683
flassh_attn on rocm if on_gfx9() and find_spec(flash_attn)
JartX Oct 22, 2025
010fa99
Merge branch 'main' into fixbug/torch_sdpa_rdna3_for_vl_models and re…
JartX Oct 24, 2025
f8aa45f
pre-commit
JartX Oct 24, 2025
0dabbf1
qwen2vl and qwen2.5vl contiguous on rocm and torch.sdpa
JartX Oct 25, 2025
c42f2d1
precommit
JartX Oct 25, 2025
1e3074b
Merge branch 'main' into fixbug/torch_sdpa_rdna3_for_vl_models
JartX Oct 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 23 additions & 8 deletions vllm/attention/layer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,12 @@
SlidingWindowSpec,
)

if current_platform.is_rocm():
from vllm.platforms.rocm import on_gfx9
else:
on_gfx9 = lambda *args, **kwargs: False


FP8_DTYPE = current_platform.fp8_dtype()
logger = init_logger(__name__)
USE_XFORMERS_OPS = None
Expand Down Expand Up @@ -93,14 +99,22 @@ def check_upstream_fa_availability(dtype: torch.dtype):

def maybe_get_vit_flash_attn_backend(
attn_backend: _Backend, use_upstream_fa: bool
) -> tuple[_Backend, Callable]:
if (
attn_backend != _Backend.FLASH_ATTN
and attn_backend != _Backend.ROCM_AITER_FA
and check_upstream_fa_availability(torch.get_default_dtype())
):
attn_backend = _Backend.FLASH_ATTN
use_upstream_fa = True
) -> tuple[_Backend, Callable | None]:
if current_platform.is_rocm():
if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9():
Comment thread
JartX marked this conversation as resolved.
attn_backend = _Backend.ROCM_AITER_FA
elif on_gfx9():

@tjtanaa tjtanaa Oct 21, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This condition was crafted for ROCm platform (on_gfx9()).

        if (
            attn_backend != _Backend.FLASH_ATTN
            and attn_backend != _Backend.ROCM_AITER_FA
            and check_upstream_fa_availability(torch.get_default_dtype())
        ):

On on_gfx9(), we will always attempt to use flash_attn, but if it does not support, then we fallback to _Backend.TORCH_SDPA

@JartX JartX Oct 21, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tjtanaa The first step was simply to check if it was rocm and if it was on_gfx1x and if so, return _Backend.TORCH_SDPA, None directly. Would you like to leave it like this with a comment indicating why? I've only seen it affect RDNA3 with the upstream flow. I think this way we could mark the required changes as resolved. If not, please tell me how the code should be :)

attn_backend = _Backend.FLASH_ATTN
else:
return _Backend.TORCH_SDPA, None
else:
if (

@tjtanaa tjtanaa Oct 21, 2025

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As mentioned above, this condition and its content is also applicable to ROCm on gfx9.

attn_backend != _Backend.FLASH_ATTN
and attn_backend != _Backend.ROCM_AITER_FA
and check_upstream_fa_availability(torch.get_default_dtype())
):
attn_backend = _Backend.FLASH_ATTN
use_upstream_fa = True
Comment thread
JartX marked this conversation as resolved.

if current_platform.is_rocm() and attn_backend == _Backend.FLASH_ATTN:
use_upstream_fa = True
Expand Down Expand Up @@ -534,6 +548,7 @@ def forward(
value = torch.repeat_interleave(value, num_repeat, dim=2)

if self.is_flash_attn_backend:
assert self._flash_attn_varlen_func is not None
cu_seqlens_q = torch.arange(
0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device
)
Expand Down