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 12, 09:45 UTC
Remove upstream fa checks (#29471)
Signed-off-by: mingyuanm <mingyuanm@nvidia.com> Signed-off-by: Roger Wang <hey@rogerw.io> Co-authored-by: Roger Wang <hey@rogerw.io>
Mingyuan Ma committed
Nov 28, 2025 at 05:52 UTC
460d8bbf2d19fc4757b3253050a0a18b61c136e3
13 files changed
+18
-148
vllm/attention/layer.py
+7
-50
@@ -56,53 +56,28 @@ FP8_DTYPE = current_platform.fp8_dtype()
56
logger = init_logger(__name__)
57
58
59
-def check_upstream_fa_availability(dtype: torch.dtype):
60
- if (
61
- dtype in (torch.float16, torch.bfloat16)
62
- and current_platform.is_cuda()
63
- and current_platform.has_device_capability(80)
64
- ):
65
- from transformers.utils import is_flash_attn_2_available
66
-
67
- return is_flash_attn_2_available()
68
- if current_platform.is_rocm():
69
- from importlib.util import find_spec
70
-
71
- return find_spec("flash_attn") is not None
72
- return False
73
-
74
-
59
def maybe_get_vit_flash_attn_backend(
60
attn_backend: AttentionBackendEnum,
77
- use_upstream_fa: bool,
61
attn_backend_override: AttentionBackendEnum | None = None,
62
) -> tuple[AttentionBackendEnum, Callable | None]:
63
if current_platform.is_rocm():
64
if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9():
65
attn_backend = AttentionBackendEnum.ROCM_AITER_FA
83
-
66
elif (
85
- check_upstream_fa_availability(torch.get_default_dtype())
67
+ attn_backend_override is None
68
and on_gfx9()
87
- and attn_backend_override is None
69
+ and attn_backend == AttentionBackendEnum.FLASH_ATTN
70
):
89
- attn_backend = AttentionBackendEnum.FLASH_ATTN
90
- use_upstream_fa = True
71
+ pass
72
else:
73
return AttentionBackendEnum.TORCH_SDPA, None
93
-
74
elif current_platform.is_cuda():
95
- if (
96
- attn_backend != AttentionBackendEnum.FLASH_ATTN
97
- and check_upstream_fa_availability(torch.get_default_dtype())
98
- ):
99
- attn_backend = AttentionBackendEnum.FLASH_ATTN
100
- use_upstream_fa = True
75
+ pass
76
elif current_platform.is_xpu():
77
assert attn_backend == AttentionBackendEnum.FLASH_ATTN, (
78
"XPU platform only supports FLASH_ATTN as vision attention backend."
79
)
105
- use_upstream_fa = False
80
+ pass
81
else:
82
return AttentionBackendEnum.TORCH_SDPA, None
83
@@ -113,10 +88,7 @@ def maybe_get_vit_flash_attn_backend(
88
if attn_backend == AttentionBackendEnum.ROCM_AITER_FA:
89
from aiter import flash_attn_varlen_func
90
else:
116
- if use_upstream_fa:
117
- from flash_attn import flash_attn_varlen_func
118
- else:
119
- from vllm.attention.utils.fa_utils import flash_attn_varlen_func
91
+ from vllm.attention.utils.fa_utils import flash_attn_varlen_func
92
else:
93
flash_attn_varlen_func = None
94
@@ -501,11 +473,6 @@ class MultiHeadAttention(nn.Module):
473
attn_backend_override=attn_backend_override,
474
)
475
504
- # Some auto-selected backends can be upgraded
505
- # to upstream flash attention if available.
506
- # If vllm native fa is selected, we use it directly.
507
- use_upstream_fa = False
508
-
476
self.attn_backend = (
477
backend
478
if backend
@@ -521,7 +488,6 @@ class MultiHeadAttention(nn.Module):
488
self.attn_backend, self._flash_attn_varlen_func = (
489
maybe_get_vit_flash_attn_backend(
490
self.attn_backend,
524
- use_upstream_fa,
491
attn_backend_override=attn_backend_override,
492
)
493
)
@@ -531,17 +497,8 @@ class MultiHeadAttention(nn.Module):
497
AttentionBackendEnum.ROCM_AITER_FA,
498
}
499
534
- # this condition is just to make sure that the
535
- # use_upstream_fa in the log is correct
536
- if (
537
- current_platform.is_rocm()
538
- and self.attn_backend == AttentionBackendEnum.FLASH_ATTN
539
- ):
540
- use_upstream_fa = True
541
-
500
logger.info_once(
543
- f"MultiHeadAttention attn_backend: {self.attn_backend}, "
544
- f"use_upstream_fa: {use_upstream_fa}"
501
+ f"Using {self.attn_backend} for MultiHeadAttention in multimodal encoder."
502
)
503
504
def forward(
vllm/attention/ops/vit_attn_wrappers.py
+2
-8
@@ -27,15 +27,11 @@ def flash_attn_maxseqlen_wrapper(
27
max_seqlen: torch.Tensor,
28
batch_size: int,
29
is_rocm_aiter: bool,
30
- use_upstream_fa: bool,
30
) -> torch.Tensor:
31
if is_rocm_aiter:
32
from aiter import flash_attn_varlen_func
33
else:
35
- if use_upstream_fa:
36
- from flash_attn import flash_attn_varlen_func
37
- else:
38
- from vllm.attention.utils.fa_utils import flash_attn_varlen_func
34
+ from vllm.attention.utils.fa_utils import flash_attn_varlen_func
35
q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
36
output = flash_attn_varlen_func(
37
q,
@@ -62,7 +58,6 @@ def flash_attn_maxseqlen_wrapper_fake(
58
max_seqlen: torch.Tensor,
59
batch_size: int,
60
is_rocm_aiter: bool,
65
- use_upstream_fa: bool,
61
) -> torch.Tensor:
62
b, s, h, d = q.shape
63
return torch.empty((s, b, h * d), dtype=q.dtype, device=q.device)
@@ -83,10 +78,9 @@ def vit_flash_attn_wrapper(
78
max_seqlen: torch.Tensor,
79
batch_size: int,
80
is_rocm_aiter: bool,
86
- use_upstream_fa: bool,
81
) -> torch.Tensor:
82
return torch.ops.vllm.flash_attn_maxseqlen_wrapper(
89
- q, k, v, cu_seqlens, max_seqlen, batch_size, is_rocm_aiter, use_upstream_fa
83
+ q, k, v, cu_seqlens, max_seqlen, batch_size, is_rocm_aiter
84
)
85
86
vllm/attention/utils/fa_utils.py
+8
@@ -18,6 +18,14 @@ elif current_platform.is_xpu():
18
reshape_and_cache_flash = ops.reshape_and_cache_flash
19
flash_attn_varlen_func = ops.flash_attn_varlen_func
20
get_scheduler_metadata = ops.get_scheduler_metadata
21
+elif current_platform.is_rocm():
22
+ try:
23
+ from flash_attn import flash_attn_varlen_func # noqa: F401
24
+ except ImportError as e:
25
+ raise ImportError(
26
+ "Rocm platform requires upstream flash-attn "
27
+ "to be installed. Please install flash-attn first."
28
+ ) from e
29
30
31
def get_flash_attn_version(requires_alibi: bool = False) -> int | None:
vllm/model_executor/models/dots_ocr.py
-8
@@ -11,7 +11,6 @@ from transformers.models.qwen2_vl import Qwen2VLProcessor
11
12
from vllm.attention.backends.registry import AttentionBackendEnum
13
from vllm.attention.layer import (
14
- check_upstream_fa_availability,
14
maybe_get_vit_flash_attn_backend,
15
)
16
from vllm.config import VllmConfig
@@ -294,12 +293,10 @@ class DotsVisionAttention(nn.Module):
293
torch.get_default_dtype(),
294
attn_backend_override=attn_backend_override,
295
)
297
- self.use_upstream_fa = False
296
297
self.attn_backend, self.flash_attn_varlen_func = (
298
maybe_get_vit_flash_attn_backend(
299
self.attn_backend,
302
- self.use_upstream_fa,
300
attn_backend_override=attn_backend_override,
301
)
302
)
@@ -569,11 +566,6 @@ class DotsVisionTransformer(nn.Module):
566
dtype=torch.get_default_dtype(),
567
attn_backend_override=attn_backend_override,
568
)
572
- if (
573
- self.attn_backend != AttentionBackendEnum.FLASH_ATTN
574
- and check_upstream_fa_availability(torch.get_default_dtype())
575
- ):
576
- self.attn_backend = AttentionBackendEnum.FLASH_ATTN
569
self.out_hidden_size = config.hidden_size
570
# Keep blocks for compatibility with other vision towers
571
num_layers = (
vllm/model_executor/models/ernie45_vl.py
-9
@@ -38,7 +38,6 @@ from transformers import BatchFeature
38
39
from vllm.attention.backends.registry import AttentionBackendEnum
40
from vllm.attention.layer import (
41
- check_upstream_fa_availability,
41
maybe_get_vit_flash_attn_backend,
42
)
43
from vllm.config import VllmConfig
@@ -201,12 +200,9 @@ class Ernie4_5_VisionAttention(nn.Module):
200
attn_backend_override=attn_backend_override,
201
)
202
204
- self.use_upstream_fa = False
205
-
203
self.attn_backend, self.flash_attn_varlen_func = (
204
maybe_get_vit_flash_attn_backend(
205
self.attn_backend,
209
- self.use_upstream_fa,
206
attn_backend_override=attn_backend_override,
207
)
208
)
@@ -498,11 +494,6 @@ class Ernie4_5_VisionTransformer(nn.Module):
494
dtype=torch.get_default_dtype(),
495
attn_backend_override=attn_backend_override,
496
)
501
- if (
502
- self.attn_backend != AttentionBackendEnum.FLASH_ATTN
503
- and check_upstream_fa_availability(torch.get_default_dtype())
504
- ):
505
- self.attn_backend = AttentionBackendEnum.FLASH_ATTN
497
498
@property
499
def dtype(self) -> torch.dtype:
vllm/model_executor/models/glm4_1v.py
+1
-11
@@ -47,10 +47,7 @@ from transformers.models.glm4v.video_processing_glm4v import Glm4vVideoProcessor
47
from transformers.video_utils import VideoMetadata
48
49
from vllm.attention.backends.registry import AttentionBackendEnum
50
-from vllm.attention.layer import (
51
- check_upstream_fa_availability,
52
- maybe_get_vit_flash_attn_backend,
53
-)
50
+from vllm.attention.layer import maybe_get_vit_flash_attn_backend
51
from vllm.config import VllmConfig
52
from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
53
from vllm.distributed import get_tensor_model_parallel_world_size, parallel_state
@@ -296,12 +293,10 @@ class Glm4vVisionAttention(nn.Module):
293
dtype=torch.get_default_dtype(),
294
attn_backend_override=attn_backend_override,
295
)
299
- self.use_upstream_fa = False
296
297
self.attn_backend, self.flash_attn_varlen_func = (
298
maybe_get_vit_flash_attn_backend(
299
self.attn_backend,
304
- self.use_upstream_fa,
300
attn_backend_override=attn_backend_override,
301
)
302
)
@@ -730,11 +725,6 @@ class Glm4vVisionTransformer(nn.Module):
725
dtype=torch.get_default_dtype(),
726
attn_backend_override=attn_backend_override,
727
)
733
- if (
734
- self.attn_backend != AttentionBackendEnum.FLASH_ATTN
735
- and check_upstream_fa_availability(torch.get_default_dtype())
736
- ):
737
- self.attn_backend = AttentionBackendEnum.FLASH_ATTN
728
729
@property
730
def dtype(self) -> torch.dtype:
vllm/model_executor/models/keye.py
-1
@@ -418,7 +418,6 @@ class KeyeSiglipAttention(nn.Module):
418
self.attn_backend, self.flash_attn_varlen_func = (
419
maybe_get_vit_flash_attn_backend(
420
self.attn_backend,
421
- use_upstream_fa=False,
421
attn_backend_override=attn_backend_override,
422
)
423
)
vllm/model_executor/models/paddleocr_vl.py
-15
@@ -33,7 +33,6 @@ from transformers.utils import torch_int
33
34
from vllm.attention.backends.registry import AttentionBackendEnum
35
from vllm.attention.layer import (
36
- check_upstream_fa_availability,
36
maybe_get_vit_flash_attn_backend,
37
)
38
from vllm.attention.ops.vit_attn_wrappers import (
@@ -582,7 +581,6 @@ class SiglipAttention(nn.Module):
581
prefix: str = "",
582
attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
583
attn_backend_override: AttentionBackendEnum | None = None,
585
- use_upstream_fa: bool = False,
584
) -> None:
585
super().__init__()
586
@@ -612,11 +610,9 @@ class SiglipAttention(nn.Module):
610
)
611
612
self.attn_backend = attn_backend
615
- self.use_upstream_fa = use_upstream_fa
613
self.attn_backend, self.flash_attn_varlen_func = (
614
maybe_get_vit_flash_attn_backend(
615
self.attn_backend,
619
- self.use_upstream_fa,
616
attn_backend_override=attn_backend_override,
617
)
618
)
@@ -680,7 +676,6 @@ class SiglipAttention(nn.Module):
676
max_seqlen,
677
batch_size,
678
self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA,
683
- self.use_upstream_fa,
679
)
680
elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
681
outputs = []
@@ -783,7 +778,6 @@ class SiglipEncoderLayer(nn.Module):
778
*,
779
attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
780
attn_backend_override: AttentionBackendEnum | None = None,
786
- use_upstream_fa: bool = False,
781
):
782
super().__init__()
783
self.embed_dim = config.hidden_size
@@ -796,7 +790,6 @@ class SiglipEncoderLayer(nn.Module):
790
prefix=f"{prefix}.self_attn",
791
attn_backend=attn_backend,
792
attn_backend_override=attn_backend_override,
799
- use_upstream_fa=use_upstream_fa,
793
)
794
self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
795
self.mlp = SiglipMLP(
@@ -852,13 +845,6 @@ class SiglipEncoder(nn.Module):
845
dtype=torch.get_default_dtype(),
846
attn_backend_override=attn_backend_override,
847
)
855
- self.use_upstream_fa = False
856
- if self.attn_backend not in {
857
- AttentionBackendEnum.FLASH_ATTN,
858
- AttentionBackendEnum.ROCM_AITER_FA,
859
- } and check_upstream_fa_availability(torch.get_default_dtype()):
860
- self.attn_backend = AttentionBackendEnum.FLASH_ATTN
861
- self.use_upstream_fa = True
848
if self.attn_backend not in {
849
AttentionBackendEnum.FLASH_ATTN,
850
AttentionBackendEnum.TORCH_SDPA,
@@ -875,7 +861,6 @@ class SiglipEncoder(nn.Module):
861
prefix=f"{prefix}.layers.{layer_idx}",
862
attn_backend=self.attn_backend,
863
attn_backend_override=attn_backend_override,
878
- use_upstream_fa=self.use_upstream_fa,
864
)
865
for layer_idx in range(config.num_hidden_layers)
866
]
vllm/model_executor/models/qwen2_5_vl.py
-18
@@ -307,7 +307,6 @@ class Qwen2_5_VisionAttention(nn.Module):
307
prefix: str = "",
308
use_data_parallel: bool = False,
309
attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
310
- use_upstream_fa: bool = False,
310
attn_backend_override: AttentionBackendEnum | None = None,
311
) -> None:
312
super().__init__()
@@ -344,24 +343,13 @@ class Qwen2_5_VisionAttention(nn.Module):
343
disable_tp=use_data_parallel,
344
)
345
self.attn_backend = attn_backend
347
- self.use_upstream_fa = use_upstream_fa
346
self.attn_backend, self.flash_attn_varlen_func = (
347
maybe_get_vit_flash_attn_backend(
348
self.attn_backend,
351
- self.use_upstream_fa,
349
attn_backend_override=attn_backend_override,
350
)
351
)
355
- # On ROCm with FLASH_ATTN backend, upstream flash_attn is used
356
- from vllm.platforms import current_platform
352
358
- if (
359
- current_platform.is_rocm()
360
- and self.attn_backend == AttentionBackendEnum.FLASH_ATTN
361
- ):
362
- self.use_upstream_fa = True
363
- if current_platform.is_xpu():
364
- self.use_upstream_fa = False
353
self.is_flash_attn_backend = self.attn_backend in {
354
AttentionBackendEnum.FLASH_ATTN,
355
AttentionBackendEnum.ROCM_AITER_FA,
@@ -415,7 +403,6 @@ class Qwen2_5_VisionAttention(nn.Module):
403
max_seqlen,
404
batch_size,
405
self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA,
418
- self.use_upstream_fa,
406
)
407
elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
408
# Execute attention entry by entry for speed & less VRAM.
@@ -459,7 +446,6 @@ class Qwen2_5_VisionBlock(nn.Module):
446
prefix: str = "",
447
use_data_parallel: bool = False,
448
attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
462
- use_upstream_fa: bool = False,
449
attn_backend_override: AttentionBackendEnum | None = None,
450
) -> None:
451
super().__init__()
@@ -475,7 +461,6 @@ class Qwen2_5_VisionBlock(nn.Module):
461
prefix=f"{prefix}.attn",
462
use_data_parallel=use_data_parallel,
463
attn_backend=attn_backend,
478
- use_upstream_fa=use_upstream_fa,
464
attn_backend_override=attn_backend_override,
465
)
466
self.mlp = Qwen2_5_VisionMLP(
@@ -644,7 +629,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
629
is_neox_style=True,
630
)
631
647
- use_upstream_fa = False
632
self.attn_backend = get_vit_attn_backend(
633
head_size=head_dim,
634
dtype=torch.get_default_dtype(),
@@ -654,7 +638,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
638
self.attn_backend, self.flash_attn_varlen_func = (
639
maybe_get_vit_flash_attn_backend(
640
self.attn_backend,
657
- use_upstream_fa,
641
attn_backend_override=attn_backend_override,
642
)
643
)
@@ -681,7 +664,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
664
prefix=f"{prefix}.blocks.{layer_idx}",
665
use_data_parallel=use_data_parallel,
666
attn_backend=self.attn_backend,
684
- use_upstream_fa=use_upstream_fa,
667
attn_backend_override=attn_backend_override,
668
)
669
for layer_idx in range(depth)
vllm/model_executor/models/qwen2_vl.py
-8
@@ -45,7 +45,6 @@ from transformers.models.qwen2_vl.video_processing_qwen2_vl import Qwen2VLVideoP
45
46
from vllm.attention.backends.registry import AttentionBackendEnum
47
from vllm.attention.layer import (
48
- check_upstream_fa_availability,
48
maybe_get_vit_flash_attn_backend,
49
)
50
from vllm.config import VllmConfig
@@ -335,12 +334,10 @@ class Qwen2VisionAttention(nn.Module):
334
dtype=torch.get_default_dtype(),
335
attn_backend_override=attn_backend_override,
336
)
338
- self.use_upstream_fa = False
337
338
self.attn_backend, self.flash_attn_varlen_func = (
339
maybe_get_vit_flash_attn_backend(
340
self.attn_backend,
343
- self.use_upstream_fa,
341
attn_backend_override=attn_backend_override,
342
)
343
)
@@ -657,11 +654,6 @@ class Qwen2VisionTransformer(nn.Module):
654
dtype=torch.get_default_dtype(),
655
attn_backend_override=attn_backend_override,
656
)
660
- if (
661
- self.attn_backend != AttentionBackendEnum.FLASH_ATTN
662
- and check_upstream_fa_availability(torch.get_default_dtype())
663
- ):
664
- self.attn_backend = AttentionBackendEnum.FLASH_ATTN
657
658
@property
659
def dtype(self) -> torch.dtype:
vllm/model_executor/models/qwen3_omni_moe_thinker.py
-6
@@ -47,7 +47,6 @@ from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import (
47
from transformers.models.whisper import WhisperFeatureExtractor
48
49
from vllm.attention.backends.registry import AttentionBackendEnum
50
-from vllm.attention.layer import check_upstream_fa_availability
50
from vllm.compilation.decorators import support_torch_compile
51
from vllm.config import VllmConfig
52
from vllm.distributed import get_pp_group
@@ -381,11 +380,6 @@ class Qwen3Omni_VisionTransformer(nn.Module):
380
dtype=torch.get_default_dtype(),
381
attn_backend_override=attn_backend_override,
382
)
384
- if (
385
- self.attn_backend != AttentionBackendEnum.FLASH_ATTN
386
- and check_upstream_fa_availability(torch.get_default_dtype())
387
- ):
388
- self.attn_backend = AttentionBackendEnum.FLASH_ATTN
383
384
@property
385
def dtype(self) -> torch.dtype:
vllm/model_executor/models/qwen3_vl.py
-12
@@ -49,7 +49,6 @@ from transformers.models.qwen3_vl.video_processing_qwen3_vl import (
49
from transformers.video_utils import VideoMetadata
50
51
from vllm.attention.backends.registry import AttentionBackendEnum
52
-from vllm.attention.layer import check_upstream_fa_availability
52
from vllm.compilation.decorators import support_torch_compile
53
from vllm.config import VllmConfig
54
from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
@@ -202,7 +201,6 @@ class Qwen3_VisionBlock(nn.Module):
201
prefix: str = "",
202
use_data_parallel: bool = False,
203
attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
205
- use_upstream_fa: bool = False,
204
) -> None:
205
super().__init__()
206
if norm_layer is None:
@@ -217,7 +215,6 @@ class Qwen3_VisionBlock(nn.Module):
215
prefix=f"{prefix}.attn",
216
use_data_parallel=use_data_parallel,
217
attn_backend=attn_backend,
220
- use_upstream_fa=use_upstream_fa,
218
)
219
self.mlp = Qwen3_VisionMLP(
220
dim,
@@ -378,14 +375,6 @@ class Qwen3_VisionTransformer(nn.Module):
375
dtype=torch.get_default_dtype(),
376
attn_backend_override=attn_backend_override,
377
)
381
- use_upstream_fa = False
382
- if (
383
- self.attn_backend != AttentionBackendEnum.FLASH_ATTN
384
- and self.attn_backend != AttentionBackendEnum.ROCM_AITER_FA
385
- and check_upstream_fa_availability(torch.get_default_dtype())
386
- ):
387
- self.attn_backend = AttentionBackendEnum.FLASH_ATTN
388
- use_upstream_fa = True
378
379
if self.attn_backend not in {
380
AttentionBackendEnum.FLASH_ATTN,
@@ -407,7 +396,6 @@ class Qwen3_VisionTransformer(nn.Module):
396
prefix=f"{prefix}.blocks.{layer_idx}",
397
use_data_parallel=use_data_parallel,
398
attn_backend=self.attn_backend,
410
- use_upstream_fa=use_upstream_fa,
399
)
400
for layer_idx in range(vision_config.depth)
401
]
vllm/model_executor/models/siglip2navit.py
-2
@@ -255,12 +255,10 @@ class Siglip2Attention(nn.Module):
255
dtype=torch.get_default_dtype(),
256
attn_backend_override=attn_backend_override,
257
)
258
- self.use_upstream_fa = False
258
259
self.attn_backend, self.flash_attn_varlen_func = (
260
maybe_get_vit_flash_attn_backend(
261
self.attn_backend,
263
- self.use_upstream_fa,
262
attn_backend_override=attn_backend_override,
263
)
264
)