@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6cbd41909e

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

Feature/vit attention unification# 23880 (#23978)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

baonudesifeizhai committed Sep 10, 2025 at 09:10 UTC 6cbd41909edaee82786b5a896ccc2fbccd934aa8

9 files changed +70 -58

tests/kernels/attention/test_mha_attn.py

+17 -5

@@ -23,6 +23,9 @@ def clear_cache():

23 """Clear lru cache to ensure each test case runs without caching.

24 """

25 _cached_get_attn_backend.cache_clear()

26 + # Clear xformers availability cache

27 + import vllm.attention.layer as layer_module

28 + layer_module.USE_XFORMERS_OPS = None

31 @pytest.mark.parametrize("device", ["cpu", "hip", "cuda"])

@@ -33,19 +36,28 @@ def test_mha_attn_platform(device: str):

36 torch.set_default_dtype(torch.float16)

38 if device == "cpu":

36 - with patch("vllm.attention.selector.current_platform", CpuPlatform()):

39 + with patch("vllm.attention.selector.current_platform",

40 + CpuPlatform()), \

41 + patch("vllm.platforms.current_platform", CpuPlatform()):

42 attn = MultiHeadAttention(16, 64, scale=1)

38 - assert attn.attn_backend == _Backend.TORCH_SDPA

43 + assert attn.attn_backend == _Backend.TORCH_SDPA_VLLM_V1

44 elif device == "hip":

40 - with patch("vllm.attention.selector.current_platform", RocmPlatform()):

45 + with patch("vllm.attention.selector.current_platform",

46 + RocmPlatform()), \

47 + patch("vllm.platforms.current_platform", RocmPlatform()), \

48 + patch("vllm.attention.layer.current_platform", RocmPlatform()):

49 attn = MultiHeadAttention(16, 64, scale=1)

50 assert attn.attn_backend == _Backend.TORCH_SDPA

51 else:

44 - with patch("vllm.attention.selector.current_platform", CudaPlatform()):

52 + with patch("vllm.attention.selector.current_platform",

53 + CudaPlatform()), \

54 + patch("vllm.platforms.current_platform", CudaPlatform()):

55 attn = MultiHeadAttention(16, 64, scale=1)

56 assert attn.attn_backend == _Backend.XFORMERS

48 - with patch("vllm.attention.selector.current_platform", CudaPlatform()):

58 + with patch("vllm.attention.selector.current_platform",

59 + CudaPlatform()), \

60 + patch("vllm.platforms.current_platform", CudaPlatform()):

61 attn = MultiHeadAttention(16, 72, scale=1)

62 assert attn.attn_backend == _Backend.XFORMERS

vllm/attention/layer.py

+19 -6

@@ -360,13 +360,13 @@ class MultiHeadAttention(nn.Module):

360 # currently, only torch_sdpa is supported on rocm

361 self.attn_backend = _Backend.TORCH_SDPA

362 else:

363 - if backend in (_Backend.FLASH_ATTN, _Backend.FLASH_ATTN_VLLM_V1,

364 - _Backend.FLEX_ATTENTION):

365 - backend = _Backend.XFORMERS

366 -

363 self.attn_backend = backend if backend in {

368 - _Backend.TORCH_SDPA, _Backend.XFORMERS, _Backend.PALLAS_VLLM_V1

369 - } else _Backend.TORCH_SDPA

364 + _Backend.TORCH_SDPA,

365 + _Backend.TORCH_SDPA_VLLM_V1,

366 + _Backend.XFORMERS,

367 + _Backend.PALLAS_VLLM_V1,

368 + _Backend.ROCM_AITER_FA,

369 + } else current_platform.get_vit_attn_backend()

370

371 if (self.attn_backend == _Backend.XFORMERS

372 and not check_xformers_availability()):

@@ -413,6 +413,19 @@ class MultiHeadAttention(nn.Module):

413 from torch_xla.experimental.custom_kernel import flash_attention

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

415 out = out.transpose(1, 2)

416 + elif self.attn_backend == _Backend.ROCM_AITER_FA:

417 + from aiter import flash_attn_varlen_func

418 +

419 + # ROCm Flash Attention expects (batch, seq, heads, head_dim)

420 + out = flash_attn_varlen_func(query,

421 + key,

422 + value,

423 + softmax_scale=self.scale)

424 + else:

425 + # ViT attention hasn't supported this backend yet

426 + raise NotImplementedError(

427 + f"ViT attention hasn't supported {self.attn_backend} "

428 + f"backend yet.")

429

430 return out.reshape(bsz, q_len, -1)

431

vllm/model_executor/models/idefics2_vision_model.py

@@ -170,6 +170,7 @@ class Idefics2VisionAttention(nn.Module):

170 quant_config=quant_config,

