[Attention][FlashInfer] Enable FP8 FlashInfer (TRTLLM) MLA decode

[MatthewBonanni]

Purpose

Enable FP8 kv cache for FLASHINFER_MLA backend.

Test Plan

Correctness

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

Performance

VLLM_ATTENTION_BACKEND=FLASHINFER_MLA vllm bench throughput --model=deepseek-ai/DeepSeek-V2-Lite-Chat --dataset-name=random --input-len=8192 --output-len=1024 --num-prompts=100 --kv-cache-dtype=<dtype>

Test Result

Correctness

with <dtype>="auto":

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

with <dtype>="fp8":

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6384|±  |0.0132|
|     |       |strict-match    |     5|exact_match|↑  |0.6262|±  |0.0133|

Performance

with <dtype>="auto": Throughput: 3.22 requests/s, 29668.09 total tokens/s, 3296.62 output tokens/s
with <dtype>="fp8": Throughput: 3.55 requests/s, 32757.46 total tokens/s, 3639.90 output tokens/s

---

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.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[mgoin]

Could you use the _float version of these?

        # We also keep q/k/v_scale on host (cpu) memory for attention
        # backends that require the scales to be on host instead of on device.
        # e.g. Flashinfer
        self._q_scale_float = 1.0
        self._k_scale_float = 1.0
        self._v_scale_float = 1.0

Also could this be calculated ahead of time rather than each forward pass?

[MatthewBonanni]

I can definitely switch to the float ones. It's possible that each layer could have a different scale though, right?

[mgoin]

Yes, definitely. Sorry I wasn't looking carefully at the difference between layer and self here, we can keep the local computation. Was just trying to avoid CPU ops if possible

[MatthewBonanni]

No problem, thanks for the review!