@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9c5ee91b2a

Fork 0

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 14, 05:30 UTC

[ROCm] [VL] [Bugfix] Fix vit flash attn dispatcher logic for ROCm (#26104)

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

TJian committed Oct 2, 2025 at 22:34 UTC 9c5ee91b2af834fb3221787d63ac025badbe0168

9 files changed +154 -141

vllm/attention/layer.py

+49 -23

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Attention layer."""

4 -from typing import List, Optional

4 +from typing import Callable, List, Optional

6 import torch

7 import torch.nn as nn

@@ -68,9 +68,39 @@ def check_upstream_fa_availability(dtype: torch.dtype):

68 ) and current_platform.has_device_capability(80):

69 from transformers.utils import is_flash_attn_2_available

70 return is_flash_attn_2_available()

71 + if current_platform.is_rocm():

72 + from importlib.util import find_spec

73 + return find_spec("flash_attn") is not None

74 return False

77 +def maybe_get_vit_flash_attn_backend(

78 + attn_backend: _Backend,

79 + use_upstream_fa: bool) -> tuple[_Backend, Callable]:

80 + if attn_backend != _Backend.FLASH_ATTN and \

81 + attn_backend != _Backend.ROCM_AITER_FA and \

82 + check_upstream_fa_availability(torch.get_default_dtype()):

83 + attn_backend = _Backend.FLASH_ATTN

84 + use_upstream_fa = True

85 +

86 + if current_platform.is_rocm() and \

87 + attn_backend == _Backend.FLASH_ATTN:

88 + use_upstream_fa = True

89 +

90 + if (attn_backend in {_Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA}):

91 + if attn_backend == _Backend.ROCM_AITER_FA:

92 + from aiter import flash_attn_varlen_func

93 + else:

94 + if use_upstream_fa:

95 + from flash_attn import flash_attn_varlen_func

96 + else:

97 + from vllm.vllm_flash_attn import flash_attn_varlen_func

98 + else:

99 + flash_attn_varlen_func = None

100 +

101 + return attn_backend, flash_attn_varlen_func

102 +

103 +

104 class Attention(nn.Module, AttentionLayerBase):

105 """Attention layer.

106

@@ -410,13 +440,9 @@ class MultiHeadAttention(nn.Module):

440 # to upstream flash attention if available.

441 # If vllm native fa is selected, we use it directly.

442 use_upstream_fa = False

413 - if backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(

414 - dtype):

415 - backend = _Backend.FLASH_ATTN

416 - use_upstream_fa = True

443

418 - if current_platform.is_rocm() or current_platform.is_xpu():

419 - # currently, only torch_sdpa is supported on rocm/xpu

444 + if current_platform.is_xpu():

445 + # currently, only torch_sdpa is supported on xpu

446 self.attn_backend = _Backend.TORCH_SDPA

447 else:

448

@@ -428,17 +454,25 @@ class MultiHeadAttention(nn.Module):

454 _Backend.FLASH_ATTN,

455 } else _Backend.TORCH_SDPA

456

457 + self.attn_backend, self._flash_attn_varlen_func \

458 + = maybe_get_vit_flash_attn_backend(

459 + self.attn_backend,

460 + use_upstream_fa,

461 + )

462 +

463 if (self.attn_backend == _Backend.XFORMERS

464 and not check_xformers_availability()):

465 self.attn_backend = _Backend.TORCH_SDPA

466

435 - if self.attn_backend == _Backend.FLASH_ATTN:

436 - if use_upstream_fa:

437 - from flash_attn import flash_attn_varlen_func

438 - self._flash_attn_varlen_func = flash_attn_varlen_func

439 - else:

440 - from vllm.vllm_flash_attn import flash_attn_varlen_func

441 - self._flash_attn_varlen_func = flash_attn_varlen_func

467 + self.is_flash_attn_backend = self.attn_backend in {

468 + _Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA

469 + }

470 +

471 + # this condition is just to make sure that the

472 + # use_upstream_fa in the log is correct

473 + if current_platform.is_rocm() \

474 + and self.attn_backend == _Backend.FLASH_ATTN:

475 + use_upstream_fa = True

476