171 prefix=f"{prefix}.out_proj",

172 )

173 + # Use unified MultiHeadAttention with Flash Attention support

174 self.attn = MultiHeadAttention(self.num_heads_per_partition,

175 self.head_dim, self.scale)

176

@@ -181,6 +182,8 @@ class Idefics2VisionAttention(nn.Module):

182 hidden_states

183 ) # batch_size, q_len, 3 * num_heads_per_partition * head_dim

184 query_states, key_states, value_states = qkv.chunk(3, dim=-1)

185 +

186 + # Use unified MultiHeadAttention implementation

187 out = self.attn(query_states, key_states, value_states)

188 attn_output, _ = self.out_proj(out)

189 return attn_output

vllm/model_executor/models/intern_vit.py

+6 -5

@@ -255,6 +255,10 @@ class InternSdpaAttention(nn.Module):

255

256 self.proj = nn.Linear(self.dummy_dim, self.embed_dim)

257

258 + # Use unified MultiHeadAttention with automatic backend selection

259 + self.attn = MultiHeadAttention(self.num_heads, self.head_dim,

260 + self.scale)

261 +

262 def forward(self, x: torch.Tensor) -> torch.Tensor:

263 B, N, C = x.shape

264 qkv = self.qkv(x)

@@ -268,12 +272,9 @@ class InternSdpaAttention(nn.Module):

272 B_, N_, H_, D_ = q.shape

273 q = self.q_norm(q.flatten(-2, -1)).view(B_, N_, H_, D_)

274 k = self.k_norm(k.flatten(-2, -1)).view(B_, N_, H_, D_)

271 - q = q.transpose(1, 2)

272 - k = k.transpose(1, 2)

273 - v = v.transpose(1, 2)

275

275 - x = F.scaled_dot_product_attention(q, k, v, scale=self.scale)

276 - x = x.transpose(1, 2).reshape(B, N, -1)

276 + # Use unified MultiHeadAttention with automatic backend selection

277 + x = self.attn(q, k, v)

278

279 x = self.proj(x)

280 return x

vllm/model_executor/models/interns1_vit.py

+7 -10

@@ -12,10 +12,10 @@ from typing import Optional

13 import torch

14 import torch.nn as nn

15 -import torch.nn.functional as F

15 from transformers import PretrainedConfig

16 from transformers.utils import torch_int

18 +from vllm.attention.layer import MultiHeadAttention

19 from vllm.model_executor.layers.activation import get_act_fn

20 from vllm.model_executor.layers.layernorm import RMSNorm

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

@@ -206,6 +206,10 @@ class InternSdpaAttention(nn.Module):

206

207 self.projection_layer = nn.Linear(self.dummy_dim, self.embed_dim)

208

209 + # Use unified MultiHeadAttention with automatic backend selection

210 + self.attn = MultiHeadAttention(self.num_heads, self.head_dim,

211 + self.scale)

212 +

213 def forward(self, x: torch.Tensor) -> torch.Tensor:

214 B, N, C = x.shape

215

@@ -213,20 +217,13 @@ class InternSdpaAttention(nn.Module):

217 k = self.k_proj(x)

218 v = self.v_proj(x)

219

216 - q = q.view(B, N, self.num_heads, self.head_dim)

217 - k = k.view(B, N, self.num_heads, self.head_dim)

218 - v = v.view(B, N, self.num_heads, self.head_dim)

219 -

220 if self.qk_normalization:

221 B_, N_, H_, D_ = q.shape

222 q = self.q_norm(q.flatten(-2, -1)).view(B_, N_, H_, D_)

223 k = self.k_norm(k.flatten(-2, -1)).view(B_, N_, H_, D_)

224 - q = q.transpose(1, 2)

225 - k = k.transpose(1, 2)

226 - v = v.transpose(1, 2)

224

228 - x = F.scaled_dot_product_attention(q, k, v, scale=self.scale)

229 - x = x.transpose(1, 2).reshape(B, N, -1)

225 + # Use unified MultiHeadAttention with automatic backend selection

226 + x = self.attn(q, k, v)

227

228 x = self.projection_layer(x)

229 return x

vllm/model_executor/models/mllama.py

+9 -15

@@ -35,6 +35,7 @@ from transformers.models.mllama.processing_mllama import (

36 import vllm.distributed.parallel_state as ps

37 from vllm.attention import Attention, AttentionMetadata, AttentionType

38 +from vllm.attention.layer import MultiHeadAttention

39 from vllm.attention.ops.paged_attn import PagedAttention

40 from vllm.attention.selector import _Backend

41 from vllm.config import VllmConfig

@@ -517,6 +518,10 @@ class MllamaVisionSdpaAttention(nn.Module):

518 prefix=f"{prefix}.o_proj",

519 )

