[MM Encoder]: Migrate legacy ViT MultiHeadAttention to new MMEncoderAttention interface

[Isotr0py]

Purpose

• Following PR for [CustomOp][MM] Extract MMEncoderAttention as CustomOp and replace the backend of QwenVisionAttention with it. #30125
• Migrate MultiHeadAttention usage to new MMEncoderAttention

Test Plan

pytest - s-v tests/kernels/attention/test_attention.py

pytest -s -v tests/kernels/attention/test_mha_attn.py

Test Result

Test should pass

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request refactors the MultiHeadAttention class, used for multimodal encoders, into a new MMEncoderAttention class, moving its definition to vllm/attention/layers/mm_encoder_attention.py and removing it from vllm/attention/layer.py. All instances and imports of MultiHeadAttention across various model implementations (e.g., AIMV2, BLIP, CLIP, GLM4V, Idefics2, InternViT, MLlama4, MoLMo, SigLIP, Step3-VL, Whisper) and their respective test files have been updated to use MMEncoderAttention. The MMEncoderAttention class now directly integrates Flash Attention backend selection logic and removes a redundant reshape_qkv_to_3d method. However, a review comment points out a critical issue in the torch_sdpa_wrapper within vllm/attention/ops/vit_attn_wrappers.py, where torch.split is incorrectly applied on the sequence length dimension (dim=1) for batched inputs, assuming packed tensors. This causes a dimension mismatch and will lead to errors, with the reviewer suggesting to split along the batch dimension (dim=0) or use an alternative approach for handling batched inputs with SDPA in variable-length attention.

[gemini-code-assist]

The logic in torch_sdpa_wrapper for handling batched inputs appears to be incorrect. The torch.split on dim=1 (lines 104-106) assumes that the input tensors q, k, v are packed (i.e., shape (1, total_tokens, ...)), but they are passed as batched tensors of shape (batch_size, seq_len, ...). This will cause torch.split to fail because sum(lens) will not match q.shape[1].

Additionally, the new block at lines 109-112 for when cu_seqlens is None is also flawed. It computes lens based on cu_seqlens which is computed for a uniform batch, but torch.split will still fail for the same reason.

To fix this, you should probably split along the batch dimension (dim=0) or use a different approach to handle batched inputs with SDPA for varlen attention.