477 logger.info_once(

478 f"MultiHeadAttention attn_backend: {self.attn_backend}, "

@@ -466,7 +500,7 @@ class MultiHeadAttention(nn.Module):

500 key = torch.repeat_interleave(key, num_repeat, dim=2)

501 value = torch.repeat_interleave(value, num_repeat, dim=2)

502

469 - if self.attn_backend == _Backend.FLASH_ATTN:

503 + if self.is_flash_attn_backend:

504 cu_seqlens_q = torch.arange(0, (bsz + 1) * q_len,

505 step=q_len,

506 dtype=torch.int32,

@@ -507,14 +541,6 @@ class MultiHeadAttention(nn.Module):

541 from torch_xla.experimental.custom_kernel import flash_attention

542 out = flash_attention(query, key, value, sm_scale=self.scale)

543 out = out.transpose(1, 2)

510 - elif self.attn_backend == _Backend.ROCM_AITER_FA:

511 - from aiter import flash_attn_varlen_func

512 -

513 - # ROCm Flash Attention expects (batch, seq, heads, head_dim)

514 - out = flash_attn_varlen_func(query,

515 - key,

516 - value,

517 - softmax_scale=self.scale)

544 else:

545 # ViT attention hasn't supported this backend yet

546 raise NotImplementedError(

vllm/model_executor/models/dots_ocr.py

+19 -22

@@ -10,7 +10,8 @@ from torch.nn import LayerNorm

10 from transformers.models.qwen2_vl import Qwen2VLProcessor

12 from vllm.attention.backends.registry import _Backend

13 -from vllm.attention.layer import check_upstream_fa_availability

13 +from vllm.attention.layer import (check_upstream_fa_availability,

14 + maybe_get_vit_flash_attn_backend)

15 from vllm.config import VllmConfig

16 from vllm.distributed import utils as dist_utils

17 from vllm.distributed.parallel_state import (

@@ -267,10 +268,12 @@ class DotsVisionAttention(nn.Module):

268 self.attn_backend = get_vit_attn_backend(

269 self.hidden_size_per_attention_head, torch.get_default_dtype())

270 self.use_upstream_fa = False

270 - if self.attn_backend != _Backend.FLASH_ATTN and \

271 - check_upstream_fa_availability(torch.get_default_dtype()):

272 - self.attn_backend = _Backend.FLASH_ATTN

273 - self.use_upstream_fa = True

271 +

272 + self.attn_backend, self.flash_attn_varlen_func \

273 + = maybe_get_vit_flash_attn_backend(

274 + self.attn_backend,

275 + self.use_upstream_fa,

276 + )

277 if self.attn_backend not in {

278 _Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS,

279 _Backend.ROCM_AITER_FA

@@ -306,25 +309,18 @@ class DotsVisionAttention(nn.Module):

309 q, k = torch.chunk(qk_rotated, 2, dim=0)

310

311 if self.is_flash_attn_backend:

309 - if self.attn_backend == _Backend.ROCM_AITER_FA:

310 - from aiter import flash_attn_varlen_func

311 - else:

312 - if self.use_upstream_fa:

313 - from flash_attn import flash_attn_varlen_func

314 - else:

315 - from vllm.vllm_flash_attn import flash_attn_varlen_func

312 q_ = q.reshape(bs * q.shape[1], q.shape[2], q.shape[3])

313 k_ = k.reshape(bs * k.shape[1], k.shape[2], k.shape[3])

314 v_ = v.reshape(bs * v.shape[1], v.shape[2], v.shape[3])

319 - output = flash_attn_varlen_func(q_,

320 - k_,

321 - v_,

322 - cu_seqlens_q=cu_seqlens,

323 - cu_seqlens_k=cu_seqlens,

324 - max_seqlen_q=max_seqlen,

325 - max_seqlen_k=max_seqlen,

326 - dropout_p=0.0,

327 - causal=False)

315 + output = self.flash_attn_varlen_func(q_,

316 + k_,

317 + v_,

318 + cu_seqlens_q=cu_seqlens,

319 + cu_seqlens_k=cu_seqlens,

320 + max_seqlen_q=max_seqlen,

321 + max_seqlen_k=max_seqlen,

322 + dropout_p=0.0,

323 + causal=False)

324 context_layer = output.view(bs, -1,

325 self.num_attention_heads_per_partition,

326 self.hidden_size_per_attention_head)

@@ -611,7 +607,8 @@ class DotsVisionTransformer(nn.Module):

607 self, cu_seqlens: torch.Tensor

608 ) -> tuple[Optional[int], Optional[list[int]]]:

609 max_seqlen, seqlens = None, None

614 - if self.attn_backend == _Backend.FLASH_ATTN:

610 + if (self.attn_backend == _Backend.FLASH_ATTN

611 + or self.attn_backend == _Backend.ROCM_AITER_FA):

612 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

613 elif self.attn_backend == _Backend.XFORMERS:

614 seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()

vllm/model_executor/models/ernie45_vl.py

+23 -26

@@ -35,7 +35,8 @@ from einops import rearrange, repeat

35 from transformers import BatchFeature

37 from vllm.attention.backends.registry import _Backend

38 -from vllm.attention.layer import check_upstream_fa_availability

38 +from vllm.attention.layer import (check_upstream_fa_availability,

39 + maybe_get_vit_flash_attn_backend)

40 from vllm.config import VllmConfig

41 from vllm.distributed import parallel_state

42 from vllm.distributed import utils as dist_utils

@@ -176,14 +177,18 @@ class Ernie4_5_VisionAttention(nn.Module):

177 dtype=torch.get_default_dtype())

178

179 self.use_upstream_fa = False

179 - if self.attn_backend != _Backend.FLASH_ATTN and \

180 - check_upstream_fa_availability(torch.get_default_dtype()):

181 - self.attn_backend = _Backend.FLASH_ATTN

182 - self.use_upstream_fa = True

180 +

181 + self.attn_backend, self.flash_attn_varlen_func \

182 + = maybe_get_vit_flash_attn_backend(

183 + self.attn_backend,

184 + self.use_upstream_fa,

185 + )

186

187 if self.attn_backend not in {

185 - _Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS,

186 - _Backend.ROCM_AITER_FA

188 + _Backend.FLASH_ATTN,

189 + _Backend.TORCH_SDPA,

190 + _Backend.XFORMERS,

191 + _Backend.ROCM_AITER_FA,

192 }:

193 raise RuntimeError(

194 f"Ernie45-VL does not support {self.attn_backend} backend now."

@@ -239,27 +244,18 @@ class Ernie4_5_VisionAttention(nn.Module):

244 q, k = torch.chunk(qk_rotated, 2, dim=0)

245

246 if self.is_flash_attn_backend:

242 - # from vllm_flash_attn.flash_attn_interface import (

243 - # flash_attn_varlen_func)

244 - if self.attn_backend == _Backend.ROCM_AITER_FA:

245 - from aiter import flash_attn_varlen_func

246 - else:

247 - if self.use_upstream_fa:

248 - from flash_attn import flash_attn_varlen_func

249 - else:

250 - from vllm.vllm_flash_attn import flash_attn_varlen_func

247

248 q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])

249

254 - output = flash_attn_varlen_func(q,

255 - k,

256 - v,

257 - cu_seqlens_q=cu_seqlens,

258 - cu_seqlens_k=cu_seqlens,

259 - max_seqlen_q=max_seqlen,

260 - max_seqlen_k=max_seqlen,

261 - dropout_p=0.0,

262 - causal=False)

250 + output = self.flash_attn_varlen_func(q,

251 + k,

252 + v,

253 + cu_seqlens_q=cu_seqlens,

254 + cu_seqlens_k=cu_seqlens,

255 + max_seqlen_q=max_seqlen,

256 + max_seqlen_k=max_seqlen,

257 + dropout_p=0.0,

258 + causal=False)

259

260 context_layer = rearrange(output,

261 "(b s) h d -> s b (h d)",

@@ -516,7 +512,8 @@ class Ernie4_5_VisionTransformer(nn.Module):

512 self, cu_seqlens: torch.Tensor

513 ) -> tuple[Optional[int], Optional[list[int]]]:

514 max_seqlen, seqlens = None, None

519 - if self.attn_backend == _Backend.FLASH_ATTN:

515 + if (self.attn_backend == _Backend.FLASH_ATTN

516 + or self.attn_backend == _Backend.ROCM_AITER_FA):

517 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

518 elif self.attn_backend == _Backend.XFORMERS:

519 seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()

vllm/model_executor/models/glm4_1v.py

+17 -14

@@ -47,7 +47,8 @@ from transformers.models.glm4v.video_processing_glm4v import (

47 from transformers.video_utils import VideoMetadata

49 from vllm.attention.backends.registry import _Backend

50 -from vllm.attention.layer import check_upstream_fa_availability

50 +from vllm.attention.layer import (check_upstream_fa_availability,

51 + maybe_get_vit_flash_attn_backend)

52 from vllm.config import VllmConfig

53 from vllm.distributed import (get_tensor_model_parallel_world_size,

54 parallel_state)

@@ -263,19 +264,26 @@ class Glm4vVisionAttention(nn.Module):

264 head_size=self.hidden_size_per_attention_head,

265 dtype=torch.get_default_dtype())

266 self.use_upstream_fa = False

266 - if self.attn_backend != _Backend.FLASH_ATTN and \

267 - check_upstream_fa_availability(torch.get_default_dtype()):

268 - self.attn_backend = _Backend.FLASH_ATTN

269 - self.use_upstream_fa = True

267 +

268 + self.attn_backend, self.flash_attn_varlen_func \

269 + = maybe_get_vit_flash_attn_backend(

270 + self.attn_backend,

271 + self.use_upstream_fa,

272 + )

273

274 if self.attn_backend not in {

275 _Backend.FLASH_ATTN,

276 _Backend.TORCH_SDPA,

277 _Backend.XFORMERS,

278 + _Backend.ROCM_AITER_FA,

279 }:

280 raise RuntimeError(

281 f"GLM-4V does not support {self.attn_backend} backend now.")

282

283 + self.is_flash_attn_backend = self.attn_backend in {

284 + _Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA

285 + }

286 +

287 def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:

288 # [s, b, 3 * head * head_dim]

289 seq_len, bs, _ = qkv.shape

@@ -316,17 +324,11 @@ class Glm4vVisionAttention(nn.Module):

324 qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)

325 q, k = torch.chunk(qk_rotated, 2, dim=0)

326

319 - if self.attn_backend == _Backend.FLASH_ATTN:

320 - # from vllm_flash_attn.flash_attn_interface import (

321 - # flash_attn_varlen_func)

322 - if self.use_upstream_fa:

323 - from flash_attn import flash_attn_varlen_func

324 - else:

325 - from vllm.vllm_flash_attn import flash_attn_varlen_func

327 + if self.is_flash_attn_backend:

328

329 q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])

