[V1] Fixed fp8 kv-cache issues for the FlashInfer attention backend #20746

chenyang78 · 2025-07-10T08:36:12Z

This PR fixed fp8 kv-cache issues for the FlashInfer attn backend.

Along with #17005, got reasonable eval results on B200:

$ VLLM_USE_V1=1 VLLM_ATTENTION_BACKEND=FLASHINFER lm_eval --model vllm --model_args pretrained=meta-llama/Llama-3.1-8B-Instruct,kv_cache_dtype=fp8 --trust_remote_code --tasks gsm8k --num_fewshot 5 --batch_size auto
...
vllm (pretrained=meta-llama/Llama-3.1-8B-Instruct,kv_cache_dtype=fp8,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.7779|±  |0.0114|
|     |       |strict-match    |     5|exact_match|↑  |0.7582|±  |0.0118|

compared with bf16 kv-cache

$ VLLM_USE_V1=1 VLLM_ATTENTION_BACKEND=FLASHINFER lm_eval --model vllm --model_args pretrained=meta-llama/Llama-3.1-8B-Instruct --trust_remote_code --tasks gsm8k --num_fewshot 5 --batch_size auto
...
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.7756|±  |0.0115|
|     |       |strict-match    |     5|exact_match|↑  |0.7498|±  |0.0119|

Tags:

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

Test Plan

Test Result

(Optional) Documentation Update

This PR fixed fp8 kv-cache issues for the FlashInfer attn backend. Along with vllm-project#17005, got reasonable eval results on B200: ``` $ VLLM_USE_V1=1 VLLM_ATTENTION_BACKEND=FLASHINFER lm_eval --model vllm --model_args pretrained=meta-llama/Llama-3.1-8B-Instruct,kv_cache_dtype=fp8 --trust_remote_code --tasks gsm8k --num_fewshot 5 --batch_size auto ... vllm (pretrained=meta-llama/Llama-3.1-8B-Instruct,kv_cache_dtype=fp8,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto |Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| |gsm8k| 3|flexible-extract| 5|exact_match|↑ |0.7779|± |0.0114| | | |strict-match | 5|exact_match|↑ |0.7582|± |0.0118| ``` compared with bf16 kv-cache ``` $ VLLM_USE_V1=1 VLLM_ATTENTION_BACKEND=FLASHINFER lm_eval --model vllm --model_args pretrained=meta-llama/Llama-3.1-8B-Instruct --trust_remote_code --tasks gsm8k --num_fewshot 5 --batch_size auto ... |Tasks|Version| Filter |n-shot| Metric | |Value | |Stderr| |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:| |gsm8k| 3|flexible-extract| 5|exact_match|↑ |0.7756|± |0.0115| | | |strict-match | 5|exact_match|↑ |0.7498|± |0.0119| ``` Tags: Signed-off-by: Yang Chen <[email protected]>

github-actions · 2025-07-10T08:36:20Z

👋 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

Summary of Changes

Hello @chenyang78, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves critical issues related to the use of FP8 KV-cache with the FlashInfer attention backend. The changes ensure that FP8 data types are correctly identified, handled, and cast throughout the KV-cache management and attention computation processes, leading to stable and accurate model evaluation results when using FP8 precision.

Highlights

FP8 Type Mapping Utility: Introduced a new static method get_fp8_dtype_for_flashinfer within FlashInferBackend to accurately map string representations of FP8 data types (e.g., "fp8", "fp8_e4m3", "fp8_e5m2") to their corresponding torch.dtype objects, ensuring correct type resolution for FlashInfer.
Correct KV-Cache Type Handling in Build Process: Modified the build method within FlashInferBackend to ensure the correct torch.dtype for FP8 KV-cache is passed to FlashInferMetadata. This change addresses an issue where the gpu_model_runner might incorrectly convert fp8 to uint8, preserving the intended FP8 precision.
Explicit KV-Cache Type Casting in Forward Pass: Added logic to the forward method of FlashInferBackend to explicitly view the KV-cache tensor to the correct torch.float8_e4m3fn or torch.float8_e5m2 dtype when FP8 KV-cache is enabled. This ensures that FlashInfer receives the KV-cache data in the precise FP8 format it expects for computation.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This PR fixes fp8 kv-cache issues for the FlashInfer attention backend. I suggested storing the fp8 dtype in self to avoid re-deriving it, improving efficiency and consistency.

gemini-code-assist · 2025-07-10T08:37:49Z

vllm/v1/attention/backends/flashinfer.py

+        if cache_dtype.startswith("fp8"):
+            kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
+                cache_dtype)


