[TPU][V1] MHA Pallas backend

[NickLucche]

The torch sdpa F.scaled_dot_product_attention impl is particularly inefficient on TPU.
It's both slow and also causes issues when XLA compiling as reported here #15051, which is a major obstacle in multimodal adoption on TPU.

This PR adds a Pallas branch to the vanilla MHA module used in most ViTs.

Some numbers:

Pre PR (TPUv6):

VLLM_USE_V1=1 vllm serve llava-hf/llava-1.5-7b-hf --max-model-len 2512 --max-num-seqs 16 --max-num-batched-tokens 512 --chat-template examples/template_llava.jinja 

============ Serving Benchmark Result ============
Successful requests:                     100       
Benchmark duration (s):                  144.89    
Total input tokens:                      1388      
Total generated tokens:                  7859      
Request throughput (req/s):              0.69      
Output token throughput (tok/s):         54.24     
Total Token throughput (tok/s):          63.82     
---------------Time to First Token----------------
Mean TTFT (ms):                          99840.40  
Median TTFT (ms):                        104402.54 
P99 TTFT (ms):                           141547.78 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          215.15    
Median TPOT (ms):                        133.66    
P99 TPOT (ms):                           930.35    
---------------Inter-token Latency----------------
Mean ITL (ms):                           206.93    
Median ITL (ms):                         42.99     
P99 ITL (ms):                            3931.00   
==================================================

Post PR (TPUv6):

VLLM_USE_V1=1 vllm serve llava-hf/llava-1.5-7b-hf --max-model-len 2512 --max-num-seqs 16 --max-num-batched-tokens 512 --chat-template examples/template_llava.jinja 

============ Serving Benchmark Result ============
Successful requests:                     100       
Benchmark duration (s):                  39.41     
Total input tokens:                      1388      
Total generated tokens:                  7749      
Request throughput (req/s):              2.54      
Output token throughput (tok/s):         196.60    
Total Token throughput (tok/s):          231.82    
---------------Time to First Token----------------
Mean TTFT (ms):                          20138.36  
Median TTFT (ms):                        21187.89  
P99 TTFT (ms):                           35919.35  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          67.93     
Median TPOT (ms):                        61.32     
P99 TPOT (ms):                           135.02    
---------------Inter-token Latency----------------
Mean ITL (ms):                           66.88     
Median ITL (ms):                         38.68     
P99 ITL (ms):                            151.16    
==================================================

Throughput boost is significant, almost 4x in this instance.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[mgoin]

LGTM nice work! FYI @Isotr0py

[mgoin]

Shouldn't this also check for the default behavior env bar as well?

[NickLucche]

I remember I discussed it with @robertgshaw2-redhat 🤔
Maybe I can check if the env was set rather than its value?

[Isotr0py]

The ViT part LGTM too!

[NickLucche]

FYI CI errors look unrelated, it's the same thing we have on main https://buildkite.com/vllm/ci/builds/16011#0195b8a1-d554-46b6-9929-798657dcabbf

[Isotr0py]

You can merge from main branch, the failing kennel tests have been fixed by #15282

[NickLucche]

Thanks! Didn't see the PR mb