[Kernel] Add FP8 support with FlashMLA backend

[MatthewBonanni]

Essential Elements of an Effective PR Description Checklist

• 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
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

Enable FP8 KV cache with MLA

Test Plan

Correctness

pytest tests/kernels/attention/test_flashmla.py
pytest tests/kernels/attention/test_cache.py::test_gather_and_maybe_dequant_cache_mla

Accuracy

With kv_cache_type = "auto":
VLLM_ATTENTION_BACKEND=FLASHMLA lm_eval --model vllm --model_args '{"pretrained": "deepseek-ai/DeepSeek-V2-Lite-Chat", "trust_remote_code": true, "kv_cache_dtype": "auto"}' --tasks gsm8k --batch_size auto

With kv_cache_type = "fp8":
VLLM_ATTENTION_BACKEND=FLASHMLA lm_eval --model vllm --model_args '{"pretrained": "deepseek-ai/DeepSeek-V2-Lite-Chat", "trust_remote_code": true, "kv_cache_dtype": "fp8"}' --tasks gsm8k --batch_size auto

Performance

With kv_cache_type = "auto":
VLLM_ATTENTION_BACKEND=FLASHMLA chg run --gpus 1 -- vllm bench throughput --model=deepseek-ai/DeepSeek-V2-Lite-Chat --dataset-name=random --input-len=512 --output-len=512 --num-prompts=10000 --kv-cache-dtype=auto

With kv_cache_type = "fp8":
VLLM_ATTENTION_BACKEND=FLASHMLA chg run --gpus 1 -- vllm bench throughput --model=deepseek-ai/DeepSeek-V2-Lite-Chat --dataset-name=random --input-len=512 --output-len=512 --num-prompts=10000 --kv-cache-dtype=fp8

Test Result

Correctness

Tests pass

Accuracy

With kv_cache_type = "auto":

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6763|±  |0.0129|
|     |       |strict-match    |     5|exact_match|↑  |0.6634|±  |0.0130|

With kv_cache_type = "fp8":

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6543|±  |0.0131|
|     |       |strict-match    |     5|exact_match|↑  |0.6452|±  |0.0132|

Performance

On 1x H100:

Here are the results for 512/512:
--kv-cache-dtype=auto: Throughput: 26.37 requests/s, 26975.81 total tokens/s, 13499.72 output tokens/s
--kv-cache-dtype=fp8: Throughput: 27.99 requests/s, 28635.39 total tokens/s, 14330.23 output tokens/s

Here are the results for 8192/1024:
--kv-cache-dtype=auto: Throughput: 2.40 requests/s, 22143.47 total tokens/s, 2460.56 output tokens/s
--kv-cache-dtype=fp8: Throughput: 3.25 requests/s, 29971.81 total tokens/s, 3330.44 output tokens/s

(Optional) Documentation Update

[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, @MatthewBonanni.

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

[gemini-code-assist]

Code Review

This pull request adds FP8 support for the FlashMLA attention backend. The changes include updating the FlashMLA sources in CMake, modifying the FlashMLA Python ops to handle FP8 scaling factors, and adding FP8 data types to the FlashMLA tests. The logic for enabling FP8 in the attention backend seems correct and is consistently applied across the codebase. I have one suggestion to improve code clarity in the test suite.

[mergify]

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

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

[MatthewBonanni]

CI build fails because FlashMLA git tag hasn't been updated yet, will update prior to merge

EDIT: Done

[alexm-redhat]

@MatthewBonanni this is great work, FP8 MLA is an important feature to have. In general, the PR LGTM, left some minor comments. Some questions about perf:

1. The benchmark you did used prompt/decode 512/512, would be interesting to see perf numbers for 8192/1024 prompt/decode. It should exercise the prefill chunking more aggressively.
2. Would be interesting to profile the MLA code to see how much improvement FP8 MLA related code is actually making vs the other MLA code. Maybe there is some low hanging fruit to improve (the code that surrounds the new FP8 MLA)
3. Did you had a chance to compare perf of F8 MLA vs plain CUTLASS MLA on B200 (fp16)? Just to know that FP8 helps vs CUTLASS MLA.

[LucasWilkinson]

Really great work! Overall seems very close to landable; left a couple comments

[MatthewBonanni]

@MatthewBonanni this is great work, FP8 MLA is an important feature to have. In general, the PR LGTM, left some minor comments. Some questions about perf:

1. The benchmark you did used prompt/decode 512/512, would be interesting to see perf numbers for 8192/1024 prompt/decode. It should exercise the prefill chunking more aggressively.
2. Would be interesting to profile the MLA code to see how much improvement FP8 MLA related code is actually making vs the other MLA code. Maybe there is some low hanging fruit to improve (the code that surrounds the new FP8 MLA)
3. Did you had a chance to compare perf of F8 MLA vs plain CUTLASS MLA on B200 (fp16)? Just to know that FP8 helps vs CUTLASS MLA.

@alexm-redhat Thanks for your review!

1. Here are the results for 8192/1024:
   --kv-cache-dtype=auto: Throughput: 2.40 requests/s, 22143.47 total tokens/s, 2460.56 output tokens/s
   --kv-cache-dtype=fp8: Throughput: 3.25 requests/s, 29971.81 total tokens/s, 3330.44 output tokens/s
2. I'll work on investigating this further! Right now this is just a first pass to get FP8 enabled
3. FlashMLA doesn't support Blackwell yet so I don't think that I can make this comparison

[LucasWilkinson]

LGTM thanks!