PR #41613: [On hold] feat(kernels): Migrate mm_encoder_attn to vLLM IR

• Repository: vllm-project/vllm
• Author: harshaljanjani
• State: open | merged: False
• Link: https://github.com/vllm-project/vllm/pull/41613

Description (problem / solution / changelog)

Status

On hold on IR workspace support (per maintainer feedback). Will rebase and complete once that lands.

Purpose

→ This PR migrates the multimodal encoder attention (MMEncoderAttention) to the new vLLM IR framework, as discussed in https://github.com/vllm-project/vllm/issues/32676 → Followed existing IR ops implementation patterns and the spec detailed in https://github.com/vllm-project/vllm/issues/32358 → Added a test file following the IR pattern, and tested inference in native, triton, and default flash_attn modes. Moved all hardware-specific routing in MMEncoderAttention into _call_ir_op and preserved the FlashInfer path to handle FP8 caching.

Test Plan

<ins>A > Existing MHA Tests (Regression)</ins> → The following test harness was executed to check that no regressions were introduced in the existing MHA and FlashInfer paths: tests/kernels/attention/test_mha_attn.py <ins>B > New IR Kernel Tests (Added)</ins> → The following test harness was added in this PR and executed to check the functional robustness of the new IR op across various data types (FP16, BF16), batch sizes, and sequence lengths: tests/kernels/ir/test_mm_encoder_attn.py <ins>C > IR Meta-Tests (Integration)</ins> → The following test harness was added in this PR and executed to check that the new op is correctly registered in the IR framework: tests/kernels/ir/test_ir_ops.py <ins>D > E2E Sanity (Inference)</ins> → Wrote a sanity check script and executed inference to check for E2E sanity across all modes: https://gist.github.com/harshaljanjani/432d9fb7aa0bfdc550a63774872ea073

Test Results

<ins>A, B and C:</ins><br> <img alt="Untitled design (3)" src="https://github.com/user-attachments/assets/985b7e78-5a92-4549-99af-c4b1bcb172ad" />

<ins>D:</ins><br> → <ins>Backend: flash_attn</ins><br> <img alt="2-1" src="https://github.com/user-attachments/assets/853e7879-2b65-4aea-ba5d-c37662191594" />

→ <ins>Backend: triton</ins><br> <img alt="2-2" src="https://github.com/user-attachments/assets/57951cb1-2646-48dc-bc2c-1c3858e94d6f" />

→ <ins>Backend: native</ins><br> <img alt="2-3" src="https://github.com/user-attachments/assets/1171fbf7-bdfd-4f33-a43b-9e7fea868118" />

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Was this discussed/approved via a Github issue?
• 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

Changed files

• tests/kernels/ir/test_mm_encoder_attn.py (added, +174/-0)
• vllm/config/kernel.py (modified, +3/-0)
• vllm/ir/ops/__init__.py (modified, +2/-1)
• vllm/ir/ops/attention.py (added, +80/-0)
• vllm/kernels/__init__.py (modified, +2/-2)
• vllm/kernels/attention_ops.py (added, +136/-0)
• vllm/model_executor/layers/attention/mm_encoder_attention.py (modified, +25/-141)
• vllm/platforms/cuda.py (modified, +6/-1)
• vllm/platforms/rocm.py (modified, +6/-1)
• vllm/platforms/xpu.py (modified, +3/-1)