330

329 - output = flash_attn_varlen_func(

331 + output = self.flash_attn_varlen_func(

332 q,

333 k,

334 v,

@@ -774,7 +776,8 @@ class Glm4vVisionTransformer(nn.Module):

776 ) -> tuple[Optional[int], Optional[list[int]]]:

777 max_seqlen, seqlens = None, None

778 seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()

777 - if self.attn_backend == _Backend.FLASH_ATTN:

779 + if (self.attn_backend == _Backend.FLASH_ATTN

780 + or self.attn_backend == _Backend.ROCM_AITER_FA):

781 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

782 return max_seqlen, seqlens

783

vllm/model_executor/models/qwen2_5_vl.py

+17 -17

@@ -39,7 +39,8 @@ from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (

39 Qwen2_5_VLConfig, Qwen2_5_VLVisionConfig)

41 from vllm.attention.backends.registry import _Backend

42 -from vllm.attention.layer import check_upstream_fa_availability

42 +from vllm.attention.layer import (check_upstream_fa_availability,

43 + maybe_get_vit_flash_attn_backend)

44 from vllm.config import VllmConfig

45 from vllm.distributed import parallel_state

46 from vllm.distributed import utils as dist_utils

@@ -302,6 +303,11 @@ class Qwen2_5_VisionAttention(nn.Module):

303 disable_tp=use_data_parallel)

