A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 05:00 UTC
[ROCm][Bugfix] Fix `fa_version` argument error in `flash_attn_maxseqlen_wrapper` for ROCm without aiter (#30909)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Dec 18, 2025 at 02:45 UTC
8da6ae49c3d99e147bc739e54fe57e43581d887b
1 file changed
+5
-4
vllm/attention/ops/vit_attn_wrappers.py
+5
-4
@@ -28,7 +28,7 @@ def flash_attn_maxseqlen_wrapper(
28
max_seqlen: torch.Tensor,
29
batch_size: int,
30
is_rocm_aiter: bool,
31
- fa_version: int,
31
+ fa_version: int | None,
32
) -> torch.Tensor:
33
kwargs = {}
34
if is_rocm_aiter:
@@ -36,7 +36,8 @@ def flash_attn_maxseqlen_wrapper(
36
else:
37
from vllm.attention.utils.fa_utils import flash_attn_varlen_func
38
39
- kwargs["fa_version"] = fa_version
39
+ if not current_platform.is_rocm() and fa_version is not None:
40
+ kwargs["fa_version"] = fa_version
41
q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
42
output = flash_attn_varlen_func(
43
q,
@@ -62,7 +63,7 @@ def flash_attn_maxseqlen_wrapper_fake(
63
max_seqlen: torch.Tensor,
64
batch_size: int,
65
is_rocm_aiter: bool,
65
- fa_version: int,
66
+ fa_version: int | None,
67
) -> torch.Tensor:
68
return torch.empty_like(q)
69
@@ -82,7 +83,7 @@ def vit_flash_attn_wrapper(
83
max_seqlen: torch.Tensor,
84
batch_size: int,
85
is_rocm_aiter: bool,
85
- fa_version: int,
86
+ fa_version: int | None,
87
) -> torch.Tensor:
88
return torch.ops.vllm.flash_attn_maxseqlen_wrapper(
89
q, k, v, cu_seqlens, max_seqlen, batch_size, is_rocm_aiter, fa_version