[RFC]: Reorganizing ViT Abstraction and Attention Selection Logic

[tjtanaa]

Motivation.

This RFC is aimed to address the following issues:

1. The ViT right now is still pretty coupled with Text backbone attention. This RFC will further the effort to decouple the ViT and the text backbone attention.

2. Another pain point is that the overriding of the ViT logic is scattered all around the places. We should avoid doing ViT logic overriding in model definition classes. The platform class should define the logic of what ViT is supported and how it should be overwritten.

• The above logic is applied to general use case. As of the time of this RFC is proposed. This single logic is applied to all of the VL model below:
  • vllm/model_executor/models/qwen2_5_vl.py (this is shared by Qwen2.5 VL and Qwen3 VL)
  • vllm/model_executor/models/dots_ocr.py
  • vllm/model_executor/models/ernie45_vl.py
  • vllm/model_executor/models/glm4_1v.py
  • vllm/model_executor/models/qwen2_vl.py
  • vllm/model_executor/models/siglip2navit.py
• If the model only supports a specific type of attention, the ViT overriding logic will be implemented explicitly in the model definition file (model.py).

3. Since the introduction of torch.compile into the ViT, currently only starting with qwen vl model in PR [Misc][qwen2_5_vl][torch.compile] Enable supports_torch_compile on generic nn.Module and demonstrate speedup on Qwen Vision model #23207 , the AMD ViT Code path are broken. New approach will try to accommodate this new feature. torch.compile has brought a lot of performance improvement and we can now consider to replace triton kernels with pytorch native implementation as there are possibilities that torch.compile code is faster than custom triton kernel code.
   One of the proven case is [Bugfix][ROCm] Fix ViT rotary embeddings for torch.compile compatibility on ROCm #27748 . The torch.compile rotary_embedding is faster than triton rotary_embedding from flash_attn .

4. Ensure ViT changes take into account the other model definition files model.py files, as current changes only involves qwen2_5_vl.py which potentially affecting all other models.py files.

• vllm/model_executor/models/dots_ocr.py
• vllm/model_executor/models/ernie45_vl.py
• vllm/model_executor/models/glm4_1v.py
• vllm/model_executor/models/qwen2_vl.py
• vllm/model_executor/models/siglip2navit.py

Proposed Change.

NOTE: More changes to the details will come in while I am writing up a version with all these changes.

For a first quick reorganization of the ViT Attention while retaining current use of --mm-encoder-atttention-backend interface, introduced in the PR is #27061 , and a bugfix PR #27124 .

1. First, we should shrink down the https://github.com/vllm-project/vllm/pull/27061/files#r2443909604 the _Backend by introducing another _MHA_Backend registry.

class _MHA_Backend(enum.Enum):
    VLLM_FLASH_ATTN = enum.auto()  # CUDA-only
    FLASH_ATTN = enum.auto()  # CUDA/ROCm
    XFORMERS = enum.auto()  # CUDA
    ROCM_AITER_FA = enum.auto()  # ROCM-only
    TORCH_SDPA = enum.auto()  # CUDA/ROCm/TPU/XPU/CPU
    PALLAS = enum.auto() # TPU only

2. Make sure that the ViT attention is a platform specific. We should determine platform interface. We also perform override in the platform interface. We should avoid doing that in the model.py files.

• get_vit_attn_backend in the platform interface has to be able to access the --mm-encoder-attn-backend.

• In the platform interface, we should only return _MHA_Backend, we should not return the functions. The functions should only be returned through maybe_get_vit_flash_attn_backend . If the model only supports a specific type of attention, the ViT overriding logic will be implemented explicitly in the model definition file (model.py).

class Platform:

    ...

    @classmethod
    def get_supported_vit_attn_backends(cls) -> list["_MHA_Backend"]:
        from vllm.attention.backends.registry import _MHA_Backend

        return [
            _MHA_Backend.TORCH_SDPA,
        ]

    @classmethod
    def get_vit_attn_backend(
        cls,
        head_size: int,
        dtype: torch.dtype,
        backend: Optional["_MHA_Backend"] = None,
    ) -> "_MHA_Backend":
        # ViT Attention should be checked and override
        # in the platform-specific implementation.
        # we should not override this in any other places,
        # like the model_executor/models/<model_name>.py

        # So the steps are:
        # 1. Check if the backend is None or not:
        #    a. If not, check if the backend is supported by the platform.
        #    b. If None, continue to the default selection logic.

        # Import _Backend here to avoid circular import.
        from vllm.attention.backends.registry import _MHA_Backend

        if backend is not None:
            assert backend in cls.get_supported_vit_attn_backends(), (
                f"Backend {backend} is not supported for vit attention"
                f"Supported backends are: {cls.get_supported_vit_attn_backends()}"
            )
            logger.info_once(f"Using backend {backend} for vit attention")
            return backend

        logger.info_once(
            f"Using default backend {_MHA_Backend.TORCH_SDPA} for vit attention"
        )
        return _MHA_Backend.TORCH_SDPA

• Honor --mm-encoder-attn-backend so that we can write unit tests to test all different backends. AMD Instinct GPU is able to test all backends. Radeon GPUs only are able to use the TORCH_SDPA code path.

• We need to deprecate this line https://github.com/vllm-project/vllm/blob/33a0ea5f3264b5b2f571b8a53357e10efcc94670/vllm/model_executor/models/vision.py#L96 it is using VLLM_ATTENTION_BACKEND which is for Text Backbone. The ViT should not use this environment variable.

