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

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[MM][Core] Decouple ViT backend from LM backend (#27061)

Signed-off-by: Roger Wang <hey@rogerw.io>

Roger Wang committed Oct 21, 2025 at 00:30 UTC c3a2c6ac5f9a13c5c407c69199dd0fa68145acad

16 files changed +230 -17

tests/config/test_multimodal_config.py new

+25

@@ -0,0 +1,25 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +

4 +import pytest

5 +

6 +from vllm.attention.backends.registry import _Backend

7 +from vllm.config.multimodal import MultiModalConfig

8 +

9 +

10 +def test_mm_encoder_attn_backend_str_conversion():

11 + config = MultiModalConfig(mm_encoder_attn_backend="FLASH_ATTN")

12 + assert config.mm_encoder_attn_backend == _Backend.FLASH_ATTN

13 +

14 +

15 +def test_mm_encoder_attn_backend_invalid():

16 + with pytest.raises(ValueError):

17 + MultiModalConfig(mm_encoder_attn_backend="not_a_backend")

18 +

19 +

20 +def test_mm_encoder_attn_backend_hash_updates():

21 + base_hash = MultiModalConfig().compute_hash()

22 + overridden_hash = MultiModalConfig(

23 + mm_encoder_attn_backend=_Backend.FLASH_ATTN

24 + ).compute_hash()

25 + assert base_hash != overridden_hash

vllm/attention/layer.py

+10 -1

@@ -16,6 +16,7 @@ from vllm.attention.backends.registry import _Backend, backend_name_to_enum

16 from vllm.attention.selector import get_attn_backend

17 from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target

18 from vllm.config import CacheConfig, get_current_vllm_config

19 +from vllm.config.multimodal import MultiModalConfig

20 from vllm.config.vllm import VllmConfig

21 from vllm.distributed.kv_transfer import (

22 get_kv_transfer_group,

@@ -443,6 +444,7 @@ class MultiHeadAttention(nn.Module):

444 # This has no effect, it is only here to make it easier to swap

445 # between Attention and MultiHeadAttention

446 prefix: str = "",

447 + multimodal_config: MultiModalConfig | None = None,

448 ) -> None:

449 super().__init__()

450 self.num_heads = num_heads

@@ -462,7 +464,14 @@ class MultiHeadAttention(nn.Module):

464 dtype = torch.get_default_dtype()

465

466 # Determine the attention backend

465 - backend = get_vit_attn_backend(head_size=head_size, dtype=dtype)

467 + attn_backend_override = None

468 + if multimodal_config is not None:

469 + attn_backend_override = multimodal_config.mm_encoder_attn_backend

470 + backend = get_vit_attn_backend(

471 + head_size=head_size,

472 + dtype=dtype,

473 + attn_backend_override=attn_backend_override,

474 + )

475

476 # Some auto-selected backends can be upgraded

477 # to upstream flash attention if available.

vllm/config/model.py

@@ -50,6 +50,7 @@ if TYPE_CHECKING:

51 import vllm.model_executor.layers.quantization as me_quant

52 import vllm.model_executor.models as me_models

53 + from vllm.attention.backends.registry import _Backend

54 from vllm.config.load import LoadConfig

55 from vllm.config.parallel import ParallelConfig

56 from vllm.model_executor.layers.quantization import QuantizationMethods

@@ -57,6 +58,7 @@ if TYPE_CHECKING:

58 else:

59 PretrainedConfig = Any

61 + _Backend = Any

62 me_quant = LazyLoader(

63 "model_executor", globals(), "vllm.model_executor.layers.quantization"

64 )

@@ -307,6 +309,7 @@ class ModelConfig:

309 mm_processor_cache_type: InitVar[MMCacheType | None] = None

310 mm_shm_cache_max_object_size_mb: InitVar[int | None] = None

311 mm_encoder_tp_mode: InitVar[MMEncoderTPMode | None] = None

312 + mm_encoder_attn_backend: InitVar[_Backend | str | None] = None

313 interleave_mm_strings: InitVar[bool | None] = None

314 skip_mm_profiling: InitVar[bool | None] = None

315 video_pruning_rate: InitVar[float | None] = None

@@ -424,6 +427,7 @@ class ModelConfig:

427 mm_processor_cache_type: MMCacheType | None,

428 mm_shm_cache_max_object_size_mb: int | None,

429 mm_encoder_tp_mode: MMEncoderTPMode | None,

430 + mm_encoder_attn_backend: _Backend | str | None,

431 interleave_mm_strings: bool | None,

432 skip_mm_profiling: bool | None,

433 video_pruning_rate: float | None,

@@ -733,6 +737,7 @@ class ModelConfig:

737 mm_processor_cache_type=mm_processor_cache_type,

738 mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,

739 mm_encoder_tp_mode=mm_encoder_tp_mode,

740 + mm_encoder_attn_backend=mm_encoder_attn_backend,

741 interleave_mm_strings=interleave_mm_strings,

742 skip_mm_profiling=skip_mm_profiling,

743 video_pruning_rate=video_pruning_rate,

vllm/config/multimodal.py

+38 -4

@@ -3,13 +3,18 @@

4 import hashlib

5 from collections.abc import Mapping

6 -from typing import Any, Literal, TypeAlias

6 +from typing import TYPE_CHECKING, Any, Literal, TypeAlias

8 from pydantic import ConfigDict, Field, field_validator, model_validator

9 from pydantic.dataclasses import dataclass

11 from vllm.config.utils import config

13 +if TYPE_CHECKING:

14 + from vllm.attention.backends.registry import _Backend

15 +else:

16 + _Backend = Any

17 +

19 @dataclass

20 class BaseDummyOptions:

@@ -112,6 +117,10 @@ class MultiModalConfig:

117 DP (which is controlled by `--data-parallel-size`).

118 This is only supported on a per-model basis and falls back to

119 `"weights"` if the encoder does not support DP."""

120 + mm_encoder_attn_backend: _Backend | None = None

121 + """Optional override for the multi-modal encoder attention backend when

122 + using vision transformers. Accepts any value from

123 + `vllm.attention.backends.registry._Backend` (e.g. `FLASH_ATTN`)."""

124 interleave_mm_strings: bool = False

125 """Enable fully interleaved support for multimodal prompts, while using

126 --chat-template-content-format=string."""

@@ -148,6 +157,29 @@ class MultiModalConfig:

157 value[k] = BaseDummyOptions(**v)

158 return value

159

160 + @field_validator("mm_encoder_attn_backend", mode="before")

161 + @classmethod

162 + def _validate_mm_encoder_attn_backend(cls, value: object) -> _Backend | None:

163 + from vllm.attention.backends.registry import (

164 + _Backend as BackendEnum,

165 + )

166 + from vllm.attention.backends.registry import (

167 + backend_name_to_enum,

168 + )

169 +

170 + if value is None or isinstance(value, BackendEnum):

171 + return value

172 +

173 + if isinstance(value, str):

174 + candidate = backend_name_to_enum(value.upper())

175 + if candidate is not None:

176 + return candidate

177 +

178 + valid_backends = ", ".join(sorted(BackendEnum.__members__.keys()))

179 + raise ValueError(

180 + f"Invalid mm encoder attention backend. Expected one of: {valid_backends}."

181 + )

182 +

183 @model_validator(mode="after")

184 def _validate_multimodal_config(self):

185 if self.mm_processor_cache_type != "shm" and (

@@ -172,9 +204,11 @@ class MultiModalConfig:

204 excluding anything before input ids/embeddings and after

205 the final hidden states.

206 """

175 - # no factors to consider.

176 - # this config will not affect the computation graph.

177 - factors: list[Any] = []

207 + factors: list[Any] = [

208 + self.mm_encoder_attn_backend.name

209 + if self.mm_encoder_attn_backend is not None

210 + else None

211 + ]

212 hash_str = hashlib.md5(str(factors).encode(), usedforsecurity=False).hexdigest()

213 return hash_str

214

vllm/engine/arg_utils.py

@@ -32,6 +32,7 @@ from pydantic.fields import FieldInfo

32 from typing_extensions import TypeIs, deprecated

34 import vllm.envs as envs

35 +from vllm.attention.backends.registry import _Backend

36 from vllm.config import (

37 CacheConfig,

38 CompilationConfig,

@@ -451,6 +452,9 @@ class EngineArgs:

452 MultiModalConfig.mm_shm_cache_max_object_size_mb

453 )

454 mm_encoder_tp_mode: MMEncoderTPMode = MultiModalConfig.mm_encoder_tp_mode

455 + mm_encoder_attn_backend: _Backend | str | None = (

456 + MultiModalConfig.mm_encoder_attn_backend

457 + )

458 io_processor_plugin: str | None = None

459 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling

460 video_pruning_rate: float = MultiModalConfig.video_pruning_rate

@@ -914,6 +918,10 @@ class EngineArgs:

918 multimodal_group.add_argument(

919 "--mm-encoder-tp-mode", **multimodal_kwargs["mm_encoder_tp_mode"]

920 )

921 + multimodal_group.add_argument(

922 + "--mm-encoder-attn-backend",

923 + **multimodal_kwargs["mm_encoder_attn_backend"],

924 + )

925 multimodal_group.add_argument(

926 "--interleave-mm-strings", **multimodal_kwargs["interleave_mm_strings"]

927 )

@@ -1160,6 +1168,7 @@ class EngineArgs:

1168 mm_processor_cache_type=self.mm_processor_cache_type,

1169 mm_shm_cache_max_object_size_mb=self.mm_shm_cache_max_object_size_mb,

1170 mm_encoder_tp_mode=self.mm_encoder_tp_mode,

1171 + mm_encoder_attn_backend=self.mm_encoder_attn_backend,

1172 pooler_config=self.pooler_config,

1173 override_pooler_config=self.override_pooler_config,

1174 logits_processor_pattern=self.logits_processor_pattern,

vllm/model_executor/models/dots_ocr.py

+17 -2

@@ -256,6 +256,7 @@ class DotsVisionAttention(nn.Module):

256 quant_config: QuantizationConfig | None = None,

257 prefix: str = "",

258 use_data_parallel: bool = False,

259 + attn_backend_override: _Backend | None = None,

260 ) -> None:

261 super().__init__()

262

@@ -288,7 +289,9 @@ class DotsVisionAttention(nn.Module):

289 )

290 # Select attention backend

291 self.attn_backend = get_vit_attn_backend(

291 - self.hidden_size_per_attention_head, torch.get_default_dtype()

292 + self.hidden_size_per_attention_head,

293 + torch.get_default_dtype(),

294 + attn_backend_override=attn_backend_override,

295 )

296 self.use_upstream_fa = False

297

@@ -510,6 +513,7 @@ class DotsVisionBlock(nn.Module):

513 quant_config: QuantizationConfig | None = None,

514 prefix: str = "",

515 use_data_parallel: bool = False,

516 + attn_backend_override: _Backend | None = None,

517 ):

518 super().__init__()

519

@@ -521,6 +525,7 @@ class DotsVisionBlock(nn.Module):

525 quant_config=quant_config,

526 prefix=f"{prefix}.attn",

527 use_data_parallel=use_data_parallel,

528 + attn_backend_override=attn_backend_override,

529 )

530 self.norm1 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)

531 self.mlp = DotsSwiGLUFFN(

@@ -561,6 +566,7 @@ class DotsVisionTransformer(nn.Module):

566 require_post_norm: bool | None = None,

567 prefix: str = "",

568 use_data_parallel: bool = False,

569 + attn_backend_override: _Backend | None = None,

570 ) -> None:

571 super().__init__()

572 self.config = config

@@ -571,7 +577,9 @@ class DotsVisionTransformer(nn.Module):

577 head_dim = config.embed_dim // config.num_attention_heads

578 self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)

579 self.attn_backend = get_vit_attn_backend(

574 - head_size=head_dim, dtype=torch.get_default_dtype()

580 + head_size=head_dim,

581 + dtype=torch.get_default_dtype(),

582 + attn_backend_override=attn_backend_override,

583 )

584 if self.attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(

585 torch.get_default_dtype()

@@ -591,6 +599,7 @@ class DotsVisionTransformer(nn.Module):

599 quant_config=quant_config,

600 prefix=f"{prefix}.blocks.{i}",

601 use_data_parallel=use_data_parallel,

602 + attn_backend_override=attn_backend_override,

603 )

604 for i in range(num_layers)

605 ]

@@ -750,11 +759,17 @@ class DotsOCRForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsLoRA

759 self.config.vision_config = vision_config

760 else:

761 vision_config = self.config.vision_config

762 + attn_backend_override = (

763 + multimodal_config.mm_encoder_attn_backend

764 + if multimodal_config is not None

765 + else None

766 + )

767 self.vision_tower = DotsVisionTransformer(

768 vision_config,

769 quant_config=self.quant_config,

770 prefix=maybe_prefix(prefix, "vision_tower"),

771 use_data_parallel=self.use_data_parallel,

772 + attn_backend_override=attn_backend_override,

773 )

774 self.language_model: Qwen2ForCausalLM = init_vllm_registered_model(

775 vllm_config=vllm_config,

vllm/model_executor/models/ernie45_vl.py

+15 -1

@@ -164,6 +164,7 @@ class Ernie4_5_VisionAttention(nn.Module):

164 projection_size: int,

165 quant_config: QuantizationConfig | None = None,

166 prefix: str = "",

167 + attn_backend_override: _Backend | None = None,

168 ) -> None:

169 super().__init__()

170 # Per attention head and per partition values.

@@ -196,6 +197,7 @@ class Ernie4_5_VisionAttention(nn.Module):

197 self.attn_backend = get_vit_attn_backend(

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.use_upstream_fa = False

@@ -367,6 +369,7 @@ class Ernie4_5_VisionBlock(nn.Module):

369 norm_layer: Callable[[int], nn.Module] | None = None,

370 quant_config: QuantizationConfig | None = None,

371 prefix: str = "",

372 + attn_backend_override: _Backend | None = None,

373 ) -> None:

374 super().__init__()

375

@@ -382,6 +385,7 @@ class Ernie4_5_VisionBlock(nn.Module):

385 projection_size=dim,

386 quant_config=quant_config,

387 prefix=f"{prefix}.attn",

388 + attn_backend_override=attn_backend_override,

389 )

390