520

521 + # Use unified MultiHeadAttention with automatic backend selection

522 + self.attn = MultiHeadAttention(self.num_local_heads, self.head_dim,

523 + 1.0 / math.sqrt(self.head_dim))

524 +

525 def forward(

526 self,

527 hidden_state: torch.Tensor,

@@ -524,21 +529,10 @@ class MllamaVisionSdpaAttention(nn.Module):

529 ) -> torch.Tensor:

530 qkv, _ = self.qkv_proj(hidden_state)

531 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

527 - q = q.view(q.shape[0], q.shape[1], self.num_local_heads,

528 - self.head_dim).transpose(1, 2)

529 - k = k.view(k.shape[0], k.shape[1], self.num_local_heads,

530 - self.head_dim).transpose(1, 2)

531 - v = v.view(v.shape[0], v.shape[1], self.num_local_heads,

532 - self.head_dim).transpose(1, 2)

533 -

534 - # TODO: remove padding in image encoder

535 - attn_output = F.scaled_dot_product_attention(q,

536 - k,

537 - v,

538 - attn_mask=attention_mask,

539 - dropout_p=0.0)

540 -

541 - attn_output = attn_output.transpose(1, 2).contiguous()

532 +

533 + # Use unified MultiHeadAttention with automatic backend selection

534 + attn_output = self.attn(q, k, v)

535 +

536 attn_output = attn_output.reshape(attn_output.shape[0],

537 attn_output.shape[1], -1)

538 output, _ = self.o_proj(attn_output)

vllm/model_executor/models/step3_vl.py

+7 -16

@@ -16,6 +16,7 @@ from torchvision import transforms

16 from torchvision.transforms.functional import InterpolationMode

17 from transformers import BatchFeature, PretrainedConfig, TensorType

19 +from vllm.attention.layer import MultiHeadAttention

20 from vllm.config import VllmConfig

21 from vllm.distributed import get_tensor_model_parallel_world_size

22 from vllm.model_executor.layers.activation import get_act_fn

@@ -682,9 +683,9 @@ class Step3VisionAttention(nn.Module):

683 prefix=f"{prefix}.out_proj",

684 disable_tp=use_data_parallel)

685

685 - def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):

686 - return tensor.view(bsz, seq_len, self.num_heads,

687 - self.head_dim).transpose(1, 2).contiguous()

686 + # Use unified MultiHeadAttention with automatic backend selection

687 + self.attn = MultiHeadAttention(self.num_heads, self.head_dim,

688 + self.scale)

689

690 def forward(

691 self,

@@ -696,19 +697,9 @@ class Step3VisionAttention(nn.Module):

697 # get query proj

698 qkv, _ = self.qkv_proj(hidden_states)

699 q, k, v = qkv.chunk(chunks=3, dim=-1)

699 - q = q.view(bsz, tgt_len, self.num_heads, self.head_dim)

700 - k = k.view(bsz, tgt_len, self.num_heads, self.head_dim)

701 - v = v.view(bsz, tgt_len, self.num_heads, self.head_dim)

702 - q = q.transpose(1, 2)

703 - k = k.transpose(1, 2)

704 - v = v.transpose(1, 2)

705 - attn_output = F.scaled_dot_product_attention(q,

706 - k,

707 - v,

708 - scale=self.scale,

709 - is_causal=False)

710 - attn_output = attn_output.transpose(1, 2).reshape(

711 - bsz, tgt_len, self.num_heads * self.head_dim)

700 +

701 + # Use unified MultiHeadAttention with automatic backend selection

702 + attn_output = self.attn(q, k, v)

703

704 attn_output, _ = self.out_proj(attn_output)

705

vllm/model_executor/models/vision.py

+1 -1

@@ -122,4 +122,4 @@ def resolve_visual_encoder_outputs(

122 uses_last_layer = feature_sample_layers[-1] in (len(hs_pool) - 1, -1)

123 if post_layer_norm is not None and uses_last_layer:

124 hs_pool[-1] = post_layer_norm(encoder_outputs)

125 - return torch.cat(hs_pool, dim=-1)

125 + return torch.cat(hs_pool, dim=-1)

\ No newline at end of file

vllm/platforms/interface.py

@@ -48,6 +48,7 @@ class _Backend(enum.Enum):

48 ROCM_AITER_MLA_VLLM_V1 = enum.auto()

49 ROCM_AITER_FA = enum.auto() # used for ViT attn backend

50 TORCH_SDPA = enum.auto()

51 + TORCH_SDPA_VLLM_V1 = enum.auto()

52 FLASHINFER = enum.auto()

53 FLASHINFER_VLLM_V1 = enum.auto()

54 TRITON_MLA = enum.auto() # Supported by V1