304 self.attn_backend = attn_backend

305 self.use_upstream_fa = use_upstream_fa

306 + self.attn_backend, self.flash_attn_varlen_func \

307 + = maybe_get_vit_flash_attn_backend(

308 + self.attn_backend,

309 + self.use_upstream_fa,

310 + )

311 self.is_flash_attn_backend = self.attn_backend in {

312 _Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA

313 }

@@ -354,25 +360,18 @@ class Qwen2_5_VisionAttention(nn.Module):

360 q, k = torch.chunk(qk_rotated, 2, dim=0)

361

362 if self.is_flash_attn_backend:

357 - if self.attn_backend == _Backend.ROCM_AITER_FA:

358 - from aiter import flash_attn_varlen_func

359 - else:

360 - if self.use_upstream_fa:

361 - from flash_attn import flash_attn_varlen_func

362 - else:

363 - from vllm.vllm_flash_attn import flash_attn_varlen_func

363

364 q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])

365

367 - output = flash_attn_varlen_func(q,

368 - k,

369 - v,

370 - cu_seqlens_q=cu_seqlens,

371 - cu_seqlens_k=cu_seqlens,

372 - max_seqlen_q=max_seqlen,

373 - max_seqlen_k=max_seqlen,

374 - dropout_p=0.0,

375 - causal=False)