391 self.mlp = Ernie4_5_VisionMLP(

@@ -458,6 +462,7 @@ class Ernie4_5_VisionTransformer(nn.Module):

462 norm_eps: float = 1e-6,

463 quant_config: QuantizationConfig | None = None,

464 prefix: str = "",

465 + attn_backend_override: _Backend | None = None,

466 ) -> None:

467 super().__init__()

468 patch_size = vision_config.patch_size

@@ -493,6 +498,7 @@ class Ernie4_5_VisionTransformer(nn.Module):

498 norm_layer=norm_layer,

499 quant_config=quant_config,

500 prefix=f"{prefix}.blocks.{layer_idx}",

501 + attn_backend_override=attn_backend_override,

502 )

503 for layer_idx in range(depth)

504 ]

@@ -504,7 +510,9 @@ class Ernie4_5_VisionTransformer(nn.Module):

510 self.ln = nn.LayerNorm(hidden_size, eps=1e-6)

511

512 self.attn_backend = get_vit_attn_backend(

507 - head_size=head_dim, dtype=torch.get_default_dtype()

513 + head_size=head_dim,

514 + dtype=torch.get_default_dtype(),

515 + attn_backend_override=attn_backend_override,

516 )

517 if self.attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(

518 torch.get_default_dtype()

@@ -1327,11 +1335,17 @@ class Ernie4_5_VLMoeForConditionalGeneration(

1335 self.config = config

1336 self.multimodal_config = multimodal_config

1337

1338 + attn_backend_override = (

1339 + multimodal_config.mm_encoder_attn_backend

1340 + if multimodal_config is not None

1341 + else None

1342 + )

1343 self.vision_model = Ernie4_5_VisionTransformer(

1344 config.vision_config,

1345 norm_eps=getattr(config, "rms_norm_eps", 1e-6),

1346 quant_config=quant_config,

1347 prefix=maybe_prefix(prefix, "vision_model"),

1348 + attn_backend_override=attn_backend_override,

1349 )

1350

1351 self.language_model = Ernie4_5_VLMoeForCausalLM(

vllm/model_executor/models/glm4_1v.py

+15 -1

@@ -247,6 +247,7 @@ class Glm4vVisionAttention(nn.Module):

247 quant_config: QuantizationConfig | None = None,

248 prefix: str = "",

249 use_data_parallel: bool = False,

250 + attn_backend_override: _Backend | None = None,

251 ) -> None:

252 super().__init__()

253 # Per attention head and per partition values.

@@ -287,6 +288,7 @@ class Glm4vVisionAttention(nn.Module):

288 self.attn_backend = get_vit_attn_backend(

289 head_size=self.hidden_size_per_attention_head,

290 dtype=torch.get_default_dtype(),

291 + attn_backend_override=attn_backend_override,

292 )

293 self.use_upstream_fa = False

294

@@ -417,6 +419,7 @@ class Glm4vVisionBlock(nn.Module):

419 quant_config: QuantizationConfig | None = None,

420 prefix: str = "",

421 use_data_parallel: bool = False,

422 + attn_backend_override: _Backend | None = None,

423 ) -> None:

424 super().__init__()

425 if norm_layer is None:

@@ -430,6 +433,7 @@ class Glm4vVisionBlock(nn.Module):

433 quant_config=quant_config,

434 prefix=f"{prefix}.attn",

435 use_data_parallel=use_data_parallel,

436 + attn_backend_override=attn_backend_override,

437 )

438 self.mlp = Glm4vVisionMLP(

439 dim,

@@ -696,6 +700,7 @@ class Glm4vVisionTransformer(nn.Module):

700 quant_config: QuantizationConfig | None = None,

701 prefix: str = "",

702 use_data_parallel: bool = False,

703 + attn_backend_override: _Backend | None = None,

704 ) -> None:

705 super().__init__()

706

@@ -731,6 +736,7 @@ class Glm4vVisionTransformer(nn.Module):

736 quant_config=quant_config,

737 prefix=f"{prefix}.blocks.{layer_idx}",

738 use_data_parallel=self.use_data_parallel,

739 + attn_backend_override=attn_backend_override,

740 )

741 for layer_idx in range(depth)

742 ]

@@ -759,7 +765,9 @@ class Glm4vVisionTransformer(nn.Module):

765 )

766

767 self.attn_backend = get_vit_attn_backend(

762 - head_size=head_dim, dtype=torch.get_default_dtype()

768 + head_size=head_dim,

769 + dtype=torch.get_default_dtype(),

770 + attn_backend_override=attn_backend_override,

771 )

