[WIP][Kernel] Flashinfer MLA support

[LucasWilkinson]

CUDA graphs still broken

Based off: #12639

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

🚀

[yzh119]

Hi @LucasWilkinson , thank you for the great work!

Would you mind trying the https://github.com/flashinfer-ai/flashinfer/releases/tag/v0.2.2 where we fixed some potential bugs for cudagraph and also add a new template for hopper (flashinfer-ai/flashinfer#887).

[mergify]

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

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

[LucasWilkinson]

Hi @LucasWilkinson , thank you for the great work!

Would you mind trying the https://github.com/flashinfer-ai/flashinfer/releases/tag/v0.2.2 where we fixed some potential bugs for cudagraph and also add a new template for hopper (flashinfer-ai/flashinfer#887).

Thank you for your amazing work, big fan of flashinfer! Very quickly tested it and im still having coda-graph issues but ill definitely play with it more tmrw and keep you up to date here 👍

[mergify]

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

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

[mergify]

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

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

[LucasWilkinson]

V1 benchmarking results 8xH200:

FLASHINFER

Data Preview:
  backend  input_tokens  output_tokens  output_toks/s     req/s  median_itl_ms  median_ttft_ms
3    vllm          1000           1000    1248.113964  1.248114      35.447731      921.937089
2    vllm          5000           1000    1001.868678  1.001869      39.285003     4376.742943
4    vllm         10000           1000     487.413100  0.487413      52.301841     7087.848108
1    vllm         32000           1000     153.426390  0.153426      54.572381   117342.499398


FLASHMLA

Data Preview:
  backend  input_tokens  output_tokens  output_toks/s     req/s  median_itl_ms  median_ttft_ms
3    vllm          1000           1000    1316.813496  1.316813      32.389333     1344.864228
2    vllm          5000           1000    1104.185700  1.104186      33.006376     4300.458221
4    vllm         10000           1000     625.310863  0.625311      34.492233     6999.176600
1    vllm         32000           1000     179.463985  0.179464      38.213399   101831.181545

[DefTruth]

@LucasWilkinson Hi~ Can you update this PR and resolve the conflicts? I can help test flashinfer_mla on an Ada device (which cannot use FlashMLA). I hope flashinfer_mla can boost the decode throughput in our product environment. I want to do the same thing, but I've found that there is already a PR integrating flashinfer. Does flashinfer_mla support Ada?

[hmellor]

Closing as stale