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, 11:15 UTC
[CustomOp][MM] Extract MMEncoderAttention as CustomOp and replace the backend of QwenVisionAttention with it. (#30125)
Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>
Shanshan Shen committed
Dec 15, 2025 at 11:13 UTC
87b4d1557dc83addf5dd8568862ee7a74882200a
24 files changed
+1264
-853
tests/models/multimodal/generation/test_vit_backend_functionality.py
new
+434
@@ -0,0 +1,434 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Consolidated test for ViT attention backend functionality across multiple models.
5
+
6
+This test validates that each multimodal model can successfully generate outputs
7
+using different ViT attention backends. Tests are parametrized by model and backend.
8
+"""
9
+
10
+from dataclasses import asdict
11
+from typing import Any
12
+
13
+import pytest
14
+from transformers import AutoProcessor
15
+
16
+from vllm import LLM, EngineArgs, SamplingParams
17
+from vllm.attention.backends.registry import AttentionBackendEnum
18
+from vllm.multimodal.utils import encode_image_base64
19
+from vllm.multimodal.video import sample_frames_from_video
20
+from vllm.platforms import current_platform
21
+
22
+from ....utils import create_new_process_for_each_test
23
+from ...utils import dummy_hf_overrides
24
+
25
+# Dots.OCR prompt from official repository
26
+# https://github.com/rednote-hilab/dots.ocr/blob/d72d1d8c5bdd0362eb264f714cdbd1e5daa7cdff/dots_ocr/utils/prompts.py#L3
27
+# ruff: noqa: E501
28
+DOTS_OCR_PROMPT = """Please output the layout information from the PDF image, including each layout element's bbox, its category, and the corresponding text content within the bbox.
29
+
30
+1. Bbox format: [x1, y1, x2, y2]
31
+
32
+2. Layout Categories: The possible categories are ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header', 'Picture', 'Section-header', 'Table', 'Text', 'Title'].
33
+
34
+3. Text Extraction & Formatting Rules:
35
+ - Picture: For the 'Picture' category, the text field should be omitted.
36
+ - Formula: Format its text as LaTeX.
37
+ - Table: Format its text as HTML.
38
+ - All Others (Text, Title, etc.): Format their text as Markdown.
39
+
40
+4. Constraints:
41
+ - The output text must be the original text from the image, with no translation.
42
+ - All layout elements must be sorted according to human reading order.
43
+
44
+5. Final Output: The entire output must be a single JSON object.
45
+"""
46
+
47
+VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>"
48
+
49
+
50
+# Model configurations
51
+MODEL_CONFIGS: dict[str, dict[str, Any]] = {
52
+ "dots_ocr": {
53
+ "model_name": "rednote-hilab/dots.ocr",
54
+ "interface": "llm_chat",
55
+ "max_model_len": 32768,
56
+ "max_num_seqs": 1,
57
+ "limit_mm_per_prompt": {"image": 1},
58
+ "sampling_params": {
59
+ "temperature": 0.1,
60
+ "max_tokens": 16384,
61
+ "top_p": 0.9,
62
+ "stop_token_ids": None,
63
+ },
64
+ "use_specific_image": "stop_sign",
65
+ "prompt_builder": "build_dots_ocr_prompt",
66
+ "output_validator": lambda x: len(x) > 10 and "stop" in x.lower(),
67
+ },
68
+ "ernie45_vl": {
69
+ "model_name": "baidu/ERNIE-4.5-VL-28B-A3B-PT",
70
+ "interface": "llm_generate",
71
+ "max_model_len": 16384,
72
+ "max_num_seqs": 2,
73
+ "sampling_params": {
74
+ "temperature": 0.0,
75
+ "max_tokens": 256,
76
+ "stop_token_ids": None,
77
+ },
78
+ "use_processor": True,
79
+ "question": "What is the content of each image?",
80
+ },
81
+ "glm4_1v": {
82
+ "model_name": "zai-org/GLM-4.1V-9B-Thinking",
83
+ "interface": "llm_generate",
84
+ "max_model_len": 32768,
85
+ "max_num_seqs": 2,
86
+ "sampling_params": {
87
+ "temperature": 0.0,
88
+ "max_tokens": 256,
89
+ "stop_token_ids": None,
90
+ },
91
+ "use_processor": True,
92
+ "question": "What is the content of each image?",
93
+ },
94
+ "keye_vl": {
95
+ "model_name": "Kwai-Keye/Keye-VL-8B-Preview",
96
+ "interface": "llm_generate",
97
+ "max_model_len": 8192,
98
+ "max_num_seqs": 5,
99
+ "sampling_params": {
100
+ "temperature": 0.0,
101
+ "max_tokens": 256,
102
+ "stop_token_ids": None,
103
+ },
104
+ "supported_backends": {
105
+ AttentionBackendEnum.FLASH_ATTN,
106
+ AttentionBackendEnum.ROCM_AITER_FA,
107
+ },
108
+ "use_processor": True,
109
+ "question": "What is the content of each image?",
110
+ },
111
+ "ovis2_5": {
112
+ "model_name": "AIDC-AI/Ovis2.5-2B",
113
+ "interface": "llm_generate",
114
+ "max_model_len": 8192,
115
+ "max_num_seqs": 2,
116
+ "sampling_params": {
117
+ "temperature": 0.0,
118
+ "max_tokens": 256,
119
+ "stop_token_ids": None,
120
+ },
121
+ "prompt_builder": "build_ovis_prompt",
122
+ "question": "What is the content of each image?",
123
+ },
124
+ "qwen2_5_vl": {
125
+ "model_name": "Qwen/Qwen2.5-VL-3B-Instruct",
126
+ "interface": "vllm_runner",
127
+ "media_type": "video",
128
+ "max_model_len": 4000,
129
+ "max_num_seqs": 1,
130
+ "limit_mm_per_prompt": {"video": 1},
131
+ "sampling_params": {
132
+ "max_tokens": 128,
133
+ },
134
+ "runner_kwargs": {
135
+ "runner": "generate",
136
+ "dtype": "bfloat16",
137
+ },
138
+ "video_params": {
139
+ "num_frames": 16,
140
+ "pruning_rates": [0.0, 0.75],
141
+ },
142
+ },
143
+ "qwen2_5_omni": {
144
+ "model_name": "Qwen/Qwen2.5-Omni-3B",
145
+ "interface": "llm_generate",
146
+ "max_model_len": 32768,
147
+ "max_num_seqs": 2,
148
+ "limit_mm_per_prompt": {"image": 3, "video": 3, "audio": 3},
149
+ "sampling_params": {
150
+ "temperature": 0.6,
151
+ "top_p": 0.95,
152
+ "top_k": 20,
153
+ "max_tokens": 16384,
154
+ },
155
+ "use_processor": True,
156
+ "question": "What is the content of each image?",
157
+ },
158
+ "qwen3_omni": {
159
+ "model_name": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
160
+ "interface": "llm_generate",
161
+ "max_model_len": 32768,
162
+ "max_num_seqs": 2,
163
+ "limit_mm_per_prompt": {"image": 3, "video": 3, "audio": 3},
164
+ "sampling_params": {
165
+ "temperature": 0.6,
166
+ "top_p": 0.95,
167
+ "top_k": 20,
168
+ "max_tokens": 16384,
169
+ },
170
+ "use_processor": True,
171
+ "question": "What is the content of each image?",
172
+ },
173
+}
174
+
175
+
176
+# Prompt builder functions
177
+def build_dots_ocr_prompt(images, config):
178
+ """Build Dots.OCR specific prompt with OCR instructions."""
179
+ # Use only stop_sign image for Dots.OCR
180
+ image = images[0] # Already filtered to stop_sign
181
+
182
+ image_url = f"data:image/jpeg;base64,{encode_image_base64(image)}"
183
+
184
+ placeholders = [{"type": "image_url", "image_url": {"url": image_url}}]
185
+ messages = [
186
+ {
187
+ "role": "user",
188
+ "content": [
189
+ *placeholders,
190
+ {
191
+ "type": "text",
192
+ "text": f"<|img|><|imgpad|><|endofimg|>{DOTS_OCR_PROMPT}",
193
+ },
194
+ ],
195
+ },
196
+ ]
197
+
198
+ return messages
199
+
200
+
201
+def build_processor_prompt(images, config):
202
+ """Build prompt using AutoProcessor.apply_chat_template()."""
203
+ processor = AutoProcessor.from_pretrained(
204
+ config["model_name"], trust_remote_code=True
205
+ )
206
+
207
+ image_urls = [
208
+ f"data:image/jpeg;base64,{encode_image_base64(img)}" for img in images
209
+ ]
210
+ placeholders = [{"type": "image", "image": url} for url in image_urls]
211
+ messages = [
212
+ {
213
+ "role": "user",
214
+ "content": [
215
+ *placeholders,
216
+ {"type": "text", "text": config["question"]},
217
+ ],
218
+ },
219
+ ]
220
+
221
+ return processor.apply_chat_template(
222
+ messages, tokenize=False, add_generation_prompt=True
223
+ )
224
+
225
+
226
+def build_ovis_prompt(images, config):
227
+ """Build Ovis2.5 specific prompt with custom format."""
228
+ image_urls = [
229
+ f"data:image/jpeg;base64,{encode_image_base64(img)}" for img in images
230
+ ]
231
+
232
+ placeholders = "\n".join(
233
+ f"Image-{i}: <image>\n" for i, _ in enumerate(image_urls, start=1)
234
+ )
235
+
236
+ return (
237
+ f"<|im_start|>user\n\n{placeholders}\n{config['question']}<|im_end|>\n"
238
+ "<|im_start|>assistant\n"
239
+ )
240
+
241
+
242
+def build_qwen2_5_video_prompt():
243
+ """Build Qwen2.5-VL video prompt with EVS placeholder."""
244
+ return (
245
+ f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
246
+ f"<|im_start|>user\n{VIDEO_PLACEHOLDER}"
247
+ "Describe this video with a short sentence (no more than 20 words)"
248
+ "<|im_end|><|im_start|>assistant\n"
249
+ )
250
+
251
+
252
+# Handler functions
253
+def run_llm_generate_test(config, mm_encoder_attn_backend, image_assets):
254
+ """Standard LLM.generate() interface handler."""
255
+ images = [asset.pil_image for asset in image_assets]
256
+
257
+ # Build prompt
258
+ if config.get("use_processor"):
259
+ prompt = build_processor_prompt(images, config)
260
+ else:
261
+ prompt_builder_name = config.get("prompt_builder", "build_ovis_prompt")
262
+ prompt_builder = globals()[prompt_builder_name]
263
+ prompt = prompt_builder(images, config)
264
+
265
+ # Determine limit_mm_per_prompt
266
+ limit_mm_per_prompt = config.get("limit_mm_per_prompt", {"image": len(images)})
267
+
268
+ # Create engine
269
+ engine_args = EngineArgs(
270
+ model=config["model_name"],
271
+ trust_remote_code=True,
272
+ max_model_len=config["max_model_len"],
273
+ max_num_seqs=config["max_num_seqs"],
274
+ limit_mm_per_prompt=limit_mm_per_prompt,
275
+ mm_encoder_attn_backend=mm_encoder_attn_backend,
276
+ hf_overrides=dummy_hf_overrides,
277
+ load_format="dummy",
278
+ )
279
+
280
+ engine_dict = asdict(engine_args) | {"seed": 42}
281
+ llm = LLM(**engine_dict)
282
+
283
+ # Generate
284
+ sampling_params = SamplingParams(**config["sampling_params"])
285
+ outputs = llm.generate(
286
+ {
287
+ "prompt": prompt,
288
+ "multi_modal_data": {"image": images},
289
+ },
290
+ sampling_params=sampling_params,
291
+ )
292
+
293
+ # Validate
294
+ for o in outputs:
295
+ generated_text = o.outputs[0].text
296
+ validator = config.get("output_validator", lambda x: len(x) > 10)
297
+ assert validator(generated_text), (
298
+ f"Validation failed for {config['model_name']}: {generated_text}"
299
+ )
300
+
301
+
302
+def run_llm_chat_test(config, mm_encoder_attn_backend, image_assets):
303
+ """LLM.chat() interface handler for Dots.OCR."""
304
+ # Filter to stop_sign image only
305
+ stop_sign_image = [
306
+ asset.pil_image for asset in image_assets if asset.name == "stop_sign"
307
+ ][0]
308
+
309
+ # Build messages
310
+ messages = build_dots_ocr_prompt([stop_sign_image], config)
311
+
312
+ # Create engine
313
+ engine_args = EngineArgs(
314
+ model=config["model_name"],
315
+ trust_remote_code=True,
316
+ max_model_len=config["max_model_len"],
317
+ max_num_seqs=config["max_num_seqs"],
318
+ limit_mm_per_prompt=config["limit_mm_per_prompt"],
319
+ mm_encoder_attn_backend=mm_encoder_attn_backend,
320
+ hf_overrides=dummy_hf_overrides,
321
+ load_format="dummy",
322
+ )
323
+
324
+ engine_dict = asdict(engine_args) | {"seed": 42}
325
+ llm = LLM(**engine_dict)
326
+
327
+ # Generate using chat
328
+ sampling_params = SamplingParams(**config["sampling_params"])
329
+ outputs = llm.chat(messages=messages, sampling_params=sampling_params)
330
+
331
+ # Validate
332
+ for o in outputs:
333
+ generated_text = o.outputs[0].text
334
+ validator = config.get("output_validator", lambda x: len(x) > 10)
335
+ assert validator(generated_text), (
336
+ f"Validation failed for {config['model_name']}: {generated_text}"
337
+ )
338
+
339
+
340
+def run_video_test(config, mm_encoder_attn_backend, video_assets, vllm_runner):
341
+ """Video test with EVS (Efficient Video Sampling) handler."""
342
+ for pruning_rate in config["video_params"]["pruning_rates"]:
343
+ num_frames = config["video_params"]["num_frames"]
344
+
345
+ # Sample frames from video
346
+ sampled_vids = [
347
+ sample_frames_from_video(asset.np_ndarrays, num_frames)
348
+ for asset in video_assets
349
+ ]
350
+
351
+ # Build prompt and prepare video
352
+ prompt = build_qwen2_5_video_prompt()
353
+ prompts = [prompt]
354
+ videos = [sampled_vids[0]]
355
+
356
+ # Run with vllm_runner context manager
357
+ with vllm_runner(
358
+ config["model_name"],
359
+ max_model_len=config["max_model_len"],
360
+ max_num_seqs=config["max_num_seqs"],
361
+ limit_mm_per_prompt=config["limit_mm_per_prompt"],
362
+ tensor_parallel_size=1,
363
+ video_pruning_rate=pruning_rate,
364
+ mm_encoder_attn_backend=mm_encoder_attn_backend,
365
+ hf_overrides=dummy_hf_overrides,
366
+ load_format="dummy",
367
+ **config["runner_kwargs"],
368
+ ) as vllm_model:
369
+ outputs = vllm_model.generate_greedy(
370
+ prompts,
371
+ config["sampling_params"]["max_tokens"],
372
+ videos=videos,
373
+ )
374
+
375
+ # Validate output
376
+ assert len(outputs) == 1, f"Expected 1 output, got {len(outputs)}"
377
+ output_ids, output_text = outputs[0]
378
+ assert len(output_ids) > 0, "Generated no output IDs"
379
+ assert len(output_text) > 0, "Generated empty text"
380
+ assert isinstance(output_text, str), (
381
+ f"Output is not string: {type(output_text)}"
382
+ )
383
+
384
+
385
+# Main test function
386
+@pytest.mark.parametrize("model_key", list(MODEL_CONFIGS.keys()))
387
+@pytest.mark.parametrize(
388
+ "mm_encoder_attn_backend",
389
+ [None] + current_platform.get_supported_vit_attn_backends(),
390
+)
391
+@create_new_process_for_each_test()
392
+def test_vit_backend_functionality(
393
+ model_key: str,
394
+ mm_encoder_attn_backend: AttentionBackendEnum | None,
395
+ image_assets,
396
+ video_assets,
397
+ vllm_runner,
398
+ request,
399
+):
400
+ """Test ViT attention backend functionality for multimodal models.
401
+
402
+ This test validates that each model can successfully generate outputs
403
+ using different ViT attention backends. The test:
404
+ 1. Filters unsupported backends per model
405
+ 2. Applies appropriate GPU marks
406
+ 3. Routes to the correct test handler based on interface
407
+ 4. Validates output meets minimum requirements
408
+ """
409
+ config = MODEL_CONFIGS[model_key]
410
+
411
+ # Step 1: Backend filtering
412
+ if (
413
+ "supported_backends" in config
414
+ and mm_encoder_attn_backend is not None
415
+ and mm_encoder_attn_backend not in config["supported_backends"]
416
+ ):
417
+ pytest.skip(
418
+ f"{model_key} does not support {mm_encoder_attn_backend} backend now."
419
+ )
420
+
421
+ # Step 2: Apply GPU marks dynamically
422
+ if "gpu_marks" in config:
423
+ for mark in config["gpu_marks"]:
424
+ request.applymarker(mark)
425
+
426
+ # Step 3: Route to appropriate handler
427
+ if config.get("media_type") == "video":
428
+ run_video_test(config, mm_encoder_attn_backend, video_assets, vllm_runner)
429
+ elif config["interface"] == "llm_chat":
430
+ run_llm_chat_test(config, mm_encoder_attn_backend, image_assets)
431
+ elif config["interface"] == "llm_generate":
432
+ run_llm_generate_test(config, mm_encoder_attn_backend, image_assets)
433
+ else:
434
+ raise ValueError(f"Unknown interface: {config['interface']}")
vllm/attention/layer.py
+5
-68
@@ -3,7 +3,6 @@
3
"""Attention layer."""
4
5
import functools
6
-from collections.abc import Callable
6
from typing import cast
7
8
import torch
@@ -17,6 +16,7 @@ from vllm.attention.backends.abstract import (
16
MLAAttentionImpl,
17
)
18
from vllm.attention.backends.registry import AttentionBackendEnum
19
+from vllm.attention.layers.mm_encoder_attention import maybe_get_vit_flash_attn_backend
20
from vllm.attention.selector import get_attn_backend
21
from vllm.attention.utils.fa_utils import get_flash_attn_version
22
from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
@@ -49,58 +49,9 @@ from vllm.v1.kv_cache_interface import (
49
SlidingWindowSpec,
50
)
51
52
-if current_platform.is_rocm():
53
- from vllm.platforms.rocm import on_gfx9
54
-else:
55
- on_gfx9 = lambda *args, **kwargs: False
56
-
57
-
58
-FP8_DTYPE = current_platform.fp8_dtype()
52
logger = init_logger(__name__)
53
54
62
-def maybe_get_vit_flash_attn_backend(
63
- attn_backend: AttentionBackendEnum,
64
- attn_backend_override: AttentionBackendEnum | None = None,
65
-) -> tuple[AttentionBackendEnum, Callable | None]:
66
- if current_platform.is_rocm():
67
- if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9():
68
- attn_backend = AttentionBackendEnum.ROCM_AITER_FA
69
- elif (
70
- attn_backend_override is None
71
- and on_gfx9()
72
- and attn_backend == AttentionBackendEnum.FLASH_ATTN
73
- ):
74
- pass
75
- else:
76
- return AttentionBackendEnum.TORCH_SDPA, None
77
- elif current_platform.is_cuda():
78
- pass
79
- elif current_platform.is_xpu():
80
- assert attn_backend == AttentionBackendEnum.FLASH_ATTN, (
81
- "XPU platform only supports FLASH_ATTN as vision attention backend."
82
- )
83
- pass
84
- else:
85
- return AttentionBackendEnum.TORCH_SDPA, None
86
-
87
- if attn_backend in {
88
- AttentionBackendEnum.FLASH_ATTN,
89
- AttentionBackendEnum.ROCM_AITER_FA,
90
- }:
91
- if attn_backend == AttentionBackendEnum.ROCM_AITER_FA:
92
- from aiter import flash_attn_varlen_func
93
- else:
94
- try:
95
- from vllm.attention.utils.fa_utils import flash_attn_varlen_func
96
- except ImportError:
97
- flash_attn_varlen_func = None
98
- else:
99
- flash_attn_varlen_func = None
100
-
101
- return attn_backend, flash_attn_varlen_func
102
-
103
-
55
def _init_kv_cache_quant(
56
layer: nn.Module,
57
quant_config: QuantizationConfig | None,
@@ -496,29 +447,15 @@ class MultiHeadAttention(nn.Module):
447
attn_backend_override = None
448
if multimodal_config is not None:
449
attn_backend_override = multimodal_config.mm_encoder_attn_backend
499
- backend = get_vit_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
505
- self.attn_backend = (
506
- backend
507
- if backend
508
- in {
509
- AttentionBackendEnum.TORCH_SDPA,
510
- AttentionBackendEnum.PALLAS,
511
- AttentionBackendEnum.ROCM_AITER_FA,
512
- AttentionBackendEnum.FLASH_ATTN,
513
- }
514
- else AttentionBackendEnum.TORCH_SDPA
515
- )
516
-
517
- self.attn_backend, self._flash_attn_varlen_func = (
518
- maybe_get_vit_flash_attn_backend(
519
- self.attn_backend,
520
- attn_backend_override=attn_backend_override,
521
- )
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 {
vllm/attention/layers/mm_encoder_attention.py
new
+284
@@ -0,0 +1,284 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Callable
5
+
6
+import torch
7
+
8
+from vllm.attention.backends.registry import AttentionBackendEnum
9
+from vllm.attention.ops.vit_attn_wrappers import (
10
+ vit_flash_attn_wrapper,
11
+ vit_torch_sdpa_wrapper,
12
+)
13
+from vllm.config import MultiModalConfig
14
+from vllm.logger import init_logger
15
+from vllm.model_executor.custom_op import CustomOp
16
+from vllm.model_executor.models.vision import get_vit_attn_backend
17
+
18
+logger = init_logger(__name__)
19
+
20
+
21
+def maybe_get_vit_flash_attn_backend(
22
+ attn_backend: AttentionBackendEnum | None,
23
+) -> Callable | None:
24
+ # At this point,
25
+ # we already have the attn_backend,
26
+ # overriding logic is done in the platform-specific implementation.
27
+ # so we don't need to override backend here.
28
+ # Just return the attn_backend and flash_attn_varlen_func.
29
+
30
+ if attn_backend == AttentionBackendEnum.FLASH_ATTN:
31
+ from vllm.attention.utils.fa_utils import flash_attn_varlen_func
32
+ elif attn_backend == AttentionBackendEnum.ROCM_AITER_FA:
33
+ from aiter import flash_attn_varlen_func
34
+ else:
35
+ flash_attn_varlen_func = None
36
+
37
+ # if attn_backend is TORCH_SDPA,
38
+ # it will reach here and the flash_attn_varlen_func will be None.
39
+ return flash_attn_varlen_func
40
+
41
+
42
+@CustomOp.register("mm_encoder_attn")
43
+class MMEncoderAttention(CustomOp):
44
+ """Multi-headed attention without any cache, used for multimodal encoder."""
45
+
46
+ def __init__(
47
+ self,
48
+ num_heads: int,
49
+ head_size: int,
50
+ scale: float | None = None,
51
+ num_kv_heads: int | None = None,
52
+ prefix: str = "",
53
+ multimodal_config: MultiModalConfig | None = None,
54
+ ) -> None:
55
+ """
56
+ Args:
57
+ num_heads: number of attention heads per partition.
58
+ head_size: hidden_size per attention head.
59
+ scale: scale factor.
60
+ num_kv_heads: number of kv heads.
61
+ prefix: This has no effect, it is only here to make it easier to
62
+ swap between Attention and MultiHeadAttention
63
+ multimodal_config: configs for multi-modal.
64
+ """
65
+ super().__init__()
66
+
67
+ self.num_heads = num_heads
68
+ self.head_size = head_size
69
+ self.scale = scale
70
+ self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
71
+ self.layer_name = prefix
72
+
73
+ assert self.num_heads % self.num_kv_heads == 0, (
74
+ f"num_heads ({self.num_heads}) is not "
75
+ f"divisible by num_kv_heads ({self.num_kv_heads})"
76
+ )
77
+ self.num_queries_per_kv = self.num_heads // self.num_kv_heads
78
+
79
+ # During model initialization, the default dtype is set as the model
80
+ # weight and activation dtype.
81
+ dtype = torch.get_default_dtype()
82
+
83
+ # Try to get vision attention backend from multimodal_config.
84
+ attn_backend_override = None
85
+ if multimodal_config is not None:
86
+ attn_backend_override = multimodal_config.mm_encoder_attn_backend
87
+
88
+ # Get device-specific vision attention backend.
89
+ self.attn_backend = get_vit_attn_backend(
90
+ head_size=head_size,
91
+ dtype=dtype,
92
+ attn_backend_override=attn_backend_override,
93
+ )
94
+
95
+ self.is_flash_attn_backend = self.attn_backend in {
96
+ AttentionBackendEnum.FLASH_ATTN,
97
+ AttentionBackendEnum.ROCM_AITER_FA,
98
+ }
99
+
100
+ self.flash_attn_varlen_func = maybe_get_vit_flash_attn_backend(
101
+ self.attn_backend,
102
+ )
103
+
104
+ logger.info_once(f"Using {self.attn_backend} for MMEncoderAttention.")
105
+
106
+ @classmethod
107
+ def enabled(cls) -> bool:
108
+ return True
109
+
110
+ def reshape_qkv_to_4d(
111
+ self,
112
+ query: torch.Tensor,
113
+ key: torch.Tensor,
114
+ value: torch.Tensor,
115
+ bsz: int,
116
+ q_len: int,
117
+ kv_len: int,
118
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
119
+ """
120
+ Reshape query, key, value to 4D tensors:
121
+ (batch_size, seq_len, num_heads, head_size)
122
+ """
123
+ query = query.view(bsz, q_len, self.num_heads, self.head_size)
124
+ key = key.view(bsz, kv_len, self.num_kv_heads, self.head_size)
125
+ value = value.view(bsz, kv_len, self.num_kv_heads, self.head_size)
126
+
127
+ if (num_repeat := self.num_queries_per_kv) > 1:
128
+ # Handle MQA and GQA
129
+ key = torch.repeat_interleave(key, num_repeat, dim=2)
130
+ value = torch.repeat_interleave(value, num_repeat, dim=2)
131
+
132
+ return query, key, value
133
+
134
+ def reshape_qkv_to_3d(
135
+ self,
136
+ query: torch.Tensor,
137
+ key: torch.Tensor,
138
+ value: torch.Tensor,
139
+ bsz: int,
140
+ q_len: int,
141
+ kv_len: int,
142
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
143
+ """
144
+ Reshape query, key, value to 3D tensors:
145
+ (batch_size * seq_len, num_heads, head_size)
146
+ """
147
+ query = query.view(bsz * q_len, self.num_heads, self.head_size)
148
+ key = key.view(bsz * kv_len, self.num_kv_heads, self.head_size)
149
+ value = value.view(bsz * kv_len, self.num_kv_heads, self.head_size)
150
+
151
+ if (num_repeat := self.num_queries_per_kv) > 1:
152
+ # Handle MQA and GQA
153
+ key = torch.repeat_interleave(key, num_repeat, dim=1)
154
+ value = torch.repeat_interleave(value, num_repeat, dim=1)
155
+
156
+ return query, key, value
157
+
158
+ def _forward_sdpa(
159
+ self,
160
+ query: torch.Tensor,
161
+ key: torch.Tensor,
162
+ value: torch.Tensor,
163
+ cu_seqlens: torch.Tensor | None = None,
164
+ ) -> torch.Tensor:
165
+ # TODO(Isotr0py): Migrate MultiHeadAttention
166
+ assert cu_seqlens is not None
167
+
168
+ bsz, q_len = query.size()[:2]
169
+ kv_len = key.size(1)
170
+
171
+ query, key, value = self.reshape_qkv_to_4d(
172
+ query, key, value, bsz, q_len, kv_len
173
+ )
174
+
175
+ output = vit_torch_sdpa_wrapper(
176
+ q=query,
177
+ k=key,
178
+ v=value,
179
+ cu_seqlens=cu_seqlens,
180
+ )
181
+ return output
182
+
183
+ def _forward_fa(
184
+ self,
185
+ query: torch.Tensor,
186
+ key: torch.Tensor,
187
+ value: torch.Tensor,
188
+ cu_seqlens: torch.Tensor | None = None,
189
+ max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
190
+ ) -> torch.Tensor:
191
+ assert self.flash_attn_varlen_func is not None, (
192
+ "Flash attention function is not set."
193
+ )
194
+ # # TODO(Isotr0py): Migrate MultiHeadAttention
195
+ assert cu_seqlens is not None and max_seqlen is not None
196
+
197
+ bsz = query.shape[0]
198
+
199
+ output = vit_flash_attn_wrapper(
200
+ q=query,
201
+ k=key,
202
+ v=value,
203
+ cu_seqlens=cu_seqlens,
204
+ max_seqlen=max_seqlen,
205
+ batch_size=bsz,
206
+ is_rocm_aiter=(self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA),
207
+ )
208
+ return output
209
+
210
+ def forward_native(
211
+ self,
212
+ query: torch.Tensor,
213
+ key: torch.Tensor,
214
+ value: torch.Tensor,
215
+ cu_seqlens: torch.Tensor | None = None,
216
+ max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
217
+ ) -> torch.Tensor:
218
+ return self._forward_sdpa(query, key, value, cu_seqlens)
219
+
220
+ def forward_cuda(
221
+ self,
222
+ query: torch.Tensor,
223
+ key: torch.Tensor,
224
+ value: torch.Tensor,
225
+ cu_seqlens: torch.Tensor | None = None,
226
+ max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
227
+ ) -> torch.Tensor:
228
+ if self.is_flash_attn_backend:
229
+ return self._forward_fa(query, key, value, cu_seqlens, max_seqlen)
230
+ elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
231
+ return self._forward_sdpa(query, key, value, cu_seqlens)
232
+ else:
233
+ raise ValueError(
234
+ f"Unsupported multi-modal encoder attention backend for CUDA: "
235
+ f"{self.attn_backend}."
236
+ )
237
+
238
+ def forward_cpu(
239
+ self,
240
+ query: torch.Tensor,
241
+ key: torch.Tensor,
242
+ value: torch.Tensor,
243
+ cu_seqlens: torch.Tensor | None = None,
244
+ max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
245
+ ) -> torch.Tensor:
246
+ return self._forward_sdpa(query, key, value, cu_seqlens)
247
+
248
+ def forward_xpu(
249
+ self,
250
+ query: torch.Tensor,
251
+ key: torch.Tensor,
252
+ value: torch.Tensor,
253
+ cu_seqlens: torch.Tensor | None = None,
254
+ max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
255
+ ) -> torch.Tensor:
256
+ assert self.is_flash_attn_backend, (
257
+ "XPU only supports FLASH_ATTN for vision attention."
258
+ )
259
+ return self._forward_fa(query, key, value, cu_seqlens, max_seqlen)
260
+
261
+ def forward_tpu(
262
+ self,
263
+ query: torch.Tensor,
264
+ key: torch.Tensor,
265
+ value: torch.Tensor,
266
+ cu_seqlens: torch.Tensor | None = None,
267
+ max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
268
+ ) -> torch.Tensor:
269
+ assert self.attn_backend == AttentionBackendEnum.PALLAS, (
270
+ f"MMEncoderAttention on TPU only supports PALLAS backend, "
271
+ f"but got {self.attn_backend}."
272
+ )
273
+ if cu_seqlens is None:
274
+ query, key, value = (x.transpose(1, 2) for x in (query, key, value))
275
+ from torch_xla.experimental.custom_kernel import flash_attention
276
+
277
+ out = flash_attention(query, key, value, sm_scale=self.scale)
278
+ out = out.transpose(1, 2)
279
+ return out
280
+ logger.warning_once(
281
+ "PALLAS backend with cu_seqlens is not supported for ViT yet. ",
282
+ "Falling back to SDPA implementation.",
283
+ )
284
+ return self._forward_sdpa(query, key, value, cu_seqlens)
vllm/attention/ops/vit_attn_wrappers.py
+3
-8
@@ -44,9 +44,7 @@ def flash_attn_maxseqlen_wrapper(
44
dropout_p=0.0,
45
causal=False,
46
)
47
- context_layer = einops.rearrange(
48
- output, "(b s) h d -> s b (h d)", b=batch_size
49
- ).contiguous()
47
+ context_layer = einops.rearrange(output, "(b s) h d -> b s h d", b=batch_size)
48
return context_layer
49
50
@@ -59,8 +57,7 @@ def flash_attn_maxseqlen_wrapper_fake(
57
batch_size: int,
58
is_rocm_aiter: bool,
59
) -> torch.Tensor:
62
- b, s, h, d = q.shape
63
- return torch.empty((s, b, h * d), dtype=q.dtype, device=q.device)
60
+ return torch.empty_like(q)
61
62
63
direct_register_custom_op(
@@ -106,7 +103,6 @@ def torch_sdpa_wrapper(
103
output_i = einops.rearrange(output_i, "b h s d -> b s h d ")
104
outputs.append(output_i)
105
context_layer = torch.cat(outputs, dim=1)
109
- context_layer = einops.rearrange(context_layer, "b s h d -> s b (h d)").contiguous()
106
return context_layer
107
108
@@ -116,8 +112,7 @@ def torch_sdpa_wrapper_fake(
112
v: torch.Tensor,
113
cu_seqlens: torch.Tensor,
114
) -> torch.Tensor:
119
- b, s, h, d = q.shape
120
- return torch.empty((s, b, h * d), dtype=q.dtype, device=q.device)
115
+ return torch.empty_like(q)
116
117
118
direct_register_custom_op(
vllm/model_executor/models/dots_ocr.py
+46
-83
@@ -5,15 +5,14 @@ from typing import Annotated, Literal, TypeAlias
5
6
import torch
7
import torch.nn as nn
8
-import torch.nn.functional as F
8
from torch.nn import LayerNorm
9
from transformers.models.qwen2_vl import Qwen2VLProcessor
10
11
from vllm.attention.backends.registry import AttentionBackendEnum
13
-from vllm.attention.layer import (
14
- maybe_get_vit_flash_attn_backend,
12
+from vllm.attention.layers.mm_encoder_attention import (
13
+ MMEncoderAttention,
14
)
16
-from vllm.config import VllmConfig
15
+from vllm.config import MultiModalConfig, VllmConfig
16
from vllm.config.multimodal import BaseDummyOptions
17
from vllm.distributed import utils as dist_utils
18
from vllm.distributed.parallel_state import (
@@ -254,11 +253,15 @@ class DotsVisionAttention(nn.Module):
253
bias: bool = True,
254
*,
255
quant_config: QuantizationConfig | None = None,
256
+ multimodal_config: MultiModalConfig | None = None,
257
prefix: str = "",
258
- use_data_parallel: bool = False,
259
- attn_backend_override: AttentionBackendEnum | None = None,
258
) -> None:
259
super().__init__()
260
+ use_data_parallel = (
261
+ multimodal_config.mm_encoder_tp_mode == "data"
262
+ if multimodal_config
263
+ else False
264
+ )
265
266
self.embed_dim = dim
267
self.tp_size = (
@@ -287,31 +290,13 @@ class DotsVisionAttention(nn.Module):
290
prefix=f"{prefix}.proj",
291
disable_tp=use_data_parallel,
292
)
290
- # Select attention backend
291
- self.attn_backend = get_vit_attn_backend(
292
- self.hidden_size_per_attention_head,
293
- torch.get_default_dtype(),
294
- attn_backend_override=attn_backend_override,
295
- )
293
297
- self.attn_backend, self.flash_attn_varlen_func = (
298
- maybe_get_vit_flash_attn_backend(
299
- self.attn_backend,
300
- attn_backend_override=attn_backend_override,
301
- )
294
+ self.attn = MMEncoderAttention(
295
+ num_heads=self.num_attention_heads_per_partition,
296
+ head_size=self.hidden_size_per_attention_head,
297
+ multimodal_config=multimodal_config,
298
+ prefix=f"{prefix}.attn",
299
)
303
- if self.attn_backend not in {
304
- AttentionBackendEnum.FLASH_ATTN,
305
- AttentionBackendEnum.TORCH_SDPA,
306
- AttentionBackendEnum.ROCM_AITER_FA,
307
- }:
308
- raise RuntimeError(
309
- f"Unsupported vision attention backend: {self.attn_backend}"
310
- )
311
- self.is_flash_attn_backend = self.attn_backend in {
312
- AttentionBackendEnum.FLASH_ATTN,
313
- AttentionBackendEnum.ROCM_AITER_FA,
314
- }
300
301
def forward(
302
self,
@@ -319,7 +304,7 @@ class DotsVisionAttention(nn.Module):
304
cu_seqlens: torch.Tensor,
305
rotary_pos_emb: torch.Tensor | None = None,
306
*,
322
- max_seqlen: int | None = None,
307
+ max_seqlen: torch.Tensor | None = None,
308
) -> torch.Tensor:
309
# [S, C] -> [S, B=1, C]
310
x = hidden_states.unsqueeze(1)
@@ -336,41 +321,13 @@ class DotsVisionAttention(nn.Module):
321
qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
322
q, k = torch.chunk(qk_rotated, 2, dim=0)
323
339
- if self.is_flash_attn_backend:
340
- q_ = q.reshape(bs * q.shape[1], q.shape[2], q.shape[3])
341
- k_ = k.reshape(bs * k.shape[1], k.shape[2], k.shape[3])
342
- v_ = v.reshape(bs * v.shape[1], v.shape[2], v.shape[3])
343
- output = self.flash_attn_varlen_func(
344
- q_,
345
- k_,
346
- v_,
347
- cu_seqlens_q=cu_seqlens,
348
- cu_seqlens_k=cu_seqlens,
349
- max_seqlen_q=max_seqlen,
350
- max_seqlen_k=max_seqlen,
351
- dropout_p=0.0,
352
- causal=False,
353
- )
354
- context_layer = output.view(
355
- bs,
356
- -1,
357
- self.num_attention_heads_per_partition,
358
- self.hidden_size_per_attention_head,
359
- )
360
- elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
361
- outputs = []
362
- for i in range(1, len(cu_seqlens)):
363
- s = int(cu_seqlens[i - 1])
364
- e = int(cu_seqlens[i])
365
- q_i = q[:, s:e].permute(0, 2, 1, 3)
366
- k_i = k[:, s:e].permute(0, 2, 1, 3)
367
- v_i = v[:, s:e].permute(0, 2, 1, 3)
368
- out_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
369
- out_i = out_i.permute(0, 2, 1, 3)
370
- outputs.append(out_i)
371
- context_layer = torch.cat(outputs, dim=1) if outputs else q[:, :0]
372
- else:
373
- raise RuntimeError("Unsupported attention backend")
324
+ context_layer = self.attn(
325
+ query=q,
326
+ key=k,
327
+ value=v,
328
+ cu_seqlens=cu_seqlens,
329
+ max_seqlen=max_seqlen,
330
+ )
331
332
# [B,S,H,D] -> [S,B,H*D] -> [S, C]
333
context_layer = context_layer.permute(1, 0, 2, 3).contiguous()
@@ -385,14 +342,19 @@ class DotsSwiGLUFFN(nn.Module):
342
config,
343
*,
344
quant_config: QuantizationConfig | None = None,
345
+ multimodal_config: MultiModalConfig | None = None,
346
prefix: str = "",
389
- use_data_parallel: bool = False,
347
):
348
super().__init__()
349
hidden_features = config.intermediate_size
350
in_features = config.embed_dim
351
bias = config.use_bias
352
353
+ use_data_parallel = (
354
+ multimodal_config.mm_encoder_tp_mode == "data"
355
+ if multimodal_config
356
+ else False
357
+ )
358
# Referenced aimv2.py AIMv2SwiGLUFFN
359
self.fc13 = MergedColumnParallelLinear(
360
in_features,
@@ -498,9 +460,8 @@ class DotsVisionBlock(nn.Module):
460
config,
461
*,
462
quant_config: QuantizationConfig | None = None,
463
+ multimodal_config: MultiModalConfig | None = None,
464
prefix: str = "",
502
- use_data_parallel: bool = False,
503
- attn_backend_override: AttentionBackendEnum | None = None,
465
):
466
super().__init__()
467
@@ -510,16 +471,15 @@ class DotsVisionBlock(nn.Module):
471
num_heads=config.num_attention_heads,
472
bias=config.use_bias,
473
quant_config=quant_config,
474
+ multimodal_config=multimodal_config,
475
prefix=f"{prefix}.attn",
514
- use_data_parallel=use_data_parallel,
515
- attn_backend_override=attn_backend_override,
476
)
477
self.norm1 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
478
self.mlp = DotsSwiGLUFFN(
479
config,
480
quant_config=quant_config,
481
+ multimodal_config=multimodal_config,
482
prefix=f"{prefix}.mlp",
522
- use_data_parallel=use_data_parallel,
483
)
484
self.norm2 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
485
@@ -546,12 +506,11 @@ class DotsVisionTransformer(nn.Module):
506
self,
507
config: DotsVisionConfig,
508
quant_config: QuantizationConfig | None = None,
509
+ multimodal_config: MultiModalConfig | None = None,
510
*,
511
num_hidden_layers_override: int | None = None,
512
require_post_norm: bool | None = None,
513
prefix: str = "",
553
- use_data_parallel: bool = False,
554
- attn_backend_override: AttentionBackendEnum | None = None,
514
) -> None:
515
super().__init__()
516
self.config = config
@@ -561,6 +520,11 @@ class DotsVisionTransformer(nn.Module):
520
521
head_dim = config.embed_dim // config.num_attention_heads
522
self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
523
+ attn_backend_override = (
524
+ multimodal_config.mm_encoder_attn_backend
525
+ if multimodal_config is not None
526
+ else None
527
+ )
528
self.attn_backend = get_vit_attn_backend(
529
head_size=head_dim,
530
dtype=torch.get_default_dtype(),
@@ -578,9 +542,8 @@ class DotsVisionTransformer(nn.Module):
542
DotsVisionBlock(
543
config,
544
quant_config=quant_config,
545
+ multimodal_config=multimodal_config,
546
prefix=f"{prefix}.blocks.{i}",
582
- use_data_parallel=use_data_parallel,
583
- attn_backend_override=attn_backend_override,
547
)
548
for i in range(num_layers)
549
]
@@ -592,6 +555,11 @@ class DotsVisionTransformer(nn.Module):
555
else:
556
self.post_trunk_norm = None
557
558
+ use_data_parallel = (
559
+ multimodal_config.mm_encoder_tp_mode == "data"
560
+ if multimodal_config
561
+ else False
562
+ )
563
self.merger = PatchMerger(
564
dim=config.hidden_size,
565
context_dim=config.embed_dim,
@@ -647,7 +615,7 @@ class DotsVisionTransformer(nn.Module):
615
self.attn_backend == AttentionBackendEnum.FLASH_ATTN
616
or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
617
):
650
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
618
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
619
return max_seqlen
620
621
def forward(
@@ -733,17 +701,12 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA
701
self.config.vision_config = vision_config
702
else:
703
vision_config = self.config.vision_config
736
- attn_backend_override = (
737
- multimodal_config.mm_encoder_attn_backend
738
- if multimodal_config is not None
739
- else None
740
- )
704
+
705
self.vision_tower = DotsVisionTransformer(
706
vision_config,
707
quant_config=self.quant_config,
708
+ multimodal_config=multimodal_config,
709
prefix=maybe_prefix(prefix, "vision_tower"),
745
- use_data_parallel=self.use_data_parallel,
746
- attn_backend_override=attn_backend_override,
710
)
711
self.language_model: Qwen2ForCausalLM = init_vllm_registered_model(
712
vllm_config=vllm_config,
vllm/model_executor/models/ernie45_vl.py
+28
-80
@@ -37,10 +37,10 @@ from einops import rearrange, repeat
37
from transformers import BatchFeature
38
39
from vllm.attention.backends.registry import AttentionBackendEnum
40
-from vllm.attention.layer import (
41
- maybe_get_vit_flash_attn_backend,
40
+from vllm.attention.layers.mm_encoder_attention import (
41
+ MMEncoderAttention,
42
)
43
-from vllm.config import VllmConfig
43
+from vllm.config import MultiModalConfig, VllmConfig
44
from vllm.config.multimodal import BaseDummyOptions
45
from vllm.distributed import parallel_state
46
from vllm.distributed import utils as dist_utils
@@ -163,8 +163,8 @@ class Ernie4_5_VisionAttention(nn.Module):
163
num_heads: int,
164
projection_size: int,
165
quant_config: QuantizationConfig | None = None,
166
+ multimodal_config: MultiModalConfig | None = None,
167
prefix: str = "",
167
- attn_backend_override: AttentionBackendEnum | None = None,
168
) -> None:
169
super().__init__()
170
# Per attention head and per partition values.
@@ -193,33 +193,13 @@ class Ernie4_5_VisionAttention(nn.Module):
193
prefix=f"{prefix}.proj",
194
)
195
196
- # Detect attention implementation.
197
- self.attn_backend = get_vit_attn_backend(
196
+ self.attn = MMEncoderAttention(
197
+ num_heads=self.num_attention_heads_per_partition,
198
head_size=self.hidden_size_per_attention_head,
199
- dtype=torch.get_default_dtype(),
200
- attn_backend_override=attn_backend_override,
201
- )
202
-
203
- self.attn_backend, self.flash_attn_varlen_func = (
204
- maybe_get_vit_flash_attn_backend(
205
- self.attn_backend,
206
- attn_backend_override=attn_backend_override,
207
- )
199
+ multimodal_config=multimodal_config,
200
+ prefix=f"{prefix}.attn",
201
)
202
210
- if self.attn_backend not in {
211
- AttentionBackendEnum.FLASH_ATTN,
212
- AttentionBackendEnum.TORCH_SDPA,
213
- AttentionBackendEnum.ROCM_AITER_FA,
214
- }:
215
- raise RuntimeError(
216
- f"Ernie45-VL does not support {self.attn_backend} backend now."
217
- )
218
- self.is_flash_attn_backend = self.attn_backend in {
219
- AttentionBackendEnum.FLASH_ATTN,
220
- AttentionBackendEnum.ROCM_AITER_FA,
221
- }
222
-
203
def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
204
# [s, b, 3 * head * head_dim]
205
seq_len, bs, _ = qkv.shape
@@ -253,14 +233,13 @@ class Ernie4_5_VisionAttention(nn.Module):
233
x: torch.Tensor,
234
cu_seqlens: torch.Tensor,
235
rotary_pos_emb: torch.Tensor,
256
- max_seqlen: int | None = None, # Only used for Flash Attention
236
+ max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
237
) -> torch.Tensor:
238
# [s, b, c] --> [s, b, head * 3 * head_dim]
239
x, _ = self.qkv(x)
240
241
# [s, b, 3 * head * head_dim] -> 3 * [s, b, head, head_dim]
242
q, k, v = self.split_qkv(x)
263
- batch_size = q.shape[1]
243
244
q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous() for x in (q, k, v))
245
if rotary_pos_emb is not None:
@@ -268,43 +247,14 @@ class Ernie4_5_VisionAttention(nn.Module):
247
qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
248
q, k = torch.chunk(qk_rotated, 2, dim=0)
249
271
- if self.is_flash_attn_backend:
272
- q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
273
-
274
- output = self.flash_attn_varlen_func(
275
- q,
276
- k,
277
- v,
278
- cu_seqlens_q=cu_seqlens,
279
- cu_seqlens_k=cu_seqlens,
280
- max_seqlen_q=max_seqlen,
281
- max_seqlen_k=max_seqlen,
282
- dropout_p=0.0,
283
- causal=False,
284
- )
285
-
286
- context_layer = rearrange(
287
- output, "(b s) h d -> s b (h d)", b=batch_size
288
- ).contiguous()
289
- elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
290
- # Execute attention entry by entry for speed & less VRAM.
291
- outputs = []
292
-
293
- lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
294
- q_chunks = torch.split(q, lens, dim=1)
295
- k_chunks = torch.split(k, lens, dim=1)
296
- v_chunks = torch.split(v, lens, dim=1)
297
- for q_i, k_i, v_i in zip(q_chunks, k_chunks, v_chunks):
298
- q_i, k_i, v_i = (
299
- rearrange(x, "b s h d -> b h s d") for x in [q_i, k_i, v_i]
300
- )
301
- output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
302
- output_i = rearrange(output_i, "b h s d -> b s h d ")
303
- outputs.append(output_i)
304
- context_layer = torch.cat(outputs, dim=1)
305
- context_layer = rearrange(
306
- context_layer, "b s h d -> s b (h d)"
307
- ).contiguous()
250
+ output = self.attn(
251
+ query=q,
252
+ key=k,
253
+ value=v,
254
+ cu_seqlens=cu_seqlens,
255
+ max_seqlen=max_seqlen,
256
+ )
257
+ context_layer = rearrange(output, "b s h d -> s b (h d)").contiguous()
258
259
output, _ = self.proj(context_layer)
260
return output
@@ -350,8 +300,8 @@ class Ernie4_5_VisionBlock(nn.Module):
300
act_layer: type[nn.Module] = QuickGELU,
301
norm_layer: Callable[[int], nn.Module] | None = None,
302
quant_config: QuantizationConfig | None = None,
303
+ multimodal_config: MultiModalConfig | None = None,
304
prefix: str = "",
354
- attn_backend_override: AttentionBackendEnum | None = None,
305
) -> None:
306
super().__init__()
307
@@ -366,8 +316,8 @@ class Ernie4_5_VisionBlock(nn.Module):
316
num_heads=num_heads,
317
projection_size=dim,
318
quant_config=quant_config,
319
+ multimodal_config=multimodal_config,
320
prefix=f"{prefix}.attn",
370
- attn_backend_override=attn_backend_override,
321
)
322
323
self.mlp = Ernie4_5_VisionMLP(
@@ -383,7 +333,7 @@ class Ernie4_5_VisionBlock(nn.Module):
333
hidden_states: torch.Tensor,
334
cu_seqlens: torch.Tensor,
335
rotary_pos_emb: torch.Tensor,
386
- max_seqlen: int | None = None, # Only used for Flash Attention
336
+ max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
337
) -> torch.Tensor:
338
hidden_states = hidden_states + self.attn(
339
self.norm1(hidden_states),
@@ -441,8 +391,8 @@ class Ernie4_5_VisionTransformer(nn.Module):
391
vision_config,
392
norm_eps: float = 1e-6,
393
quant_config: QuantizationConfig | None = None,
394
+ multimodal_config: MultiModalConfig | None = None,
395
prefix: str = "",
445
- attn_backend_override: AttentionBackendEnum | None = None,
396
) -> None:
397
super().__init__()
398
patch_size = vision_config.patch_size
@@ -477,8 +427,8 @@ class Ernie4_5_VisionTransformer(nn.Module):
427
mlp_ratio=mlp_ratio,
428
norm_layer=norm_layer,
429
quant_config=quant_config,
430
+ multimodal_config=multimodal_config,
431
prefix=f"{prefix}.blocks.{layer_idx}",
481
- attn_backend_override=attn_backend_override,
432
)
433
for layer_idx in range(depth)
434
]
@@ -489,6 +439,9 @@ class Ernie4_5_VisionTransformer(nn.Module):
439
)
440
self.ln = nn.LayerNorm(hidden_size, eps=1e-6)
441
442
+ attn_backend_override = (
443
+ multimodal_config.mm_encoder_attn_backend if multimodal_config else None
444
+ )
445
self.attn_backend = get_vit_attn_backend(
446
head_size=head_dim,
447
dtype=torch.get_default_dtype(),
@@ -535,13 +488,13 @@ class Ernie4_5_VisionTransformer(nn.Module):
488
rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
489
return rotary_pos_emb
490
538
- def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> int | None:
491
+ def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> torch.Tensor | None:
492
max_seqlen = None
493
if (
494
self.attn_backend == AttentionBackendEnum.FLASH_ATTN
495
or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
496
):
544
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
497
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
498
return max_seqlen
499
500
def forward(
@@ -1304,17 +1257,12 @@ class Ernie4_5_VLMoeForConditionalGeneration(
1257
self.config = config
1258
self.multimodal_config = multimodal_config
1259
1307
- attn_backend_override = (
1308
- multimodal_config.mm_encoder_attn_backend
1309
- if multimodal_config is not None
1310
- else None
1311
- )
1260
self.vision_model = Ernie4_5_VisionTransformer(
1261
config.vision_config,
1262
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
1263
quant_config=quant_config,
1264
+ multimodal_config=multimodal_config,
1265
prefix=maybe_prefix(prefix, "vision_model"),
1317
- attn_backend_override=attn_backend_override,
1266
)
1267
1268
self.language_model = Ernie4_5_VLMoeForCausalLM(
vllm/model_executor/models/glm4_1v.py
+46
-91
@@ -47,8 +47,10 @@ 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 maybe_get_vit_flash_attn_backend
51
-from vllm.config import VllmConfig
50
+from vllm.attention.layers.mm_encoder_attention import (
51
+ MMEncoderAttention,
52
+)
53
+from vllm.config import MultiModalConfig, VllmConfig
54
from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
55
from vllm.distributed import get_tensor_model_parallel_world_size, parallel_state
56
from vllm.distributed import utils as dist_utils
@@ -191,10 +193,15 @@ class Glm4vVisionMLP(nn.Module):
193
hidden_features: int,
194
bias: bool = False,
195
quant_config: QuantizationConfig | None = None,
196
+ multimodal_config: MultiModalConfig | None = None,
197
prefix: str = "",
195
- use_data_parallel: bool = False,
198
):
199
super().__init__()
200
+ use_data_parallel = (
201
+ multimodal_config.mm_encoder_tp_mode == "data"
202
+ if multimodal_config
203
+ else False
204
+ )
205
self.gate_up_proj = MergedColumnParallelLinear(
206
input_size=in_features,
207
output_sizes=[hidden_features] * 2,
@@ -248,12 +255,16 @@ class Glm4vVisionAttention(nn.Module):
255
num_heads: int,
256
projection_size: int,
257
quant_config: QuantizationConfig | None = None,
258
+ multimodal_config: MultiModalConfig | None = None,
259
prefix: str = "",
252
- use_data_parallel: bool = False,
253
- attn_backend_override: AttentionBackendEnum | None = None,
260
) -> None:
261
super().__init__()
262
# Per attention head and per partition values.
263
+ use_data_parallel = (
264
+ multimodal_config.mm_encoder_tp_mode == "data"
265
+ if multimodal_config
266
+ else False
267
+ )
268
self.tp_size = (
269
1 if use_data_parallel else get_tensor_model_parallel_world_size()
270
)
@@ -287,34 +298,12 @@ class Glm4vVisionAttention(nn.Module):
298
disable_tp=use_data_parallel,
299
)
300
290
- # Detect attention implementation.
291
- self.attn_backend = get_vit_attn_backend(
301
+ self.attn = MMEncoderAttention(
302
+ num_heads=self.num_attention_heads_per_partition,
303
head_size=self.hidden_size_per_attention_head,
293
- dtype=torch.get_default_dtype(),
294
- attn_backend_override=attn_backend_override,
295
- )
296
-
297
- self.attn_backend, self.flash_attn_varlen_func = (
298
- maybe_get_vit_flash_attn_backend(
299
- self.attn_backend,
300
- attn_backend_override=attn_backend_override,
301
- )
304
+ multimodal_config=multimodal_config,
305
)
306
304
- if self.attn_backend not in {
305
- AttentionBackendEnum.FLASH_ATTN,
306
- AttentionBackendEnum.TORCH_SDPA,
307
- AttentionBackendEnum.ROCM_AITER_FA,
308
- }:
309
- raise RuntimeError(
310
- f"GLM-4V does not support {self.attn_backend} backend now."
311
- )
312
-
313
- self.is_flash_attn_backend = self.attn_backend in {
314
- AttentionBackendEnum.FLASH_ATTN,
315
- AttentionBackendEnum.ROCM_AITER_FA,
316
- }
317
-
307
def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
308
# [s, b, 3 * head * head_dim]
309
seq_len, bs, _ = qkv.shape
@@ -338,14 +327,13 @@ class Glm4vVisionAttention(nn.Module):
327
cu_seqlens: torch.Tensor,
328
rotary_pos_emb_cos: torch.Tensor,
329
rotary_pos_emb_sin: torch.Tensor,
341
- max_seqlen: int | None = None, # Only used for Flash Attention
330
+ max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
331
) -> torch.Tensor:
332
# [s, b, c] --> [s, b, head * 3 * head_dim]
333
x, _ = self.qkv(x)
334
335
# [s, b, 3 * head * head_dim] -> 3 * [s, b, head, head_dim]
336
q, k, v = self.split_qkv(x)
348
- batch_size = q.shape[1]
337
338
q, k, v = (rearrange(x, "s b ... -> b s ...").contiguous() for x in (q, k, v))
339
if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:
@@ -356,43 +344,14 @@ class Glm4vVisionAttention(nn.Module):
344
)
345
q, k = torch.chunk(qk_rotated, 2, dim=0)
346
359
- if self.is_flash_attn_backend:
360
- q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
361
-
362
- output = self.flash_attn_varlen_func(
363
- q,
364
- k,
365
- v,
366
- cu_seqlens_q=cu_seqlens,
367
- cu_seqlens_k=cu_seqlens,
368
- max_seqlen_q=max_seqlen,
369
- max_seqlen_k=max_seqlen,
370
- dropout_p=0.0,
371
- causal=False,
372
- )
373
-
374
- context_layer = rearrange(
375
- output, "(b s) h d -> s b (h d)", b=batch_size
376
- ).contiguous()
377
- elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
378
- # Execute attention entry by entry for speed & less VRAM.
379
- outputs = []
380
-
381
- lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
382
- q_chunks = torch.split(q, lens, dim=1)
383
- k_chunks = torch.split(k, lens, dim=1)
384
- v_chunks = torch.split(v, lens, dim=1)
385
- for q_i, k_i, v_i in zip(q_chunks, k_chunks, v_chunks):
386
- q_i, k_i, v_i = (
387
- rearrange(x, "b s h d -> b h s d") for x in [q_i, k_i, v_i]
388
- )
389
- output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
390
- output_i = rearrange(output_i, "b h s d -> b s h d ")
391
- outputs.append(output_i)
392
- context_layer = torch.cat(outputs, dim=1)
393
- context_layer = rearrange(
394
- context_layer, "b s h d -> s b (h d)"
395
- ).contiguous()
347
+ context_layer = self.attn(
348
+ query=q,
349
+ key=k,
350
+ value=v,
351
+ cu_seqlens=cu_seqlens,
352
+ max_seqlen=max_seqlen,
353
+ )
354
+ context_layer = rearrange(context_layer, "b s h d -> s b (h d)").contiguous()
355
356
output, _ = self.proj(context_layer)
357
return output
@@ -406,9 +365,8 @@ class Glm4vVisionBlock(nn.Module):
365
mlp_hidden_dim: int,
366
norm_layer: Callable[[int], nn.Module] | None = None,
367
quant_config: QuantizationConfig | None = None,
368
+ multimodal_config: MultiModalConfig | None = None,
369
prefix: str = "",
410
- use_data_parallel: bool = False,
411
- attn_backend_override: AttentionBackendEnum | None = None,
370
) -> None:
371
super().__init__()
372
if norm_layer is None:
@@ -420,17 +378,16 @@ class Glm4vVisionBlock(nn.Module):
378
num_heads=num_heads,
379
projection_size=dim,
380
quant_config=quant_config,
381
+ multimodal_config=multimodal_config,
382
prefix=f"{prefix}.attn",
424
- use_data_parallel=use_data_parallel,
425
- attn_backend_override=attn_backend_override,
383
)
384
self.mlp = Glm4vVisionMLP(
385
dim,
386
mlp_hidden_dim,
387
bias=False,
388
quant_config=quant_config,
389
+ multimodal_config=multimodal_config,
390
prefix=f"{prefix}.mlp",
433
- use_data_parallel=use_data_parallel,
391
)
392
393
def forward(
@@ -489,11 +446,16 @@ class Glm4vPatchMerger(nn.Module):
446
d_model: int,
447
context_dim: int,
448
quant_config: QuantizationConfig | None = None,
449
+ multimodal_config: MultiModalConfig | None = None,
450
bias: bool = False,
451
prefix: str = "",
494
- use_data_parallel: bool = False,
452
) -> None:
453
super().__init__()
454
+ use_data_parallel = (
455
+ multimodal_config.mm_encoder_tp_mode == "data"
456
+ if multimodal_config
457
+ else False
458
+ )
459
self.hidden_size = d_model
460
self.proj = ColumnParallelLinear(
461
self.hidden_size,
@@ -649,19 +611,19 @@ class Glm4vVisionTransformer(nn.Module):
611
vision_config: Glm4vVisionConfig,
612
norm_eps: float = 1e-6,
613
quant_config: QuantizationConfig | None = None,
614
+ multimodal_config: MultiModalConfig | None = None,
615
prefix: str = "",
653
- use_data_parallel: bool = False,
654
- attn_backend_override: AttentionBackendEnum | None = None,
616
) -> None:
617
super().__init__()
618
619
+ assert multimodal_config is not None, "multimodal_config must be provided"
620
+
621
patch_size = vision_config.patch_size
622
temporal_patch_size = vision_config.temporal_patch_size
623
in_channels = vision_config.in_channels
624
depth = vision_config.depth
625
self.hidden_size = vision_config.hidden_size
626
self.num_heads = vision_config.num_heads
664
- self.use_data_parallel = use_data_parallel
627
628
self.patch_size = vision_config.patch_size
629
self.spatial_merge_size = vision_config.spatial_merge_size
@@ -690,9 +652,8 @@ class Glm4vVisionTransformer(nn.Module):
652
mlp_hidden_dim=vision_config.out_hidden_size,
653
norm_layer=norm_layer,
654
quant_config=quant_config,
655
+ multimodal_config=multimodal_config,
656
prefix=f"{prefix}.blocks.{layer_idx}",
694
- use_data_parallel=self.use_data_parallel,
695
- attn_backend_override=attn_backend_override,
657
)
658
for layer_idx in range(depth)
659
]
@@ -701,9 +662,9 @@ class Glm4vVisionTransformer(nn.Module):
662
d_model=vision_config.out_hidden_size,
663
context_dim=vision_config.intermediate_size,
664
quant_config=quant_config,
665
+ multimodal_config=multimodal_config,
666
bias=False,
667
prefix=f"{prefix}.merger",
706
- use_data_parallel=self.use_data_parallel,
668
)
669
self.embeddings = Glm4vVisionEmbeddings(vision_config)
670
@@ -723,7 +684,7 @@ class Glm4vVisionTransformer(nn.Module):
684
self.attn_backend = get_vit_attn_backend(
685
head_size=head_dim,
686
dtype=torch.get_default_dtype(),
726
- attn_backend_override=attn_backend_override,
687
+ attn_backend_override=multimodal_config.mm_encoder_attn_backend,
688
)
689
690
@property
@@ -775,13 +736,13 @@ class Glm4vVisionTransformer(nn.Module):
736
def compute_attn_mask_seqlen(
737
self,
738
cu_seqlens: torch.Tensor,
778
- ) -> int | None:
739
+ ) -> torch.Tensor | None:
740
max_seqlen = None
741
if (
742
self.attn_backend == AttentionBackendEnum.FLASH_ATTN
743
or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA
744
):
784
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
745
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
746
return max_seqlen
747
748
def forward(
@@ -1465,18 +1426,12 @@ class Glm4vForConditionalGeneration(
1426
self.multimodal_config = multimodal_config
1427
self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
1428
1468
- attn_backend_override = (
1469
- multimodal_config.mm_encoder_attn_backend
1470
- if multimodal_config is not None
1471
- else None
1472
- )
1429
self.visual = Glm4vVisionTransformer(
1430
config.vision_config,
1431
norm_eps=getattr(config, "rms_norm_eps", 1e-5),
1432
quant_config=quant_config,
1433
+ multimodal_config=multimodal_config,
1434
prefix=maybe_prefix(prefix, "visual"),
1478
- use_data_parallel=self.use_data_parallel,
1479
- attn_backend_override=attn_backend_override,
1435
)
1436
1437
if config.model_type == "glm4v":
vllm/model_executor/models/keye.py
+27
-80
@@ -9,7 +9,6 @@ from typing import Annotated, Any, Literal, TypeAlias, TypeVar
9
import numpy as np
10
import torch
11
import torch.nn as nn
12
-import torch.nn.functional as F
12
from einops import rearrange
13
from transformers import PretrainedConfig
14
from transformers.activations import GELUActivation
@@ -17,11 +16,10 @@ from transformers.feature_extraction_utils import BatchFeature
16
from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
17
from transformers.utils import torch_int
18
20
-from vllm.attention.backends.registry import AttentionBackendEnum
21
-from vllm.attention.layer import (
22
- maybe_get_vit_flash_attn_backend,
19
+from vllm.attention.layers.mm_encoder_attention import (
20
+ MMEncoderAttention,
21
)
24
-from vllm.config import VllmConfig
22
+from vllm.config import MultiModalConfig, VllmConfig
23
from vllm.config.multimodal import BaseDummyOptions
24
from vllm.distributed import get_tensor_model_parallel_world_size
25
from vllm.logger import init_logger
@@ -80,7 +78,6 @@ from .utils import (
78
is_pp_missing_parameter,
79
maybe_prefix,
80
)
83
-from .vision import get_vit_attn_backend
81
82
logger = init_logger(__name__)
83
@@ -369,8 +366,8 @@ class KeyeSiglipAttention(nn.Module):
366
self,
367
config: PretrainedConfig,
368
quant_config: QuantizationConfig | None = None,
369
+ multimodal_config: MultiModalConfig | None = None,
370
prefix: str = "",
373
- attn_backend_override: AttentionBackendEnum | None = None,
371
):
372
super().__init__()
373
self.config = config
@@ -408,34 +405,14 @@ class KeyeSiglipAttention(nn.Module):
405
prefix=f"{prefix}.out_proj",
406
)
407
411
- # Detect attention implementation.
412
- self.attn_backend = get_vit_attn_backend(
408
+ self.attn = MMEncoderAttention(
409
+ num_heads=self.num_heads,
410
head_size=self.head_dim,
414
- dtype=torch.get_default_dtype(),
415
- attn_backend_override=attn_backend_override,
411
+ num_kv_heads=self.num_kv_heads,
412
+ prefix=f"{prefix}.attn",
413
+ multimodal_config=multimodal_config,
414
)
415
418
- self.attn_backend, self.flash_attn_varlen_func = (
419
- maybe_get_vit_flash_attn_backend(
420
- self.attn_backend,
421
- attn_backend_override=attn_backend_override,
422
- )
423
- )
424
-
425
- if self.attn_backend not in {
426
- AttentionBackendEnum.FLASH_ATTN,
427
- AttentionBackendEnum.TORCH_SDPA,
428
- AttentionBackendEnum.ROCM_AITER_FA,
429
- }:
430
- raise RuntimeError(
431
- f"Keye-VL does not support {self.attn_backend} backend now."
432
- )
433
-
434
- self.is_flash_attn_backend = self.attn_backend in {
435
- AttentionBackendEnum.FLASH_ATTN,
436
- AttentionBackendEnum.ROCM_AITER_FA,
437
- }
438
-
416
def forward(
417
self,
418
hidden_states: torch.Tensor,
@@ -450,8 +427,7 @@ class KeyeSiglipAttention(nn.Module):
427
dim=-1,
428
)
429
453
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
454
- batch_size = q.shape[0]
430
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
431
432
if rope_emb is None:
433
q = q.view(*q.shape[:-1], self.num_heads, self.head_dim)
@@ -482,38 +458,14 @@ class KeyeSiglipAttention(nn.Module):
458
self.head_dim,
459
)
460
485
- if self.is_flash_attn_backend:
486
- q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
487
-
488
- output = self.flash_attn_varlen_func(
489
- q,
490
- k,
491
- v,
492
- cu_seqlens_q=cu_seqlens,
493
- cu_seqlens_k=cu_seqlens,
494
- max_seqlen_q=max_seqlen,
495
- max_seqlen_k=max_seqlen,
496
- causal=False,
497
- softmax_scale=self.scale,
498
- )
499
- context_layer = rearrange(output, "(b s) ... -> b s ...", b=batch_size)
500
- elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
501
- outputs = []
502
- for i in range(1, len(cu_seqlens)):
503
- start_idx = cu_seqlens[i - 1]
504
- end_idx = cu_seqlens[i]
505
- q_i = q[:, start_idx:end_idx]
506
- k_i = k[:, start_idx:end_idx]
507
- v_i = v[:, start_idx:end_idx]
508
- q_i, k_i, v_i = (
509
- rearrange(x, "b s h d -> b h s d") for x in (q_i, k_i, v_i)
510
- )
511
- output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
512
- output_i = rearrange(output_i, "b h s d -> b s h d ")
513
- outputs.append(output_i)
514
- context_layer = torch.cat(outputs, dim=1) if outputs else q[:, :0]
515
-
516
- context_layer = rearrange(context_layer, "b s h d -> b s (h d)").contiguous()
461
+ context_layer = self.attn(
462
+ query=q,
463
+ key=k,
464
+ value=v,
465
+ cu_seqlens=cu_seqlens,
466
+ max_seqlen=max_seqlen,
467
+ )
468
+ context_layer = rearrange(context_layer, "b s h d -> b s (h d)")
469
470
output, _ = self.out_proj(context_layer)
471
return output
@@ -547,8 +499,8 @@ class KeyeSiglipEncoderLayer(nn.Module):
499
self,
500
config: PretrainedConfig,
501
quant_config: QuantizationConfig | None = None,
502
+ multimodal_config: MultiModalConfig | None = None,
503
prefix: str = "",
551
- attn_backend_override: AttentionBackendEnum | None = None,
504
):
505
super().__init__()
506
self.embed_dim = config.hidden_size
@@ -556,8 +508,8 @@ class KeyeSiglipEncoderLayer(nn.Module):
508
self.self_attn = KeyeSiglipAttention(
509
config,
510
quant_config=quant_config,
511
+ multimodal_config=multimodal_config,
512
prefix=f"{prefix}.self_attn",
560
- attn_backend_override=attn_backend_override,
513
)
514
self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
515
self.mlp = SiglipMLP(
@@ -601,8 +553,8 @@ class KeyeSiglipEncoder(nn.Module):
553
self,
554
config: PretrainedConfig,
555
quant_config: QuantizationConfig | None = None,
556
+ multimodal_config: MultiModalConfig | None = None,
557
prefix: str = "",
605
- attn_backend_override: AttentionBackendEnum | None = None,
558
):
559
super().__init__()
560
self.config = config
@@ -614,8 +566,8 @@ class KeyeSiglipEncoder(nn.Module):
566
KeyeSiglipEncoderLayer(
567
config,
568
quant_config=quant_config,
569
+ multimodal_config=multimodal_config,
570
prefix=f"{prefix}.layers.{layer_idx}",
618
- attn_backend_override=attn_backend_override,
571
)
572
for layer_idx in range(config.num_hidden_layers)
573
]
@@ -696,8 +648,8 @@ class KeyeSiglipVisionTransformer(nn.Module):
648
self,
649
config: PretrainedConfig,
650
quant_config: QuantizationConfig | None = None,
651
+ multimodal_config: MultiModalConfig | None = None,
652
prefix: str = "",
700
- attn_backend_override: AttentionBackendEnum | None = None,
653
):
654
super().__init__()
655
self.config = config
@@ -707,8 +659,8 @@ class KeyeSiglipVisionTransformer(nn.Module):
659
self.encoder = KeyeSiglipEncoder(
660
config,
661
quant_config=quant_config,
662
+ multimodal_config=multimodal_config,
663
prefix=f"{prefix}.encoder",
711
- attn_backend_override=attn_backend_override,
664
)
665
self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
666
@@ -779,16 +731,16 @@ class KeyeSiglipVisionModel(nn.Module):
731
self,
732
config: PretrainedConfig,
733
quant_config: QuantizationConfig | None = None,
734
+ multimodal_config: MultiModalConfig | None = None,
735
prefix: str = "",
783
- attn_backend_override: AttentionBackendEnum | None = None,
736
):
737
super().__init__()
738
739
self.vision_model = KeyeSiglipVisionTransformer(
740
config,
741
quant_config=quant_config,
742
+ multimodal_config=multimodal_config,
743
prefix=f"{prefix}.vision_model",
791
- attn_backend_override=attn_backend_override,
744
)
745
self.quant_config = quant_config
746
@@ -1329,16 +1281,11 @@ class BaseKeyeModule(nn.Module):
1281
self.config = config
1282
self.multimodal_config = multimodal_config
1283
1332
- attn_backend_override = (
1333
- multimodal_config.mm_encoder_attn_backend
1334
- if multimodal_config is not None
1335
- else None
1336
- )
1284
self.visual = KeyeSiglipVisionModel(
1285
config.vision_config,
1286
quant_config=quant_config,
1287
+ multimodal_config=multimodal_config,
1288
prefix=maybe_prefix(prefix, "visual"),
1341
- attn_backend_override=attn_backend_override,
1289
)
1290
1291
self.mlp_AR = self._build_projector(
vllm/model_executor/models/opencua.py
+1
-7
@@ -240,18 +240,12 @@ class OpenCUAForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):
240
)
241
242
if multimodal_config.get_limit_per_prompt("image"):
243
- attn_backend_override = (
244
- multimodal_config.mm_encoder_attn_backend
245
- if multimodal_config is not None
246
- else None
247
- )
243
self.visual = OpenCUAVisionTransformer(
244
vision_config=config.vision_config,
245
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
246
quant_config=self.quant_config,
247
+ multimodal_config=self.multimodal_config,
248
prefix=maybe_prefix(prefix, "visual"),
253
- use_data_parallel=self.use_data_parallel,
254
- attn_backend_override=attn_backend_override,
249
)
250
else:
251
self.visual = None
vllm/model_executor/models/ovis2_5.py
+6
-16
@@ -10,8 +10,7 @@ import torch
10
import torch.nn as nn
11
from transformers import BaseImageProcessor, BatchFeature, PretrainedConfig
12
13
-from vllm.attention.backends.registry import AttentionBackendEnum
14
-from vllm.config import VllmConfig
13
+from vllm.config import MultiModalConfig, VllmConfig
14
from vllm.config.multimodal import BaseDummyOptions
15
from vllm.model_executor.layers.linear import ReplicatedLinear
16
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -104,18 +103,16 @@ class VisualTokenizer(torch.nn.Module):
103
config: PretrainedConfig,
104
visual_vocab_size: int,
105
quant_config: QuantizationConfig | None = None,
106
+ multimodal_config: MultiModalConfig | None = None,
107
prefix: str = "",
108
- use_data_parallel: bool = False,
109
- attn_backend_override: AttentionBackendEnum | None = None,
108
):
109
super().__init__()
110
self.config = config
111
self.vit = self._init_backbone(
112
config=config,
113
quant_config=quant_config,
114
+ multimodal_config=multimodal_config,
115
prefix=f"{prefix}.vit",
117
- use_data_parallel=use_data_parallel,
118
- attn_backend_override=attn_backend_override,
116
)
117
# reserved tokens for INDICATOR_IDS
118
head_dim = visual_vocab_size - len(INDICATOR_IDS)
@@ -133,18 +130,16 @@ class VisualTokenizer(torch.nn.Module):
130
self,
131
config: PretrainedConfig,
132
quant_config: QuantizationConfig | None = None,
133
+ multimodal_config: QuantizationConfig | None = None,
134
prefix: str = "",
137
- use_data_parallel: bool = False,
138
- attn_backend_override: AttentionBackendEnum | None = None,
135
):
136
model_type = config.model_type
137
if model_type == "siglip2_navit":
138
return Siglip2NavitModel(
139
config=config,
140
quant_config=quant_config,
141
+ multimodal_config=multimodal_config,
142
prefix=prefix,
146
- use_data_parallel=use_data_parallel,
147
- attn_backend_override=attn_backend_override,
143
)
144
raise ValueError(f"Unsupported visual tokenizer model_type: {model_type}")
145
@@ -468,17 +463,12 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):
463
prefix=maybe_prefix(prefix, "llm"),
464
)
465
471
- attn_backend_override = (
472
- multimodal_config.mm_encoder_attn_backend
473
- if multimodal_config is not None
474
- else None
475
- )
466
self.visual_tokenizer = VisualTokenizer(
467
config=config.vit_config,
468
visual_vocab_size=config.visual_vocab_size,
469
+ multimodal_config=multimodal_config,
470
quant_config=quant_config,
471
prefix=f"{prefix}.visual_tokenizer",
481
- attn_backend_override=attn_backend_override,
472
)
473
474
self.vte = VisualEmbedding(config.visual_vocab_size, config.hidden_size)
vllm/model_executor/models/paddleocr_vl.py
+30
-75
@@ -22,7 +22,6 @@ from typing import Annotated, Literal
22
import numpy as np
23
import torch
24
import torch.nn as nn
25
-import torch.nn.functional as F
25
from einops import rearrange, repeat
26
from transformers import BatchFeature, PretrainedConfig
27
from transformers.activations import GELUActivation
@@ -32,13 +31,10 @@ from transformers.modeling_outputs import (
31
from transformers.utils import torch_int
32
33
from vllm.attention.backends.registry import AttentionBackendEnum
35
-from vllm.attention.layer import (
36
- maybe_get_vit_flash_attn_backend,
34
+from vllm.attention.layers.mm_encoder_attention import (
35
+ MMEncoderAttention,
36
)
38
-from vllm.attention.ops.vit_attn_wrappers import (
39
- vit_flash_attn_wrapper,
40
-)
41
-from vllm.config import VllmConfig
37
+from vllm.config import MultiModalConfig, VllmConfig
38
from vllm.config.multimodal import BaseDummyOptions
39
from vllm.distributed import parallel_state
40
from vllm.distributed import utils as dist_utils
@@ -578,9 +574,8 @@ class SiglipAttention(nn.Module):
574
num_heads: int,
575
projection_size: int,
576
quant_config: QuantizationConfig | None = None,
577
+ multimodal_config: MultiModalConfig | None = None,
578
prefix: str = "",
582
- attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
583
- attn_backend_override: AttentionBackendEnum | None = None,
579
) -> None:
580
super().__init__()
581
@@ -608,18 +603,12 @@ class SiglipAttention(nn.Module):
603
quant_config=quant_config,
604
prefix=f"{prefix}.out_proj",
605
)
611
-
612
- self.attn_backend = attn_backend
613
- self.attn_backend, self.flash_attn_varlen_func = (
614
- maybe_get_vit_flash_attn_backend(
615
- self.attn_backend,
616
- attn_backend_override=attn_backend_override,
617
- )
606
+ self.attn = MMEncoderAttention(
607
+ num_heads=self.num_attention_heads_per_partition,
608
+ head_size=self.hidden_size_per_attention_head,
609
+ multimodal_config=multimodal_config,
610
+ prefix=f"{prefix}.attn",
611
)
619
- self.is_flash_attn_backend = self.attn_backend in {
620
- AttentionBackendEnum.FLASH_ATTN,
621
- AttentionBackendEnum.ROCM_AITER_FA,
622
- }
612
613
def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
614
seq_len, bs, _ = qkv.shape
@@ -665,44 +654,16 @@ class SiglipAttention(nn.Module):
654
qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
655
q, k = torch.chunk(qk_rotated, 2, dim=0)
656
668
- if self.is_flash_attn_backend:
669
- if max_seqlen is None:
670
- raise ValueError("Flash attention backend requires max_seqlen.")
671
- context_layer = vit_flash_attn_wrapper(
672
- q,
673
- k,
674
- v,
675
- cu_seqlens,
676
- max_seqlen,
677
- batch_size,
678
- self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA,
679
- )
680
- elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
681
- outputs = []
682
- for i in range(1, len(cu_seqlens)):
683
- start_idx = cu_seqlens[i - 1]
684
- end_idx = cu_seqlens[i]
685
- q_i = q[:, start_idx:end_idx]
686
- k_i = k[:, start_idx:end_idx]
687
- v_i = v[:, start_idx:end_idx]
688
- q_i, k_i, v_i = (
689
- rearrange(tensor, "b s h d -> b h s d")
690
- for tensor in (q_i, k_i, v_i)
691
- )
692
- output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
693
- output_i = rearrange(output_i, "b h s d -> b s h d")
694
- outputs.append(output_i)
695
- context_layer = torch.cat(outputs, dim=1)
696
- context_layer = rearrange(
697
- context_layer, "b s h d -> s b (h d)"
698
- ).contiguous()
699
- else:
700
- raise RuntimeError(
701
- f"PaddleOCR-VL does not support {self.attn_backend} backend now."
702
- )
657
+ context_layer = self.attn(
658
+ query=q,
659
+ key=k,
660
+ value=v,
661
+ cu_seqlens=cu_seqlens,
662
+ max_seqlen=max_seqlen,
663
+ )
664
+ context_layer = rearrange(context_layer, "b s h d -> b s (h d)")
665
666
output, _ = self.out_proj(context_layer)
705
- output = rearrange(output, "s b d -> b s d")
667
return output
668
669
@@ -774,10 +735,8 @@ class SiglipEncoderLayer(nn.Module):
735
self,
736
config: PretrainedConfig,
737
quant_config: QuantizationConfig | None = None,
738
+ multimodal_config: MultiModalConfig | None = None,
739
prefix: str = "",
778
- *,
779
- attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
780
- attn_backend_override: AttentionBackendEnum | None = None,
740
):
741
super().__init__()
742
self.embed_dim = config.hidden_size
@@ -787,9 +746,8 @@ class SiglipEncoderLayer(nn.Module):
746
num_heads=config.num_attention_heads,
747
projection_size=config.hidden_size,
748
quant_config=quant_config,
749
+ multimodal_config=multimodal_config,
750
prefix=f"{prefix}.self_attn",
791
- attn_backend=attn_backend,
792
- attn_backend_override=attn_backend_override,
751
)
752
self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
753
self.mlp = SiglipMLP(
@@ -832,14 +790,18 @@ class SiglipEncoder(nn.Module):
790
self,
791
config: PretrainedConfig,
792
quant_config: QuantizationConfig | None = None,
793
+ multimodal_config: MultiModalConfig | None = None,
794
prefix: str = "",
836
- attn_backend_override: AttentionBackendEnum | None = None,
795
):
796
super().__init__()
797
self.config = config
798
embed_dim = config.hidden_size
799
num_heads = config.num_attention_heads
800
head_dim = embed_dim // num_heads
801
+
802
+ attn_backend_override = (
803
+ multimodal_config.mm_encoder_attn_backend if multimodal_config else None
804
+ )
805
self.attn_backend = get_vit_attn_backend(
806
head_size=head_dim,
807
dtype=torch.get_default_dtype(),
@@ -858,9 +820,8 @@ class SiglipEncoder(nn.Module):
820
SiglipEncoderLayer(
821
config,
822
quant_config=quant_config,
823
+ multimodal_config=multimodal_config,
824
prefix=f"{prefix}.layers.{layer_idx}",
862
- attn_backend=self.attn_backend,
863
- attn_backend_override=attn_backend_override,
825
)
826
for layer_idx in range(config.num_hidden_layers)
827
]
@@ -941,8 +902,8 @@ class SiglipVisionTransformer(nn.Module):
902
self,
903
config: PretrainedConfig,
904
quant_config: QuantizationConfig | None = None,
905
+ multimodal_config: MultiModalConfig | None = None,
906
prefix: str = "",
945
- attn_backend_override: AttentionBackendEnum | None = None,
907
):
908
super().__init__()
909
self.config = config
@@ -952,8 +913,8 @@ class SiglipVisionTransformer(nn.Module):
913
self.encoder = SiglipEncoder(
914
config,
915
quant_config=quant_config,
916
+ multimodal_config=multimodal_config,
917
prefix=f"{prefix}.encoder",
956
- attn_backend_override=attn_backend_override,
918
)
919
self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
920
@@ -991,16 +952,16 @@ class SiglipVisionModel(nn.Module):
952
self,
953
config,
954
quant_config: QuantizationConfig | None = None,
955
+ multimodal_config: MultiModalConfig | None = None,
956
prefix: str = "",
995
- attn_backend_override: AttentionBackendEnum | None = None,
957
):
958
super().__init__()
959
960
self.vision_model = SiglipVisionTransformer(
961
config,
962
quant_config=quant_config,
963
+ multimodal_config=multimodal_config,
964
prefix=f"{prefix}.vision_model",
1003
- attn_backend_override=attn_backend_override,
965
)
966
self.quant_config = quant_config
967
@@ -1119,17 +1080,11 @@ class PaddleOCRVLForConditionalGeneration(nn.Module, SupportsMultiModal, Support
1080
self.config = config
1081
self.multimodal_config = multimodal_config
1082
1122
- attn_backend_override = (
1123
- multimodal_config.mm_encoder_attn_backend
1124
- if multimodal_config is not None
1125
- else None
1126
- )
1127
-
1083
self.visual = SiglipVisionModel(
1084
config=config.vision_config,
1085
quant_config=quant_config,
1086
+ multimodal_config=multimodal_config,
1087
prefix=maybe_prefix(prefix, "visual"),
1132
- attn_backend_override=attn_backend_override,
1088
)
1089
self.mlp_AR = Projector(config, config.vision_config)
1090
vllm/model_executor/models/qwen2_5_omni_thinker.py
+1
@@ -845,6 +845,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
845
norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
846
quant_config=quant_config,
847
prefix=maybe_prefix(prefix, "visual"),
848
+ multimodal_config=multimodal_config,
849
)
850
else:
851
self.visual = None
vllm/model_executor/models/qwen2_5_vl.py
+48
-76
@@ -42,13 +42,9 @@ from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
42
)
43
44
from vllm.attention.backends.registry import AttentionBackendEnum
45
-from vllm.attention.layer import maybe_get_vit_flash_attn_backend
46
-from vllm.attention.ops.vit_attn_wrappers import (
47
- vit_flash_attn_wrapper,
48
- vit_torch_sdpa_wrapper,
49
-)
45
+from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
46
from vllm.compilation.decorators import support_torch_compile
51
-from vllm.config import VllmConfig
47
+from vllm.config import MultiModalConfig, VllmConfig
48
from vllm.distributed import parallel_state
49
from vllm.distributed import utils as dist_utils
50
from vllm.forward_context import set_forward_context
@@ -267,10 +263,15 @@ class Qwen2_5_VisionMLP(nn.Module):
263
bias: bool = False,
264
act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,
265
quant_config: QuantizationConfig | None = None,
266
+ multimodal_config: MultiModalConfig | None = None,
267
prefix: str = "",
271
- use_data_parallel: bool = False,
268
):
269
super().__init__()
270
+ use_data_parallel = (
271
+ multimodal_config.mm_encoder_tp_mode == "data"
272
+ if multimodal_config
273
+ else False
274
+ )
275
self.gate_up_proj = MergedColumnParallelLinear(
276
input_size=in_features,
277
output_sizes=[hidden_features] * 2, # [gate_proj, up_proj]
@@ -304,13 +305,16 @@ class Qwen2_5_VisionAttention(nn.Module):
305
num_heads: int,
306
projection_size: int,
307
quant_config: QuantizationConfig | None = None,
308
+ multimodal_config: MultiModalConfig | None = None,
309
prefix: str = "",
308
- use_data_parallel: bool = False,
309
- attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
310
- attn_backend_override: AttentionBackendEnum | None = None,
310
) -> None:
311
super().__init__()
312
# Per attention head and per partition values.
313
+ use_data_parallel = (
314
+ multimodal_config.mm_encoder_tp_mode == "data"
315
+ if multimodal_config
316
+ else False
317
+ )
318
self.tp_size = (
319
1
320
if use_data_parallel
@@ -342,18 +346,12 @@ class Qwen2_5_VisionAttention(nn.Module):
346
prefix=f"{prefix}.proj",
347
disable_tp=use_data_parallel,
348
)
345
- self.attn_backend = attn_backend
346
- self.attn_backend, self.flash_attn_varlen_func = (
347
- maybe_get_vit_flash_attn_backend(
348
- self.attn_backend,
349
- attn_backend_override=attn_backend_override,
350
- )
351
- )
349
353
- self.is_flash_attn_backend = self.attn_backend in {
354
- AttentionBackendEnum.FLASH_ATTN,
355
- AttentionBackendEnum.ROCM_AITER_FA,
356
- }
350
+ self.attn = MMEncoderAttention(
351
+ num_heads=self.num_attention_heads_per_partition,
352
+ head_size=self.hidden_size_per_attention_head,
353
+ multimodal_config=multimodal_config,
354
+ )
355
356
def forward(
357
self,
@@ -394,32 +392,17 @@ class Qwen2_5_VisionAttention(nn.Module):
392
else:
393
q, k, v = qkv.unbind(dim=2)
394
397
- if self.is_flash_attn_backend:
398
- context_layer = vit_flash_attn_wrapper(
399
- q,
400
- k,
401
- v,
402
- cu_seqlens,
403
- max_seqlen,
404
- batch_size,
405
- self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA,
406
- )
407
- elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
408
- # Execute attention entry by entry for speed & less VRAM.
409
- from vllm.platforms import current_platform
410
-
411
- # Never remove the next contiguous logic
412
- # Without it, hallucinations occur with the backend
413
- if current_platform.is_rocm():
414
- q = q.contiguous()
415
- k = k.contiguous()
416
- v = v.contiguous()
417
- context_layer = vit_torch_sdpa_wrapper(
418
- q,
419
- k,
420
- v,
421
- cu_seqlens,
422
- )
395
+ context_layer = self.attn(
396
+ query=q,
397
+ key=k,
398
+ value=v,
399
+ cu_seqlens=cu_seqlens,
400
+ max_seqlen=max_seqlen,
401
+ )
402
+
403
+ context_layer = einops.rearrange(
404
+ context_layer, "b s h d -> s b (h d)", b=batch_size
405
+ ).contiguous()
406
407
output, _ = self.proj(context_layer)
408
return output
@@ -443,10 +426,8 @@ class Qwen2_5_VisionBlock(nn.Module):
426
act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,
427
norm_layer: Callable[[int], nn.Module] | None = None,
428
quant_config: QuantizationConfig | None = None,
429
+ multimodal_config: MultiModalConfig | None = None,
430
prefix: str = "",
447
- use_data_parallel: bool = False,
448
- attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
449
- attn_backend_override: AttentionBackendEnum | None = None,
431
) -> None:
432
super().__init__()
433
if norm_layer is None:
@@ -458,10 +439,8 @@ class Qwen2_5_VisionBlock(nn.Module):
439
num_heads=num_heads,
440
projection_size=dim,
441
quant_config=quant_config,
442
+ multimodal_config=multimodal_config,
443
prefix=f"{prefix}.attn",
462
- use_data_parallel=use_data_parallel,
463
- attn_backend=attn_backend,
464
- attn_backend_override=attn_backend_override,
444
)
445
self.mlp = Qwen2_5_VisionMLP(
446
dim,
@@ -469,8 +448,8 @@ class Qwen2_5_VisionBlock(nn.Module):
448
act_fn=act_fn,
449
bias=True,
450
quant_config=quant_config,
451
+ multimodal_config=multimodal_config,
452
prefix=f"{prefix}.mlp",
473
- use_data_parallel=use_data_parallel,
453
)
454
455
def forward(
@@ -542,10 +521,15 @@ class Qwen2_5_VisionPatchMerger(nn.Module):
521
norm_layer: Callable[[int], nn.Module] | None = None,
522
spatial_merge_size: int = 2,
523
quant_config: QuantizationConfig | None = None,
524
+ multimodal_config: MultiModalConfig | None = None,
525
prefix: str = "",
546
- use_data_parallel: bool = False,
526
) -> None:
527
super().__init__()
528
+ use_data_parallel = (
529
+ multimodal_config.mm_encoder_tp_mode == "data"
530
+ if multimodal_config
531
+ else False
532
+ )
533
self.hidden_size = context_dim * (spatial_merge_size**2)
534
if norm_layer is None:
535
norm_layer = partial(nn.LayerNorm, eps=1e-6)
@@ -586,9 +570,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
570
vision_config: Qwen2_5_VLVisionConfig,
571
norm_eps: float = 1e-6,
572
quant_config: QuantizationConfig | None = None,
573
+ multimodal_config: MultiModalConfig | None = None,
574
prefix: str = "",
590
- use_data_parallel: bool = False,
591
- attn_backend_override: AttentionBackendEnum | None = None,
575
) -> None:
576
super().__init__()
577
@@ -598,7 +581,6 @@ class Qwen2_5_VisionTransformer(nn.Module):
581
depth = vision_config.depth
582
self.hidden_size = vision_config.hidden_size
583
self.num_heads = vision_config.num_heads
601
- self.use_data_parallel = use_data_parallel
584
self.out_hidden_size = vision_config.out_hidden_size
585
586
# args for get_window_index_thw
@@ -629,19 +611,17 @@ class Qwen2_5_VisionTransformer(nn.Module):
611
rope_parameters={"partial_rotary_factor": 0.5},
612
)
613
614
+ attn_backend_override = (
615
+ multimodal_config.mm_encoder_attn_backend
616
+ if multimodal_config is not None
617
+ else None
618
+ )
619
self.attn_backend = get_vit_attn_backend(
620
head_size=head_dim,
621
dtype=torch.get_default_dtype(),
622
attn_backend_override=attn_backend_override,
623
)
624
638
- self.attn_backend, self.flash_attn_varlen_func = (
639
- maybe_get_vit_flash_attn_backend(
640
- self.attn_backend,
641
- attn_backend_override=attn_backend_override,
642
- )
643
- )
644
-
625
if self.attn_backend not in {
626
AttentionBackendEnum.FLASH_ATTN,
627
AttentionBackendEnum.TORCH_SDPA,
@@ -661,10 +641,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
641
act_fn=get_act_and_mul_fn(vision_config.hidden_act),
642
norm_layer=norm_layer,
643
quant_config=quant_config,
644
+ multimodal_config=multimodal_config,
645
prefix=f"{prefix}.blocks.{layer_idx}",
665
- use_data_parallel=use_data_parallel,
666
- attn_backend=self.attn_backend,
667
- attn_backend_override=attn_backend_override,
646
)
647
for layer_idx in range(depth)
648
]
@@ -677,8 +655,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
655
norm_layer=norm_layer,
656
spatial_merge_size=self.spatial_merge_size,
657
quant_config=quant_config,
658
+ multimodal_config=multimodal_config,
659
prefix=f"{prefix}.merger",
681
- use_data_parallel=use_data_parallel,
660
)
661
662
@property
@@ -1200,18 +1178,12 @@ class Qwen2_5_VLForConditionalGeneration(
1178
if multimodal_config.get_limit_per_prompt(
1179
"image"
1180
) or multimodal_config.get_limit_per_prompt("video"):
1203
- attn_backend_override = (
1204
- multimodal_config.mm_encoder_attn_backend
1205
- if multimodal_config is not None
1206
- else None
1207
- )
1181
self.visual = Qwen2_5_VisionTransformer(
1182
vision_config=config.vision_config,
1183
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
1184
quant_config=self.quant_config,
1185
prefix=maybe_prefix(prefix, "visual"),
1213
- use_data_parallel=self.use_data_parallel,
1214
- attn_backend_override=attn_backend_override,
1186
+ multimodal_config=multimodal_config,
1187
)
1188
else:
1189
self.visual = None
vllm/model_executor/models/qwen2_vl.py
+47
-96
@@ -33,7 +33,6 @@ from typing import Annotated, Any, Literal, TypeAlias
33
import numpy as np
34
import torch
35
import torch.nn as nn
36
-import torch.nn.functional as F
36
from einops import rearrange
37
from transformers import BatchFeature
38
from transformers.models.qwen2_vl import Qwen2VLImageProcessor, Qwen2VLProcessor
@@ -45,10 +44,8 @@ from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
44
from transformers.models.qwen2_vl.video_processing_qwen2_vl import Qwen2VLVideoProcessor
45
46
from vllm.attention.backends.registry import AttentionBackendEnum
48
-from vllm.attention.layer import (
49
- maybe_get_vit_flash_attn_backend,
50
-)
51
-from vllm.config import VllmConfig
47
+from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
48
+from vllm.config import MultiModalConfig, VllmConfig
49
from vllm.config.multimodal import BaseDummyOptions
50
from vllm.distributed import parallel_state, tensor_model_parallel_all_gather
51
from vllm.distributed import utils as dist_utils
@@ -251,10 +248,15 @@ class Qwen2VisionMLP(nn.Module):
248
hidden_features: int,
249
act_layer: type[nn.Module] = QuickGELU,
250
quant_config: QuantizationConfig | None = None,
251
+ multimodal_config: MultiModalConfig | None = None,
252
prefix: str = "",
255
- use_data_parallel: bool = False,
253
):
254
super().__init__()
255
+ use_data_parallel = (
256
+ multimodal_config.mm_encoder_tp_mode == "data"
257
+ if multimodal_config
258
+ else False
259
+ )
260
self.fc1 = ColumnParallelLinear(
261
in_features,
262
hidden_features,
@@ -295,12 +297,16 @@ class Qwen2VisionAttention(nn.Module):
297
num_heads: int,
298
projection_size: int,
299
quant_config: QuantizationConfig | None = None,
300
+ multimodal_config: MultiModalConfig | None = None,
301
prefix: str = "",
299
- use_data_parallel: bool = False,
300
- attn_backend_override: AttentionBackendEnum | None = None,
302
) -> None:
303
super().__init__()
304
# Per attention head and per partition values.
305
+ use_data_parallel = (
306
+ multimodal_config.mm_encoder_tp_mode == "data"
307
+ if multimodal_config
308
+ else False
309
+ )
310
self.tp_size = (
311
1
312
if use_data_parallel
@@ -329,34 +335,12 @@ class Qwen2VisionAttention(nn.Module):
335
disable_tp=use_data_parallel,
336
)
337
332
- # Detect attention implementation.
333
- self.attn_backend = get_vit_attn_backend(
338
+ self.attn = MMEncoderAttention(
339
+ num_heads=self.num_attention_heads_per_partition,
340
head_size=self.hidden_size_per_attention_head,
335
- dtype=torch.get_default_dtype(),
336
- attn_backend_override=attn_backend_override,
337
- )
338
-
339
- self.attn_backend, self.flash_attn_varlen_func = (
340
- maybe_get_vit_flash_attn_backend(
341
- self.attn_backend,
342
- attn_backend_override=attn_backend_override,
343
- )
341
+ multimodal_config=multimodal_config,
342
)
343
346
- if self.attn_backend not in {
347
- AttentionBackendEnum.FLASH_ATTN,
348
- AttentionBackendEnum.TORCH_SDPA,
349
- AttentionBackendEnum.ROCM_AITER_FA,
350
- }:
351
- raise RuntimeError(
352
- f"Qwen2-VL does not support {self.attn_backend} backend now."
353
- )
354
-
355
- self.is_flash_attn_backend = self.attn_backend in {
356
- AttentionBackendEnum.FLASH_ATTN,
357
- AttentionBackendEnum.ROCM_AITER_FA,
358
- }
359
-
344
def split_qkv(self, qkv: torch.Tensor) -> tuple[torch.Tensor, ...]:
345
# [s, b, 3 * head * head_dim]
346
seq_len, bs, _ = qkv.shape
@@ -398,7 +382,6 @@ class Qwen2VisionAttention(nn.Module):
382
383
# [s, b, 3 * head * head_dim] -> 3 * [s, b, head, head_dim]
384
q, k, v = self.split_qkv(x)
401
- batch_size = q.shape[1]
385
386
q, k, v = (rearrange(x, "s b ... -> b s ...") for x in (q, k, v))
387
@@ -409,49 +392,15 @@ class Qwen2VisionAttention(nn.Module):
392
)
393
q, k = torch.chunk(qk_rotated, 2, dim=0)
394
412
- if self.is_flash_attn_backend:
413
- q, k, v = (rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])
414
-
415
- output = self.flash_attn_varlen_func(
416
- q,
417
- k,
418
- v,
419
- cu_seqlens_q=cu_seqlens,
420
- cu_seqlens_k=cu_seqlens,
421
- max_seqlen_q=max_seqlen,
422
- max_seqlen_k=max_seqlen,
423
- dropout_p=0.0,
424
- causal=False,
425
- )
395
+ context_layer = self.attn(
396
+ query=q,
397
+ key=k,
398
+ value=v,
399
+ cu_seqlens=cu_seqlens,
400
+ max_seqlen=max_seqlen,
401
+ )
402
427
- context_layer = rearrange(
428
- output, "(b s) h d -> s b (h d)", b=batch_size
429
- ).contiguous()
430
- elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
431
- # Execute attention entry by entry for speed & less VRAM.
432
- from vllm.platforms import current_platform
433
-
434
- if current_platform.is_rocm():
435
- q = q.contiguous()
436
- k = k.contiguous()
437
- v = v.contiguous()
438
- outputs = []
439
-
440
- lens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
441
- q_chunks = torch.split(q, lens, dim=1)
442
- k_chunks = torch.split(k, lens, dim=1)
443
- v_chunks = torch.split(v, lens, dim=1)
444
- for q_i, k_i, v_i in zip(q_chunks, k_chunks, v_chunks):
445
- q_i, k_i, v_i = (
446
- rearrange(x, "b s h d -> b h s d") for x in [q_i, k_i, v_i]
447
- )
448
- output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
449
- output_i = rearrange(output_i, "b h s d -> b s h d ")
450
- outputs.append(output_i)
451
- context_layer = torch.cat(outputs, dim=1)
452
- context_layer = rearrange(
453
- context_layer, "b s h d -> s b (h d)"
454
- ).contiguous()
403
+ context_layer = rearrange(context_layer, "b s h d -> s b (h d)").contiguous()
404
405
output, _ = self.proj(context_layer)
406
return output
@@ -466,9 +415,8 @@ class Qwen2VisionBlock(nn.Module):
415
act_layer: type[nn.Module] = QuickGELU,
416
norm_layer: Callable[[int], nn.Module] | None = None,
417
quant_config: QuantizationConfig | None = None,
418
+ multimodal_config: MultiModalConfig | None = None,
419
prefix: str = "",
470
- use_data_parallel: bool = False,
471
- attn_backend_override: AttentionBackendEnum | None = None,
420
) -> None:
421
super().__init__()
422
if norm_layer is None:
@@ -482,17 +430,16 @@ class Qwen2VisionBlock(nn.Module):
430
num_heads=num_heads,
431
projection_size=dim,
432
quant_config=quant_config,
433
+ multimodal_config=multimodal_config,
434
prefix=f"{prefix}.attn",
486
- use_data_parallel=use_data_parallel,
487
- attn_backend_override=attn_backend_override,
435
)
436
self.mlp = Qwen2VisionMLP(
437
dim,
438
mlp_hidden_dim,
439
act_layer=act_layer,
440
quant_config=quant_config,
441
+ multimodal_config=multimodal_config,
442
prefix=f"{prefix}.mlp",
495
- use_data_parallel=use_data_parallel,
443
)
444
445
def forward(
@@ -552,10 +499,15 @@ class Qwen2VisionPatchMerger(nn.Module):
499
norm_layer: Callable[[int], nn.Module] | None = None,
500
spatial_merge_size: int = 2,
501
quant_config: QuantizationConfig | None = None,
502
+ multimodal_config: MultiModalConfig | None = None,
503
prefix: str = "",
556
- use_data_parallel: bool = False,
504
) -> None:
505
super().__init__()
506
+ use_data_parallel = (
507
+ multimodal_config.mm_encoder_tp_mode == "data"
508
+ if multimodal_config
509
+ else False
510
+ )
511
self.hidden_size = context_dim * (spatial_merge_size**2)
512
if norm_layer is None:
513
norm_layer = partial(nn.LayerNorm, eps=1e-6)
@@ -599,9 +551,8 @@ class Qwen2VisionTransformer(nn.Module):
551
vision_config: Qwen2VLVisionConfig,
552
norm_eps: float = 1e-6,
553
quant_config: QuantizationConfig | None = None,
554
+ multimodal_config: MultiModalConfig | None = None,
555
prefix: str = "",
603
- use_data_parallel: bool = False,
604
- attn_backend_override: AttentionBackendEnum | None = None,
556
) -> None:
557
super().__init__()
558
@@ -615,7 +566,11 @@ class Qwen2VisionTransformer(nn.Module):
566
num_heads = vision_config.num_heads
567
mlp_ratio = vision_config.mlp_ratio
568
618
- self.use_data_parallel = use_data_parallel
569
+ self.use_data_parallel = (
570
+ multimodal_config.mm_encoder_tp_mode == "data"
571
+ if multimodal_config
572
+ else False
573
+ )
574
self.out_hidden_size = vision_config.hidden_size
575
576
self.spatial_merge_size = spatial_merge_size
@@ -647,8 +602,7 @@ class Qwen2VisionTransformer(nn.Module):
602
norm_layer=norm_layer,
603
quant_config=quant_config,
604
prefix=f"{prefix}.blocks.{layer_idx}",
650
- use_data_parallel=use_data_parallel,
651
- attn_backend_override=attn_backend_override,
605
+ multimodal_config=multimodal_config,
606
)
607
for layer_idx in range(depth)
608
]
@@ -659,7 +613,10 @@ class Qwen2VisionTransformer(nn.Module):
613
norm_layer=norm_layer,
614
quant_config=quant_config,
615
prefix=f"{prefix}.merger",
662
- use_data_parallel=use_data_parallel,
616
+ multimodal_config=multimodal_config,
617
+ )
618
+ attn_backend_override = (
619
+ multimodal_config.mm_encoder_attn_backend if multimodal_config else None
620
)
621
self.attn_backend = get_vit_attn_backend(
622
head_size=head_dim,
@@ -720,7 +677,7 @@ class Qwen2VisionTransformer(nn.Module):
677
AttentionBackendEnum.FLASH_ATTN,
678
AttentionBackendEnum.ROCM_AITER_FA,
679
}:
723
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
680
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
681
return max_seqlen
682
683
def forward(
@@ -1324,18 +1281,12 @@ class Qwen2VLForConditionalGeneration(
1281
if multimodal_config.get_limit_per_prompt(
1282
"image"
1283
) or multimodal_config.get_limit_per_prompt("video"):
1327
- attn_backend_override = (
1328
- multimodal_config.mm_encoder_attn_backend
1329
- if multimodal_config is not None
1330
- else None
1331
- )
1284
self.visual = Qwen2VisionTransformer(
1285
config.vision_config,
1286
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
1287
quant_config=quant_config,
1288
+ multimodal_config=multimodal_config,
1289
prefix=maybe_prefix(prefix, "visual"),
1337
- use_data_parallel=self.use_data_parallel,
1338
- attn_backend_override=attn_backend_override,
1290
)
1291
else:
1292
self.visual = None
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+12
-8
@@ -48,7 +48,7 @@ from transformers.models.whisper import WhisperFeatureExtractor
48
49
from vllm.attention.backends.registry import AttentionBackendEnum
50
from vllm.compilation.decorators import support_torch_compile
51
-from vllm.config import VllmConfig
51
+from vllm.config import MultiModalConfig, VllmConfig
52
from vllm.distributed import get_pp_group
53
from vllm.logger import init_logger
54
from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY
@@ -192,6 +192,7 @@ class Qwen3_VisionBlock(nn.Module):
192
mlp_hidden_dim: int,
193
act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,
194
norm_layer: Callable[[int], nn.Module] | None = None,
195
+ multimodal_config: MultiModalConfig | None = None,
196
quant_config: QuantizationConfig | None = None,
197
prefix: str = "",
198
) -> None:
@@ -205,6 +206,7 @@ class Qwen3_VisionBlock(nn.Module):
206
num_heads=num_heads,
207
projection_size=dim,
208
quant_config=quant_config,
209
+ multimodal_config=multimodal_config,
210
prefix=f"{prefix}.attn",
211
)
212
self.mlp = Qwen3_VisionMLP(
@@ -299,8 +301,8 @@ class Qwen3Omni_VisionTransformer(nn.Module):
301
vision_config,
302
norm_eps: float = 1e-6,
303
quant_config: QuantizationConfig | None = None,
304
+ multimodal_config: MultiModalConfig | None = None,
305
prefix: str = "",
303
- attn_backend_override: AttentionBackendEnum | None = None,
306
) -> None:
307
super().__init__()
308
self.hidden_size = vision_config.hidden_size
@@ -347,6 +349,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):
349
act_fn=_ACTIVATION_REGISTRY[vision_config.hidden_act],
350
norm_layer=norm_layer,
351
quant_config=quant_config,
352
+ multimodal_config=multimodal_config,
353
prefix=f"{prefix}.blocks.{layer_idx}",
354
)
355
for layer_idx in range(vision_config.depth)
@@ -376,6 +379,12 @@ class Qwen3Omni_VisionTransformer(nn.Module):
379
]
380
)
381
382
+ attn_backend_override = (
383
+ multimodal_config.mm_encoder_attn_backend
384
+ if multimodal_config is not None
385
+ else None
386
+ )
387
+
388
self.attn_backend = get_vit_attn_backend(
389
head_size=head_dim,
390
dtype=torch.get_default_dtype(),
@@ -1188,17 +1197,12 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1197
1198
self.audio_tower = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config)
1199
1191
- attn_backend_override = (
1192
- multimodal_config.mm_encoder_attn_backend
1193
- if multimodal_config is not None
1194
- else None
1195
- )
1200
self.visual = Qwen3Omni_VisionTransformer(
1201
vision_config=thinker_config.vision_config,
1202
norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),
1203
quant_config=quant_config,
1204
prefix=maybe_prefix(prefix, "visual"),
1201
- attn_backend_override=attn_backend_override,
1205
+ multimodal_config=multimodal_config,
1206
)
1207
self.quant_config = quant_config
1208
vllm/model_executor/models/qwen3_vl.py
+24
-22
@@ -50,7 +50,7 @@ from transformers.video_utils import VideoMetadata
50
51
from vllm.attention.backends.registry import AttentionBackendEnum
52
from vllm.compilation.decorators import support_torch_compile
53
-from vllm.config import VllmConfig
53
+from vllm.config import MultiModalConfig, VllmConfig
54
from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
55
from vllm.distributed import get_pp_group
56
from vllm.logger import init_logger
@@ -169,10 +169,15 @@ class Qwen3_VisionMLP(nn.Module):
169
bias: bool = False,
170
act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,
171
quant_config: QuantizationConfig | None = None,
172
+ multimodal_config: MultiModalConfig | None = None,
173
prefix: str = "",
173
- use_data_parallel: bool = False,
174
):
175
super().__init__()
176
+ use_data_parallel = (
177
+ multimodal_config.mm_encoder_tp_mode == "data"
178
+ if multimodal_config
179
+ else False
180
+ )
181
self.linear_fc1 = ColumnParallelLinear(
182
in_features,
183
hidden_features,
@@ -206,10 +211,9 @@ class Qwen3_VisionBlock(nn.Module):
211
mlp_hidden_dim: int,
212
act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,
213
norm_layer: Callable[[int], nn.Module] | None = None,
214
+ multimodal_config: MultiModalConfig | None = None,
215
quant_config: QuantizationConfig | None = None,
216
prefix: str = "",
211
- use_data_parallel: bool = False,
212
- attn_backend: AttentionBackendEnum = AttentionBackendEnum.TORCH_SDPA,
217
) -> None:
218
super().__init__()
219
if norm_layer is None:
@@ -221,9 +225,8 @@ class Qwen3_VisionBlock(nn.Module):
225
num_heads=num_heads,
226
projection_size=dim,
227
quant_config=quant_config,
228
+ multimodal_config=multimodal_config,
229
prefix=f"{prefix}.attn",
225
- use_data_parallel=use_data_parallel,
226
- attn_backend=attn_backend,
230
)
231
self.mlp = Qwen3_VisionMLP(
232
dim,
@@ -231,8 +234,8 @@ class Qwen3_VisionBlock(nn.Module):
234
act_fn=act_fn,
235
bias=True,
236
quant_config=quant_config,
237
+ multimodal_config=multimodal_config,
238
prefix=f"{prefix}.mlp",
235
- use_data_parallel=use_data_parallel,
239
)
240
241
def forward(
@@ -264,10 +267,15 @@ class Qwen3_VisionPatchMerger(nn.Module):
267
spatial_merge_size: int = 2,
268
use_postshuffle_norm: bool = False,
269
quant_config: QuantizationConfig | None = None,
270
+ multimodal_config: MultiModalConfig | None = None,
271
prefix: str = "",
268
- use_data_parallel: bool = False,
272
) -> None:
273
super().__init__()
274
+ use_data_parallel = (
275
+ multimodal_config.mm_encoder_tp_mode == "data"
276
+ if multimodal_config
277
+ else False
278
+ )
279
self.hidden_size = context_dim * (spatial_merge_size**2)
280
281
self.use_postshuffle_norm = use_postshuffle_norm
@@ -313,9 +321,8 @@ class Qwen3_VisionTransformer(nn.Module):
321
vision_config: Qwen3VLVisionConfig,
322
norm_eps: float = 1e-6,
323
quant_config: QuantizationConfig | None = None,
324
+ multimodal_config: MultiModalConfig | None = None,
325
prefix: str = "",
317
- use_data_parallel: bool = False,
318
- attn_backend_override: AttentionBackendEnum | None = None,
326
) -> None:
327
super().__init__()
328
self.hidden_size = vision_config.hidden_size
@@ -326,7 +333,6 @@ class Qwen3_VisionTransformer(nn.Module):
333
self.spatial_merge_unit = self.spatial_merge_size**2
334
self.temporal_patch_size = vision_config.temporal_patch_size
335
self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes
329
- self.use_data_parallel = use_data_parallel
336
self.num_grid_per_side = int(self.num_position_embeddings**0.5)
337
338
# NOTE: This is used for creating empty tensor for all_gather for
@@ -359,8 +365,8 @@ class Qwen3_VisionTransformer(nn.Module):
365
norm_layer=norm_layer,
366
spatial_merge_size=self.spatial_merge_size,
367
quant_config=quant_config,
368
+ multimodal_config=multimodal_config,
369
prefix=f"{prefix}.merger",
363
- use_data_parallel=use_data_parallel,
370
)
371
372
self.deepstack_merger_list = nn.ModuleList(
@@ -372,13 +378,16 @@ class Qwen3_VisionTransformer(nn.Module):
378
use_postshuffle_norm=True,
379
norm_layer=norm_layer,
380
quant_config=quant_config,
381
+ multimodal_config=multimodal_config,
382
prefix=f"{prefix}.deepstack_merger_list.{layer_idx}",
376
- use_data_parallel=use_data_parallel,
383
)
384
for layer_idx in range(len(self.deepstack_visual_indexes))
385
]
386
)
387
388
+ attn_backend_override = (
389
+ multimodal_config.mm_encoder_attn_backend if multimodal_config else None
390
+ )
391
self.attn_backend = get_vit_attn_backend(
392
head_size=head_dim,
393
dtype=torch.get_default_dtype(),
@@ -402,9 +411,8 @@ class Qwen3_VisionTransformer(nn.Module):
411
act_fn=_ACTIVATION_REGISTRY[vision_config.hidden_act],
412
norm_layer=norm_layer,
413
quant_config=quant_config,
414
+ multimodal_config=multimodal_config,
415
prefix=f"{prefix}.blocks.{layer_idx}",
406
- use_data_parallel=use_data_parallel,
407
- attn_backend=self.attn_backend,
416
)
417
for layer_idx in range(vision_config.depth)
418
]
@@ -1277,18 +1285,12 @@ class Qwen3VLForConditionalGeneration(
1285
) and not multimodal_config.get_limit_per_prompt("video"):
1286
self.visual = None
1287
else:
1280
- attn_backend_override = (
1281
- multimodal_config.mm_encoder_attn_backend
1282
- if multimodal_config is not None
1283
- else None
1284
- )
1288
self.visual = Qwen3_VisionTransformer(
1289
config.vision_config,
1290
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
1291
quant_config=quant_config,
1292
+ multimodal_config=multimodal_config,
1293
prefix=maybe_prefix(prefix, "visual"),
1290
- use_data_parallel=self.use_data_parallel,
1291
- attn_backend_override=attn_backend_override,
1294
)
1295
1296
self.language_model = Qwen3LLMForCausalLM(
vllm/model_executor/models/qwen3_vl_moe.py
+1
-2
@@ -418,7 +418,6 @@ class Qwen3VLMoeForConditionalGeneration(
418
419
self.config = config
420
self.multimodal_config = multimodal_config
421
- self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
421
422
if not multimodal_config.get_limit_per_prompt(
423
"image"
@@ -429,8 +428,8 @@ class Qwen3VLMoeForConditionalGeneration(
428
config.vision_config,
429
norm_eps=getattr(config, "rms_norm_eps", 1e-6),
430
quant_config=quant_config,
431
+ multimodal_config=multimodal_config,
432
prefix=maybe_prefix(prefix, "visual"),
433
- use_data_parallel=self.use_data_parallel,
433
)
434
435
self.language_model = Qwen3MoeLLMForCausalLM(
vllm/model_executor/models/siglip2navit.py
+45
-84
@@ -13,7 +13,8 @@ from transformers import Siglip2VisionConfig
13
from transformers.configuration_utils import PretrainedConfig
14
15
from vllm.attention.backends.registry import AttentionBackendEnum
16
-from vllm.attention.layer import maybe_get_vit_flash_attn_backend
16
+from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
17
+from vllm.config import MultiModalConfig
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.conv import Conv2dLayer
@@ -28,8 +29,6 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
29
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
30
from vllm.platforms import current_platform
31
31
-from .vision import get_vit_attn_backend
32
-
32
33
class VisionRotaryEmbedding(nn.Module):
34
def __init__(self, dim: int, theta: float = 10000.0) -> None:
@@ -190,7 +189,7 @@ def apply_rotary_pos_emb(
189
) -> tuple[torch.Tensor, torch.Tensor]:
190
cos = cos.chunk(2, dim=-1)[0].contiguous()
191
sin = sin.chunk(2, dim=-1)[0].contiguous()
193
- if is_flash_attn_backend and not current_platform.is_xpu():
192
+ if is_flash_attn_backend and current_platform.is_cuda():
193
from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
194
195
apply_rotary_emb_func = apply_rotary_emb
@@ -208,6 +207,7 @@ class Siglip2Attention(nn.Module):
207
self,
208
config: Siglip2VisionConfig,
209
quant_config: QuantizationConfig | None = None,
210
+ multimodal_config: MultiModalConfig | None = None,
211
prefix: str = "",
212
use_data_parallel: bool = False,
213
attn_backend_override: AttentionBackendEnum | None = None,
@@ -227,20 +227,25 @@ class Siglip2Attention(nn.Module):
227
self.dropout = config.attention_dropout
228
self.is_causal = False
229
230
- # TODO(Isotr0py): Enable data parallel after we support
231
- # disabling TP on parallel linear layer
230
+ use_data_parallel = (
231
+ multimodal_config.mm_encoder_tp_mode == "data"
232
+ if multimodal_config
233
+ else False
234
+ )
235
self.qkv_proj = QKVParallelLinear(
236
hidden_size=self.embed_dim,
237
head_size=self.head_dim,
238
total_num_heads=self.num_heads,
239
quant_config=quant_config,
240
prefix=f"{prefix}.qkv_proj",
241
+ disable_tp=use_data_parallel,
242
)
243
self.out_proj = RowParallelLinear(
244
input_size=self.embed_dim,
245
output_size=self.embed_dim,
246
quant_config=quant_config,
247
prefix=f"{prefix}.out_proj",
248
+ disable_tp=use_data_parallel,
249
)
250
251
self.tp_size = (
@@ -249,31 +254,13 @@ class Siglip2Attention(nn.Module):
254
self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
255
self.use_rope = config.use_rope
256
252
- # Detect attention implementation.
253
- self.attn_backend = get_vit_attn_backend(
257
+ self.attn = MMEncoderAttention(
258
+ num_heads=self.num_heads_per_partition,
259
head_size=self.head_dim,
255
- dtype=torch.get_default_dtype(),
256
- attn_backend_override=attn_backend_override,
257
- )
258
-
259
- self.attn_backend, self.flash_attn_varlen_func = (
260
- maybe_get_vit_flash_attn_backend(
261
- self.attn_backend,
262
- attn_backend_override=attn_backend_override,
263
- )
260
+ prefix=f"{prefix}.attn",
261
+ multimodal_config=multimodal_config,
262
)
263
266
- if self.attn_backend not in {
267
- AttentionBackendEnum.FLASH_ATTN,
268
- AttentionBackendEnum.TORCH_SDPA,
269
- AttentionBackendEnum.ROCM_AITER_FA,
270
- }:
271
- self.attn_backend = AttentionBackendEnum.TORCH_SDPA
272
- self.is_flash_attn_backend = self.attn_backend in {
273
- AttentionBackendEnum.FLASH_ATTN,
274
- AttentionBackendEnum.ROCM_AITER_FA,
275
- }
276
-
264
def forward(
265
self,
266
hidden_states: torch.Tensor,
@@ -298,46 +285,23 @@ class Siglip2Attention(nn.Module):
285
keys.unsqueeze(0),
286
cos,
287
sin,
301
- self.is_flash_attn_backend,
288
+ self.attn.is_flash_attn_backend,
289
)
290
queries = queries.squeeze(0)
291
keys = keys.squeeze(0)
292
306
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
307
- if self.is_flash_attn_backend:
308
- attn_output = self.flash_attn_varlen_func(
309
- queries,
310
- keys,
311
- values,
312
- cu_seqlens_q=cu_seqlens,
313
- cu_seqlens_k=cu_seqlens,
314
- max_seqlen_q=max_seqlen,
315
- max_seqlen_k=max_seqlen,
316
- ).reshape(seq_length, -1)
317
- elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
318
- # Execute attention entry by entry for speed & less VRAM.
319
- batch_size = cu_seqlens.shape[0] - 1
320
- outputs = []
321
- cu = cu_seqlens.tolist()
322
- for i in range(batch_size):
323
- start_idx = cu[i]
324
- end_idx = cu[i + 1]
325
-
326
- # Each sequence is processed independently.
327
- q_i = queries[start_idx:end_idx].unsqueeze(0)
328
- k_i = keys[start_idx:end_idx].unsqueeze(0)
329
- v_i = values[start_idx:end_idx].unsqueeze(0)
330
-
331
- # (1, seq_len, num_heads, head_dim) ->
332
- # (1, num_heads, seq_len, head_dim)
333
- q_i, k_i, v_i = [x.transpose(1, 2) for x in (q_i, k_i, v_i)]
334
-
335
- output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
336
- # (1, num_heads, seq_len, head_dim) -> (seq_len, embed_dim)
337
- output_i = output_i.transpose(1, 2).reshape(end_idx - start_idx, -1)
338
- outputs.append(output_i)
339
-
340
- attn_output = torch.cat(outputs, dim=0)
293
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
294
+ attn_output = self.attn(
295
+ query=queries.unsqueeze(0),
296
+ key=keys.unsqueeze(0),
297
+ value=values.unsqueeze(0),
298
+ cu_seqlens=cu_seqlens,
299
+ max_seqlen=max_seqlen,
300
+ )
301
+ attn_output = attn_output.reshape(
302
+ seq_length, self.num_heads_per_partition * self.head_dim
303
+ )
304
+
305
attn_output, _ = self.out_proj(attn_output)
306
return attn_output
307
@@ -347,25 +311,30 @@ class Siglip2MLP(nn.Module):
311
self,
312
config: Siglip2VisionConfig,
313
quant_config: QuantizationConfig | None = None,
314
+ multimodal_config: MultiModalConfig | None = None,
315
prefix: str = "",
351
- use_data_parallel: bool = False,
316
):
317
super().__init__()
318
self.config = config
319
+ use_data_parallel = (
320
+ multimodal_config.mm_encoder_tp_mode == "data"
321
+ if multimodal_config
322
+ else False
323
+ )
324
self.activation_fn = get_act_fn(config.hidden_act)
356
- # TODO(Isotr0py): Enable data parallel after we support
357
- # disabling TP on parallel linear layer
325
self.fc1 = ColumnParallelLinear(
326
config.hidden_size,
327
config.intermediate_size,
328
quant_config=quant_config,
329
prefix=f"{prefix}.fc1",
330
+ disable_tp=use_data_parallel,
331
)
332
self.fc2 = RowParallelLinear(
333
config.intermediate_size,
334
config.hidden_size,
335
quant_config=quant_config,
336
prefix=f"{prefix}.fc2",
337
+ disable_tp=use_data_parallel,
338
)
339
340
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -380,9 +349,8 @@ class Siglip2EncoderLayer(nn.Module):
349
self,
350
config: Siglip2VisionConfig,
351
quant_config: QuantizationConfig | None = None,
352
+ multimodal_config: MultiModalConfig | None = None,
353
prefix: str = "",
384
- use_data_parallel: bool = False,
385
- attn_backend_override: AttentionBackendEnum | None = None,
354
):
355
super().__init__()
356
self.embed_dim = config.hidden_size
@@ -390,16 +358,15 @@ class Siglip2EncoderLayer(nn.Module):
358
self.self_attn = Siglip2Attention(
359
config,
360
quant_config=quant_config,
361
+ multimodal_config=multimodal_config,
362
prefix=f"{prefix}.self_attn",
394
- use_data_parallel=use_data_parallel,
395
- attn_backend_override=attn_backend_override,
363
)
364
self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
365
self.mlp = Siglip2MLP(
366
config,
367
quant_config=quant_config,
368
+ multimodal_config=multimodal_config,
369
prefix=f"{prefix}.mlp",
402
- use_data_parallel=use_data_parallel,
370
)
371
372
def forward(
@@ -444,9 +411,8 @@ class Siglip2Encoder(nn.Module):
411
self,
412
config: Siglip2VisionConfig,
413
quant_config: QuantizationConfig | None = None,
414
+ multimodal_config: MultiModalConfig | None = None,
415
prefix: str = "",
448
- use_data_parallel: bool = False,
449
- attn_backend_override: AttentionBackendEnum | None = None,
416
):
417
super().__init__()
418
self.config = config
@@ -455,9 +421,8 @@ class Siglip2Encoder(nn.Module):
421
Siglip2EncoderLayer(
422
config,
423
quant_config=quant_config,
424
+ multimodal_config=multimodal_config,
425
prefix=f"{prefix}.layers.{idx}",
459
- use_data_parallel=use_data_parallel,
460
- attn_backend_override=attn_backend_override,
426
)
427
for idx in range(config.num_hidden_layers)
428
]
@@ -630,9 +595,8 @@ class Siglip2VisionTransformer(nn.Module):
595
self,
596
config: Siglip2VisionConfig,
597
quant_config: QuantizationConfig | None = None,
598
+ multimodal_config: MultiModalConfig | None = None,
599
prefix: str = "",
634
- use_data_parallel: bool = False,
635
- attn_backend_override: AttentionBackendEnum | None = None,
600
):
601
super().__init__()
602
self.config = config
@@ -642,9 +606,8 @@ class Siglip2VisionTransformer(nn.Module):
606
self.encoder = Siglip2Encoder(
607
config,
608
quant_config=quant_config,
609
+ multimodal_config=multimodal_config,
610
prefix=f"{prefix}.encoder",
646
- use_data_parallel=use_data_parallel,
647
- attn_backend_override=attn_backend_override,
611
)
612
self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
613
@@ -671,18 +634,16 @@ class Siglip2NavitModel(torch.nn.Module):
634
self,
635
config: Siglip2VisionConfig,
636
quant_config: QuantizationConfig | None = None,
637
+ multimodal_config: MultiModalConfig | None = None,
638
prefix: str = "",
675
- use_data_parallel: bool = False,
676
- attn_backend_override: AttentionBackendEnum | None = None,
639
):
640
super().__init__()
641
642
self.vision_model = Siglip2VisionTransformer(
643
config,
644
quant_config=quant_config,
645
+ multimodal_config=multimodal_config,
646
prefix=f"{prefix}.vision_model",
684
- use_data_parallel=use_data_parallel,
685
- attn_backend_override=attn_backend_override,
647
)
648
649
def forward(
vllm/model_executor/models/vision.py
+8
-5
@@ -88,14 +88,17 @@ def get_vit_attn_backend(
88
"""
89
Get the available attention backend for Vision Transformer.
90
"""
91
- if attn_backend_override is not None:
92
- return attn_backend_override
91
+ attn_backend = attn_backend_override
92
93
selected_backend = get_current_vllm_config().attention_config.backend
95
- if selected_backend is not None:
96
- return selected_backend
94
+ if attn_backend is None:
95
+ attn_backend = selected_backend
96
98
- return current_platform.get_vit_attn_backend(head_size, dtype)
97
+ return current_platform.get_vit_attn_backend(
98
+ head_size,
99
+ dtype,
100
+ backend=attn_backend,
101
+ )
102
103
104
def should_torch_compile_mm_vit(vllm_config: VllmConfig) -> bool:
vllm/platforms/cuda.py
+36
-18
@@ -7,7 +7,7 @@ pynvml. However, it should not initialize cuda context.
7
import os
8
from collections.abc import Callable
9
from functools import cache, wraps
10
-from typing import TYPE_CHECKING, TypeVar
10
+from typing import TYPE_CHECKING, Optional, TypeVar
11
12
import torch
13
from typing_extensions import ParamSpec
@@ -255,23 +255,6 @@ class CudaPlatformBase(Platform):
255
torch.cuda.reset_peak_memory_stats(device)
256
return torch.cuda.max_memory_allocated(device)
257
258
- @classmethod
259
- def get_vit_attn_backend(
260
- cls, head_size: int, dtype: torch.dtype
261
- ) -> "AttentionBackendEnum":
262
- # Try FlashAttention first
263
- if (cc := cls.get_device_capability()) and cc.major >= 8:
264
- try:
265
- backend_class = AttentionBackendEnum.FLASH_ATTN.get_class()
266
- if backend_class.supports_head_size(
267
- head_size
268
- ) and backend_class.supports_dtype(dtype):
269
- return AttentionBackendEnum.FLASH_ATTN
270
- except ImportError:
271
- pass
272
-
273
- return AttentionBackendEnum.TORCH_SDPA
274
-
258
@classmethod
259
def get_valid_backends(
260
cls,
@@ -418,6 +401,41 @@ class CudaPlatformBase(Platform):
401
402
return selected_backend.get_path()
403
404
+ @classmethod
405
+ def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
406
+ return [
407
+ AttentionBackendEnum.TORCH_SDPA,
408
+ AttentionBackendEnum.FLASH_ATTN,
409
+ ]
410
+
411
+ @classmethod
412
+ def get_vit_attn_backend(
413
+ cls,
414
+ head_size: int,
415
+ dtype: torch.dtype,
416
+ backend: Optional["AttentionBackendEnum"] = None,
417
+ ) -> "AttentionBackendEnum":
418
+ if backend is not None:
419
+ assert backend in cls.get_supported_vit_attn_backends(), (
420
+ f"Backend {backend} is not supported for vit attention. "
421
+ f"Supported backends are: {cls.get_supported_vit_attn_backends()}"
422
+ )
423
+ logger.info_once(f"Using backend {backend} for vit attention")
424
+ return backend
425
+
426
+ # Try FlashAttention first
427
+ if (cc := cls.get_device_capability()) and cc.major >= 8:
428
+ try:
429
+ backend_class = AttentionBackendEnum.FLASH_ATTN.get_class()
430
+ if backend_class.supports_head_size(
431
+ head_size
432
+ ) and backend_class.supports_dtype(dtype):
433
+ return AttentionBackendEnum.FLASH_ATTN
434
+ except ImportError:
435
+ pass
436
+
437
+ return AttentionBackendEnum.TORCH_SDPA
438
+
439
@classmethod
440
def get_punica_wrapper(cls) -> str:
441
return "vllm.lora.punica_wrapper.punica_gpu.PunicaWrapperGPU"
vllm/platforms/interface.py
+38
-7
@@ -7,7 +7,7 @@ import platform
7
import random
8
import sys
9
from datetime import timedelta
10
-from typing import TYPE_CHECKING, Any, NamedTuple
10
+from typing import TYPE_CHECKING, Any, NamedTuple, Optional
11
12
import numpy as np
13
import torch
@@ -222,12 +222,6 @@ class Platform:
222
with contextlib.suppress(ImportError):
223
import vllm._moe_C # noqa: F401
224
225
- @classmethod
226
- def get_vit_attn_backend(
227
- cls, head_size: int, dtype: torch.dtype
228
- ) -> "AttentionBackendEnum":
229
- return AttentionBackendEnum.TORCH_SDPA
230
-
225
@classmethod
226
def get_attn_backend_cls(
227
cls,
@@ -245,6 +239,43 @@ class Platform:
239
"""Get the attention backend class of a device."""
240
return ""
241
242
+ @classmethod
243
+ def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
244
+ return [
245
+ AttentionBackendEnum.TORCH_SDPA,
246
+ ]
247
+
248
+ @classmethod
249
+ def get_vit_attn_backend(
250
+ cls,
251
+ head_size: int,
252
+ dtype: torch.dtype,
253
+ backend: Optional["AttentionBackendEnum"] = None,
254
+ ) -> "AttentionBackendEnum":
255
+ """
256
+ Get the vision attention backend class of a device.
257
+
258
+ NOTE: ViT Attention should be checked and override in the platform-specific
259
+ implementation. we should not override this in any other places, like
260
+ the model_executor/models/<model_name>.py.
261
+
262
+ We check if the backend is None or not:
263
+ 1. If not, check if the backend is supported by the platform.
264
+ 2. If None, continue to the default selection logic.
265
+ """
266
+ if backend is not None:
267
+ assert backend in cls.get_supported_vit_attn_backends(), (
268
+ f"Backend {backend} is not supported for vit attention"
269
+ f"Supported backends are: {cls.get_supported_vit_attn_backends()}"
270
+ )
271
+ logger.info_once(f"Using backend {backend} for vit attention")
272
+ return backend
273
+
274
+ logger.info_once(
275
+ f"Using default backend {AttentionBackendEnum.TORCH_SDPA} for vit attention"
276
+ )
277
+ return AttentionBackendEnum.TORCH_SDPA
278
+
279
@classmethod
280
def get_device_capability(
281
cls,
vllm/platforms/rocm.py
+38
-19
@@ -3,7 +3,7 @@
3
4
import os
5
from functools import cache, lru_cache, wraps
6
-from typing import TYPE_CHECKING
6
+from typing import TYPE_CHECKING, Optional
7
8
import torch
9
@@ -187,24 +187,6 @@ class RocmPlatform(Platform):
187
if not on_gfx9():
188
supported_quantization += ["bitsandbytes"]
189
190
- @classmethod
191
- def get_vit_attn_backend(
192
- cls, head_size: int, dtype: torch.dtype
193
- ) -> AttentionBackendEnum:
194
- from importlib.util import find_spec
195
-
196
- from vllm._aiter_ops import rocm_aiter_ops
197
-
198
- if rocm_aiter_ops.is_mha_enabled():
199
- # Note: AITER FA is only supported for Qwen-VL models.
200
- # TODO: Add support for other VL models in their model class.
201
- return AttentionBackendEnum.ROCM_AITER_FA
202
-
203
- if on_gfx9() and find_spec("flash_attn") is not None:
204
- return AttentionBackendEnum.FLASH_ATTN
205
-
206
- return AttentionBackendEnum.TORCH_SDPA
207
-
190
@classmethod
191
def get_attn_backend_cls(
192
cls,
@@ -322,6 +304,43 @@ class RocmPlatform(Platform):
304
"ROCm. Note that V0 attention backends have been removed."
305
)
306
307
+ @classmethod
308
+ def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
309
+ return [
310
+ AttentionBackendEnum.FLASH_ATTN,
311
+ AttentionBackendEnum.ROCM_AITER_FA,
312
+ AttentionBackendEnum.TORCH_SDPA,
313
+ ]
314
+
315
+ @classmethod
316
+ def get_vit_attn_backend(
317
+ cls,
318
+ head_size: int,
319
+ dtype: torch.dtype,
320
+ backend: Optional["AttentionBackendEnum"] = None,
321
+ ) -> "AttentionBackendEnum":
322
+ if backend is not None:
323
+ assert backend in cls.get_supported_vit_attn_backends(), (
324
+ f"Backend {backend} is not supported for vit attention. "
325
+ f"Supported backends are: {cls.get_supported_vit_attn_backends()}"
326
+ )
327
+ logger.info_once(f"Using backend {backend} for vit attention")
328
+ return backend
329
+
330
+ from importlib.util import find_spec
331
+
332
+ from vllm._aiter_ops import rocm_aiter_ops
333
+
334
+ if rocm_aiter_ops.is_mha_enabled():
335
+ # Note: AITER FA is only supported for Qwen-VL models.
336
+ # TODO: Add support for other VL models in their model class.
337
+ return AttentionBackendEnum.ROCM_AITER_FA
338
+
339
+ if on_gfx9() and find_spec("flash_attn") is not None:
340
+ return AttentionBackendEnum.FLASH_ATTN
341
+
342
+ return AttentionBackendEnum.TORCH_SDPA
343
+
344
@classmethod
345
def set_device(cls, device: torch.device) -> None:
346
"""
vllm/platforms/tpu.py
+27
-1
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import contextlib
5
-from typing import TYPE_CHECKING, cast
5
+from typing import TYPE_CHECKING, Optional, cast
6
7
import torch
8
from tpu_info import device
@@ -75,6 +75,32 @@ class TpuPlatform(Platform):
75
logger.info("Using Pallas V1 backend.")
76
return AttentionBackendEnum.PALLAS.get_path()
77
78
+ @classmethod
79
+ def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
80
+ return [
81
+ AttentionBackendEnum.PALLAS,
82
+ ]
83
+
84
+ @classmethod
85
+ def get_vit_attn_backend(
86
+ cls,
87
+ head_size: int,
88
+ dtype: torch.dtype,
89
+ backend: Optional["AttentionBackendEnum"] = None,
90
+ ) -> "AttentionBackendEnum":
91
+ if backend is not None:
92
+ assert backend in cls.get_supported_vit_attn_backends(), (
93
+ f"Backend {backend} is not supported for vit attention"
94
+ f"Supported backends are: {cls.get_supported_vit_attn_backends()}."
95
+ )
96
+ logger.info_once(f"Using backend {backend} for vit attention.")
97
+ return backend
98
+
99
+ logger.info_once(
100
+ f"Using default backend {AttentionBackendEnum.PALLAS} for vit attention."
101
+ )
102
+ return AttentionBackendEnum.PALLAS
103
+
104
@classmethod
105
def set_device(cls, device: torch.device) -> None:
106
"""
vllm/platforms/xpu.py
+29
-7
@@ -3,7 +3,7 @@
3
4
import contextlib
5
import os
6
-from typing import TYPE_CHECKING
6
+from typing import TYPE_CHECKING, Optional
7
8
import torch
9
@@ -77,6 +77,34 @@ class XPUPlatform(Platform):
77
logger.info("Using Flash Attention backend.")
78
return AttentionBackendEnum.FLASH_ATTN.get_path()
79
80
+ @classmethod
81
+ def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
82
+ # XPU only supports FLASH_ATTN for vision attention.
83
+ return [
84
+ AttentionBackendEnum.FLASH_ATTN,
85
+ ]
86
+
87
+ @classmethod
88
+ def get_vit_attn_backend(
89
+ cls,
90
+ head_size: int,
91
+ dtype: torch.dtype,
92
+ backend: Optional["AttentionBackendEnum"] = None,
93
+ ) -> "AttentionBackendEnum":
94
+ if backend is not None:
95
+ assert backend in cls.get_supported_vit_attn_backends(), (
96
+ f"Backend {backend} is not supported for vit attention. "
97
+ f"Supported backends are: "
98
+ f"{cls.get_supported_vit_attn_backends()}."
99
+ )
100
+ logger.info_once(f"Using backend {backend} for vit attention")
101
+ return backend
102
+
103
+ logger.info_once(
104
+ f"Using backend {AttentionBackendEnum.FLASH_ATTN} for vit attention"
105
+ )
106
+ return AttentionBackendEnum.FLASH_ATTN
107
+
108
@classmethod
109
def set_device(cls, device: torch.device) -> None:
110
"""
@@ -110,12 +138,6 @@ class XPUPlatform(Platform):
138
device_props = torch.xpu.get_device_properties(device_id)
139
return device_props.total_memory
140
113
- @classmethod
114
- def get_vit_attn_backend(
115
- cls, head_size: int, dtype: torch.dtype
116
- ) -> "AttentionBackendEnum":
117
- return AttentionBackendEnum.FLASH_ATTN
118
-
141
@classmethod
142
def inference_mode(cls):
143
return torch.no_grad()