772 if self.attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(

773 torch.get_default_dtype()

@@ -1437,12 +1445,18 @@ class Glm4vForConditionalGeneration(

1445 self.multimodal_config = multimodal_config

1446 self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

1447

1448 + attn_backend_override = (

1449 + multimodal_config.mm_encoder_attn_backend

1450 + if multimodal_config is not None

1451 + else None

1452 + )

1453 self.visual = Glm4vVisionTransformer(

1454 config.vision_config,

1455 norm_eps=getattr(config, "rms_norm_eps", 1e-5),

1456 quant_config=quant_config,

1457 prefix=maybe_prefix(prefix, "visual"),

1458 use_data_parallel=self.use_data_parallel,

1459 + attn_backend_override=attn_backend_override,

1460 )

1461

1462 if config.model_type == "glm4v":

vllm/model_executor/models/keye.py

+18 -1

@@ -353,6 +353,7 @@ class KeyeSiglipAttention(nn.Module):

353 config: PretrainedConfig,

354 quant_config: QuantizationConfig | None = None,

355 prefix: str = "",

356 + attn_backend_override: _Backend | None = None,

357 ):

358 super().__init__()

359 self.config = config

@@ -392,7 +393,9 @@ class KeyeSiglipAttention(nn.Module):

393

394 # Detect attention implementation.

395 self.attn_backend = get_vit_attn_backend(

395 - head_size=self.head_dim, dtype=torch.get_default_dtype()

396 + head_size=self.head_dim,

397 + dtype=torch.get_default_dtype(),

398 + attn_backend_override=attn_backend_override,

399 )

400

401 self.use_upstream_fa = False

@@ -521,6 +524,7 @@ class KeyeSiglipEncoderLayer(nn.Module):

524 config: PretrainedConfig,

525 quant_config: QuantizationConfig | None = None,

526 prefix: str = "",

527 + attn_backend_override: _Backend | None = None,

528 ):

529 super().__init__()

530 self.embed_dim = config.hidden_size

@@ -529,6 +533,7 @@ class KeyeSiglipEncoderLayer(nn.Module):

533 config,

534 quant_config=quant_config,

535 prefix=f"{prefix}.self_attn",

536 + attn_backend_override=attn_backend_override,

537 )

538 self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

539 self.mlp = SiglipMLP(

@@ -573,6 +578,7 @@ class KeyeSiglipEncoder(nn.Module):

578 config: PretrainedConfig,

579 quant_config: QuantizationConfig | None = None,

580 prefix: str = "",

581 + attn_backend_override: _Backend | None = None,

582 ):

583 super().__init__()

584 self.config = config

@@ -585,6 +591,7 @@ class KeyeSiglipEncoder(nn.Module):

591 config,

592 quant_config=quant_config,

593 prefix=f"{prefix}.layers.{layer_idx}",

594 + attn_backend_override=attn_backend_override,

595 )

596 for layer_idx in range(config.num_hidden_layers)

597 ]

@@ -666,6 +673,7 @@ class KeyeSiglipVisionTransformer(nn.Module):

673 config: PretrainedConfig,

674 quant_config: QuantizationConfig | None = None,

675 prefix: str = "",

676 + attn_backend_override: _Backend | None = None,

677 ):

678 super().__init__()

679 self.config = config

@@ -676,6 +684,7 @@ class KeyeSiglipVisionTransformer(nn.Module):

684 config,

685 quant_config=quant_config,

686 prefix=f"{prefix}.encoder",

687 + attn_backend_override=attn_backend_override,

688 )

689 self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

690

@@ -747,6 +756,7 @@ class KeyeSiglipVisionModel(nn.Module):

756 config: PretrainedConfig,

757 quant_config: QuantizationConfig | None = None,

758 prefix: str = "",

759 + attn_backend_override: _Backend | None = None,

760 ):

761 super().__init__()

762

@@ -754,6 +764,7 @@ class KeyeSiglipVisionModel(nn.Module):

764 config,

765 quant_config=quant_config,

766 prefix=f"{prefix}.vision_model",

767 + attn_backend_override=attn_backend_override,

768 )

769 self.quant_config = quant_config

770

@@ -1296,10 +1307,16 @@ class BaseKeyeModule(nn.Module):

1307 self.config = config

1308 self.multimodal_config = multimodal_config

1309

1310 + attn_backend_override = (

1311 + multimodal_config.mm_encoder_attn_backend

1312 + if multimodal_config is not None

1313 + else None

1314 + )

1315 self.visual = KeyeSiglipVisionModel(

1316 config.vision_config,

1317 quant_config=quant_config,

1318 prefix=maybe_prefix(prefix, "visual"),

1319 + attn_backend_override=attn_backend_override,

1320 )

1321

1322 self.mlp_AR = self._build_projector(

vllm/model_executor/models/ovis2_5.py

+12

@@ -10,6 +10,7 @@ import torch

10 import torch.nn as nn

11 from transformers import BaseImageProcessor, BatchFeature, PretrainedConfig

13 +from vllm.attention.backends.registry import _Backend

14 from vllm.config import VllmConfig

15 from vllm.config.multimodal import BaseDummyOptions

16 from vllm.model_executor.layers.linear import ReplicatedLinear

@@ -105,6 +106,7 @@ class VisualTokenizer(torch.nn.Module):

106 quant_config: QuantizationConfig | None = None,

107 prefix: str = "",

108 use_data_parallel: bool = False,

109 + attn_backend_override: _Backend | None = None,

110 ):

111 super().__init__()

112 self.config = config

@@ -113,6 +115,7 @@ class VisualTokenizer(torch.nn.Module):

115 quant_config=quant_config,

116 prefix=f"{prefix}.vit",

117 use_data_parallel=use_data_parallel,

118 + attn_backend_override=attn_backend_override,

119 )

120 # reserved tokens for INDICATOR_IDS

121 head_dim = visual_vocab_size - len(INDICATOR_IDS)

@@ -132,6 +135,7 @@ class VisualTokenizer(torch.nn.Module):

135 quant_config: QuantizationConfig | None = None,

136 prefix: str = "",

137 use_data_parallel: bool = False,

138 + attn_backend_override: _Backend | None = None,

139 ):

140 model_type = config.model_type

141 if model_type == "siglip2_navit":

@@ -140,6 +144,7 @@ class VisualTokenizer(torch.nn.Module):

144 quant_config=quant_config,

145 prefix=prefix,

146 use_data_parallel=use_data_parallel,

147 + attn_backend_override=attn_backend_override,

148 )

149 raise ValueError(f"Unsupported visual tokenizer model_type: {model_type}")

150

@@ -457,6 +462,7 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

462 super().__init__()

463 config = vllm_config.model_config.hf_config

464 quant_config = vllm_config.quant_config

465 + multimodal_config = vllm_config.model_config.multimodal_config

466

467 self.config: PretrainedConfig = config

468 self.llm = init_vllm_registered_model(

@@ -464,11 +470,17 @@ class Ovis2_5(nn.Module, SupportsMultiModal, SupportsPP):

470 prefix=maybe_prefix(prefix, "llm"),

471 )

472

473 + attn_backend_override = (

474 + multimodal_config.mm_encoder_attn_backend

475 + if multimodal_config is not None

476 + else None

477 + )

478 self.visual_tokenizer = VisualTokenizer(

479 config=config.vit_config,

480 visual_vocab_size=config.visual_vocab_size,

481 quant_config=quant_config,

482 prefix=f"{prefix}.visual_tokenizer",

483 + attn_backend_override=attn_backend_override,

484 )

485

486 self.vte = VisualEmbedding(config.visual_vocab_size, config.hidden_size)

vllm/model_executor/models/qwen2_5_vl.py

+10 -1

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

637 quant_config: QuantizationConfig | None = None,

638 prefix: str = "",

639 use_data_parallel: bool = False,

640 + attn_backend_override: _Backend | None = None,

641 ) -> None:

642 super().__init__()

643

@@ -669,7 +670,9 @@ class Qwen2_5_VisionTransformer(nn.Module):

670

671 use_upstream_fa = False

672 self.attn_backend = get_vit_attn_backend(

672 - head_size=head_dim, dtype=torch.get_default_dtype()

673 + head_size=head_dim,

674 + dtype=torch.get_default_dtype(),

675 + attn_backend_override=attn_backend_override,

676 )

677 if (

678 self.attn_backend != _Backend.FLASH_ATTN

@@ -1226,12 +1229,18 @@ class Qwen2_5_VLForConditionalGeneration(

1229 if multimodal_config.get_limit_per_prompt(

1230 "image"

1231 ) or multimodal_config.get_limit_per_prompt("video"):

1232 + attn_backend_override = (

1233 + multimodal_config.mm_encoder_attn_backend

1234 + if multimodal_config is not None

1235 + else None

1236 + )

1237 self.visual = Qwen2_5_VisionTransformer(

1238 config.vision_config,

1239 norm_eps=getattr(config, "rms_norm_eps", 1e-6),

1240 quant_config=self.quant_config,

1241 prefix=maybe_prefix(prefix, "visual"),

1242 use_data_parallel=self.use_data_parallel,

1243 + attn_backend_override=attn_backend_override,

1244 )

1245 else:

1246 self.visual = None

vllm/model_executor/models/qwen2_vl.py

+15 -1

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

320 quant_config: QuantizationConfig | None = None,

321 prefix: str = "",

322 use_data_parallel: bool = False,

323 + attn_backend_override: _Backend | None = None,

324 ) -> None:

325 super().__init__()

326 # Per attention head and per partition values.

@@ -355,6 +356,7 @@ class Qwen2VisionAttention(nn.Module):

356 self.attn_backend = get_vit_attn_backend(

357 head_size=self.hidden_size_per_attention_head,

358 dtype=torch.get_default_dtype(),

359 + attn_backend_override=attn_backend_override,

360 )

361 self.use_upstream_fa = False

362

@@ -497,6 +499,7 @@ class Qwen2VisionBlock(nn.Module):

499 quant_config: QuantizationConfig | None = None,

500 prefix: str = "",

501 use_data_parallel: bool = False,

502 + attn_backend_override: _Backend | None = None,

503 ) -> None:

504 super().__init__()