• Added a logger.info_once so that users know which _MHA_Backend is selected in the end.

• Clean up cuda code path. Since vllm.vllm_flash_attn is just a wrapper for flash_attn library, on cuda, we always use vllm.vllm_flash_attn instead of flash_attn.

• vllm/vllm/attention/layer.py

  Lines 120 to 125 in ba33e88

  	elif current_platform.is_cuda():
  	if attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(
  	torch.get_default_dtype()
  	):
  	attn_backend = _Backend.FLASH_ATTN
  	use_upstream_fa = True

  .

3. Enable unit tests to test all different backends. Since there are large model sizes, we will check the VRAM size, if it is large enough, we run it. We provide such a unit test so that developers can run locally.

Feedback Period.

Changes

CC List.

@ywang96 @DarkLight1337 @Isotr0py

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[robertgshaw2-redhat]

Thanks for proposing this. Im strongly supportive of this

[tjtanaa]

Drafting for Phase 2:

1) Consolidate the ViT forward to use MHA Attention class

The following snippet is found in all of the model definition files listed below.

vllm/vllm/model_executor/models/qwen2_5_vl.py

Lines 402 to 466 in e675118

	def forward(
	self,
	x: torch.Tensor,
	cu_seqlens: torch.Tensor,
	rotary_pos_emb: torch.Tensor,
	max_seqlen: torch.Tensor, # Only used for Flash Attention
	seqlens: torch.Tensor, # Only used for xFormers
	) -> torch.Tensor:
	# [s, b, c] --> [s, b, head * 3 * head_dim]
	x, _ = self.qkv(x)
	# [s, b, 3 * head * head_dim] -> 3 * [s, b, head, head_dim]
	q, k, v = self.split_qkv(x)
	batch_size = q.shape[1]
	q, k, v = (einops.rearrange(x, "s b ... -> b s ...") for x in (q, k, v))
	if rotary_pos_emb is not None:
	# [2 * b, s, heads, head_dim]
	qk_concat = torch.cat([q, k], dim=0)
	qk_rotated = apply_rotary_pos_emb_vision(qk_concat, rotary_pos_emb)
	q, k = torch.chunk(qk_rotated, 2, dim=0)
	if self.is_flash_attn_backend:
	context_layer = vit_flash_attn_wrapper(
	q,
	k,
	v,
	cu_seqlens,
	max_seqlen,
	batch_size,
	self.attn_backend == _Backend.ROCM_AITER_FA,
	self.use_upstream_fa,
	)
	elif self.attn_backend == _Backend.TORCH_SDPA:
	# Execute attention entry by entry for speed & less VRAM.
	from vllm.platforms import current_platform
	# Never remove the next contiguous logic
	# Without it, hallucinations occur with the backend
	if current_platform.is_rocm():
	q = q.contiguous()
	k = k.contiguous()
	v = v.contiguous()
	outputs = []
	for i in range(1, len(cu_seqlens)):
	start_idx = cu_seqlens[i - 1]
	end_idx = cu_seqlens[i]
	q_i = q[:, start_idx:end_idx]
	k_i = k[:, start_idx:end_idx]
	v_i = v[:, start_idx:end_idx]
	q_i, k_i, v_i = (
	einops.rearrange(x, "b s h d -> b h s d") for x in [q_i, k_i, v_i]
	)
	output_i = F.scaled_dot_product_attention(q_i, k_i, v_i, dropout_p=0.0)
	output_i = einops.rearrange(output_i, "b h s d -> b s h d ")
	outputs.append(output_i)
	context_layer = torch.cat(outputs, dim=1)
	context_layer = einops.rearrange(
	context_layer, "b s h d -> s b (h d)"
	).contiguous()
	elif self.attn_backend == _Backend.XFORMERS:
	context_layer = vit_xformers_attn_wrapper(q, k, v, seqlens)
	output, _ = self.proj(context_layer)
	return output

• vllm/model_executor/models/qwen2_5_vl.py
• vllm/model_executor/models/dots_ocr.py
• vllm/model_executor/models/keye.py
• vllm/model_executor/models/siglip2navit.py
• vllm/model_executor/models/ernie45_vl.py
• vllm/model_executor/models/glm4_1v.py
• vllm/model_executor/models/qwen2_vl.py

2) Decide whether to unify the use of flash attention on CUDA platform to only use from vllm.vllm_flash_attn import flash_attn_varlen_func, deprecating from flash_attn import flash_attn_varlen_func.

This PR #24347 introduced the use of from vllm.vllm_flash_attn import flash_attn_varlen_func so that ViT uses FA3. This is because in the upstream flash_attn repository, the FA3 API is not the same as FA2 API. Moreover, the from flash_attn import flash_attn_varlen_func only imports the use of FA2 API. vLLM has maintained its own wrapper in repo https://github.com/vllm-project/flash-attention.git. vLLM has implemented auto-selection logic in the wrapper repo to use FA3 whenever possible, else use FA2. This could mean we can deprecate the from flash_attn import flash_attn_varlen_func and always use from vllm.vllm_flash_attn import flash_attn_varlen_func.

If we agree to do that, then we can deprecate _MHA_Backend.VLLM_FLASH_ATTN, and also manage all the flash_attn_varlen_func through this utility file vllm/attention/utils/fa_utils.py