366 + output = self.flash_attn_varlen_func(q,

367 + k,

368 + v,

369 + cu_seqlens_q=cu_seqlens,

370 + cu_seqlens_k=cu_seqlens,

371 + max_seqlen_q=max_seqlen,

372 + max_seqlen_k=max_seqlen,

373 + dropout_p=0.0,

374 + causal=False)

375

376 context_layer = rearrange(output,

377 "(b s) h d -> s b (h d)",

@@ -618,6 +617,7 @@ class Qwen2_5_VisionTransformer(nn.Module):

617 self.attn_backend = get_vit_attn_backend(

618 head_size=head_dim, dtype=torch.get_default_dtype())

619 if self.attn_backend != _Backend.FLASH_ATTN and \

620 + self.attn_backend != _Backend.ROCM_AITER_FA and \

621 check_upstream_fa_availability(

622 torch.get_default_dtype()):

623 self.attn_backend = _Backend.FLASH_ATTN

vllm/model_executor/models/qwen2_vl.py

+18 -22

@@ -42,7 +42,8 @@ from transformers.models.qwen2_vl.video_processing_qwen2_vl import (

42 Qwen2VLVideoProcessor)

44 from vllm.attention.backends.registry import _Backend

45 -from vllm.attention.layer import check_upstream_fa_availability

45 +from vllm.attention.layer import (check_upstream_fa_availability,

46 + maybe_get_vit_flash_attn_backend)

47 from vllm.config import VllmConfig

48 from vllm.distributed import parallel_state, tensor_model_parallel_all_gather

49 from vllm.distributed import utils as dist_utils

@@ -319,11 +320,12 @@ class Qwen2VisionAttention(nn.Module):

320 head_size=self.hidden_size_per_attention_head,

321 dtype=torch.get_default_dtype())

322 self.use_upstream_fa = False

322 - if self.attn_backend != _Backend.FLASH_ATTN and \

323 - check_upstream_fa_availability(

324 - torch.get_default_dtype()):

325 - self.attn_backend = _Backend.FLASH_ATTN

326 - self.use_upstream_fa = True

323 +

324 + self.attn_backend, self.flash_attn_varlen_func \

325 + = maybe_get_vit_flash_attn_backend(

326 + self.attn_backend,

327 + self.use_upstream_fa,

328 + )

329

330 if self.attn_backend not in {

331 _Backend.FLASH_ATTN, _Backend.TORCH_SDPA, _Backend.XFORMERS,

@@ -331,6 +333,7 @@ class Qwen2VisionAttention(nn.Module):

333 }:

334 raise RuntimeError(

335 f"Qwen2-VL does not support {self.attn_backend} backend now.")

336 +

337 self.is_flash_attn_backend = self.attn_backend in {

338 _Backend.FLASH_ATTN, _Backend.ROCM_AITER_FA

339 }

@@ -383,25 +386,18 @@ class Qwen2VisionAttention(nn.Module):

386 q, k = torch.chunk(qk_rotated, 2, dim=0)

387

388 if self.is_flash_attn_backend:

386 - if self.attn_backend == _Backend.ROCM_AITER_FA:

387 - from aiter import flash_attn_varlen_func

388 - else:

389 - if self.use_upstream_fa:

390 - from flash_attn import flash_attn_varlen_func

391 - else:

392 - from vllm.vllm_flash_attn import flash_attn_varlen_func

389

390 q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])

391

396 - output = flash_attn_varlen_func(q,

397 - k,

398 - v,

399 - cu_seqlens_q=cu_seqlens,

400 - cu_seqlens_k=cu_seqlens,

401 - max_seqlen_q=max_seqlen,

402 - max_seqlen_k=max_seqlen,

403 - dropout_p=0.0,

404 - causal=False)

