[V1][Core] FlashInfer attention backend for V1

[aurickq]

Implement FlashInfer attention backend for V1.

Outstanding items that need feedback:

• Need a better way to abstract out the AttentionMetadata creation interface, currently it just handles it case-by-case in the GPU runner for each attention backend. Solved after rebase
• Is the use_cascade_attention logic still valid for FlashInfer? This was copied from the FlashAttention implementation.
• Cascade attention seems to have different outputs from normal attention. Sometimes this difference is significant, but still coherent. Is this an expected behavior of cascade attention? Seems to match now after fixing a bug

Current limitations:

• It appears we cannot set the kv cache dtype using the MultiLevelCascadeAttentionWrapper interface in FlashInfer. This means we only support cascade attention when the kv cache dtype is the same as the query dtype.
• Only tested on simple cases so far, need more comprehensive testing.

[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.

🚀

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aurickq.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[aarnphm]

Very much early on, but from the looks of things seems like a lot of code moved from v0.

tests would be awesome here as well 😃

[aurickq]

PR is ready. But there are some performance issues that seem to be due to using BatchPrefillWithPagedKVCacheWrapper for all prefill and decode tokens. Maybe this wrapper is not well-optimized for decoding sequences with short query lengths.

V1 FlashAttention:

Input token throughput (tok/s):          21448.32  
Output token throughput (tok/s):         2745.38

V1 FlashInfer (This PR):

Input token throughput (tok/s):          11372.04  
Output token throughput (tok/s):         1455.62

V0 FlashInfer:

Input token throughput (tok/s):          15157.36  
Output token throughput (tok/s):         1940.14

V0 FlashInfer with BatchDecodeWithPagedKVCacheWrapper replaced with BatchPrefillWithPagedKVCacheWrapper:

Input token throughput (tok/s):          9479.15   
Output token throughput (tok/s):         1213.33

There is a big performance degradation from using BatchPrefillWithPagedKVCacheWrapper for decode tokens in V0. In V1, we use a single BatchPrefillWithPagedKVCacheWrapper for mixed prefill/decode batches, which may be causing the poor performance.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aurickq.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork