@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 700a5ad6c6
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 11, 05:00 UTC

[MM Encoder]: Migrate legacy ViT `MultiHeadAttention` to new `MMEncoderAttention` interface (#30684)

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

Isotr0py committed Dec 19, 2025 at 02:04 UTC 700a5ad6c616358f42db7d9b55e8bc9caa140ca5
20 files changed +182 -266
tests/kernels/attention/test_attention.py
+3 -2
@@ -9,7 +9,8 @@ import torch
9 from tests.kernels.allclose_default import get_default_atol, get_default_rtol
10 from tests.kernels.utils import opcheck
11 from vllm import _custom_ops as ops
12 -from vllm.attention.layer import Attention, MultiHeadAttention
12 +from vllm.attention.layer import Attention
13 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
14 from vllm.platforms import current_platform
15 from vllm.utils.mem_utils import get_max_shared_memory_bytes
16
@@ -442,7 +443,7 @@ def ref_multi_query_kv_attention(
443 return torch.cat(ref_outputs, dim=0)
444
445
445 -@pytest.mark.parametrize("attention_cls", [Attention, MultiHeadAttention])
446 +@pytest.mark.parametrize("attention_cls", [Attention, MMEncoderAttention])
447 def test_num_heads_not_divisble_by_num_kv_heads(attention_cls: type) -> None:
448 head_size = 64
449 scale = float(1.0 / (head_size**0.5))
tests/kernels/attention/test_mha_attn.py
+67 -11
@@ -3,16 +3,17 @@
3 """
4 Test:
5
6 -* Tests for MultiHeadAttention layer
6 +* Tests for MMEncoderAttention layer
7 """
8
9 +import itertools
10 from unittest.mock import patch
11
12 import pytest
13 import torch
14
15 from vllm.attention.backends.registry import AttentionBackendEnum
15 -from vllm.attention.layer import MultiHeadAttention
16 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
17 from vllm.attention.selector import _cached_get_attn_backend
18 from vllm.platforms import current_platform
19 from vllm.platforms.cpu import CpuPlatform
@@ -42,35 +43,31 @@ def test_mha_attn_platform(device: str):
43
44 if device == "cpu":
45 with (
45 - patch("vllm.attention.layer.current_platform", CpuPlatform()),
46 patch("vllm.model_executor.models.vision.current_platform", CpuPlatform()),
47 ):
48 - attn = MultiHeadAttention(16, 64, scale=1)
48 + attn = MMEncoderAttention(16, 64, scale=1)
49 assert attn.attn_backend == AttentionBackendEnum.TORCH_SDPA
50 elif device == "hip":
51 with (
52 - patch("vllm.attention.layer.current_platform", RocmPlatform()),
52 patch("vllm.model_executor.models.vision.current_platform", RocmPlatform()),
53 ):
55 - attn = MultiHeadAttention(16, 64, scale=1)
54 + attn = MMEncoderAttention(16, 64, scale=1)
55 assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
56 else:
57 # Test CUDA with head_size=64 (divisible by 32)
58 # - should use vLLM's FlashAttention
59 with (
61 - patch("vllm.attention.layer.current_platform", CudaPlatform()),
60 patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
61 ):
64 - attn = MultiHeadAttention(16, 64, scale=1)
62 + attn = MMEncoderAttention(16, 64, scale=1)
63 assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
64
65 # Test CUDA with head_size=72 (not divisible by 32)
66 # - should use vLLM's FlashAttention
67 with (
70 - patch("vllm.attention.layer.current_platform", CudaPlatform()),
68 patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
69 ):
73 - attn = MultiHeadAttention(16, 72, scale=1)
70 + attn = MMEncoderAttention(16, 72, scale=1)
71 assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
72
73
@@ -94,6 +91,10 @@ def ref_attention(
91
92 BATCH_SIZES = [1, 16]
93 SEQ_LENS = [1]
94 +VAR_SEQ_LENS = [
95 + [2, 2],
96 + [2, 3, 4],
97 +]
98 NUM_HEADS = [1, 16]
99 NUM_KV_HEADS = [1]
100 HEAD_SIZES = [64, 80]
@@ -130,7 +131,7 @@ def test_mha_attn_forward(
131 k = torch.randn(batch_size, seq_len, num_kv_heads * head_size)
132 v = torch.randn(batch_size, seq_len, num_kv_heads * head_size)
133 scale = 1.0 / head_size**0.5
133 - attn = MultiHeadAttention(
134 + attn = MMEncoderAttention(
135 num_heads, head_size, scale=scale, num_kv_heads=num_kv_heads
136 )
137 output = attn(q, k, v)
@@ -151,3 +152,58 @@ def test_mha_attn_forward(
152 scale=scale,
153 ).reshape(batch_size, seq_len, num_heads * head_size)
154 torch.testing.assert_close(output, ref_output)
155 +
156 +
157 +@pytest.mark.parametrize("var_seq_len", VAR_SEQ_LENS)
158 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
159 +@pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS)
160 +@pytest.mark.parametrize("head_size", HEAD_SIZES)
161 +@pytest.mark.parametrize("dtype", DTYPES)
162 +@pytest.mark.parametrize("device", CUDA_DEVICES)
163 +def test_mha_attn_varlen_forward(
164 + var_seq_len: list[int],
165 + num_heads: int,
166 + num_kv_heads: int,
167 + head_size: int,
168 + dtype: torch.dtype,
169 + device: str,
170 +):
171 + current_platform.seed_everything(0)
172 + torch.set_default_device(device)
173 + torch.set_default_dtype(dtype)
174 +
175 + q = torch.randn(1, sum(var_seq_len), num_heads, head_size)
176 + k = torch.randn(1, sum(var_seq_len), num_kv_heads, head_size)
177 + v = torch.randn(1, sum(var_seq_len), num_kv_heads, head_size)
178 + cu_seqlens = torch.tensor(
179 + [0] + list(itertools.accumulate(var_seq_len)), dtype=torch.int32
180 + )
181 + scale = 1.0 / head_size**0.5
182 + attn = MMEncoderAttention(
183 + num_heads, head_size, scale=scale, num_kv_heads=num_kv_heads
184 + )
185 + output = attn(
186 + q, k, v, cu_seqlens=cu_seqlens, max_seqlen=torch.tensor(max(var_seq_len))
187 + )
188 +
189 + assert num_heads % num_kv_heads == 0
190 + num_queries_per_kv = num_heads // num_kv_heads
191 + if num_queries_per_kv > 1:
192 + k = torch.repeat_interleave(k, num_queries_per_kv, dim=2)
193 + v = torch.repeat_interleave(v, num_queries_per_kv, dim=2)
194 +
195 + ref_output = []
196 + for q_i, k_i, v_i in zip(
197 + torch.split(q, var_seq_len, dim=1),
198 + torch.split(k, var_seq_len, dim=1),
199 + torch.split(v, var_seq_len, dim=1),
200 + ):
201 + output_i = ref_attention(
202 + q_i,
203 + k_i,
204 + v_i,
205 + scale=scale,
206 + )
207 + ref_output.append(output_i)
208 + ref_output = torch.cat(ref_output, dim=1)
209 + torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)
tests/v1/tpu/test_mha_attn.py
+3 -3
@@ -3,7 +3,7 @@
3 """
4 Test:
5
6 -* Tests for MultiHeadAttention layer
6 +* Tests for MMEncoderAttention layer
7 """
8
9 import pytest
@@ -12,7 +12,7 @@ import torch_xla
12 import torch_xla.core
13 import torch_xla.core.xla_model
14
15 -from vllm.attention.layer import MultiHeadAttention
15 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
16 from vllm.attention.selector import _cached_get_attn_backend
17 from vllm.platforms import current_platform
18
@@ -69,7 +69,7 @@ def test_mha_attn_forward(
69 k = torch.randn(batch_size, seq_len, num_kv_heads * head_size, device=device)
70 v = torch.randn(batch_size, seq_len, num_kv_heads * head_size, device=device)
71 scale = 1.0 / head_size**0.5
72 - attn = MultiHeadAttention(
72 + attn = MMEncoderAttention(
73 num_heads, head_size, scale=scale, num_kv_heads=num_kv_heads
74 )
75 output = attn(q, k, v)
vllm/attention/layer.py
-132
@@ -2,12 +2,10 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Attention layer."""
4
5 -import functools
5 from typing import cast
6
7 import torch
8 import torch.nn as nn
10 -import torch.nn.functional as F
9
10 import vllm.envs as envs
11 from vllm.attention.backends.abstract import (
@@ -16,13 +14,10 @@ from vllm.attention.backends.abstract import (
14 MLAAttentionImpl,
15 )
16 from vllm.attention.backends.registry import AttentionBackendEnum
19 -from vllm.attention.layers.mm_encoder_attention import maybe_get_vit_flash_attn_backend
17 from vllm.attention.selector import get_attn_backend
21 -from vllm.attention.utils.fa_utils import get_flash_attn_version
18 from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
19 from vllm.attention.utils.kv_transfer_utils import maybe_transfer_kv_layer
20 from vllm.config import CacheConfig, get_current_vllm_config
25 -from vllm.config.multimodal import MultiModalConfig
21 from vllm.config.vllm import VllmConfig
22 from vllm.forward_context import ForwardContext, get_forward_context
23 from vllm.logger import init_logger
@@ -36,7 +31,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
31 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
32 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
33 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
39 -from vllm.model_executor.models.vision import get_vit_attn_backend
34 from vllm.platforms import current_platform
35 from vllm.utils.torch_utils import (
36 direct_register_custom_op,
@@ -412,132 +406,6 @@ class Attention(nn.Module, AttentionLayerBase):
406 )
407
408
415 -class MultiHeadAttention(nn.Module):
416 - """Multi-headed attention without any cache, used for ViT."""
417 -
418 - def __init__(
419 - self,
420 - num_heads: int,
421 - head_size: int,
422 - scale: float,
423 - num_kv_heads: int | None = None,
424 - # This has no effect, it is only here to make it easier to swap
425 - # between Attention and MultiHeadAttention
426 - prefix: str = "",
427 - multimodal_config: MultiModalConfig | None = None,
428 - ) -> None:
429 - super().__init__()
430 - self.num_heads = num_heads
431 - self.head_size = head_size
432 - self.scale = scale
433 - self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
434 - self.layer_name = prefix
435 -
436 - assert self.num_heads % self.num_kv_heads == 0, (
437 - f"num_heads ({self.num_heads}) is not "
438 - f"divisible by num_kv_heads ({self.num_kv_heads})"
439 - )
440 - self.num_queries_per_kv = self.num_heads // self.num_kv_heads
441 -
442 - # During model initialization, the default dtype is set as the model
443 - # weight and activation dtype.
444 - dtype = torch.get_default_dtype()
445 -
446 - # Determine the attention backend
447 - attn_backend_override = None
448 - if multimodal_config is not None:
449 - attn_backend_override = multimodal_config.mm_encoder_attn_backend
450 -
451 - self.attn_backend = get_vit_attn_backend(
452 - head_size=head_size,
453 - dtype=dtype,
454 - attn_backend_override=attn_backend_override,
455 - )
456 -
457 - self._flash_attn_varlen_func = maybe_get_vit_flash_attn_backend(
458 - self.attn_backend,
459 - )
460 -
461 - self.is_flash_attn_backend = self.attn_backend in {
462 - AttentionBackendEnum.FLASH_ATTN,
463 - AttentionBackendEnum.ROCM_AITER_FA,
464 - }
465 -
466 - self.fa_version = None
467 - if (
468 - self.attn_backend == AttentionBackendEnum.FLASH_ATTN
469 - and current_platform.is_cuda()
470 - ):
471 - self.fa_version = get_flash_attn_version()
472 - assert self._flash_attn_varlen_func is not None
473 - self._flash_attn_varlen_func = functools.partial(
474 - self._flash_attn_varlen_func, fa_version=self.fa_version
475 - )
476 -
477 - logger.info_once(
478 - f"Using {self.attn_backend} for MultiHeadAttention in multimodal encoder."
479 - )
480 -
481 - def forward(
482 - self,
483 - query: torch.Tensor,
484 - key: torch.Tensor,
485 - value: torch.Tensor,
486 - ) -> torch.Tensor:
487 - """Input shape:
488 - (batch_size x seq_len x hidden_size) or
489 - (batch_size x seq_len x num_heads x head_size)
490 - """
491 - bsz, q_len = query.size()[:2]
492 - kv_len = key.size(1)
493 -
494 - query = query.view(bsz, q_len, self.num_heads, self.head_size)
495 - key = key.view(bsz, kv_len, self.num_kv_heads, self.head_size)
496 - value = value.view(bsz, kv_len, self.num_kv_heads, self.head_size)
497 -
498 - if (num_repeat := self.num_queries_per_kv) > 1:
499 - # Handle MQA and GQA
500 - key = torch.repeat_interleave(key, num_repeat, dim=2)
501 - value = torch.repeat_interleave(value, num_repeat, dim=2)
502 -
503 - if self.is_flash_attn_backend:
504 - assert self._flash_attn_varlen_func is not None
505 - cu_seqlens_q = torch.arange(
506 - 0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device
507 - )
508 - cu_seqlens_k = torch.arange(
509 - 0, (bsz + 1) * kv_len, step=kv_len, dtype=torch.int32, device=key.device
510 - )
511 -
512 - out = self._flash_attn_varlen_func(
513 - query.flatten(0, 1),
514 - key.flatten(0, 1),
515 - value.flatten(0, 1),
516 - cu_seqlens_q=cu_seqlens_q,
517 - cu_seqlens_k=cu_seqlens_k,
518 - max_seqlen_q=q_len,
519 - max_seqlen_k=kv_len,
520 - softmax_scale=self.scale,
521 - )
522 - elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
523 - query, key, value = (x.transpose(1, 2) for x in (query, key, value))
524 - out = F.scaled_dot_product_attention(query, key, value, scale=self.scale)
525 - out = out.transpose(1, 2)
526 - elif self.attn_backend == AttentionBackendEnum.PALLAS:
527 - query, key, value = (x.transpose(1, 2) for x in (query, key, value))
528 - from torch_xla.experimental.custom_kernel import flash_attention
529 -
530 - out = flash_attention(query, key, value, sm_scale=self.scale)
531 - out = out.transpose(1, 2)
532 - else:
533 - # ViT attention hasn't supported this backend yet
534 - raise NotImplementedError(
535 - f"ViT attention hasn't supported {self.attn_backend} backend yet."
536 - )
537 -
538 - return out.reshape(bsz, q_len, -1)
539 -
540 -
409 class MLAAttention(nn.Module, AttentionLayerBase):
410 """Multi-Head Latent Attention layer.
411
vllm/attention/layers/mm_encoder_attention.py
+27 -63
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from collections.abc import Callable
4
5 import torch
6
@@ -19,27 +18,6 @@ from vllm.model_executor.models.vision import get_vit_attn_backend
18 logger = init_logger(__name__)
19
20
22 -def maybe_get_vit_flash_attn_backend(
23 - attn_backend: AttentionBackendEnum | None,
24 -) -> Callable | None:
25 - # At this point,
26 - # we already have the attn_backend,
27 - # overriding logic is done in the platform-specific implementation.
28 - # so we don't need to override backend here.
29 - # Just return the attn_backend and flash_attn_varlen_func.
30 -
31 - if attn_backend == AttentionBackendEnum.FLASH_ATTN:
32 - from vllm.attention.utils.fa_utils import flash_attn_varlen_func
33 - elif attn_backend == AttentionBackendEnum.ROCM_AITER_FA:
34 - from aiter import flash_attn_varlen_func
35 - else:
36 - flash_attn_varlen_func = None
37 -
38 - # if attn_backend is TORCH_SDPA,
39 - # it will reach here and the flash_attn_varlen_func will be None.
40 - return flash_attn_varlen_func
41 -
42 -
21 @CustomOp.register("mm_encoder_attn")
22 class MMEncoderAttention(CustomOp):
23 """Multi-headed attention without any cache, used for multimodal encoder."""
@@ -98,21 +76,17 @@ class MMEncoderAttention(CustomOp):
76 AttentionBackendEnum.ROCM_AITER_FA,
77 }
78
101 - self.flash_attn_varlen_func = maybe_get_vit_flash_attn_backend(
102 - self.attn_backend,
79 + self._fa_version = (
80 + get_flash_attn_version() if self.is_flash_attn_backend else None
81 )
82
105 - if self.is_flash_attn_backend:
106 - assert self.flash_attn_varlen_func is not None
107 - self._fa_version = get_flash_attn_version()
108 -
83 logger.info_once(f"Using {self.attn_backend} for MMEncoderAttention.")
84
85 @classmethod
86 def enabled(cls) -> bool:
87 return True
88
115 - def reshape_qkv_to_4d(
89 + def maybe_reshape_qkv_to_4d(
90 self,
91 query: torch.Tensor,
92 key: torch.Tensor,
@@ -136,30 +110,6 @@ class MMEncoderAttention(CustomOp):
110
111 return query, key, value
112
139 - def reshape_qkv_to_3d(
140 - self,
141 - query: torch.Tensor,
142 - key: torch.Tensor,
143 - value: torch.Tensor,
144 - bsz: int,
145 - q_len: int,
146 - kv_len: int,
147 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
148 - """
149 - Reshape query, key, value to 3D tensors:
150 - (batch_size * seq_len, num_heads, head_size)
151 - """
152 - query = query.view(bsz * q_len, self.num_heads, self.head_size)
153 - key = key.view(bsz * kv_len, self.num_kv_heads, self.head_size)
154 - value = value.view(bsz * kv_len, self.num_kv_heads, self.head_size)
155 -
156 - if (num_repeat := self.num_queries_per_kv) > 1:
157 - # Handle MQA and GQA
158 - key = torch.repeat_interleave(key, num_repeat, dim=1)
159 - value = torch.repeat_interleave(value, num_repeat, dim=1)
160 -
161 - return query, key, value
162 -
113 def _forward_sdpa(
114 self,
115 query: torch.Tensor,
@@ -167,13 +117,15 @@ class MMEncoderAttention(CustomOp):
117 value: torch.Tensor,
118 cu_seqlens: torch.Tensor | None = None,
119 ) -> torch.Tensor:
170 - # TODO(Isotr0py): Migrate MultiHeadAttention
171 - assert cu_seqlens is not None
172 -
120 + """Input shape:
121 + (batch_size x seq_len x hidden_size) or
122 + (batch_size x seq_len x num_heads x head_size)
123 + """
124 bsz, q_len = query.size()[:2]
125 kv_len = key.size(1)
126 + is_reshaped = query.dim() != 4
127
176 - query, key, value = self.reshape_qkv_to_4d(
128 + query, key, value = self.maybe_reshape_qkv_to_4d(
129 query, key, value, bsz, q_len, kv_len
130 )
131
@@ -183,6 +135,8 @@ class MMEncoderAttention(CustomOp):
135 v=value,
136 cu_seqlens=cu_seqlens,
137 )
138 + if is_reshaped:
139 + output = output.view(bsz, q_len, -1)
140 return output
141
142 def _forward_fa(
@@ -193,13 +147,21 @@ class MMEncoderAttention(CustomOp):
147 cu_seqlens: torch.Tensor | None = None,
148 max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
149 ) -> torch.Tensor:
196 - assert self.flash_attn_varlen_func is not None, (
197 - "Flash attention function is not set."
198 - )
199 - # # TODO(Isotr0py): Migrate MultiHeadAttention
200 - assert cu_seqlens is not None and max_seqlen is not None
150 + """Input shape:
151 + (batch_size x seq_len x hidden_size) or
152 + (batch_size x seq_len x num_heads x head_size)
153 + """
154 + assert (cu_seqlens is not None and max_seqlen is not None) or (
155 + cu_seqlens is None and max_seqlen is None
156 + ), "cu_seqlens and max_seqlen should be both set or both None."
157
202 - bsz = query.shape[0]
158 + bsz, q_len = query.size()[:2]
159 + kv_len = key.size(1)
160 + is_reshaped = query.dim() != 4
161 +
162 + query, key, value = self.maybe_reshape_qkv_to_4d(
163 + query, key, value, bsz, q_len, kv_len
164 + )
165
166 output = vit_flash_attn_wrapper(
167 q=query,
@@ -211,6 +173,8 @@ class MMEncoderAttention(CustomOp):
173 is_rocm_aiter=(self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA),
174 fa_version=self._fa_version,
175 )
176 + if is_reshaped:
177 + output = output.view(bsz, q_len, -1)
178 return output
179
180 def forward_native(
vllm/attention/ops/vit_attn_wrappers.py
+39 -14
@@ -24,11 +24,11 @@ def flash_attn_maxseqlen_wrapper(
24 q: torch.Tensor,
25 k: torch.Tensor,
26 v: torch.Tensor,
27 - cu_seqlens: torch.Tensor,
28 - max_seqlen: torch.Tensor,
27 batch_size: int,
28 is_rocm_aiter: bool,
29 fa_version: int | None,
30 + cu_seqlens: torch.Tensor | None = None,
31 + max_seqlen: torch.Tensor | None = None,
32 ) -> torch.Tensor:
33 kwargs = {}
34 if is_rocm_aiter:
@@ -38,6 +38,14 @@ def flash_attn_maxseqlen_wrapper(
38
39 if not current_platform.is_rocm() and fa_version is not None:
40 kwargs["fa_version"] = fa_version
41 +
42 + q_len = q.size(1)
43 + if cu_seqlens is None:
44 + cu_seqlens = torch.arange(
45 + 0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=q.device
46 + )
47 + max_seqlen = q_len if max_seqlen is None else max_seqlen.item()
48 +
49 q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
50 output = flash_attn_varlen_func(
51 q,
@@ -45,8 +53,8 @@ def flash_attn_maxseqlen_wrapper(
53 v,
54 cu_seqlens_q=cu_seqlens,
55 cu_seqlens_k=cu_seqlens,
48 - max_seqlen_q=max_seqlen.item(),
49 - max_seqlen_k=max_seqlen.item(),
56 + max_seqlen_q=max_seqlen,
57 + max_seqlen_k=max_seqlen,
58 dropout_p=0.0,
59 causal=False,
60 **kwargs,
@@ -79,24 +87,42 @@ def vit_flash_attn_wrapper(
87 q: torch.Tensor,
88 k: torch.Tensor,
89 v: torch.Tensor,
82 - cu_seqlens: torch.Tensor,
83 - max_seqlen: torch.Tensor,
90 batch_size: int,
91 is_rocm_aiter: bool,
92 fa_version: int | None,
93 + cu_seqlens: torch.Tensor | None = None,
94 + max_seqlen: torch.Tensor | None = None,
95 ) -> torch.Tensor:
96 return torch.ops.vllm.flash_attn_maxseqlen_wrapper(
89 - q, k, v, cu_seqlens, max_seqlen, batch_size, is_rocm_aiter, fa_version
97 + q,
98 + k,
99 + v,
100 + batch_size,
101 + is_rocm_aiter,
102 + fa_version,
103 + cu_seqlens,
104 + max_seqlen,
105 )
106
107
108 +def apply_sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
109 + """
110 + Input shape:
111 + (batch_size x seq_len x num_heads x head_size)
112 + """
113 + q, k, v = (einops.rearrange(x, "b s h d -> b h s d") for x in [q, k, v])
114 + output = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
115 + output = einops.rearrange(output, "b h s d -> b s h d ")
116 + return output
117 +
118 +
119 # TODO: Once we have a torch 2.10, we can use tensor slices
120 # so we won't need to wrap this in custom ops
121 def torch_sdpa_wrapper(
122 q: torch.Tensor,
123 k: torch.Tensor,
124 v: torch.Tensor,
99 - cu_seqlens: torch.Tensor,
125 + cu_seqlens: torch.Tensor | None = None,
126 ) -> torch.Tensor:
127 # Never remove the contiguous logic for ROCm
128 # Without it, hallucinations occur with the backend
@@ -105,6 +131,9 @@ def torch_sdpa_wrapper(
131 k = k.contiguous()
132 v = v.contiguous()
133
134 + if cu_seqlens is None:
135 + return apply_sdpa(q, k, v)
136 +
137 outputs = []
138
139 lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
@@ -112,11 +141,7 @@ def torch_sdpa_wrapper(
141 k_chunks = torch.split(k, lens, dim=1)
142 v_chunks = torch.split(v, lens, dim=1)
143 for q_i, k_i, v_i in zip(q_chunks, k_chunks, v_chunks):
115 - q_i, k_i, v_i = (
116 - einops.rearrange(x, "b s h d -> b h s d") for x in [q_i, k_i, v_i]
117 - )
118 - output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
119 - output_i = einops.rearrange(output_i, "b h s d -> b s h d ")
144 + output_i = apply_sdpa(q_i, k_i, v_i)
145 outputs.append(output_i)
146 context_layer = torch.cat(outputs, dim=1)
147 return context_layer
@@ -142,6 +167,6 @@ def vit_torch_sdpa_wrapper(
167 q: torch.Tensor,
168 k: torch.Tensor,
169 v: torch.Tensor,
145 - cu_seqlens: torch.Tensor,
170 + cu_seqlens: torch.Tensor | None = None,
171 ) -> torch.Tensor:
172 return torch.ops.vllm.torch_sdpa_wrapper(q, k, v, cu_seqlens)
vllm/model_executor/models/aimv2.py
+2 -2
@@ -8,7 +8,7 @@ from collections.abc import Iterable
8 import torch
9 import torch.nn as nn
10
11 -from vllm.attention.layer import MultiHeadAttention
11 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
12 from vllm.distributed import get_tensor_model_parallel_world_size
13 from vllm.distributed.utils import divide
14 from vllm.model_executor.layers.activation import SiluAndMul
@@ -126,7 +126,7 @@ class AIMv2Attention(nn.Module):
126 self.tp_size = get_tensor_model_parallel_world_size()
127 self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
128
129 - self.attn = MultiHeadAttention(
129 + self.attn = MMEncoderAttention(
130 self.num_heads_per_partition, self.head_dim, self.scale
131 )
132
vllm/model_executor/models/blip.py
+2 -2
@@ -9,7 +9,7 @@ import torch
9 import torch.nn as nn
10 from transformers import Blip2VisionConfig, BlipVisionConfig
11
12 -from vllm.attention.layer import MultiHeadAttention
12 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
13 from vllm.distributed import divide, get_tensor_model_parallel_world_size
14 from vllm.model_executor.layers.activation import get_act_fn
15 from vllm.model_executor.layers.conv import Conv2dLayer
@@ -122,7 +122,7 @@ class BlipAttention(nn.Module):
122 self.tp_size = get_tensor_model_parallel_world_size()
123 self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
124
125 - self.attn = MultiHeadAttention(
125 + self.attn = MMEncoderAttention(
126 self.num_heads_per_partition, self.head_dim, self.scale
127 )
128
vllm/model_executor/models/clip.py
+6 -5
@@ -14,7 +14,8 @@ from transformers import (
14 CLIPVisionConfig,
15 )
16
17 -from vllm.attention.layer import Attention, MultiHeadAttention
17 +from vllm.attention.layer import Attention
18 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
19 from vllm.config import VllmConfig
20 from vllm.config.multimodal import BaseDummyOptions
21 from vllm.distributed import divide, get_tensor_model_parallel_world_size
@@ -354,7 +355,7 @@ class CLIPAttention(nn.Module):
355 quant_config: QuantizationConfig | None = None,
356 *,
357 prefix: str = "",
357 - attn_cls: type[Attention] | type[MultiHeadAttention],
358 + attn_cls: type[Attention] | type[MMEncoderAttention],
359 ) -> None:
360 super().__init__()
361
@@ -449,7 +450,7 @@ class CLIPEncoderLayer(nn.Module):
450 quant_config: QuantizationConfig | None = None,
451 *,
452 prefix: str = "",
452 - attn_cls: type[Attention] | type[MultiHeadAttention],
453 + attn_cls: type[Attention] | type[MMEncoderAttention],
454 ) -> None:
455 super().__init__()
456 self.self_attn = CLIPAttention(
@@ -493,7 +494,7 @@ class CLIPEncoder(nn.Module):
494 num_hidden_layers_override: int | None = None,
495 *,
496 prefix: str = "",
496 - attn_cls: type[Attention] | type[MultiHeadAttention],
497 + attn_cls: type[Attention] | type[MMEncoderAttention],
498 ) -> None:
499 super().__init__()
500
@@ -638,7 +639,7 @@ class CLIPVisionTransformer(nn.Module):
639 quant_config=quant_config,
640 num_hidden_layers_override=num_hidden_layers_override,
641 prefix=f"{prefix}.encoder",
641 - attn_cls=MultiHeadAttention,
642 + attn_cls=MMEncoderAttention,
643 )
644
645 num_hidden_layers = config.num_hidden_layers
vllm/model_executor/models/deepencoder.py
+2 -2
@@ -18,7 +18,7 @@ import torch.nn as nn
18 import torch.nn.functional as F
19 from transformers import CLIPVisionConfig
20
21 -from vllm.attention.layer import MultiHeadAttention
21 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
22 from vllm.model_executor.layers.conv import Conv2dLayer
23 from vllm.model_executor.layers.quantization import QuantizationConfig
24 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -628,7 +628,7 @@ class DeepCLIPVisionTransformer(nn.Module):
628 quant_config=quant_config,
629 num_hidden_layers_override=num_hidden_layers_override,
630 prefix=f"{prefix}.encoder",
631 - attn_cls=MultiHeadAttention,
631 + attn_cls=MMEncoderAttention,
632 )
633
634 num_hidden_layers = config.num_hidden_layers
vllm/model_executor/models/glm4v.py
+2 -2
@@ -19,7 +19,7 @@ from transformers import BatchFeature, PreTrainedTokenizer, TensorType
19 from transformers.image_utils import ImageInput
20 from transformers.tokenization_utils_base import TextInput
21
22 -from vllm.attention.layer import MultiHeadAttention
22 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
23 from vllm.config import VllmConfig
24 from vllm.config.multimodal import BaseDummyOptions
25 from vllm.distributed import get_tensor_model_parallel_world_size
@@ -135,7 +135,7 @@ class EVA2CLIPAttention(nn.Module):
135 prefix=f"{prefix}.dense",
136 )
137
138 - self.attn = MultiHeadAttention(
138 + self.attn = MMEncoderAttention(
139 self.num_heads_per_rank, self.head_dim, self.scale
140 )
141 self.output_dropout = torch.nn.Dropout(config.dropout_prob)
vllm/model_executor/models/hunyuan_vision.py
+2 -2
@@ -34,7 +34,7 @@ import torch.nn.functional as F
34 from transformers import BatchFeature
35
36 from vllm.attention.backends.registry import AttentionBackendEnum
37 -from vllm.attention.layer import MultiHeadAttention
37 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
38 from vllm.config import MultiModalConfig, VllmConfig
39 from vllm.config.multimodal import BaseDummyOptions
40 from vllm.distributed import parallel_state
@@ -232,7 +232,7 @@ class HunYuanVisionAttention(nn.Module):
232 )
233
234 self.scale = self.hidden_size_per_attention_head**-0.5
235 - self.attn = MultiHeadAttention(
235 + self.attn = MMEncoderAttention(
236 self.num_attention_heads_per_partition,
237 self.hidden_size_per_attention_head,
238 self.scale,
vllm/model_executor/models/idefics2_vision_model.py
+4 -4
@@ -27,7 +27,7 @@ from transformers.models.idefics2.configuration_idefics2 import (
27 Idefics2VisionConfig,
28 )
29
30 -from vllm.attention.layer import MultiHeadAttention
30 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
31 from vllm.distributed import get_tensor_model_parallel_world_size
32 from vllm.model_executor.layers.activation import get_act_fn
33 from vllm.model_executor.layers.conv import Conv2dLayer
@@ -161,8 +161,8 @@ class Idefics2VisionAttention(nn.Module):
161 prefix=f"{prefix}.out_proj",
162 disable_tp=use_data_parallel,
163 )
164 - # Use unified MultiHeadAttention with Flash Attention support
165 - self.attn = MultiHeadAttention(
164 + # Use unified MMEncoderAttention with Flash Attention support
165 + self.attn = MMEncoderAttention(
166 self.num_heads_per_partition, self.head_dim, self.scale
167 )
168
@@ -175,7 +175,7 @@ class Idefics2VisionAttention(nn.Module):
175 ) # batch_size, q_len, 3 * num_heads_per_partition * head_dim
176 query_states, key_states, value_states = qkv.chunk(3, dim=-1)
177
178 - # Use unified MultiHeadAttention implementation
178 + # Use unified MMEncoderAttention implementation
179 out = self.attn(query_states, key_states, value_states)
180 attn_output, _ = self.out_proj(out)
181 return attn_output
vllm/model_executor/models/intern_vit.py
+2 -2
@@ -15,7 +15,7 @@ import torch.nn as nn
15 import torch.nn.functional as F
16 from transformers import PretrainedConfig
17
18 -from vllm.attention.layer import MultiHeadAttention
18 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
19 from vllm.distributed import (
20 divide,
21 get_tensor_model_parallel_rank,
@@ -207,7 +207,7 @@ class InternParallelAttention(nn.Module):
207 disable_tp=use_data_parallel,
208 )
209
210 - self.attn = MultiHeadAttention(
210 + self.attn = MMEncoderAttention(
211 self.num_heads_per_partition, self.head_dim, self.scale
212 )
213
vllm/model_executor/models/interns1_vit.py
+4 -4
@@ -14,7 +14,7 @@ import torch.nn as nn
14 from transformers import PretrainedConfig
15 from transformers.utils import torch_int
16
17 -from vllm.attention.layer import MultiHeadAttention
17 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
18 from vllm.model_executor.layers.activation import get_act_fn
19 from vllm.model_executor.layers.conv import Conv2dLayer
20 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -214,8 +214,8 @@ class InternSdpaAttention(nn.Module):
214
215 self.projection_layer = nn.Linear(self.dummy_dim, self.embed_dim)
216
217 - # Use unified MultiHeadAttention with automatic backend selection
218 - self.attn = MultiHeadAttention(self.num_heads, self.head_dim, self.scale)
217 + # Use unified MMEncoderAttention with automatic backend selection
218 + self.attn = MMEncoderAttention(self.num_heads, self.head_dim, self.scale)
219
220 def forward(self, x: torch.Tensor) -> torch.Tensor:
221 """x shape: (B, N, C)"""
@@ -228,7 +228,7 @@ class InternSdpaAttention(nn.Module):
228 q = self.q_norm(q)
229 k = self.k_norm(k)
230
231 - # Use unified MultiHeadAttention with automatic backend selection
231 + # Use unified MMEncoderAttention with automatic backend selection
232 x = self.attn(q, k, v)
233
234 x = self.projection_layer(x)
vllm/model_executor/models/mllama4.py
+2 -2
@@ -31,7 +31,7 @@ from transformers.models.llama4.image_processing_llama4_fast import (
31 get_best_fit,
32 )
33
34 -from vllm.attention.layer import MultiHeadAttention
34 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
35 from vllm.config import VllmConfig
36 from vllm.config.multimodal import BaseDummyOptions
37 from vllm.distributed import get_tensor_model_parallel_world_size
@@ -255,7 +255,7 @@ class Llama4VisionAttention(nn.Module):
255 self.attention_dropout = config.attention_dropout
256 self.scaling = self.head_dim**-0.5
257
258 - self.attn = MultiHeadAttention(
258 + self.attn = MMEncoderAttention(
259 self.num_local_heads, self.head_dim, self.scaling
260 )
261
vllm/model_executor/models/molmo.py
+3 -2
@@ -17,7 +17,8 @@ from transformers import BatchFeature, PretrainedConfig, ProcessorMixin, TensorT
17 from transformers.image_utils import ImageInput
18 from transformers.tokenization_utils_base import TextInput
19
20 -from vllm.attention.layer import Attention, MultiHeadAttention
20 +from vllm.attention.layer import Attention
21 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
22 from vllm.compilation.decorators import support_torch_compile
23 from vllm.config import CacheConfig, VllmConfig
24 from vllm.config.multimodal import BaseDummyOptions
@@ -222,7 +223,7 @@ class MultiHeadDotProductAttention(nn.Module):
223 )
224
225 self.scale = self.head_dim**-0.5
225 - self.attn = MultiHeadAttention(
226 + self.attn = MMEncoderAttention(
227 self.num_heads, self.head_dim, self.scale, num_kv_heads=self.num_kv_heads
228 )
229
vllm/model_executor/models/siglip.py
+5 -5
@@ -16,8 +16,8 @@ from transformers import (
16 SiglipVisionConfig,
17 )
18
19 -from vllm.attention.layer import MultiHeadAttention
19 from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
20 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
21 from vllm.config import VllmConfig
22 from vllm.config.multimodal import BaseDummyOptions
23 from vllm.distributed import divide, get_tensor_model_parallel_world_size
@@ -379,7 +379,7 @@ class SiglipAttention(nn.Module):
379 quant_config: QuantizationConfig | None = None,
380 *,
381 prefix: str = "",
382 - attn_cls: type[EncoderOnlyAttention] | type[MultiHeadAttention],
382 + attn_cls: type[EncoderOnlyAttention] | type[MMEncoderAttention],
383 ) -> None:
384 super().__init__()
385
@@ -481,7 +481,7 @@ class SiglipEncoderLayer(nn.Module):
481 quant_config: QuantizationConfig | None = None,
482 *,
483 prefix: str = "",
484 - attn_cls: type[EncoderOnlyAttention] | type[MultiHeadAttention],
484 + attn_cls: type[EncoderOnlyAttention] | type[MMEncoderAttention],
485 ) -> None:
486 super().__init__()
487
@@ -527,7 +527,7 @@ class SiglipEncoder(nn.Module):
527 num_hidden_layers_override: int | None = None,
528 *,
529 prefix: str = "",
530 - attn_cls: type[EncoderOnlyAttention] | type[MultiHeadAttention],
530 + attn_cls: type[EncoderOnlyAttention] | type[MMEncoderAttention],
531 ) -> None:
532 super().__init__()
533
@@ -700,7 +700,7 @@ class SiglipVisionTransformer(nn.Module):
700 quant_config=quant_config,
701 num_hidden_layers_override=num_hidden_layers_override,
702 prefix=f"{prefix}.encoder",
703 - attn_cls=MultiHeadAttention,
703 + attn_cls=MMEncoderAttention,
704 )
705
706 num_hidden_layers = config.num_hidden_layers
vllm/model_executor/models/step3_vl.py
+4 -4
@@ -15,7 +15,7 @@ from torchvision import transforms
15 from torchvision.transforms.functional import InterpolationMode
16 from transformers import BatchFeature, PretrainedConfig, TensorType
17
18 -from vllm.attention.layer import MultiHeadAttention
18 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
19 from vllm.config import VllmConfig
20 from vllm.config.multimodal import BaseDummyOptions
21 from vllm.distributed import get_tensor_model_parallel_world_size
@@ -753,8 +753,8 @@ class Step3VisionAttention(nn.Module):
753 disable_tp=use_data_parallel,
754 )
755
756 - # Use unified MultiHeadAttention with automatic backend selection
757 - self.attn = MultiHeadAttention(self.num_heads, self.head_dim, self.scale)
756 + # Use unified MMEncoderAttention with automatic backend selection
757 + self.attn = MMEncoderAttention(self.num_heads, self.head_dim, self.scale)
758
759 def forward(
760 self,
@@ -767,7 +767,7 @@ class Step3VisionAttention(nn.Module):
767 qkv, _ = self.qkv_proj(hidden_states)
768 q, k, v = qkv.chunk(chunks=3, dim=-1)
769
770 - # Use unified MultiHeadAttention with automatic backend selection
770 + # Use unified MMEncoderAttention with automatic backend selection
771 attn_output = self.attn(q, k, v)
772
773 attn_output, _ = self.out_proj(attn_output)
vllm/model_executor/models/whisper.py
+3 -3
@@ -16,9 +16,9 @@ from transformers import (
16 )
17 from transformers.models.whisper.modeling_whisper import sinusoids
18
19 -from vllm.attention.backends.abstract import AttentionType
20 -from vllm.attention.layer import Attention, MultiHeadAttention
19 +from vllm.attention.layer import Attention, AttentionType
20 from vllm.attention.layers.cross_attention import CrossAttention
21 +from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
22 from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig
23 from vllm.config.multimodal import BaseDummyOptions
24 from vllm.distributed import get_tensor_model_parallel_world_size
@@ -141,7 +141,7 @@ class WhisperAudioInputs(TensorSchema):
141 ]
142
143
144 -class WhisperEncoderAttention(MultiHeadAttention):
144 +class WhisperEncoderAttention(MMEncoderAttention):
145 """Multi-headed attention for Whisper encoder with 2D tensor support."""
146
147 def forward(