392 + output = self.flash_attn_varlen_func(q,

393 + k,

394 + v,

395 + cu_seqlens_q=cu_seqlens,

396 + cu_seqlens_k=cu_seqlens,

397 + max_seqlen_q=max_seqlen,

398 + max_seqlen_k=max_seqlen,

399 + dropout_p=0.0,

400 + causal=False)

401

402 context_layer = rearrange(output,

403 "(b s) h d -> s b (h d)",

vllm/model_executor/models/qwen3_vl.py

+3 -1

@@ -323,6 +323,7 @@ class Qwen3_VisionTransformer(nn.Module):

323 head_size=head_dim, dtype=torch.get_default_dtype())

324 use_upstream_fa = False

325 if self.attn_backend != _Backend.FLASH_ATTN and \

326 + self.attn_backend != _Backend.ROCM_AITER_FA and \

327 check_upstream_fa_availability(

328 torch.get_default_dtype()):

329 self.attn_backend = _Backend.FLASH_ATTN

@@ -476,7 +477,8 @@ class Qwen3_VisionTransformer(nn.Module):

477 cu_seqlens: torch.Tensor,

478 ) -> tuple[Optional[int], Optional[list[int]]]:

479 max_seqlen, seqlens = None, None

479 - if self.attn_backend == _Backend.FLASH_ATTN:

480 + if (self.attn_backend == _Backend.FLASH_ATTN

481 + or self.attn_backend == _Backend.ROCM_AITER_FA):

482 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

483 elif self.attn_backend == _Backend.XFORMERS:

484 seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()

vllm/model_executor/models/siglip2navit.py

+8 -14

@@ -14,7 +14,7 @@ from transformers import Siglip2VisionConfig

14 from transformers.configuration_utils import PretrainedConfig

16 from vllm.attention.backends.registry import _Backend

17 -from vllm.attention.layer import check_upstream_fa_availability

17 +from vllm.attention.layer import maybe_get_vit_flash_attn_backend

18 from vllm.distributed import divide, get_tensor_model_parallel_world_size

19 from vllm.model_executor.layers.activation import get_act_fn

20 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

@@ -240,11 +240,12 @@ class Siglip2Attention(nn.Module):

240 self.attn_backend = get_vit_attn_backend(

241 head_size=self.head_dim, dtype=torch.get_default_dtype())

242 self.use_upstream_fa = False

243 - if self.attn_backend != _Backend.FLASH_ATTN and \

244 - check_upstream_fa_availability(

245 - torch.get_default_dtype()):

246 - self.attn_backend = _Backend.FLASH_ATTN

247 - self.use_upstream_fa = True

243 +

244 + self.attn_backend, self.flash_attn_varlen_func \

245 + = maybe_get_vit_flash_attn_backend(

246 + self.attn_backend,

247 + self.use_upstream_fa,

248 + )

249

250 if self.attn_backend not in {

251 _Backend.FLASH_ATTN, _Backend.TORCH_SDPA,

@@ -286,14 +287,7 @@ class Siglip2Attention(nn.Module):

287

288 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

289 if self.is_flash_attn_backend:

289 - if self.attn_backend == _Backend.ROCM_AITER_FA:

290 - from aiter import flash_attn_varlen_func

291 - else:

292 - if self.use_upstream_fa:

293 - from flash_attn import flash_attn_varlen_func

294 - else:

295 - from vllm.vllm_flash_attn import flash_attn_varlen_func

296 - attn_output = flash_attn_varlen_func(

290 + attn_output = self.flash_attn_varlen_func(

291 queries, keys, values, cu_seqlens, cu_seqlens, max_seqlen,

292 max_seqlen).reshape(seq_length, -1)

293 elif self.attn_backend == _Backend.TORCH_SDPA:

vllm/platforms/rocm.py

-2

@@ -189,8 +189,6 @@ class RocmPlatform(Platform):

189 from vllm.attention.backends.registry import _Backend

190 if (envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA

191 and on_gfx9()):

192 - # Note: AITER FA is only supported for Qwen-VL models.

193 - # TODO: Add support for other VL models in their model class.

192 return _Backend.ROCM_AITER_FA

193 if on_gfx9():

194 return _Backend.FLASH_ATTN