The torch.dtype for fp8 is derived using FlashInferBackend.get_fp8_dtype_for_flashinfer.[^1] Consider storing this dtype in self during initialization to avoid re-deriving it on each call to build. This can improve efficiency.

self.kv_cache_dtype_fp8 = FlashInferBackend.get_fp8_dtype_for_flashinfer( cache_dtype)

gemini-code-assist · 2025-07-10T08:37:50Z

vllm/v1/attention/backends/flashinfer.py

            head_dim=self.kv_cache_spec.head_size,
            page_size=page_size,
-            data_type=self.kv_cache_spec.dtype,
+            data_type=kv_cache_dtype,


The torch.dtype for fp8 is already computed and stored in self.kv_cache_dtype_fp8. Instead of using kv_cache_dtype here, use self.kv_cache_dtype_fp8 to ensure consistency and avoid potential errors if kv_cache_dtype is modified elsewhere.

data_type=self.kv_cache_dtype_fp8,

gemini-code-assist · 2025-07-10T08:37:50Z

vllm/v1/attention/backends/flashinfer.py

+            torch_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
+                self.kv_cache_dtype)
+            kv_cache = kv_cache.view(torch_dtype)


The torch.dtype for fp8 is derived using FlashInferBackend.get_fp8_dtype_for_flashinfer.[^1] Consider storing this dtype in self during initialization to avoid re-deriving it here. This can improve efficiency.

kv_cache = kv_cache.view(self.kv_cache_dtype_fp8)

mgoin · 2025-07-10T18:47:36Z

Thanks for finding this fix! Pavani seemed to uncover the same in her TRTLLM decode integration PR #19825, and since that is high priority I think we'll try to merge it in there today

chenyang78 · 2025-07-10T19:03:48Z

Thanks for finding this fix! Pavani seemed to uncover the same in her TRTLLM decode integration PR #19825, and since that is high priority I think we'll try to merge it in there today

Sounds good. Thanks for the heads-up.

chenyang78 · 2025-07-10T20:17:45Z

Closing this one per discussion.

chenyang78 requested review from WoosukKwon, alexm-redhat, comaniac, njhill, robertgshaw2-redhat and ywang96 as code owners July 10, 2025 08:36

gemini-code-assist bot reviewed Jul 10, 2025

View reviewed changes

mergify bot added the v1 label Jul 10, 2025

gemini-code-assist bot reviewed Jul 10, 2025

View reviewed changes

chenyang78 mentioned this pull request Jul 10, 2025

Enable FlashInfer V1 FP8 kv cache #17005

Closed

pavanimajety mentioned this pull request Jul 10, 2025

[Core] Add Flashinfer TRTLLM Backend for Flashinfer decode path (SM100). #19825

Merged

3 tasks

chenyang78 closed this Jul 10, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

[V1] Fixed fp8 kv-cache issues for the FlashInfer attention backend #20746

[V1] Fixed fp8 kv-cache issues for the FlashInfer attention backend #20746

Uh oh!

chenyang78 commented Jul 10, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Jul 10, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jul 10, 2025

Uh oh!

gemini-code-assist bot Jul 10, 2025

Uh oh!

gemini-code-assist bot Jul 10, 2025

Uh oh!

mgoin commented Jul 10, 2025

Uh oh!

chenyang78 commented Jul 10, 2025

Uh oh!

chenyang78 commented Jul 10, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

[V1] Fixed fp8 kv-cache issues for the FlashInfer attention backend #20746

[V1] Fixed fp8 kv-cache issues for the FlashInfer attention backend #20746

Uh oh!

Conversation

chenyang78 commented Jul 10, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

github-actions bot commented Jul 10, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 10, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 10, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 10, 2025

Choose a reason for hiding this comment

Uh oh!

mgoin commented Jul 10, 2025

Uh oh!

chenyang78 commented Jul 10, 2025

Uh oh!

chenyang78 commented Jul 10, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

chenyang78 commented Jul 10, 2025 •

edited by github-actions bot

Loading