[Kernel] Flashinfer MLA (trtllm-gen) decode kernel integration

[hjjq]

E2E Benchmark:

VLLM_ATTENTION_BACKEND=FLASHINFER_MLA vllm serve deepseek-ai/DeepSeek-R1 --block-size 32 --tensor-parallel-size 8 --max-num-seqs 512

python benchmarks/benchmark_serving.py  --model deepseek-ai/DeepSeek-R1 --dataset-name random --ignore-eos --num-prompts 1024 --max-concurrency 512 --random-input-len 1024 --random-output-len 2048

TRTLLM-gen (via Flashinfer) MLA:

============ Serving Benchmark Result ============
Successful requests:                     1024      
Maximum request concurrency:             512       
Benchmark duration (s):                  434.37    
Total input tokens:                      1045957   
Total generated tokens:                  2097152   
Request throughput (req/s):              2.36      
Output token throughput (tok/s):         4827.98   
Total Token throughput (tok/s):          7235.94   
---------------Time to First Token----------------
Mean TTFT (ms):                          13770.84  
Median TTFT (ms):                        18015.33  
P99 TTFT (ms):                           26645.67  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          87.76     
Median TPOT (ms):                        84.81     
P99 TPOT (ms):                           120.83    
---------------Inter-token Latency----------------
Mean ITL (ms):                           87.79     
Median ITL (ms):                         74.03     
P99 ITL (ms):                            413.67    
=====================================================

cutlass MLA (VLLM_ATTENTION_BACKEND=CUTLASS_MLA):

============ Serving Benchmark Result ============
Successful requests:                     1024      
Maximum request concurrency:             512       
Benchmark duration (s):                  455.36    
Total input tokens:                      1045957   
Total generated tokens:                  2097152   
Request throughput (req/s):              2.25      
Output token throughput (tok/s):         4605.44   
Total Token throughput (tok/s):          6902.40   
---------------Time to First Token----------------
Mean TTFT (ms):                          13788.19  
Median TTFT (ms):                        18140.64  
P99 TTFT (ms):                           26975.27  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          91.38     
Median TPOT (ms):                        88.12     
P99 TPOT (ms):                           126.08    
---------------Inter-token Latency----------------
Mean ITL (ms):                           91.41     
Median ITL (ms):                         77.04     
P99 ITL (ms):                            417.73    
==================================================

Kernel-only microbenchmark:

bs=512:
trtllm-gen MLA kernel latency: 0.501ms
cutlass MLA kernel latency:  0.626ms
bs=1:
trtllm-gen: 0.023ms
cutlass: 0.025ms

Under high-concurrency, trtllm-gen is 25% faster than cutlass. But E2E isn't bottlenecked by MLA, therefore E2E speedup is not as significant.

Accuracy tests:

	Cutlass MLA	trtllm-gen MLA
gsm8k	0.95	0.97
gpqa	0.47	0.47
mmlu	0.85	0.85

commands:
gsm8k:

lm_eval --model local-chat-completions --tasks gsm8k --limit 100 --model_args model=deepseek-ai/DeepSeek-R1,base_url=http://0.0.0.0:8000/v1/chat/completions,max_gen_toks=16384,num_concurrent=64 --batch_size auto --fewshot_as_multiturn --apply_chat_template

gpqa:

lm_eval --model local-completions --tasks gpqa_extended_zeroshot --model_args model=deepseek-ai/DeepSeek-R1,base_url=http://0.0.0.0:8000/v1/completions,max_gen_toks=65535,num_concurrent=64,trust_remote_code=true,max_length=128000 --batch_size auto

mmlu:

lm_eval --model local-completions --tasks mmlu --model_args model=deepseek-ai/DeepSeek-R1,base_url=http://0.0.0.0:8000/v1/completions,max_gen_toks=65535,num_concurrent=64,trust_remote_code=true,max_length=128000  --batch_size auto

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

🚀

[gemini-code-assist]

Code Review

This pull request introduces a new FlashInfer MLA (Multi-LoRA Attention) decode kernel for the vLLM V1 engine. There are critical inconsistencies between the backend implementation and its corresponding test file regarding the shape of the kv_cache tensor and the function signature of the FlashInfer kernel being called. These issues need to be addressed.

[mergify]

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

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

[kushanam]

@LucasWilkinson this is ready for review. Thanks.

[mgoin]

Moving to draft while investigating performance

[mergify]

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

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

[kushanam]

@farazkh80 for review.

[MatthewBonanni]

@hjjq You might need to implement a quick fix similar to #24453 because of the changes of #23734

[kushanam]

@LucasWilkinson could you please review?

[mergify]

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

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

[hjjq]

Thanks @MatthewBonanni and @LucasWilkinson, I've made the changes and verified correctness with gsm8k. I've also reverted @mgoin 's changes so that test_flex_attention and test_triton_flash_attention are not run anymore. The current failures seem to be unrelated and exist on main.

[benchislett]

CI on main that confirms these test failures are not caused by this PR:
https://buildkite.com/vllm/ci/builds/30011/steps/canvas?sid=01992fa8-939a-4ba4-8585-1e5074bcd92d

[MatthewBonanni]

@hjjq @LucasWilkinson the remaining CI failures are all failed on main

[mergify]

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

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

[LucasWilkinson]

LGTM! Thanks!

[mgoin]

Thanks!