505 if norm_layer is None:

@@ -512,6 +515,7 @@ class Qwen2VisionBlock(nn.Module):

515 quant_config=quant_config,

516 prefix=f"{prefix}.attn",

517 use_data_parallel=use_data_parallel,

518 + attn_backend_override=attn_backend_override,

519 )

520 self.mlp = Qwen2VisionMLP(

521 dim,

@@ -662,6 +666,7 @@ class Qwen2VisionTransformer(nn.Module):

666 quant_config: QuantizationConfig | None = None,

667 prefix: str = "",

668 use_data_parallel: bool = False,

669 + attn_backend_override: _Backend | None = None,

670 ) -> None:

671 super().__init__()

672

@@ -703,6 +708,7 @@ class Qwen2VisionTransformer(nn.Module):

708 quant_config=quant_config,

709 prefix=f"{prefix}.blocks.{layer_idx}",

710 use_data_parallel=use_data_parallel,

711 + attn_backend_override=attn_backend_override,

712 )

713 for layer_idx in range(depth)

714 ]

@@ -716,7 +722,9 @@ class Qwen2VisionTransformer(nn.Module):

722 use_data_parallel=use_data_parallel,

723 )

724 self.attn_backend = get_vit_attn_backend(

719 - head_size=head_dim, dtype=torch.get_default_dtype()

725 + head_size=head_dim,

726 + dtype=torch.get_default_dtype(),

727 + attn_backend_override=attn_backend_override,

728 )

729 if self.attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(

730 torch.get_default_dtype()

@@ -1356,12 +1364,18 @@ class Qwen2VLForConditionalGeneration(

1364 if multimodal_config.get_limit_per_prompt(

1365 "image"

1366 ) or multimodal_config.get_limit_per_prompt("video"):

1367 + attn_backend_override = (

1368 + multimodal_config.mm_encoder_attn_backend

1369 + if multimodal_config is not None

1370 + else None

1371 + )

1372 self.visual = Qwen2VisionTransformer(

1373 config.vision_config,

1374 norm_eps=getattr(config, "rms_norm_eps", 1e-6),

1375 quant_config=quant_config,

1376 prefix=maybe_prefix(prefix, "visual"),

1377 use_data_parallel=self.use_data_parallel,

1378 + attn_backend_override=attn_backend_override,

1379 )

1380 else:

1381 self.visual = None

vllm/model_executor/models/qwen3_omni_moe_thinker.py

+10 -1

@@ -296,6 +296,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):

296 norm_eps: float = 1e-6,

297 quant_config: QuantizationConfig | None = None,

298 prefix: str = "",

299 + attn_backend_override: _Backend | None = None,

300 ) -> None:

301 super().__init__()

302 self.hidden_size = vision_config.hidden_size

@@ -367,7 +368,9 @@ class Qwen3Omni_VisionTransformer(nn.Module):

368 )

369

370 self.attn_backend = get_vit_attn_backend(

370 - head_size=head_dim, dtype=torch.get_default_dtype()

371 + head_size=head_dim,

372 + dtype=torch.get_default_dtype(),

373 + attn_backend_override=attn_backend_override,

374 )

375 if self.attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(

376 torch.get_default_dtype()

@@ -1144,11 +1147,17 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(

1147

1148 self.audio_tower = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config)

1149

1150 + attn_backend_override = (

1151 + multimodal_config.mm_encoder_attn_backend

1152 + if multimodal_config is not None

1153 + else None

1154 + )

1155 self.visual = Qwen3Omni_VisionTransformer(

1156 vision_config=thinker_config.vision_config,

1157 norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),

1158 quant_config=quant_config,

1159 prefix=maybe_prefix(prefix, "visual"),

1160 + attn_backend_override=attn_backend_override,

1161 )

1162 self.quant_config = quant_config

1163

vllm/model_executor/models/qwen3_vl.py

+10 -2

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

300 quant_config: QuantizationConfig | None = None,

301 prefix: str = "",

302 use_data_parallel: bool = False,

303 + attn_backend_override: _Backend | None = None,

304 ) -> None:

305 super().__init__()

306 self.hidden_size = vision_config.hidden_size

@@ -359,7 +360,9 @@ class Qwen3_VisionTransformer(nn.Module):

360 )

361

362 self.attn_backend = get_vit_attn_backend(

362 - head_size=head_dim, dtype=torch.get_default_dtype()

363 + head_size=head_dim,

364 + dtype=torch.get_default_dtype(),

365 + attn_backend_override=attn_backend_override,

366 )

367 use_upstream_fa = False

368 if (

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

382 raise RuntimeError(

383 f"Qwen3-VL does not support {self.attn_backend} backend now."

384 )

382 -

385 self.blocks = nn.ModuleList(

386 [

387 Qwen3_VisionBlock(

@@ -1214,12 +1216,18 @@ class Qwen3VLForConditionalGeneration(

1216 ) and not multimodal_config.get_limit_per_prompt("video"):

1217 self.visual = None

1218 else:

1219 + attn_backend_override = (

1220 + multimodal_config.mm_encoder_attn_backend

1221 + if multimodal_config is not None

1222 + else None

1223 + )

1224 self.visual = Qwen3_VisionTransformer(

1225 config.vision_config,

1226 norm_eps=getattr(config, "rms_norm_eps", 1e-6),

1227 quant_config=quant_config,

1228 prefix=maybe_prefix(prefix, "visual"),

1229 use_data_parallel=self.use_data_parallel,

1230 + attn_backend_override=attn_backend_override,

1231 )

1232

1233 self.language_model = Qwen3LLMForCausalLM(

vllm/model_executor/models/siglip2navit.py

+12 -1

@@ -208,6 +208,7 @@ class Siglip2Attention(nn.Module):

208 quant_config: QuantizationConfig | None = None,

209 prefix: str = "",

210 use_data_parallel: bool = False,

211 + attn_backend_override: _Backend | None = None,

212 ):

213 super().__init__()

214 self.config = config

@@ -248,7 +249,9 @@ class Siglip2Attention(nn.Module):

249

250 # Detect attention implementation.

251 self.attn_backend = get_vit_attn_backend(

251 - head_size=self.head_dim, dtype=torch.get_default_dtype()

252 + head_size=self.head_dim,

253 + dtype=torch.get_default_dtype(),

254 + attn_backend_override=attn_backend_override,

255 )

256 self.use_upstream_fa = False

257

@@ -372,6 +375,7 @@ class Siglip2EncoderLayer(nn.Module):

375 quant_config: QuantizationConfig | None = None,

376 prefix: str = "",

377 use_data_parallel: bool = False,

378 + attn_backend_override: _Backend | None = None,

379 ):

380 super().__init__()

381 self.embed_dim = config.hidden_size

@@ -381,6 +385,7 @@ class Siglip2EncoderLayer(nn.Module):

385 quant_config=quant_config,

386 prefix=f"{prefix}.self_attn",

387 use_data_parallel=use_data_parallel,

388 + attn_backend_override=attn_backend_override,

389 )

390 self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

391 self.mlp = Siglip2MLP(

@@ -434,6 +439,7 @@ class Siglip2Encoder(nn.Module):

439 quant_config: QuantizationConfig | None = None,

440 prefix: str = "",

441 use_data_parallel: bool = False,

442 + attn_backend_override: _Backend | None = None,

443 ):

444 super().__init__()

445 self.config = config

@@ -444,6 +450,7 @@ class Siglip2Encoder(nn.Module):

450 quant_config=quant_config,

451 prefix=f"{prefix}.layers.{idx}",

452 use_data_parallel=use_data_parallel,

453 + attn_backend_override=attn_backend_override,

454 )

455 for idx in range(config.num_hidden_layers)

456 ]

@@ -618,6 +625,7 @@ class Siglip2VisionTransformer(nn.Module):

625 quant_config: QuantizationConfig | None = None,

626 prefix: str = "",

627 use_data_parallel: bool = False,

628 + attn_backend_override: _Backend | None = None,

629 ):

630 super().__init__()

631 self.config = config

@@ -629,6 +637,7 @@ class Siglip2VisionTransformer(nn.Module):

637 quant_config=quant_config,

638 prefix=f"{prefix}.encoder",

639 use_data_parallel=use_data_parallel,

640 + attn_backend_override=attn_backend_override,

641 )

642 self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

643

@@ -657,6 +666,7 @@ class Siglip2NavitModel(torch.nn.Module):

666 quant_config: QuantizationConfig | None = None,

667 prefix: str = "",

668 use_data_parallel: bool = False,

669 + attn_backend_override: _Backend | None = None,

670 ):

671 super().__init__()

672

@@ -665,6 +675,7 @@ class Siglip2NavitModel(torch.nn.Module):

675 quant_config=quant_config,

676 prefix=f"{prefix}.vision_model",

677 use_data_parallel=use_data_parallel,

678 + attn_backend_override=attn_backend_override,

679 )

680

681 def forward(

vllm/model_executor/models/vision.py

+9 -1

@@ -78,10 +78,18 @@ def get_vision_encoder_info(hf_config: VisionLanguageConfig) -> VisionEncoderInf

78 raise NotImplementedError(msg)

81 -def get_vit_attn_backend(head_size: int, dtype: torch.dtype) -> _Backend:

81 +def get_vit_attn_backend(

82 + head_size: int,

83 + dtype: torch.dtype,

84 + *,

85 + attn_backend_override: _Backend | None = None,

86 +) -> _Backend:

87 """

88 Get the available attention backend for Vision Transformer.

89 """

90 + if attn_backend_override is not None:

91 + return attn_backend_override

92 +

93 # Lazy import to avoid circular dependency

94 from vllm.attention.selector import get_env_variable_attn_backend